[Proposal] Jacobian Lens (J-lens): load published lenses, fit natively, read out and intervene

[danielxmed]

Proposal

Add first-class support for the Jacobian lens (J-lens) — the interpretabilitytechnique from Anthropic's July 2026 paper Verbalizable Representations Form aGlobal Workspace in Language Models— as a tool in transformer_lens/tools/analysis/, working with bothTransformerBridge and HookedTransformer.

Concretely, one JacobianLens component that can:

1. Load published lens artifacts from the HF Hub (neuronpedia/jacobian-lens ships pre-fitted lenses for 38 open models, e.g. Gemma 2/3/4, Qwen 2.5/3/3.5/3.6, Llama 3.x, GPT-2, Pythia, OLMo 3, gpt-oss) in the official artifact format.
2. Read out per-layer, per-position vocabulary distributions: lens(h_l) = softmax(W_U · norm(J_l · h_l)), using the model's own final norm and unembedding (plus the config's logit softcap where the architecture has one, e.g. Gemma-2), with the logit lens available as the J_l = I baseline through the same code path.
3. Fit lenses natively on any hooked model with TransformerLens backward hooks, reproducing the official estimator (anthropics/jacobian-lens); fit parity would be verified numerically against the official implementation on small configs as part of the PR. Plus merge() for parallelized fitting.
4. Intervene in J-lens coordinates: steering along a J-lens vector, ablation (projecting a vector out), and the paper's pseudoinverse coordinate swap (V = [v_s, v_t], c = V⁺h, h ← h + V(σ(c) − c)).

Motivation

The J-lens is, in the paper's framing, the principled correction of the logit lens: J_l is the corpus-averaged, causal linear map relating layer-l residual directions to their final-layer counterparts (the logit lens is the special case J_l = I; the tuned lens is trained on output-matching, which is correlational and tends to "skip ahead" past intermediate computation). It surfaces interpretable content at depths where the logit lens reads out noise, and its vectors support causal interventions (steering, ablation, concept swaps). Beyond the paper's workspace results, it is already being used for alignment auditing readouts.

Right now there is no maintained, general implementation anyone can build on:

• Anthropic's companion repo (anthropics/jacobian-lens, Apache-2.0) is explicitly frozen: "Not maintained and not accepting contributions." It also contains no intervention code (readout + fitting only), and at least one community offer to contribute there is stranded (Interest in a contributed swap-intervention implementation? anthropics/jacobian-lens#2).
• Neuronpedia serves J-lens readouts at neuronpedia.org/jlens — and their inference server internally runs them on top of transformer_lens.HookedTransformer — but that lives inside their platform, not as a reusable library API.
• In the nine days since the code release, roughly forty standalone replications/ports (MLX, visualizers, single-model demos) have appeared, none of them a maintained, tested library integration.

TransformerLens already has the right substrate (residual-stream hooks, backward hooks, ActivationCache, HF Hub utilities) and an established home for exactly this kind of feature (tools/analysis/, cf. #1369). J-lens would be TL's first fitted-lens tool and the first to consume published lens artifacts from the Hub.

Pitch

Proposed API (names/placement happily adjusted to maintainer preference):

from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.tools.analysis import JacobianLens

model = TransformerBridge.boot_transformers("google/gemma-2-2b", dtype=torch.bfloat16)

# 1) load a published lens (fp16 artifact, fp32 in memory; validated against model)
lens = JacobianLens.from_pretrained(
    "neuronpedia/jacobian-lens",
    filename="gemma-2-2b/jlens/Salesforce-wikitext/gemma-2-2b_jacobian_lens.pt",
    model=model,  # validates d_model / layer count; errors on processed weights
)

# 2) per-layer readout on a prompt (JacobianLensResult dataclass)
result = lens.readout(model, "The capital of France is", layers=[6, 13, 19], positions=[-1])
result.top_tokens(k=5)                  # J-lens top-k per (layer, position)
result_ll = lens.readout(model, ..., use_jacobian=False)  # logit-lens baseline

# 3) fit natively (backward hooks; parity with the official estimator)
lens = JacobianLens.fit(model, prompts, dim_batch=8, max_seq_len=128)
merged = JacobianLens.merge([lens_a, lens_b])   # parallelize across prompt slices

# 4) interventions (returned as hook fns usable with model.hooks / run_with_hooks)
band = range(10, 24)  # an intermediate-layer band
hooks = lens.steer(" Paris", layers=band, alpha=4.0)
hooks = lens.ablate(" Paris", layers=band)
hooks = lens.swap(" France", " China", layers=band)   # pseudoinverse coordinate swap

Coordinate-basis compatibility plan (the main correctness risk). Published lenses are fitted on raw HF activations: J[l] maps the output of block l (blocks.{l}.hook_resid_post) to the final block output, pre-ln_final, and the readout applies the model's own final norm + unembed. TL's classic weight processing (fold_ln, center_writing_weights, center_unembed) changes the residual basis, so:

• TransformerBridge.boot_transformers(...) (raw HF weights by default) and HookedTransformer.from_pretrained_no_processing(...) are the supported paths.
• JacobianLens checks the model's processing state on attach and raises (not warns) when lens artifacts meet a fold_ln/centered model, with an actionable message. Same treatment for d_model/layer-count mismatches.
• Artifacts saved by TL carry metadata (model name, TL version, processing flags, hook convention, corpus, n_prompts, dtype) while staying loadable by the official package (official keys preserved; extras additive). Loading official artifacts — which carry only {J, n_prompts, source_layers, d_model} — stays fully supported.
• Readout reuses the model's own ln_final + unembed modules, so RMSNorm vs LayerNorm and Gemma-2's final_logit_softcapping are handled by construction (cf. the logit-lens normalization fixes in fixed the logit lens implementation inside ActivationCache.accumulated_resid to match the standard definition in literature and the expected and defined behavior as per the documentation in the docstring and in the docs #1077/Fix/1076 logit lens layer norm #1180 and the processing discussion in Make TransformerBridge match HookedTransformer preprocessing defaults (fold_ln, center_writing_weights, center_unembed), or document the divergence #1272).

Validation plan (all numbers to be posted in the PR):

1. Artifact parity: loading the published gemma-2-2b lens reproduces reference readouts (official jlens package on HF transformers as the numerical oracle, with top-k overlap and rank correlation thresholds on fixed prompts, per layer; Neuronpedia's public lens API as a secondary qualitative sanity check — their serving stack resolves bf16 near-ties differently, so it isn't a bitwise reference).
2. Late-layer agreement: J-lens ≈ logit lens top-k in the final layers (J_l → I).
3. Kurtosis profile: excess kurtosis of readout logits ≈ 0 in the first ~third of depth, rising through the workspace band (paper §4 structure result).
4. Fit convergence: native TL fit at small n on wikitext trends toward the published readouts (official repo: "~100 prompts is usable"; published artifacts ship per-prompt convergence CSVs to compare against).
5. Causal smoke test: the France→China coordinate swap on "The capital of France is", following the paper's flexible-generalization protocol (swap clamped at every position across an intermediate-layer band), with observed success rates reported honestly — the paper itself reports 42/48 for country swaps at α=1 on Sonnet 4.5, and no open-weights numbers exist yet.

Scope of the first PR (deliberately small; a working vertical slice):

• transformer_lens/tools/analysis/jacobian_lens.py (load/save/fit/merge/readout/ interventions, result dataclass), exported like the existing analysis tools.
• Unit tests (synthetic tiny model, no HF downloads: estimator orientation on a linear toy model where J is known in closed form, save/load round-trip, merge weighting, basis-mismatch errors) + integration tests on CI-cached models (gpt2: fit + readout invariants; the published gpt2-small lens exists for parity) + @pytest.mark.slow gemma-2-2b parity tests.
• One demo notebook (demos/Jacobian_Lens_Demo.ipynb): J-lens vs logit lens on a two-hop prompt with an unspoken intermediate, and the France→China swap; gemma-2-2b (+ a Bridge-only modern model, see below).

Follow-up work (PR2+, separate issue once this lands): artifact registry/aliases so JacobianLens.from_pretrained("gemma-2-2b") just works; converters for community formats (fit-checkpoint files, other lens zoos); gradient-pursuit sparse decomposition (J-space coordinates, occupancy); a community fitting guide (the fit parallelizes cleanly with merge, so lens coverage of new models is a natural community contribution surface).

Model coverage in the first PR: gemma-2-2b as the correctness reference (supported by both HookedTransformer and TransformerBridge; published artifact; CI already caches the -it variant, which also has a published lens, for any cached-tier tests) and Qwen/Qwen3.5-4B as the Bridge-only modern-architecture demo (the model used in the official walkthrough, with a published n=1000 artifact). I've smoke-tested both: the Qwen3.5 text-only Bridge adapter loads the model, its blocks.{i}.hook_out dims match the published artifact (32 layers, d_model 2560), and backward hooks deliver gradients — so both readout and native fitting are feasible on it.

Alternatives

• Standalone package (like the MLX ports): works, but fragments the ecosystem — TL users would still lack lens-aware hooks/tests, and the technique pairs naturally with TL's existing cache/hook idioms. The official repo being frozen makes the canonical mech-interp library the right long-term home.
• Depend on the official jlens package: it's frozen, pins transformers>=5.5, and its model adapter duplicates what TL already abstracts (and it would still leave interventions unimplemented).
• Readout-only integration (no fitting): smaller, but native fitting is what makes the feature self-sufficient for models without published artifacts, and it reuses TL's backward-hook machinery — the marginal complexity is contained.

Additional context

• Paper: Gurnee et al., Verbalizable Representations Form a Global Workspace in Language Models, Transformer Circuits Thread, 2026. https://transformer-circuits.pub/2026/workspace/index.html
• Official reference implementation (Apache-2.0, frozen): https://github.com/anthropics/jacobian-lens — the fitting estimator and artifact format to match; attribution will be included in module docstrings.
• Published artifacts: https://huggingface.co/neuronpedia/jacobian-lens (MIT card; Gemma-derived lenses remain subject to the Gemma Terms of Use, so they are always downloaded at runtime, never vendored).
• Interactive reference: https://www.neuronpedia.org/jlens
• Estimator note for reviewers: the official fit is deterministic — for each output dimension, a one-hot cotangent is planted at every valid target position at once and backpropagated (causal masking turns the per-source-position gradient into the sum over targets t' ≥ t for free); sum over targets, mean over source positions (first 16 and the final position excluded), unweighted mean over prompts; cost 1 forward + ceil(d_model/dim_batch) backwards per prompt. This maps directly onto model.hooks(bwd_hooks=...).

I have a working spike of the loader/readout path: loading the published gemma-2-2b lens through a raw-weights TransformerBridge and reading out via the model's own ln_final/unembed reproduces the official implementation's per-layer top-8 readouts (worst case 7/8 top-8 overlap across 75 layer×prompt cells on 3 fixed prompts, top-1 identical in 70/75; the residue is bf16 near-ties). If this direction looks right, I'll follow up with the PR against dev shortly. Happy to adjust API names, placement, scope, and test tiers to whatever the maintainers prefer.

Checklist

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

[jlarson4]

Hi @danielxmed! I am on board with this idea, on the whole. You've already submitted your initial PR, so I'll add more detailed feedback in that location. My initial thought on the plan overall: drop HookedTransformer support. We plan to deprecate that system in the next major release, and we don't want to bog down this implementation with the need to support two parallel systems.

Additionally, it would be greatly appreciated if you could please fix the formatting before submitting large issues like this in the future. The unnecessary returns make it incredibly difficult to read. I cleaned them up from the main body of this issue, just a note for the future.

I should have a full review up on the PR itself within the next 24 hours.