Docs: Jlens Qwen3.5-4b Demo

[KYinXu]

Description

Added a new walkthrough section using Qwen3.5-4b in demos/Jacobian_Lens_Demo.ipynb.

Key additions:

• Demonstrates support for bridge-only model architectures within the Jacobian Lens pipeline.
• Reproduces the multi-hop (twohop) reasoning experiment previously demonstrated on gemma-2-2b.

Fixes #1539 Tier 2 - second model demo

Type of change

• Documentation update

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

[KYinXu]

When loading Qwen-3.5 I received the following warning (x24+):

[/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:494](https://file+.vscode-resource.vscode-cdn.net/Users/kyleyinxu/Documents/Projects/TransformerLens/transformer_lens/model_bridge/bridge.py:494): UserWarning: Hook alias 'hook_attn_out' -> 'attn.hook_out' on BlockBridge(name='model.language_model.layers.2') did not resolve; this hook will not be accessible.
  getattr(module, "_register_aliases")()

The warning is harmless since we're not performing any hook actions on the hybrid layers, but it is something of note for clarity in the notebook cell as it becomes clutter and also as a potential ticket to implement hook support on Gated DeltaNets.

[KYinXu]

Also sidenote, notebook cell runs were done locally without a CUDA GPU or Colab's TPUs, so if default cell output is a desired clarity fix then I can look into rerunning remote.

[emerardd]

Hi @KYinXu — thanks for adding the Qwen3.5 walkthrough. I tested it locally in a combined tree with current dev and #1545. The full notebook passed nbval (21 passed in 659.80s), and the Qwen readout worked with all 32 block-output hooks available.

I noticed two reproducibility/documentation points:

1. The notebook describes QWEN_LAYERS as full-attention layers, but the live Bridge reports layers 8, 16, 24, and 30 as GatedDeltaNet linear-attention layers; only 31 is full attention. The full-attention layers are [3, 7, 11, 15, 19, 23, 27, 31]. Would it make sense either to update the wording or select layers such as [7, 15, 23, 31]?

2. Could the Qwen model revision also be pinned? My run resolved to 851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a. The lens revision is pinned, but its metadata does not contain a model revision, so it cannot catch future model-repository drift.

One integration note: #1545 and this PR both modify the same notebook from the same earlier base, so whichever lands second will need a cell-level conflict resolution and another full nbval run.

Overall, the Qwen path worked successfully in my local validation.

[jlarson4]

Thanks for putting this together @KYinXu! And thank you @emerardd for mentioning the conflict potential with your PR, valid point that I was going to address.

We will want the default cell output, the goal is to add this demo to CI as validation, so we will need the output to match.

#1545 has already merged we will want to rebase this branch onto dev to avoid conflicts. In that pass please leave the earlier gemma cells' outputs, the base branch recorded them.

A few asks for the rebase:

1. The final comparison cell has no committed output at all, as you noted. Ideally, we'd like every cell executed and nbval-clean.
2. +1 to both of @emerardd's points: pin the model revision, and fix the full-attention wording / QWEN_LAYERS choice.
3. Add a one-line comment on why the lens load bypasses the registry short name (the registry entry points at the non-n1000 artifact).
4. Update the header memory claim. Qwen3.5-4B bf16 weights alone are ~8 GB, so it won't fit an 8 GB card; say ~10 GB GPU or a CPU run.
5. Two small requests for the comparison cell: there's trailing whitespace inside the QWEN_LAYERS list literal (first line), and the final print hardcodes top_j[31]. Let's use model.cfg.n_layers - 1 (or reuse the last entry of QWEN_LAYERS) so the cell doesn't silently break if this gets reused for a different model.

The GatedDeltaNet hook-support ticket is a good idea, would you be willing to open an issue?

[KYinXu]

Got it, thanks for the feedback! Working on changes now.

[KYinXu]

Addressed all changes and rebased, thanks guys!

@emerardd

1. The notebook describes QWEN_LAYERS as full-attention layers, but the live Bridge reports layers 8, 16, 24, and 30 as GatedDeltaNet linear-attention layers; only 31 is full attention. The full-attention layers are [3, 7, 11, 15, 19, 23, 27, 31]. Would it make sense either to update the wording or select layers such as [7, 15, 23, 31]?

Fixed here, was originall using layers from Anthropic's jlens guide but it seems they just used evenly spaced intermediate layers without considering which were full-attention, so I shifted all selected layers over by one.

2. Could the Qwen model revision also be pinned? My run resolved to 851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a. The lens revision is pinned, but its metadata does not contain a model revision, so it cannot catch future model-repository drift.

Pinned the same model revision for Qwen3.5, confirmed demo still works locally.

@jlarson4

1. The final comparison cell has no committed output at all, as you noted. Ideally, we'd like every cell executed and nbval-clean.

The cell was in the original notebook as a markdown file so I kept it that way. It seems to be for if users want to try the feature themselves with separate models and doesn't contribute to the walkthrough. I kept it as markdown but let me know if we should make this executable for clarity.

3. Add a one-line comment on why the lens load bypasses the registry short name (the registry entry points at the non-n1000 artifact).

4. Update the header memory claim. Qwen3.5-4B bf16 weights alone are ~8 GB, so it won't fit an 8 GB card; say ~10 GB GPU or a CPU run.

Added clarifications here for both.

5. Two small requests for the comparison cell: there's trailing whitespace inside the QWEN_LAYERS list literal (first line), and the final print hardcodes top_j[31]. Let's use model.cfg.n_layers - 1 (or reuse the last entry of QWEN_LAYERS) so the cell doesn't silently break if this gets reused for a different model.

Fixed

The GatedDeltaNet hook-support ticket is a good idea, would you be willing to open an issue?

I'd love to, I'll open a ticket about it tomorrow and look into implementing as well!