[Proposal] Natural Language Autoencoders: explain and reconstruct Qwen activations through TransformerBridge

Start by reading the proposed API and correctness contracts in this issue, then inspect transformer_lens.tools.analysis and the existing caching, embedding-hook, and generation entry points. The scope depends on validating pinned Qwen AV/AR artifacts and reference numerical behavior before implementation. Done means usable inference with tests, documentation, and empirical compatibility evidence; the issue names no specific test files.

---

Proposal

Add a server-free Natural Language Autoencoder (NLA) inference tool to
transformer_lens.tools.analysis, initially supporting only the released
Qwen2.5-7B-Instruct/block-20 Activation Verbalizer (AV) and Activation
Reconstructor (AR) pair.

A researcher can extract one raw residual-stream activation, generate a
natural-language explanation, reconstruct a vector from that explanation, and
measure directional reconstruction similarity. The models run locally through
TransformerBridge hooks without a SGLang server.

This is a feasibility-gated feature proposal. The exact released AV and custom
AR packages must pass reference numerical checks before support is advertised.
The issue closes with usable inference, tests, documentation, and empirical
compatibility evidence, not with a compatibility report alone.

---

Motivation

The authors' Natural Language Autoencoders Produce Unsupervised Explanations of
LLM Activations releases trained models for an activation-to-text-to-activation
round trip. Its AV client uses SGLang embedding injection. Its AR uses a
truncated backbone and a separately released linear reconstruction head.

A Bridge-native implementation makes individual activation inspection practical
in ordinary TL scripts and notebooks, without operating a model server. TL
already provides caching, embedding hooks, and hook-aware generation. It does
not currently assemble these primitives into an NLA inference tool.

This consumes released inference artifacts, not the authors' RL/GRPO training
infrastructure. Reconstruction fidelity measures preservation of a vector; it
is not proof that the explanation is faithful to a human interpretation.

Related existing work is not a duplicate: TL's
Patchscopes notebook
demonstrates activation patching and generation, but not trained AV/AR
reconstruction. Jacobian Lens provides artifact and intervention precedent, not
this activation-to-text protocol. Refresh the duplicate-work search before filing.

---

User-facing scope

Supported initially

• Target: Qwen/Qwen2.5-7B-Instruct, raw HF residual basis.
• AV: kitft/nla-qwen2.5-7b-L20-av.
• AR: kitft/nla-qwen2.5-7b-L20-ar and its value_head.safetensors.
• Artifact-fixed zero-based block 20; dimension 3584.
• One unpadded example and one selected prompt token per call.
• Local stock HF architecture loading without remote code.
• A documented single-device execution configuration, established by the spike.
• Lazy AV/AR loading and independent unloading for sequential execution.

Immutable target, AV, AR, tokenizer, and reference revisions are recorded during
compatibility validation. Repository names above are artifact candidates, not
claims of completed verification. Do not ship mutable main revisions in the
support registry.

Non-goals

• Edited-explanation steering or causal validation.
• Corpus scheduling, persistence/resume, or dataset FVE.
• Gemma/Llama pairs, cross-model transfer, or arbitrary artifact layers.
• Training, SGLang serving, multimodal inputs, or quantization.
• Arbitrary device maps, automatic offload planning, concurrent calls on one
  model instance, or a general-purpose intervention framework.

---

Proposed API

from transformer_lens.tools.analysis import NaturalLanguageAutoencoder

with NaturalLanguageAutoencoder.from_pretrained(
    "qwen2.5-7b-instruct-L20", device="cuda"
) as nla:
    target = nla.load_target()
    activation = nla.extract(target, prompt="The auditor asked about reasoning.", position=-1)
    explanation = nla.verbalize(activation, max_new_tokens=512)
    if explanation.text is None:
        raise ValueError(f"Explanation parse failed: {explanation.parse_status}")
    reconstruction = nla.reconstruct(explanation.text)
    score = nla.score(activation, reconstruction.vector)
    print(explanation.text, score.cosine, score.normalized_mse)
    del target

load_target() is a convenience that loads the registry's pinned raw target and
records loading provenance. The returned target is caller-owned: NLA must not
keep a strong reference to it. extract() also accepts an externally loaded raw
Bridge target, subject to the provenance policy below.

The example above prioritizes convenience and may retain several models. The
sequential low-memory workflow is:

with NaturalLanguageAutoencoder.from_pretrained(
    "qwen2.5-7b-instruct-L20", device="cuda"
) as nla:
    target = nla.load_target()
    activation = nla.extract(
        target, prompt="The auditor asked about reasoning.", position=-1,
        output_device="cpu",
    )
    del target
    explanation = nla.verbalize(activation)
    nla.unload("av")
    if explanation.text is None:
        raise ValueError(f"Explanation parse failed: {explanation.parse_status}")
    reconstruction = nla.reconstruct(explanation.text)
    nla.unload("ar")
    score = nla.score(activation, reconstruction.vector)

Other caller references must also be released before target memory can be freed.
The implementation must not retain the target through a result, full cache,
closure, provenance record, or model ownership registry. Context exit unloads
NLA-owned AV/AR resources on both normal and exceptional exit; it does not own the
returned target. Unloading is idempotent and does not promise to free externally
referenced allocations. Document normal device memory cleanup separately.

Methods and defaults

• from_pretrained(name, device, dtype=...): resolve a verified bundled entry and
  validate local/downloaded metadata and tokenizer contracts; do not materialize
  full AV/AR weights until their operation is needed. Preserve existing HF cache
  behavior and an offline/local-files-only option. Metadata/tokenizer access is
  expected here, but importing analysis tools itself has no network side effect.
• load_target(): pinned raw target loading, eager attention, no remote code;
  returned model ownership belongs to the caller.
• extract(target, prompt, position=-1, output_device="cpu", ...): return a
  detached one-dimensional activation and provenance. CPU output is the safe
  default. No free layer override. Resolve positions against exact original
  tokens; reject batches, padding, empty token input, and out-of-range positions.
• verbalize(activation, max_new_tokens=512, temperature=1.0, use_past_kv_cache=True, ...): return parsed and raw continuation results.
  Expose explicit bounded sampling parameters, not arbitrary unchecked HF kwargs.
  Greedy decoding uses an explicit sampling flag. The 512-token default is a
  bounded convenience choice, not a claim about the paper's generation length.
• explain(target, prompt, position=-1, ...): extraction plus verbalization with
  the same guards; does not automatically unload a caller-owned target.
• reconstruct(text, output_device="cpu"): reconstruct canonical content only;
  return a detached vector and AR provenance.
• score(original, reconstructed): cosine and normalized directional MSE.
• unload(role), close(), context management: independent AV/AR owned-resource
  cleanup. An operation can lazily reload an unloaded role; concurrent use is
  unsupported rather than silently sharing mutable hook state.

Finalize the exact precision default from measured feasibility. Reference
comparison uses float32 where resources permit; only advertise additional
precision/device configurations after their checks pass. Do not silently
quantize or select an unvalidated precision to make a model fit.

---

Correctness contracts

Explanation text and AR tokenization

NLAExplanation.text is stripped inner content of exactly one complete
<explanation>...</explanation> region. Raw continuation text and generated
continuation token IDs are separate. Parse status distinguishes success from
missing, multiple, nested, incomplete, or empty/whitespace-only regions. Failures
set text=None; they do not substitute raw text.

reconstruct() accepts nonempty untagged explanation content. Reject explicit
explanation wrappers and malformed wrapper syntax instead of silently stripping
it. Do not claim semantic detection of all untagged full responses: callers
supplying strings are responsible for providing explanation content. Human edits
use the same untagged convention. Preserve authored interior whitespace; stripping
is the AV extraction convention, not a license to rewrite arbitrary AR input.

AR formats its recorded raw template once, with special tokens enabled as in the
reference. It is not a chat prompt. Exact fixed-input tokenization tests protect
against duplicate explanation tags, double BOS, and suffix changes.

Artifact metadata

Use safe YAML parsing and safetensors. Check recognized schema/kind and role,
artifact target identity, block and hidden dimension, reserved-character single
encoding and non-UNK identity, exactly one placeholder, valid position/neighbors,
finite positive actor injection norm, AR suffix, AR depth K+1, and value-head
state keys/shape with no bias. Raise clear errors rather than optional warnings.

The released Qwen actor's recorded injection norm is 150. Read it from the pinned
sidecar; do not derive or replace it with a generic default.

Check finite positive AR mse_scale and the supported direction-normalized mode.
Initially require sqrt(d) using a documented relative/absolute tolerance justified
by the pinned sidecar's serialized precision. Reject null/raw-MSE mode and other
unsupported scales. Do not use a loose tolerance to accept a different mode.

Target provenance and raw basis

Require raw Bridge execution without compatibility weight processing. Check
architecture, block availability, hidden dimension, supported precision, and
model identity/revision evidence. Reject known wrong identities/revisions and
explicitly unsupported processed/quantized targets.

load_target() supplies recorded pinned-loading provenance. This establishes
where weights were loaded from, not continuous proof that live weights were
never edited afterward. Do not call model-name/config checks weight verification
or hash a 7B model on every extraction.

For externally loaded targets, accept matching immutable loading evidence where
Bridge exposes it. If provenance cannot be established (including local files
without a matching manifest), fail by default. Permit an explicit
allow_unverified_target=True attestation for otherwise structurally compatible
raw targets, with source_verified=False in results and a warning. This does not
bypass known incompatibilities. Users are responsible for declaring later weight
edits; such runs are outside the verified artifact regime. Avoid strong target
references in provenance tracking.

Bare vector inputs to verbalize may be accepted with dimension/finite/nonzero
checks and explicit unverified-origin status. Structured extraction results
retain source identity, resolved position, tokenization/BOS policy, artifact ID,
and basis/precision provenance.

Activation and generation mechanics

Extract blocks.20.hook_out, the raw output of zero-based block 20, corresponding
to HF hidden_states[21]. Cache only the required hook, detach/copy the chosen
vector, and discard the full cache. Tokenize once without silent truncation;
record the target BOS policy and never auto-chat-template arbitrary target text.

AV renders its canonical chat template once with
apply_chat_template(..., tokenize=True, add_generation_prompt=True). Rescale
in a stable accumulation dtype, then cast to the actual embedding hook tensor
for the single-position overwrite. Only that position changes.

With KV caching, inject on prefill and leave later decode chunks alone. Without
caching, reapply at the original prompt position on every recomputed prefix.
Do not identify prefill solely through length greater than one. Reset state on
each call and remove hooks on success or failure. Existing caller hooks must not
be deleted by NLA cleanup.

Validate target length, AV prompt plus generation budget, and rendered AR length
against the supported context limits. Raise before the relevant full-model
forward on oversized input; never silently truncate. Reject invalid budgets and
sampling settings. EOS can stop generation early; incomplete/empty explanations
remain parse failures rather than successful round trips.

AR loading

The AR config contains 21 blocks and a separate bias-free Linear(d,d) head.
Its weight index omits final norm and LM-head weights. Do not use those initialized
output modules as part of reconstruction or suppress arbitrary missing keys.

Use a hook-preserving truncated execution path verified by the compatibility
spike. Prefer a backbone-only load wrapped for Bridge if supported; otherwise
validate the causal wrapper's expected missing keys and bypass the final output
path explicitly. Raw block-20 final-token output feeds the released head.

TL's stop-before-block mechanism cannot be assumed to bypass the output path
using stop_at_layer=21: the truncated stack has no block 21 to trigger it.
No new generic model loader or adapter refactor is justified without a demonstrated
need. Resolve the chosen loading strategy before implementing the AR public path.

Scoring

In a stable accumulation dtype, normalize both nonzero finite vectors to norm
sqrt(d). Per-coordinate MSE equals 2*(1-cosine) with no extra dimension
multiplier. Validate matching dimensions and handle harmless cosine range
overshoot through a documented numerical policy.

These are directional diagnostics, not raw vector MSE, dataset FVE, or scientific
validation of explanation content. Reconstruction can preserve a vector while
its text still misleads a reader.

---

Validation and acceptance

Normal CI

• Local synthetic metadata/tokenizer fixtures and tensor tests cover every guard,
  parsing outcome, normalization mode, context bound, and similarity identity.
• Existing cached tiny HF models exercise actual Bridge hook registration,
  original-position overwrite, cached/uncached generation, repeated calls,
  exception cleanup, caller-hook preservation, and applicable reconstruction
  mechanics. No mocked model/Hub loads and no 7B artifact requirement in normal CI.
• Public API composition and result detachment/ownership are tested with local
  synthetic model objects where legitimate and real cached-model composition
  where the behavior depends on Bridge execution.

Released-artifact gate

• Pin reference/client and artifacts; compare Bridge AV prefill logits against
  direct HF inputs_embeds under identical tokens, scaled vector, dtype, and
  eager attention.
• Compare Bridge AR block output plus value head against the authors' NLACritic
  for fixed canonical explanation text and exact tokenization.
• Test a real target extraction and complete sampled round trip. Sampled text
  equality and a favorable cosine threshold are not deterministic loading oracles.
• Run reference/Bridge models sequentially. Record tensor error distributions and
  empirically justified tolerances; investigate drift, not label it "fp noise".
• Verify every advertised precision/device configuration. No automatic platform
  skip or xfail to hide an incompatible package.

Definition of done

• All public methods work with one verified Qwen registry entry.
• Text, provenance, metric, loading, and memory contracts are implemented.
• Normal-CI composition and explicit slow artifact gates pass.
• Registry and required resources exist in the built wheel.
• Format, type checks, standard test tiers, docs build, and package build pass.
• User guide includes convenience and sequential examples, resource requirements,
  parsing failures, unverified-origin limitations, and faithfulness caveats.
• No model-server or training dependency is added.

---

Resource requirements and risks

Three separate 7B-class artifacts are involved across target, AV, and AR. Historical
inspection puts AR alone at roughly 10.9 GB plus a 25.7 MB head. Sequential loading
reduces peak model residency, not total download/storage or validation time.
Exact peak RAM/VRAM and timings must be measured, not promised here.

Main risks: silent scale/placeholder drift, wrong AR output basis, tokenizer/BOS
mismatch, provenance overstated as live-weight identity, and hooks that behave
correctly in a standalone forward but fail during generation. The acceptance
gates target these risks directly.

---

References and filing checklist

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

Historical audit citation/artifact checks are not newly executed compatibility
evidence. Refresh citation metadata, actual artifacts, and duplicate issue/PR
search before filing. No live-weight hashes, immutable artifact revisions, or
numerical results are invented in this proposal.

• Refresh duplicate-work and citation searches.
• Record immutable artifacts/reference and confirm loading feasibility.
• Confirm maintainer agreement on a tools/analysis inference integration.
• Implement and validate the complete round trip with the contracts above.

---

Checklist

• I have checked that there is no similar issue in the repo (required)

---