[janmenjayap]

---

Description

Adds demos/Natural_Language_Autoencoders_Demo.ipynb. It runs the released NaturalLanguage Autoencoder (NLA) round trip for Qwen2.5-7B-Instruct at block 20 throughTransformerBridge, with no model server:

1. extract one raw residual-stream activation from the target model;
2. turn it into an explanation with the Activation Verbalizer (AV), injected as one input embedding;
3. turn the explanation back into a vector with the Activation Reconstructor (AR);
4. score the reconstruction against the original with cosine and direction-normalized MSE.

This is the demo-only reframe of #1884 that @jlarson4 approved (comment), following the pattern from #1907. It includes the three requested adjustments: the released section runs in bf16 and reports measured memory, its cells are # NBVAL_SKIP with committed outputs, and the scope is this demo only.

No package code changes. Nothing changes in transformer_lens/, pyproject.toml, or uv.lock. The notebook uses only existing public Bridge methods: run_with_cache, hooks, generate, and boot_transformers(hf_model=...). pyyaml and safetensors are already dependencies through transformers.

Fixes #1884

---

What the notebook contains

Helpers, defined inline. Each is short and is followed by a cell that uses it.

Helper	What it does
check_sidecar	Validates an artifact's nla_meta.yaml against its role and the live tokenizer: schema, layer, width, placeholder token and neighbors, injection norm, templates, and the AR's prompt suffix.
extract	run_with_cache filtered to blocks.{L}.hook_out only. Returns a detached CPU float32 vector and rejects compatibility-mode Bridges.
make_injection_hook	Rescales to the sidecar norm in float64, casts to the embedding dtype, and overwrites only the placeholder row at hook_embed. Injects on every forward pass that reaches the placeholder: the prefill, and each full-prefix pass when generation runs uncached, including when generate disables the cache itself. Cached decode steps embed only the newest token and are left unchanged.
render_av_prompt, verbalize	Renders the AV chat template once, checks the placeholder's neighbors, and calls generate inside bridge.hooks(...) with explicit, bounded sampling settings.
parse_explanation	Accepts exactly one complete <explanation> region. Otherwise it returns None with a named status: missing, multiple, nested, incomplete, or empty. The notebook says that this is stricter than the reference client.
reconstruct	Formats the raw AR template, checks the suffix tokens, reads the final token's block-$L$ output, and applies the bias-free value head. It never touches the AR's absent final norm or LM head.
score	Computed in float64; asserts that $\mathrm{nMSE} = 2(1 - \cos)$.

Mechanics section, run in CI. It runs on trl-internal-testing/tiny-Qwen2ForCausalLM-2.5, pinned to revision ce8d0bf270b28c8fab026ced69fe8aa14b0f0eda and shares Qwen2.5's tokenizer and chat template. Synthetic sidecars follow the released schema. It asserts:

• (a) extract equals HF hidden_states[L + 1], and Bridge prefill logits with the injection equal HF logits from inputs_embeds with the same row replaced. Both are exact (torch.equal).
• (b) Greedy generation gives the same tokens with and without the KV cache. The hook's call log confirms the injection fires on the prefill only with the cache, and on every pass without it.
• (c) The truncated AR read-out equals the full model's hidden_states[L + 1]. It also shows two shortcuts that give the wrong vector: HF's last hidden_states entry has the final norm applied, and stop_at_layer=L + 1 doesn't stop early on a truncated stack.
• (d) Every parse outcome, plus the rejection of tagged input in reconstruct.
• (e) No hooks remain, including after a caller hook raises during verbalize. The caller's own hook survives until the caller removes it.

Released section, run locally. All of it is # NBVAL_SKIP, opt-in through TL_RUN_RELEASED_NLA=1, with outputs committed. It uses pinned revisions:

Role	Repository	Revision
target	Qwen/Qwen2.5-7B-Instruct	a09a35458c702b33eeacc393d103063234e8bc28
AV	kitft/nla-qwen2.5-7b-L20-av	b88469162777ae6553bc14208eb0cb579336f8f4
AR	kitft/nla-qwen2.5-7b-L20-ar (+ value_head.safetensors)	e2c9e57eac213d37a31612087f645ab6332c1bb6
tiny mechanics model	trl-internal-testing/tiny-Qwen2ForCausalLM-2.5	ce8d0bf270b28c8fab026ced69fe8aa14b0f0eda

The section loads one model at a time in bf16. Each stage compares the Bridge with plain Hugging Face at the same dtype, device, tokens, and eager attention: the HF reference runs first, then the same module is wrapped with boot_transformers(hf_model=...). The AR reference follows the authors' NLACritic read-out. Every prompt is reported, including any parse failure.

Docs and CI registration:

• docs/make_docs.py and docs/source/index.md register the notebook.
• docs/source/content/analysis_tools.md gets a short pointer with the faithfulness caveat, and docs/source/content/tutorials.md mentions it.
• .github/workflows/checks.yml adds the notebook to the notebook-checks matrix.
• tests/unit/test_make_docs.py adds it to the expected list of copied demos.

---

Measured results (released section)

Apple M3 Max, 48 GiB, CPU only, bf16, torch 2.11.0. Checkpoints were already cached.

Comparison, same dtype and device	Max abs difference
Target block-20 output vs HF hidden_states[21]	0.0
AV prefill logits with injection vs HF inputs_embeds (argmax agrees at every position)	0.0
AR block-20 read-out vs NLACritic-style HF	0.0

The sidecar checks all pass: injection norm 150, placeholder U+320E (token 149705), mse_scale = $\sqrt{3584}$, AR suffix [1318, 29, 366, 1708, 29], and a 3584×3584 bf16 value head with no bias. HF reports exactly model.norm.weight and lm_head.weight as missing from the AR.

Prompt (final token)	Parse status	Cosine	Normalized MSE
Travel narrative, 39 tokens (" of")	ok	0.902	0.195
Python fibonacci, 48 tokens (" return")	ok	0.785	0.430
Narrative sentence, 34 tokens (" into")	ok	0.909	0.183

Stage	Wall time	Peak RSS
Sidecars and value head	3.7 s	1.16 GiB
Target: load and extract	21.0 s	13.64 GiB
AV: load, parity, verbalize	134.6 s	13.78 GiB
AR: load, parity, reconstruct	29.9 s	13.78 GiB
Total	3.2 min	13.8 GiB

Peak RSS stays flat after the target stage, so each 7B model is released before the next one loads. Downloads total 38.6 GiB and are not included in the times.

The notebook also shows the faithfulness caveat in practice. The explanations get the prompt's format and the final token's role right, but they also contain claims that aren't in the prompts: Brussels and Belgium for the travel prompt, a nonexistent math.fib for the code prompt, and a pregnancy for the narrative. These appear alongside cosines of 0.90 and 0.91. "Reading these results" and the interpretation checklist state that a high reconstruction score does not certify individual claims.

---

Notes for review

• Parity scope. On the released models, parity covers extraction, the AV prefill, and the AR read-out. Sampled generation is not compared with HF generation. The mechanics section compares cached against uncached generation instead.
• Platform. The released section was run on CPU only. The code moves the HF reference to the same device as the Bridge on CUDA, but that path has not been run.
• Parser. parse_explanation is deliberately stricter than the reference client, which takes the first match or falls back to the raw text. The notebook states this, because parse-failure rates from the two are not directly comparable.
• Numerical margin in check (b). On the tiny model, cached and uncached logits differ by at most 3.4e-8. The smallest top-1 vs top-2 gap is 2.5e-4, so a greedy-token flip between platforms is unlikely.
• Out of scope, per [Proposal] Natural Language Autoencoders demo: explain and reconstruct a Qwen activation with TransformerBridge #1884: explanation-edit steering, corpus benchmarks or FVE, other released NLA pairs (Gemma, Llama), training, SGLang, and quantization.
• Related bug. get_tokenizer_with_bos raises KeyError for in-memory tokenizers. It is filed separately as [Bug Report] get_tokenizer_with_bos raises KeyError for in-memory tokenizers #1961, with a fix in fix: handle in-memory tokenizers in get_tokenizer_with_bos #1962. This demo does not depend on it.

---

Commits

1. docs(demos): add Natural Language Autoencoders demo with tiny Qwen2 mechanics checks: background, helpers, the CI-executed mechanics checks, and the interpretation checklist.
2. docs(demos): add released Qwen2.5-7B NLA round trip in bf16 with HF parity: the skipped, opt-in released section with committed outputs and measured resources.
3. docs: register NLA demo in docs build and notebook checks: docs build, index, analysis-tools pointer, tutorials mention, notebook-checks matrix, and the test_make_docs expectation.
4. docs(demos): pin the tiny Qwen2 checkpoint revision in the NLA demo, and
5. docs(demos): inject the NLA placeholder on every pass that covers it

Each commit passes nbval on its own.

---

Validation

Command	Result
uv run pytest --nbval-sanitize-with demos/doc_sanitize.cfg demos/Natural_Language_Autoencoders_Demo.ipynb	16 passed, 7 skipped (the released cells). It also passed at commit 1 alone (16 passed) and twice in a row, so the outputs are deterministic.
TL_RUN_RELEASED_NLA=1 uv run jupyter nbconvert --to notebook --execute --inplace demos/Natural_Language_Autoencoders_Demo.ipynb	Completed. These are the committed outputs.
uv run build-docs	Succeeded. The demo page renders the committed 7B outputs, the new links resolve, and the changed pages add no warnings.
make check-format	Passed
uv run mypy .	Passed (no issues in 390 source files)
uv run pytest tests/unit/test_make_docs.py	2 passed
make test-pr, unit tier	7966 passed, 1 failed. The failure was test_copy_demos_creates_generated_dir_when_absent, which pins the list of copied demos; it is fixed in commit 3 and now passes.
make docstring-test	13 passed
make acceptance-test	163 passed
make integration-test	1850 passed, 2 failed on the first run: test_get_params[google/gemma-2-2b-it] and test_TransformerBridge_gemma2_forward raised GatedRepoError (401) because no HF token was set. Rerun with a token: uv run pytest tests/integration/model_bridge/test_bridge_integration.py -k gemma gives 2 passed.
nbval at commits 4 and 5	16 passed, 7 skipped; black and isort clean on every code cell

The released cells were not re-executed for commits 4 and 5: the only change there removes an argument whose injection decision is identical, so the saved outputs are unchanged.

make test-pr ran before the final notebook-only revision. After that revision, only the notebook changed, and nbval, test_make_docs, check-format, and build-docs were rerun and passed.

---

Type of change

• This change requires a documentation update

---

Screenshots

Not applicable. The rendered notebook is in the docs build artifact.

---

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

All unit, docstring, acceptance, and integration tests pass locally. The feature's tests are the notebook's CI-executed mechanics section, plus the updated test_make_docs expectation.

---

[jlarson4]

Great work on the demo @janmenjayap, just a couple small procedural edits, nothing to crazy. Let me know when you've had a chance to look them over.

[jlarson4]

The hook infers the prefill from its own call count and the use_past_kv_cache argument, but generate can drop caching on its own (and does for adapters without KV support). In that scenario, only the first pass is injected while every later pass recomputes the prefix with the raw placeholder. If we adjust to decide based on whether the pass covers the placeholder, embeddings.shape[1] > position, we can remove the flag and the call-count state while still producing the two call-log patterns check (b) prints.

[janmenjayap]

Thanks, good catch. Fixed in 1924bfe: the hook now injects on every pass whose embeddings reach the placeholder (embeddings.shape[1] > position), and the use_past_kv_cache argument and call-count state are gone. Besides adapters without KV support, this also covers generate's cached branch when no cache was captured, which recomputes the full prefix. Check (b) prints the same two call-log patterns, and since the hook is no longer told the caching mode, the uncached run also shows the case where generate disables the cache itself. To keep a guard, verbalize now asserts that the prefill injected.

[jlarson4]

The released models are pinned to commits, but the tiny checkpoint loads its moving head. The outputs print token counts and the placeholder position derived from them, so a Hub update would cause this cell to error in CI unexpectedly. This checkpoint has been updated in the past, so to avoid this issue, could you please pin a revision= for the tiny checkpoint as well?

[janmenjayap]

Agreed, thanks. Pinned in 3328f10 to ce8d0bf270b28c8fab026ced69fe8aa14b0f0eda, the current head and the revision that produced the committed outputs, so nothing printed changes.

[janmenjayap]

Thanks for the review, @jlarson4! I've pushed one commit for each comment: 3328f10 pins the tiny checkpoint's revision, and 1924bfe makes the injection hook decide per pass. Only the notebook changed, and nbval passes. Please let me know if you'd like anything else adjusted.

[janmenjayap]

Hi @jlarson4, one more update. After my previous push, the NLA notebook check started failing on cell 26. That's because #1964 landed on dev and changed stop_at_layer=n_layers to stop after the last block instead of returning logits, and the notebook had relied on the old behaviour.

To fix it, I merged the latest dev into the branch (283bd2b). Then in d71407e I updated check (c): it now asserts that stop_at_layer=L + 1 equals the block-L read-out instead of printing the logits width, and the two markdown cells that described it are updated.

I re-executed the whole notebook, including the released 7B section, on the merged code. Parity with plain Hugging Face is still exact (max abs difference 0.0), and the parse statuses, explanations, and cosines are identical to before. Only the timings changed, with peak memory now 14.1 GiB. nbval passes twice in a row (16 passed, 7 skipped).

Thank you for your patience!

[jlarson4]

Thanks for putting those fixes together @janmenjayap, looks great, merging now