anthropics/jacobian-lens

anthropics/jacobian-lens

Description: Companion code for the global workspace interpretability paper

Language: Python

License: Apache-2.0

Stars: 236

Forks: 45

Open issues: 1

Created: 2026-07-02T20:30:04Z

Pushed: 2026-07-02T20:52:50Z

Default branch: main

Fork: no

Archived: no

README:

jlens — Jacobian lens

> Reference implementation. Not maintained and not accepting contributions.

Companion code for **Verbalizable Representations Form a Global Workspace in Language Models**.

The Jacobian lens reads out what an internal activation is disposed to make the model say. It linearly transports a residual-stream vector at any layer and position into the final-layer basis, then decodes it with the model's own unembedding into a ranked list of vocabulary tokens.

The transport is the average input–output Jacobian over a text corpus:

lens_l(h) = unembed( J_l @ h ), J_l = E[∂h_final / ∂h_l]

The expectation is over prompts, source positions, and all current-and-future target positions in a generic web-text corpus; the precise estimator (cotangents summed over target positions, then averaged over source positions) is documented in the [jlens.fitting](jlens/fitting.py) module docstring.

This repo fits the lens on open-weights decoder transformers, applies it, and renders the interactive layer × position view shown below. Examples use Qwen; other HuggingFace decoders adapt cleanly.


*The ASCII-face example: selecting the ^ (nose) position shows the lens reading out "nose" at mid layers, although the word never appears in the prompt.*

Install

pip install -e .

Usage

Apply

To apply a pre-fitted lens:

import transformers, jlens

hf = transformers.AutoModelForCausalLM.from_pretrained("org/model").cuda()
tok = transformers.AutoTokenizer.from_pretrained("org/model")
model = jlens.from_hf(hf, tok)

lens = jlens.JacobianLens.from_pretrained("org/lens-repo", filename="model/lens.pt")
lens_logits, model_logits, _ = lens.apply(
model, "Fact: The currency used in the country shaped like a boot is",
positions=[-2])
for layer, logits in sorted(lens_logits.items()):
print(layer, [tok.decode([t]) for t in logits[0].topk(5).indices])

Fit

To fit a lens on your own model:

lens = jlens.fit(model, prompts=my_prompts, checkpoint_path="out/ckpt.pt")
lens.save("out/jacobian_lens.pt")

The paper's lenses use 1000 sequences of 128 tokens from a pretraining-like corpus. Quality saturates quickly (§9.3); ~100 prompts is usable. This is a reference implementation and is not optimized; fitting time is dominated by the model's own backward pass. Parallelize by running fit() on disjoint slices and combining with JacobianLens.merge().

Walkthrough

[walkthrough.ipynb](walkthrough.ipynb) is the end-to-end notebook: load a model, load (or fit) a lens, apply it at a few layers, and render a slice page like the one above.

Reading a slice page:

• Each cell shows the lens top-1 word at that (position, layer); the

superscript is its rank over the full vocabulary.

• Click a cell to select a (position, layer) and pin its top-1 token; pinned

tokens get rank-tracking charts and a rank heatmap.

• The bottom row (L = n_layers − 1) is the model's actual output.

License and data

Code is released under the Apache License 2.0 — see [LICENSE](LICENSE).

The replication and lens-eval prompt sets in [data/](data/) are synthetic, authored by Anthropic, and released under the same Apache License 2.0 as the code. See the READMEs in [data/experiments/](data/experiments/) and [data/evaluations/](data/evaluations/) for what each set contains.

The slice-vis pages use d3 (ISC license), loaded from the jsDelivr CDN with subresource integrity or inlined into self-contained pages.

No model weights or text corpora are bundled; models and datasets downloaded at run time are subject to their own licenses.