gpt-oss-20b-tuned-lens

Tuned Lens for openai/gpt-oss-20b

A tuned lens (Belrose et al. 2023, Eliciting Latent Predictions from Transformers with the Tuned Lens) for openai/gpt-oss-20b. As of training there was no published tuned lens for any gpt-oss model, and AlignmentResearch/tuned-lens's model_surgery does not recognise GptOssForCausalLM (its layer/norm/unembed lookup is a hardcoded isinstance cascade). This lens was trained with a self-contained trainer (../tuned_lens_gptoss.py) that reuses this project's MXFP4 loader, but the artifact is saved in tuned-lens's native format so it is a drop-in TunedLens.

What a tuned lens is

For each layer idx, a learned affine map corrects the residual stream so the model's own final norm + unembedding decodes it into a faithful next-token distribution:

transform_hidden(h, idx) = h + translator[idx](h)      # residual delta, zero-init == logit lens
lens_logits(h, idx)      = lm_head(final_norm(transform_hidden(h, idx)))

Translators (idx 0..23) decode hidden_states[0..23]; hidden_states[24] is the model's own output. Objective: KL(model_final_dist || lens_dist) — i.e. distil the model's output distribution (NOT cross-entropy to the true next token). This is the defining tuned-lens objective.

Files

Training

• Data: HuggingFaceFW/fineweb (sample-10BT), streamed, packed into 1024-token sequences.
• Tokens: 20M (2441 steps, batch 1 × accum 8 = 8192 tok/step).
• Optimiser: AdamW, lr 1e-3, cosine decay + 3% warmup, grad-clip 1.0. Model frozen (MXFP4); only the 24 affine maps (~200M params total) are trained, in fp32.
• Hardware: single 24GB RTX 3090, ~13.1h, ~420 tok/s. Model forward under no_grad; per-layer backward+free keeps the 201k-vocab logits for at most one layer live at a time.
• Final training meanKL/layer ≈ 0.53.

Evaluation (held-out)

mean KL-to-model and top-1 agreement with the model, tuned vs the logit-lens baseline, on held-out sequences. Two distributions:

• On-distribution (80 held-out FineWeb seqs, disjoint shuffle seed): mean KL 0.55, top-1 48.7%.
• Off-distribution (60 Wikipedia extracts): mean KL 1.04 — the extra ~0.5 nats is distribution shift, not lens error; it generalises off-distribution and the improvement there is understated.

Biggest gains in the early/mid layers where the logit lens is blind (L0: KL 18.2→1.08, top-1 0%→23%; L8: 7.2→0.59, 0%→46%). Smooth monotonic; deep layers excellent (L23: KL 0.07, top-1 82%). The lens is well-fit on-distribution — the mid-layer softness seen off-distribution is largely distribution shift, so further training buys little.

Functional test (predict)

Per-layer top-1 trajectory for prompts with obvious continuations shows the tuned lens decodes the model's actual next-token earlier than the logit lens and emits coherent tokens mid-stack where the logit lens emits artifacts. E.g. "Romeo and Juliet was written by William ___": tuned lens locks onto "Shakespeare" at layer 4 and holds it; logit lens shows garbage until layer 21. Paris/oxygen/cold/ Jupiter resolve ~layer 19–20 (tuned) vs 21–22 (logit).

python ../tuned_lens_gptoss.py predict --lens . --prompt "The capital of France is"

Reproduce:

python ../tuned_lens_gptoss.py eval --lens . --data-file <held-out.jsonl> --seqs 60 --split-skip 0

Loading

Self-contained (no tuned-lens install needed) — see ../tuned_lens_gptoss.py (load_lens + lens_logits).

To load as an upstream TunedLens (requires adding a GptOssForCausalLM branch to tuned-lens's model_surgery, since it is otherwise unsupported):

from tuned_lens.nn.lenses import TunedLens          # after patching model_surgery for gpt-oss
lens = TunedLens.from_unembed_and_pretrained(unembed, lens_resource_id="<this dir>")

Status / provenance

Trained 2026-07-13 on a single RTX 3090. Benign interpretability artifact, released for research use.

Note on loading via upstream tuned-lens: as of release, AlignmentResearch/tuned-lens's model_surgery does not yet recognise GptOssForCausalLM, so from_model_and_pretrained needs a small model_surgery branch for gpt-oss (mapping to model.model.layers / model.model.norm / lm_head). Until that lands upstream, load with the self-contained helpers in tuned_lens_gptoss.py (load_lens + lens_logits), which reproduce the exact tuned-lens convention. A model_surgery PR is planned.