Fit a Jacobian lens for any HuggingFace decoder and read it out, from a browser page or the command line, entirely on your own machine.
A Jacobian lens is a per-layer linear map J_l that transports the residual
stream at layer l to the model's output layer, so unembed(J_l @ h) shows
what the model is heading towards at that depth. The estimation comes from
anthropics/jacobian-lens
(jlens); lensforge wraps it with the checks, corpus handling, evaluation,
export and UI you need to actually use it.
git clone --recursive https://github.com/<you>/lensforge
cd lensforge
python -m venv .venv && . .venv/bin/activate # Windows: .venv\Scripts\activate
pip install -e ./jacobian-lens # jlens first: it is not on PyPI
pip install -e ".[all]" # lensforge + UI + datasetsIf you cloned without --recursive, run git submodule update --init first.
On Windows pip install torch gives you the CPU build; install the CUDA wheel
from pytorch.org if you have a GPU.
lensforge ui # http://127.0.0.1:7860(python -m lensforge.cli ui is the same thing. Flags: --port, --host 0.0.0.0 for a remote GPU box, --no-browser.)
The page has four steps, in the order you use them. Everything runs locally: the model is downloaded from the Hub to the HF cache, never uploaded.
| field | what to put |
|---|---|
| model | A HuggingFace link (https://huggingface.co/org/model), an org/model id, or a local folder with the weights. Any causal decoder with a vocabulary head works; models with a classification head instead do not. |
| dtype | bf16 (default) on a GPU. fp16 if bf16 is unsupported. fp32 on CPU. |
| device_map | Blank loads on one device (GPU if available). auto spreads a large model across GPUs and CPU. |
Load model downloads the weights on first use and keeps the model in memory for the next steps.
Checks that the model can be lensed at all and measures what the fit will cost on this machine. Run it before every fit.
| field | what to put |
|---|---|
| prompts you plan to fit on | Only used for the cost estimate. Put the number you intend to use in step 3. |
| dim_batch | Directions of the Jacobian computed per backward pass. Backward passes per prompt = d_model / dim_batch. Higher is faster and uses more memory: 16 is a safe start, 64 on a large GPU. |
| source layers | How many layers get a lens, spread evenly across depth. 25 is the paper's convention. 0 fits every layer (storage grows as layers × d_model²). |
The result is a table of checks and a cost card:
| check | what a fail means |
|---|---|
layout |
Residual blocks, final norm or LM head are not where jlens guessed. |
probe_length |
Prompts are too short: the first 16 tokens are skipped and nothing is left. |
unembed_fidelity |
The detected norm + head do not reproduce the model's own logits (wrong layout). |
batch_determinism |
Dropout is on, or routing/kernels are non-deterministic. The estimator needs identical rows across the batch. |
grad_paths |
A source layer receives no gradient, or gradients are non-finite under fp16. |
| estimate | meaning |
|---|---|
backward_passes_per_prompt |
d_model / dim_batch |
measured_forward_s, measured_backward_s |
Timed on this machine with a probe prompt. |
est_per_prompt_s, est_total_h |
Extrapolated wall-clock for the number of prompts you entered. |
checkpoint_size_gb, exported_lens_size_gb |
Disk used during the fit and by the exported lens.pt (fp16). |
A fail blocks the fit: the lens would load fine and read garbage.
The corpus decides what the lens is calibrated for. J_l is one linear map
averaged over the fit prompts, so use mostly generic text plus a minority of
prompts that look like what you will inspect.
| field | what to put |
|---|---|
| corpus A / weight A | Main source and its share of the mix. wikitext streams WikiText-103. |
| corpus B / weight B | Optional second source. file:path.jsonl (one {"text": …} per line), file:path.txt (one prompt per line), or hf:repo/name:config:split:field for any HF dataset. examples/corpus.jsonl ships 400 wikitext paragraphs for offline use. |
| fit prompts | Prompts the Jacobian is averaged over. 100–300 is typical; more is smoother and slower. |
| held-out eval | Prompts kept aside to score the lens against the logit lens. 20 is enough. |
| dim_batch, source layers | Same meaning as in preflight. |
| max_seq_len | Tokens per prompt. Longer gives more positions per prompt and uses more memory. 128 is the default. |
| output directory | Where lens.pt, fit_config.json, the model card and the checkpoint go. |
| wrap prompts in the chat template | On for chat/instruct models, so the fit distribution matches how you will prompt them. Off for base models. |
Fit lens streams the progress log. The fit is resumable: run it again with
the same output directory and it continues from ckpt.pt. At the end you get
the evaluation table:
| column | meaning |
|---|---|
| J-lens top-1 | Fraction of positions where the lens at that layer already predicts the model's final top-1 token. |
| logit-lens top-1 | Same for the plain logit lens (unembedding the residual directly): the baseline to beat. |
| J-lens rank / logit-lens rank | Median rank of the model's final token in each readout. Lower is better. |
A good lens beats the logit lens in the early layers and converges to the model's output at the end. If it is already near-perfect at the first layer it is skipping ahead (describing the transport, not the layer); the verdict line says so.
Already have a lens? loads an existing lens.pt (path or HF repo id) for
the model currently loaded, skipping the fit.
| field | what to put |
|---|---|
| prompt | The text to read out. |
| tokens to track | Comma-separated tokens whose rank you want to follow across depth, e.g. yes, no or a classifier's label tokens. A multi-token entry uses its first token (a warning says so). |
| positions shown | How many trailing token positions the grid displays. The forward pass still uses the whole prompt. |
| layer stride | Show every Nth fitted layer. 2–3 keeps the grid short on deep models. |
| chat template | Wrap the prompt the same way as during the fit. |
Read out produces two things:
- Grid — one row per fitted layer, one column per position, the top-1 token the lens reads at each cell. The bottom row (model) is the model's own output. A cell in full colour matches that row; a pale cell is a token the model emits at some other position. A working lens fills in from the bottom up.
- Rank vs depth — full-vocabulary rank of each tracked token at the last position, log scale, lower is better. When a decision hinges on specific tokens, this curve is the feature to threshold on, not the top-1.
Build the full d3 page writes a self-contained interactive HTML page for the prompt (click a cell to pin a token and get rank charts) and links it.
Same operations, for scripting and for sharded fits on several machines:
lensforge preflight org/model --n-prompts 200
lensforge fit org/model --out-dir out/model --corpus "wikitext:0.7" \
--corpus "file:examples/corpus.jsonl:0.3" --chat-template \
--n-prompts 200 --n-eval 20 --dim-batch 16 --track yes --track no
lensforge eval org/model out/model/lens.pt --track yes
lensforge inspect org/model out/model/lens.pt --repl --chat-template --pin yes --pin no
lensforge inspect org/model out/model/lens.pt --prompt-file questions.txt \
--chat-template --pin yes --out-dir out/slices # HTML pages + index
lensforge fit org/model --out-dir out/s0 --shards 4 --shard 0 # on each node
lensforge merge out/s*/lens.pt --out out/merged/lens.ptlensforge <command> --help lists every flag. --push-to-hub me/my-lens
uploads an export.
- The lens is correlational: a token ranking high is not evidence the model will emit it. Testing that needs interventions in lens coordinates.
- A lens fitted on one distribution is extrapolating on another. Fit on what you will inspect.
- Models with a classification head instead of a vocabulary unembedding have nothing to decode into.
lensforge/ cli.py, server.py, static/ (the UI), preflight.py, corpus.py,
evaluate.py, export.py, readout.py
jacobian-lens/ git submodule: the jlens reference implementation
examples/ corpus.jsonl
scripts/ download_model.py, dump_corpus.py
screenshots/
tests/ python -m pytest tests (CPU, no network)
Weights (models/), fitted lenses (out/, *.pt) and generated pages are
git-ignored.
jacobian-lens is © Anthropic PBC, Apache-2.0, companion code for
Verbalizable Representations Form a Global Workspace in Language Models (2026).