anthropics/jacobian-lens

TL;DR: Anthropic's companion code for their paper "Verbalizable Representations Form a Global Workspace in Language Models" — a new lens for reading out what any intermediate activation, at any layer and position, is disposed to make the model eventually say. Instead of fitting a probe against final logits (tuned lens) or decoding raw activations directly (logit lens), it linearly transports activations into the final-layer basis using the average input–output Jacobian over a corpus, then decodes with the model's own unembedding. The global-workspace framing and the cross-position readout make this more interesting than yet another lens variant.

How it works

The core object is a per-layer linear map J_l = E[∂h_final / ∂h_l] — the Jacobian of the final residual state with respect to the layer-l activation, averaged over prompts, source positions, and (crucially) all current-and-future target positions in a generic web-text corpus. Applying the lens is then just unembed(J_l @ h): no learned probe, no optimization against the output distribution, just the model's own local linearization, averaged.

Two things distinguish this from the tuned lens family:

1. The transport is derived, not fitted. A tuned lens probe is trained to minimize divergence from the final logits, which risks the probe doing the work. Here the map is the model's actual (averaged) sensitivity structure — a more principled claim about what the activation causes downstream.
2. Averaging over future target positions means the lens reads out what an activation is disposed to make the model say later, not just at that position. The README's ASCII-face example is the pitch: select the ^ (nose) position in an ASCII face, and mid-layer activations decode to "nose" — a word that never appears in the prompt. The model has a verbalizable concept there before any verbalization.

The connection to your interests: the paper's framing is explicitly global-workspace-theory-flavored (Baars/Dehaene) — the claim being that representations which are verbalizable via this lens occupy a shared, broadcast-like subspace. The README itself doesn't summarize the paper's findings beyond the title and the demo, so read the paper for the actual evidence behind that claim.

What's actually here

• Code (Apache-2.0): fitting (jlens.fit), applying (JacobianLens.apply), merging shards fit in parallel, and an interactive layer × position slice viewer (d3-based, self-contained HTML pages).
• Data: synthetic Anthropic-authored replication and lens-eval prompt sets, Apache-licensed. No corpora or weights bundled.
• Walkthrough notebook (walkthrough.ipynb): end-to-end from model load to rendered slice page. Examples use Qwen; any HuggingFace decoder should adapt.
• No numbers in the README. No faithfulness metrics, no comparison to tuned lens, no benchmarks — the evidence lives in the paper. The interactive slice view (top-1 token per cell with vocabulary rank as superscript, rank-tracking charts on pinned tokens) is the demo.

Expectations and caveats

This is explicitly a frozen reference implementation — not maintained, not accepting contributions. Fitting is unoptimized, dominated by the model's backward pass; the paper's lenses use 1000 × 128-token sequences, but quality reportedly saturates fast (~100 prompts usable). You need enough GPU to backprop through your target model.

The main scientific caveat is the one worth testing yourself: the lens assumes an averaged linearization is a faithful transport. Averaging Jacobians over a corpus washes out input-dependence; whatever the lens reads out is filtered through that fixed linear map, so apparent "verbalizability at layer l" could partly be an artifact of the transport rather than the representation. The fitting estimator is documented in the jlens.fitting docstring — worth reading before drawing conclusions on your own models.

Try it

git clone https://github.com/anthropics/jacobian-lens && cd jacobian-lens
pip install -e .

import transformers, jlens

hf = transformers.AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B").cuda()
tok = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B")
model = jlens.from_hf(hf, tok)

# fit your own (parallelize with jlens.JacobianLens.merge)
lens = jlens.fit(model, prompts=my_prompts, checkpoint_path="out/ckpt.pt")

lens_logits, model_logits, _ = lens.apply(
    model, "Fact: The currency used in the country shaped like a boot is",
    positions=[-2])

Then open walkthrough.ipynb for the interactive layer × position view — that's the fastest way to build intuition for what this lens sees that logit/tuned lenses don't.