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 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 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.

The replication and lens-eval prompt sets in data/ are synthetic, authored by Anthropic, and released under the same Apache License 2.0 as the code. See the READMEs in data/experiments/ and 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.