senku21x/olmo3-32b-jacobian-lens-checkpoints

Jacobian lenses for OLMo-3-32B, one per training checkpoint

These are per-checkpoint Jacobian lenses (J-lenses) for OLMo-3-1125-32B. They cover stage-1 pretraining checkpoints, the released base model (after mid-training) and the three Think post-training models. Each checkpoint has its own lens, fitted on that checkpoint's weights, so you can follow how the lens and its "J-space" form over training.

Contents

One folder per checkpoint.

Folder	Model / revision
stage1-step5000	allenai/Olmo-3-1125-32B @ stage1-step5000
stage1-step10000	@ stage1-step10000
stage1-step30000	@ stage1-step30000
stage1-step100000	@ stage1-step100000
stage1-step300000	@ stage1-step300000
stage1-step656000	@ stage1-step656000
main	@ main (base: pretraining + mid-training)
think-sft	allenai/Olmo-3-32B-Think-SFT
think-dpo	allenai/Olmo-3-32B-Think-DPO
think	allenai/Olmo-3-32B-Think

Each folder holds:

• lens_J.pt (≈ 3.1 GB): a dict {layer: fp16 tensor [5120, 5120]} for source layers 0–62. J[l] is indexed [target_dim, source_dim] and maps the residual stream after block l to the residual stream after the final block (63).
• e7b.json:
  • spectra of the readable operator (k50, PR) per layer;
  • 63 × 63 lens-vector CKA, CKA vs the logit lens, and cos(J-vector, logit-vector);
  • token-centred consistency;
  • bridge-entity readout on 2-hop prompts.
• e7d_l16_check.json (where present): top eigen-axis shares of JᵀMJ at layers 0/4/8/16/24.
• readout_items.json: per-item bridge readout ranks (J-lens vs logit lens) with a control entity.
• log.txt: the fitting log.

How the lenses were fitted

The lenses were fitted with Anthropic's reference implementation jacobian-lens, calling jlens.fitting.jacobian_for_prompt unmodified.

• Estimator: the future-summed Jacobian of the final-block output with respect to layer l. Cotangents are summed over all target positions p′ ≥ p, then averaged over source positions p and over prompts.
• Positions: the first 16 are skipped (skip_first = 16, the repo default; attention sinks) and the final position is excluded. That leaves 111 of 128 positions.
• Prompts: 10 wikitext-103-raw validation paragraphs of ≥ 600 characters (skipping the first 700 such), truncated to 128 tokens.
  • This is far fewer than the paper's 1000 sequences. Treat each lens as a single noisy estimate.
• Precision and tokenization: bf16 weights and autograd, with fp32 accumulation. No BOS is prepended at any checkpoint (force_bos=False). The base checkpoints have no BOS token; the post-training tokenizers do, and it was left off so that every checkpoint is tokenized the same way.

Usage

import torch
from transformers import AutoModelForCausalLM
J = torch.load("main/lens_J.pt", weights_only=True)          # {layer: [5120, 5120] fp16}
model = AutoModelForCausalLM.from_pretrained("allenai/Olmo-3-1125-32B", revision="main", dtype=torch.bfloat16)
# h: residual stream after block l, shape [..., 5120]
def jlens_logits(h, l):
    z = h.float() @ J[l].float().T                          # transport to the final-block basis
    return model.lm_head(model.model.norm(z.to(torch.bfloat16)))

Use the lens that matches the checkpoint (revision) whose activations you are reading.

Observations (from the accompanying study; single fits, post hoc, reviewed by CPU skeptic rounds)

Note on e7b.json. Its lens_cka, lens_cka_vs_logit and lens_vs_logit_cos fields used a flawed token set: 2294 of the 3000 "frequent" tokens were zero-count filler. Corrected CKA (top-3000 word tokens of wikitext-103 validation) is in cka_freqset.json at the repo root. The lenses themselves are unaffected.

• No sharp onset. Adjacent-layer lens CKA is ≥ 0.98 at every checkpoint.
• The lens leaves the output format in a front that sweeps up the stack during pretraining. The first layer whose J-lens resembles the logit lens (CKA ≥ 0.8) is 14 / 17 / 23 / 26 / 30 / 39 / 41 at 5k / 10k / 30k / 100k / 300k / 656k / main. The strongest CKA boundary is an early/mid edge that migrates upward with it (L4 → L22).
• Early-layer readout concentrates. From 30k–100k on, L0 readout energy concentrates in a few directions, and L16's top readout axis grows from ≈ 2% to 46% of readout energy. That axis is not one fixed residual direction across checkpoints.
• Think-SFT moves partway back (first logit-like layer 30; looser early layers). This is measured on raw wikitext without a chat template.
• Noise: each lens is a single 10-prompt fit, and no replicate fit exists. Early- and mid-layer differences under ~0.07 CKA between late checkpoints are not resolvable.

Why the lens moves away from the output format (cross-position vs within-position terms) has not been tested yet.