uspace-lenses

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.

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.