GitHub - anthropics/jacobian-lens: Companion code for the global workspace interpretability paper · GitHub
/" data-turbo-transient="true" />
Skip to content
Search or jump to...
Search code, repositories, users, issues, pull requests...
-->
Search
Clear
Search syntax tips
Provide feedback
--><br>We read every piece of feedback, and take your input very seriously.
Include my email address so I can be contacted
Cancel
Submit feedback
Saved searches
Use saved searches to filter your results more quickly
-->
Name
Query
To see all available qualifiers, see our documentation.
Cancel
Create saved search
Sign in
/;ref_cta:Sign up;ref_loc:header logged out"}"<br>Sign up
Appearance settings
Resetting focus
You signed in with another tab or window. Reload to refresh your session.<br>You signed out in another tab or window. Reload to refresh your session.<br>You switched accounts on another tab or window. Reload to refresh your session.
Dismiss alert
{{ message }}
Uh oh!
There was an error while loading. Please reload this page.
anthropics
jacobian-lens
Public
Notifications<br>You must be signed in to change notification settings
Fork<br>64
Star<br>377
main
BranchesTags
Go to file
CodeOpen more actions menu
Folders and files<br>NameNameLast commit message<br>Last commit date<br>Latest commit
History<br>1 Commit<br>1 Commit
assets
assets
data
data
jlens
jlens
tests
tests
.gitignore
.gitignore
LICENSE
LICENSE
README.md
README.md
pyproject.toml
pyproject.toml
uv.lock
uv.lock
walkthrough.ipynb
walkthrough.ipynb
View all files
Repository files navigation
jlens — Jacobian lens
Reference implementation. Not maintained and not accepting contributions.
Companion code for Verbalizable Representations Form a Global Workspace in<br>Language Models .
The Jacobian lens reads out what an internal activation is disposed to make the<br>model say. It linearly transports a residual-stream vector at any layer and<br>position into the final-layer basis, then decodes it with the model's own<br>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<br>target positions in a generic web-text corpus; the precise estimator<br>(cotangents summed over target positions, then averaged over source positions)<br>is documented in the jlens.fitting module docstring.
This repo fits the lens on open-weights decoder transformers, applies it, and<br>renders the interactive layer × position view shown below. Examples use Qwen;<br>other HuggingFace decoders adapt cleanly.
The ASCII-face example: selecting the ^ (nose) position shows the lens<br>reading out "nose" at mid layers, although the word never appears in the<br>prompt.
Install
pip install -e .
Usage
Apply
To apply a pre-fitted lens:
import transformers, jlens
hf = transformers.AutoModelForCausalLM.from_pretrained("org/model").cuda()<br>tok = transformers.AutoTokenizer.from_pretrained("org/model")<br>model = jlens.from_hf(hf, tok)
lens = jlens.JacobianLens.from_pretrained("org/lens-repo", filename="model/lens.pt")<br>lens_logits, model_logits, _ = lens.apply(<br>model, "Fact: The currency used in the country shaped like a boot is",<br>positions=[-2])<br>for layer, logits in sorted(lens_logits.items()):<br>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")<br>lens.save("out/jacobian_lens.pt")
The paper's lenses use 1000 sequences of 128 tokens from a pretraining-like<br>corpus. Quality saturates quickly (§9.3); ~100 prompts is usable. This is a<br>reference implementation and is not optimized; fitting time is dominated by<br>the model's own backward pass. Parallelize by running fit() on disjoint<br>slices and combining with JacobianLens.merge().
Walkthrough
walkthrough.ipynb is the end-to-end notebook: load a<br>model, load (or fit) a lens, apply it at a few layers, and render a slice page<br>like the one above.
Reading a slice page:
Each cell shows the lens top-1 word at that (position, layer); the<br>superscript is its rank over the full vocabulary.
Click a cell to select a (position, layer) and pin its top-1 token; pinned<br>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.
The replication and lens-eval prompt sets in data/ are synthetic,<br>authored by Anthropic, and released under the same Apache License 2.0 as the<br>code. See the READMEs in data/experiments/ and<br>data/evaluations/ for what each set contains.
The slice-vis pages use d3 (ISC license), loaded<br>from the jsDelivr CDN with subresource integrity or inlined into<br>self-contained pages.
No model weights or text corpora are bundled; models and datasets downloaded<br>at run time are subject to their own...