You can deploy this model and start earning money today!

DrOwzer/
gpt-oss-20b-tuned-lens

Name: gpt-oss-20b-tuned-lens
Availability: InStock

tuned-lens

interpretability

mechanistic-interpretability

Views

No views yet

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

file	contents
`config.json`	`TunedLensConfig` fields (`d_model=2880`, `num_hidden_layers=24`, `lens_type="linear_tuned_lens"`, ...)
`params.pt`	`layer_translators.state_dict()` — keys `"0.weight"`,`"0.bias"`,...,`"23.weight"`,`"23.bias"`

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:

metric	logit lens	tuned (FineWeb, on-dist)	tuned (Wikipedia, off-dist)
mean KL (nats)	~7	0.55	1.04
mean top-1 agreement	~10%	48.7%	45.8%

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):

python

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.

DrOwzer/gpt-oss-20b-tuned-lens