Add Jacobian lens (J-lens): published-artifact loading, readout, native fitting, and interventions

[danielxmed]

Description

Implements the Jacobian lens (J-lens) as a tools/analysis component, per the
design proposal in #1505.

The J-lens (Gurnee et al., Verbalizable Representations Form a Global
Workspace in Language Models,
Transformer Circuits 2026) characterizes intermediate residual-stream
activations by their corpus-averaged, first-order causal effect on the model's
output: one d_model x d_model transport matrix per layer, read through the
model's own final norm and unembedding. The logit lens is the J = I special
case and is available through the identical code path.

What's included (transformer_lens/tools/analysis/jacobian_lens.py):

• JacobianLens.from_pretrained / load / save: the official
  anthropics/jacobian-lens
  artifact format (fp16 storage / fp32 in memory), including the 38-model zoo
  published at neuronpedia/jacobian-lens.
  Artifacts saved here add a provenance metadata key while staying loadable
  by the official package.
• readout(): per-layer, per-position vocabulary logits via the model's own
  ln_final + unembed plus the config's output_logits_soft_cap where
  configured (e.g. Gemma-2), returning a JacobianLensReadout dataclass.
• fit(): native fitting with TL backward hooks, reproducing the official
  estimator (one-hot cotangents planted at every valid target position;
  causal masking turns the per-source gradient into the sum over targets
  t' >= t; mean over valid source positions; unweighted mean over prompts).
  merge() combines shard fits exactly (n_prompts-weighted).
• Interventions implemented from the paper's Methods (the official repo ships
  none): steering_hooks (norm-matched variant, documented as such),
  ablation_hooks, and the pseudoinverse coordinate swap_hooks, all
  returning (hook_name, fn) pairs for model.hooks(...).
• Loud basis guards: published lenses are fitted on raw HF activations, so
  validate_model/fit/readout refuse (a) fold_ln'd HookedTransformers
  (normalization_type ending in Pre), (b) any compatibility-mode
  TransformerBridge — enable_compatibility_mode() processes weights by
  default and a no_processing=True call can't be reliably distinguished
  afterwards (the mirror image of DLA's requirement), and (c) models whose
  W_U is exactly mean-centered, the signature of center_unembed — which
  from_pretrained applies even with fold_ln=False. Centering of writing
  weights alone leaves no post-hoc signature, so the docs state plainly that
  from_pretrained_no_processing / raw boot_transformers are the supported
  paths. Open to a different contract here if preferred.

Works with both TransformerBridge (raw boot_transformers default — the
recommended path) and HookedTransformer (from_pretrained_no_processing).

Fixes #1505.

Type of change

• New feature (non-breaking change which adds functionality)

Validation

All numbers from real runs (local RTX GPU + a rented 24GB GPU; harness and raw
results retained by the author). Produced at the branch head; the only commit
after the main validation pass is a device-indexing robustness fix re-verified
by the full unit/integration suite.

1. Fit parity vs the official implementation: TL fit() vs official
   jlens.fit() on gpt2 (fp32 CPU, same 3 wikitext prompts, dim_batch 16,
   seq 64): worst per-layer relative Frobenius error 7.7e-07 across all
   11 layers — numerically identical estimator.
2. Artifact parity (official-jlens-on-HF as oracle, bf16 cuda, 3 fixed
   prompts, final position, every fitted layer):
   • gemma-2-2b (n=454 artifact): top-8 overlap worst 7/8 over 75 cells,
     top-1 identical 70/75 (Bridge) / 71/75 (HT no_processing); model rows
     top-1 identical. Residue is bf16 near-ties.
   • Qwen3.5-4B (n=1000 artifact, Bridge-only, hybrid GatedDeltaNet): top-8
     overlap 8/8 in 85 of 93 cells and at least 7/8 in all; top-1 identical
     91/93; tokenization and model rows identical. No architecture
     special-casing needed in the tool.
3. Late-layer agreement (J -> I): gpt2 rank of the model's top-1 token under
   the J-lens collapses 8577 (L0) -> 12 (L10) -> 0 (L11); gemma-2-2b
   J-lens/logit-lens top-10 overlap goes 0/10 (L13) -> 7/10 (L24) -> 10/10 (L25).
4. Fit convergence toward published lenses: gpt2 2-prompt fit vs published
   n=277: cosine 0.51/0.66/0.91 at L0/L5/L10. gemma-2-2b (seq 128, official
   wikitext corpus) vs published n=454: n=25 cosine from 0.868 (L0, minimum
   across all layers) to 0.998 (L24); n=100 from 0.941 to 1.000 — rising with
   depth apart from a shallow dip around L2-L5. n=100 readouts reach 8/8 top-8
   agreement with the published lens at late band layers. Qwen3.5-4B: even a
   5-prompt fit reaches cosine 0.645 (L0) to 0.998 (L30) vs the published
   n=1000 lens, crossing 0.90 by L17.
5. Kurtosis profile (paper §4 structure, reported as measured): the rise
   through the mid-to-late band reproduces on both models (gemma-2-2b:
   ~0.17 -> 0.73 over L20 -> L24; Qwen3.5-4B: ~1.0 -> ~2.2 over L12 -> L27),
   but the paper's near-zero early-layer plateau does not — early layers show
   elevated kurtosis on these small open models, consistent with their known
   noisy early readouts (the same heavy-tailed junk-token attractors are
   visible in Neuronpedia's readouts; the paper's numbers are from
   Claude-scale models). Not encoded as a test threshold.
6. Causal smoke test (paper's flexible-generalization protocol, gemma-2-2b,
   swap " France" -> " China" at every position, layers 10-24): top-1 flips
   ' French' -> ' Chinese' at alpha=1 on "Most people in France speak";
   ' Beijing' moves rank 968 -> 8 on "The capital of France is" (whose
   unperturbed top-1 on this base model is the filler ' a'); ' Asia' 8 -> 1
   on the continent template. The paper reports alpha=2 recovering some
   alpha=1 failures on Sonnet 4.5; on this small model alpha=2 instead
   overshoots into verbalizing ' China' directly, so the demo uses alpha=1.

Tests: 20 unit (a closed-form linear toy model pins the estimator and
transport orientation exactly; guards incl. the centered-unembed signature)

• 9 integration on gpt2 (Hub artifact, rank collapse, HT/Bridge agreement,
  steering, exact fit==merge identity, small-fit convergence; 2 marked slow).
  make format, mypy, and make docstring-test pass locally. Full
  make unit-test: 3249 passed with 4 failures confined to
  tests/unit/test_generate_no_tokenizer.py, which fail identically on a clean
  origin/dev checkout in the same environment (py3.12, torch 2.10,
  transformers 5.8.1): generate() concatenates a cuda:0 tensor with cpu
  tensors at HookedTransformer.py:2187 for tokenizer-free models on
  CUDA-visible machines, which CPU-only CI never hits. Unrelated to this
  additive change; I can open a separate bug report/fix for it.

The issue proposed @pytest.mark.slow gemma-2-2b parity tests; since base
gemma-2-2b is HF-gated, that coverage shipped instead as the offline
validation above, with gpt2 (CI-cached, published lens available) carrying
the in-repo parity tests.

Demo: demos/Jacobian_Lens_Demo.ipynb (gemma-2-2b, ~6GB GPU): the two-hop
boot -> Italy -> euro readout vs the logit lens, rank trajectories, the
France -> China swap, and steering. Executed end-to-end with baked outputs;
nbval passes locally. Not added to the CI notebook matrix (gated model + GPU);
can add it if preferred.

Limitations / follow-ups (PR2 roadmap in #1505)

• No gradient-pursuit sparse decomposition (J-space coordinates/occupancy),
  no artifact registry/aliases, no fit checkpoint/resume (merge() covers
  parallelism; official checkpoint files are rejected with a clear error).
• steering_hooks uses a norm-matched scale (unit lens vector times the
  activation's median per-position residual norm), following the reference
  implementation's experiment protocols rather than the paper's minimal
  h += alpha * v_t; the docstring spells this out and raw vectors are
  available via lens_vectors for the minimal form.
• Interventions operate on residual hooks at every position by default;
  position-targeted protocols are exposed via positions=.
• The kurtosis/swap phenomenology on small base models is noisier than the
  paper's Claude-scale results; the numbers above are what gemma-2-2b and
  Qwen3.5-4B actually produce.

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (docstrings; API
  docs are generated)
• 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