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

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

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