U-Space self-fitted Jacobian lenses
Jacobian lenses fitted for checkpoints that have none published, for the U-Space
uncertainty work (github.com/nils-loose/vfield-mi, branch iclr27-artifact).
Fitted with Anthropic's reference implementation (anthropics/jacobian-lens, Apache-2.0) — one-hot cotangents at every valid target position, summed over targets, averaged over source positions and prompts, first 16 positions skipped as attention sinks.
| file | model | blocks fitted | relative depth | prompts | d_model |
|---|---|---|---|---|---|
olmo_think_jacobian_lens.pt |
allenai/Olmo-3-32B-Think (64 blocks) |
38, 42, 47, 50, 54 | 0.609, 0.672, 0.750, 0.797, 0.859 | 100 wikitext-103 | 5120 |
gemma4it_selffit_jacobian_lens.pt |
google/gemma-4-31B-it (60 blocks) |
36, 39, 44, 47, 50 | 0.617, 0.667, 0.750, 0.800, 0.850 | 100 wikitext-103 | 5376 |
olmo_base100_jacobian_lens.pt |
allenai/Olmo-3-1125-32B (base) (64 blocks) |
42, 47, 50 | 0.672, 0.750, 0.797 | 100 wikitext-103 | 5120 |
magistral_100_jacobian_lens.pt |
mistralai/Magistral-Small-2507 (40 blocks) |
22, 25, 29, 33, 37 | 0.575, 0.650, 0.750, 0.850, 0.950 | 100 wikitext-103 | 5120 |
Only five blocks per model are present, not a contiguous band. Fitting is dominated by
the number of layers (each is a d_model² fp32 accumulator touched on every backward
pass), so only the depths actually read were fitted. Relative depth 0.750 is present in
both; any other layer raises a KeyError. Refit with
fit_lens.py
if you need more.
The reference implementation suggests ~1000 prompts and notes quality saturates near 100; these use 100, so they are usable but at the low end.
olmo_think exists because AI2's published lenses are fitted on base checkpoints and the
reasoning checkpoint is -Think; the lens must match the checkpoint the states come from.
gemma4it_selffit was fitted as a comparison against the published gemma-4-31B lens,
which was subsequently verified correct — so it is a control, not a replacement.
Use
import torch
from huggingface_hub import hf_hub_download
z = torch.load(hf_hub_download("tbrx/uspace-lenses",
"olmo_think_jacobian_lens.pt"), map_location="cpu",
weights_only=False)
J = z["J"] # {block -> [d_model, d_model]}
A lens transports the output of block b (hidden-state index b+1), so relative depth
is (b+1)/n_blocks. The token direction for vocabulary item v is
norm((W_U[v] * gamma) @ J[b]).
Which Olmo lens to use
Two are here and they are fitted on different checkpoints, which is the whole point:
olmo_think_jacobian_lens.pt—Olmo-3-32B-Think, the reasoning checkpoint our states come from. Agreement with AI2's published base lens is cosine 0.962.olmo_base100_jacobian_lens.pt—Olmo-3-1125-32B, the base checkpoint, i.e. a reproduction of what AI2 published. Agreement is cosine 0.996, which is what says our fitting procedure is correct rather than the lens being special.
On our end-of-think readout the base-checkpoint lens scores slightly higher than the Think self-fit (macro AUROC 0.675 / length-matched 0.615, against 0.638 / 0.607), so the matched-checkpoint argument does not survive contact with the data and either can be used. Report whichever you use, and note that Gemma is the outlier of the three families here: our Gemma self-fit agrees with the published lens at only cosine 0.747, against 0.997 for Qwen and 0.996 for Olmo.