The official walkthrough from anthropics/jacobian-lens
(Apache-2.0, license), run on Qwen3.5-4B with the neuronpedia n=1000 lens on one H200.
Sections 5 (local file server) and 6 (fitting a new lens) were not run.
Jacobian lens — walkthrough
Load a model, load a pre-fitted Jacobian lens from the Hub, apply it to a prompt, and render the interactive slice visualisation.
[1]import jlens
jlens.configure_logging()
MODEL_NAME = "Qwen/Qwen3.5-4B"
# MODEL_NAME = "Qwen/Qwen3.6-27B"
LENS_REPO = "neuronpedia/jacobian-lens"
LENS_REVISION = "qwen-n1000"
LENS_FILE = {
"Qwen/Qwen3.5-4B": "qwen3.5-4b/jlens/Salesforce-wikitext/Qwen3.5-4B_jacobian_lens_n1000.pt",
"Qwen/Qwen3.6-27B": "qwen3.6-27b/jlens/Salesforce-wikitext/Qwen3.6-27B_jacobian_lens_n1000.pt",
}[MODEL_NAME] 1. Load the model
jlens.from_hf wraps an already-loaded HuggingFace model into LensModel interface
[2]import torch
import transformers
hf_model = transformers.AutoModelForCausalLM.from_pretrained(
MODEL_NAME, dtype=torch.bfloat16
).cuda()
tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
model = jlens.from_hf(hf_model, tokenizer)
model [ERROR] `loss` is part of Qwen3_5CausalLMOutputWithPast.__init__'s signature, but not documented. Make sure to add it to the docstring of the function in /data/agastyas/Miniconda3/envs/ml/lib/python3.13/site-packages/transformers/models/qwen3_5/modeling_qwen3_5.py.
[ERROR] `logits` is part of Qwen3_5CausalLMOutputWithPast.__init__'s signature, but not documented. Make sure to add it to the docstring of the function in /data/agastyas/Miniconda3/envs/ml/lib/python3.13/site-packages/transformers/models/qwen3_5/modeling_qwen3_5.py.
HFLensModel(Qwen3_5ForCausalLM, n_layers=32, d_model=2560)
2. Load a pre-fitted lens
JacobianLens.from_pretrained pulls a .pt from the Hub (or a local path). The lens holds one [d_model, d_model] matrix per layer.
[3]lens = jlens.JacobianLens.from_pretrained(
LENS_REPO, filename=LENS_FILE, revision=LENS_REVISION
)
lens JacobianLens(d_model=2560, n_prompts=1000, source_layers=[0..30] (31 layers))
3. Apply: J-lens vs logit lens
lens.apply(model, prompt, positions=...) runs one forward pass, transports each layer's residual into the final-layer basis with J_l, and decodes through the model's own unembedding. use_jacobian=False skips the transport — that's the vanilla logit lens.
Below: a two-hop factual question, read out at the boot token. The J-lens surfaces interpretable tokens at layers where the logit lens is still noise.
[4]prompt = "Fact: The currency used in the country shaped like a boot is"
layers = [
model.n_layers // 4,
model.n_layers // 2,
model.n_layers // 4 * 3,
model.n_layers - 2,
]
jlens_logits, model_logits, _ = lens.apply(model, prompt, layers=layers, positions=[-2])
logit_lens, _, _ = lens.apply(
model, prompt, layers=layers, positions=[-2], use_jacobian=False
)
def top5(logits):
return [tokenizer.decode([t]) for t in logits.topk(5).indices]
for layer in layers:
print(f"L{layer:>3} logit-lens: {top5(logit_lens[layer][0])}")
print(f"L{layer:>3} J-lens: {top5(jlens_logits[layer][0])}")
print(f"model: {top5(model_logits[0])}") L 8 logit-lens: ['oman', 'edom', 'ולי', 'GPC', ' Urlaubs']
L 8 J-lens: [' `', ' boots', ' *', ' `\\', ' heel']
L 16 logit-lens: ['shaw', 'วย', 'amaz', 'REA', '举世']
L 16 J-lens: ['?', '？', "'?", '____', ' Italy']
L 24 logit-lens: ['的形状', '形状的', 'shape', '形状', '-shaped']
L 24 J-lens: ['-shaped', ' shape', ' shaped', 'shape', '形状']
L 30 logit-lens: [' is', ' shape', '-shaped', ' heel', ' shaped']
L 30 J-lens: [' is', ' shape', ' shaped', '-shaped', ' heel']
model: [' is', ' in', ' on', '.', ' with']
4. Render a slice page (inline)
compute_slice + build_page produce an interactive position × layer view of the lens's token ranks (the ? in the corner explains the controls). mode="embed" inlines everything so the page is self-contained.
[5]import gzip
import json
from jlens.examples import EXAMPLES, resolve_prompt
from jlens.vis import build_page, compute_slice, notebook_iframe
# English gloss for Qwen's Chinese/Japanese/Korean vocab tokens (machine-
# generated, best-effort), shown next to
# the token in the page (alt_token=).
gloss = {
int(k): v for k, v in json.load(gzip.open("assets/qwen_gloss.json.gz")).items()
}
example = next(e for e in EXAMPLES if e.slug == "multihop")
prompt = resolve_prompt(example, tokenizer)
slice_data = compute_slice(
model,
lens,
prompt,
layer_stride=2,
# Empirically on Qwen, the interesting word tokens trail punctuation and
# single-character tokens in the raw top-K; mask to word-like tokens only.
mask_display=True,
)
page, _, _ = build_page(
slice_data,
prompt,
title=example.section,
description=example.description,
alt_token=gloss,
)
notebook_iframe(page) [interactive slice page: saved as slice.html]
Output: open the interactive slice page (position × layer view of the lens readouts).
5. Render a slice page (served)
For longer prompts prefer mode="fetch": build_page writes the data as sidecar files to out_dir and the page fetches rank files lazily on pin, so it stays small regardless of how many tokens are tracked.
[ ]import os
import threading
from functools import partial
from http.server import HTTPServer, SimpleHTTPRequestHandler
from pathlib import Path
example = next(e for e in EXAMPLES if e.slug == "ascii-face")
prompt = resolve_prompt(example, tokenizer)
slice_data = compute_slice(model, lens, prompt, mask_display=True)
out_dir = Path("slices") / example.slug
page, _, _ = build_page(
slice_data,
prompt,
title=example.section,
description=example.description,
alt_token=gloss,
mode="fetch",
out_dir=out_dir,
)
(out_dir / "index.html").write_text(page)
if "_jlens_httpd" not in globals():
_handler = partial(SimpleHTTPRequestHandler, directory=os.path.abspath("slices"))
_jlens_httpd = HTTPServer(("127.0.0.1", 0), _handler)
threading.Thread(target=_jlens_httpd.serve_forever, daemon=True).start()
print(f"-> http://localhost:{_jlens_httpd.server_address[1]}/{example.slug}/") Not run here.
More to explore
A few more prompts are bundled in jlens.examples.EXAMPLES — change the slug above and see what surfaces, or try a prompt of your own.
[6]for e in EXAMPLES:
print(f"{e.slug:>24} {e.section}") multihop Multi-hop reasoning
modulation-topic Voluntary modulation: topic
modulation-arithmetic Voluntary modulation: arithmetic
ascii-face ASCII face
off-by-one Bug in code
overdose-flag Overdose flag
greatest-fear Greatest fear (don't say it)
blackmail Agentic Misalignment (blackmail honeypot)
6. Fitting
fit(model, prompts) computes J_l over the supplied prompts. 100 prompts is enough for a usable lens; the released lenses use 1000. dim_batch is the memory knob — each prompt does ceil(d_model / dim_batch) backward passes on a retained graph.
[ ]from jlens.examples import load_wikitext_prompts
prompts = load_wikitext_prompts(n_prompts=100)
lens = jlens.fit(
model, prompts, dim_batch=32, max_seq_len=128, checkpoint_path="ckpt.pt"
)
lens.save("jacobian_lens.pt") Not run here.