feat(nla): add Bridge-native Qwen activation verbalization and reconstruction

[janmenjayap]

---

Description

Add local Natural Language Autoencoder (NLA) inference infrastructure on
TransformerBridge: raw activation extraction, Activation Verbalizer (AV)
generation, Activation Reconstructor (AR) execution with a separate value head,
and directional reconstruction scoring.

Fixes #1884

The implementation remains private. Released-Qwen numerical compatibility has
not been established, the support registry is empty, and the class is not
exported from transformer_lens.tools.analysis. This PR does not currently
provide a verified public pretrained model name. The issue's complete acceptance
criteria remain outstanding as listed below.

---

Motivation

The authors' reference AV client uses SGLang embedding injection. This work
provides hook-native inference without operating a model server, consuming
released artifacts rather than implementing their training infrastructure.

---

Changes

• Typed immutable artifact references and schema-v2 sidecar validation.
• Tokenizer/template drift guards, mandatory injection scale, AR suffix/depth
  checks, normalization-mode checks, and safe bias-free value-head loading.
• Pinned raw target loading and detached single-token activation extraction.
• Loading provenance and explicit unverified-origin handling; known structured
  target, revision, layer, and basis mismatches are rejected before AV loading.
• Scoped AV embedding injection: cached prefill only and uncached reapplication
  at the original prompt position, with bounded generation and continuation-only
  explanation parsing.
• Truncated AR backbone execution under Bridge hooks, bypassing final
  normalization and avoiding execution of the unused LM head.
• Cosine and normalized MSE, independent role unloading/reloading, context
  cleanup, and rejection of concurrent operations on one instance.
• Unit, integration, reference-oracle, wheel-resource, and documentation coverage.
• In-memory tokenizer BOS handling fix exposed by composition tests.

Explicit runtime dependencies include PyYAML and safetensors; YAML typing stubs
are included in the dev group. No training/server dependencies are added. The
lockfile adds the stubs without changing existing package versions; uv also
regenerated resolution-marker ordering.

Excluded: training, edited steering, corpus execution/FVE, SGLang serving,
additional model pairs, multimodal inputs, quantization, arbitrary device maps,
and cross-model/layer transfer.

---

Correctness and ownership

The facade lives in transformer_lens/tools/analysis/_nla.py. Experimental use
requires an explicit private pair configuration. Only raw fp32/eager execution
is implemented; that restriction is not a released-artifact compatibility claim.

AV normalization matches the reference client's fp32 norm and division order
on ordinary fp32 inputs. A scale-safe fallback handles extreme inputs but is
not claimed as reference-client equivalence in that regime.

Canonical explanation text is the inner content of exactly one complete region.
Parse failures preserve the raw continuation and return text=None; malformed
responses are not silently passed to AR. AR uses its recorded raw prompt and
special-token convention, not a chat template.

The AR loader accepts exactly the omitted final norm and LM-head weights,
wraps an owned model, bypasses normalization, and executes only its backbone.
It does not rely on an absent block triggering early-stop behavior.

Results own detached vectors; returned targets belong to callers. Unloading
releases NLA-owned references, not externally held allocations. Loading provenance
is not continuous live-weight integrity. Reconstruction similarity is not proof
of human-readable faithfulness or causal validity. Normalized MSE is
2*(1-cosine), not raw vector MSE or dataset FVE.

---

Validation performed

After review fixes:

• NLA-focused unit/integration tests: 262 passed, 3 slow tests deselected.
• Full-project mypy: passed for 397 source files, with the declared YAML stubs.
• make format and make check-format: passed.
• Sphinx incremental build: succeeded with 5 warnings; no warning-free claim.
• Diff whitespace and changed-file hygiene checks: passed.
• Wheel packaging regression: passed with a fresh uv cache and actual extracted-
  wheel package/analysis imports in an isolated subprocess with network blocked.
  Build dependency resolution may access the network; package import may not.

Local tests cover cached/uncached hooks, caller-hook preservation, injected-logit
comparison with direct HF, a local truncated-Qwen checkpoint, raw AR output and
LM-head bypass, metadata/parser errors, metric identities, and lifecycle guards.
A forced deterministic continuation passes through real generation, canonical
parsing, and AR reconstruction with sequential model release checks. That tests
the software boundary, not explanation quality or released-model compatibility.

Reproduction

From the implementation worktree, using the project environment:

uv sync
set -a; source .env; set +a
uv run pytest tests/unit/tools/test_nla*.py \
  tests/integration/test_nla_bridge_execution.py \
  tests/integration/test_nla_facade_execution.py \
  tests/integration/test_nla_oracle_parity.py \
  tests/integration/test_nla_packaging.py -m "not slow" -q
make check-format
uv run mypy .
uv run build-docs

Local checks reused an existing uv-managed environment with --no-sync and
--with types-PyYAML for mypy. The commands above are normal reproduction
equivalents. Sphinx was invoked through docs.make_docs.build_docs to ensure
it built this worktree's sources when reusing an editable environment.

---

Outstanding validation and completion

Released-Qwen parity

All three slow checks were attempted and stopped at their storage prerequisite,
before model loading or tensor comparison. The full workflow required 38.5 GiB
of missing weights plus 2 GiB reserve, versus 11.8 GiB free at that run.

Use a larger HF cache volume, setting HF_HOME before starting Python, and
sufficient memory for one fp32 model plus loading/forward overhead. Execute
models sequentially:

set -a; source .env; set +a
uv run pytest tests/integration/test_nla_oracle_parity.py -m slow -s

The CPU fp32 tests initially require exact equality. Any nonzero drift must be
investigated and measured before introducing justified tolerances. Released-model
peak RAM/VRAM, timing, and calibrated tolerances have not been established.

Pinned candidate inputs:

• Target Qwen/Qwen2.5-7B-Instruct:
  a09a35458c702b33eeacc393d103063234e8bc28.
• AV kitft/nla-qwen2.5-7b-L20-av:
  b88469162777ae6553bc14208eb0cb579336f8f4.
• AR kitft/nla-qwen2.5-7b-L20-ar:
  e2c9e57eac213d37a31612087f645ab6332c1bb6.
• Inspected reference client:
  0577769b55ad4fdd96d159e983361b97fa4e7331.

These are test inputs, not registered support claims. The oracle executes local
transparent direct-HF code, not downloaded reference Python.

Full standard suite

make test-pr timed out after 200 seconds during unit tests at approximately 44%.
No complete standard-suite pass is claimed. Run it to completion and investigate
all failures before recording final acceptance results.

Public support

Public exports and the pinned Qwen registry entry remain withheld pending parity.
They still need actual from_pretrained and public-workflow tests, measured
resource documentation, and final-source format/type/full-suite/docs/build checks.
The completed inference infrastructure alone does not satisfy every criterion
of issue #1884.

---

Type of change

• New non-breaking inference infrastructure
• Documentation update
• Bug fix for in-memory tokenizer BOS handling
• Breaking change

---

Checklist

• Runtime, tests, and documentation added together
• Focused NLA tests pass after review fixes
• Format/type checks pass
• Wheel resources and actual package-import regression pass
• No training/server dependencies or remote-code execution added
• Issue [Proposal] Natural Language Autoencoders: explain and reconstruct Qwen activations through TransformerBridge #1884 linked
• Complete standard test tiers pass
• Released AV-logit and AR-vector parity pass
• Released target-to-AV-to-AR workflow completes
• Numerical tolerances and resource requirements measured
• Public API and verified Qwen registry entry delivered
• All final-source acceptance checks rerun
• Maintainer review completed

---

References

• Paper: https://transformer-circuits.pub/2026/nla/
• Authors' repository: https://github.com/kitft/natural_language_autoencoders

---

[jlarson4]

Closing this to wait for community interest on #1884. If you are interested in a feature like this, please give that issue a read!