Eliciting Latent Predictions from Transformers with the Tuned Lens

In plain words

A language model writes its answer in many steps, one per layer. The logit lens is a cheap way to peek at the answer it is building at each step: read the intermediate state as if it were the final one. For many models that peek is misleading, because the intermediate states live in a different coordinate system from the final output, and some models hide their guesses behind the input word.

The tuned lens fixes this with a small learned correction for each layer. For each layer it learns a simple linear map that turns that layer's state into the coordinates the output uses, and it trains these maps so that the guessed next-word distribution matches what the full model finally says. The base model is never changed. The result is a sequence of guesses that is more readable and more reliable than the plain peek.

The authors also use the guess sequence for two things: to find which directions inside the model matter most for its guesses, and to spot prompts that were tampered with, because the guesses follow an unusual path for those prompts.

TL;DR

• Trains one affine translator per layer of a frozen transformer, so each hidden state decodes through the model's own unembedding into a next-token distribution §3 eq. 8.
• Training minimises KL divergence from the final-layer distribution to the lens output, a distillation loss that avoids teaching the probe new facts §3 eq. 9.
• Author-reported: lower perplexity and lower bias than the logit lens across the models tested, with models up to 20B parameters §2 §3 fig. 4 fig. 5 app. A.
• Author-reported: prompt-injection detection from the latent prediction trajectory reaches AUROC 0.98 to 1.00 on five of nine tasks with Pythia 12B §5.3 tab. 1.
• The code (MIT) is the official tuned-lens library; it matches the paper on most hyperparameters, but its learning-rate convention, translator parametrisation and the final-layer handling differ from the text (see Paper vs code).

Problem

A transformer's hidden states can be read as successive refinements of a guess for the next token §1 §2. The logit lens decodes each hidden state with the final layer norm and unembedding matrix, and this works for some models. For BLOOM and OPT 125M, the authors report that the logit lens often puts its highest probability on the input token rather than on a plausible continuation, in more than half of the layers §2 app. B.1. For GPT-Neo-2.7B its bias relative to the final output is about 4 to 5 bits for most layers, compared with 0.0068 bits for the final layer of Pythia 160M against Pythia 12B §2 fig. 3. The paper argues that a biased estimator cannot be read as belief updating, because a predictable direction of change could be exploited §2.

The authors also note that intermediate layers may drift in basis from the final layer, and that some hidden states contain a few very large-variance "rogue" dimensions §3. Earlier early-exit methods change the training process, so they cannot be applied to a pretrained model without retraining §1.

Method

Overview. Each layer of a pre-LayerNorm transformer adds a residual update to its hidden state eq. 1 §2. Applying this recursively gives the output as a function of any intermediate state, with the remaining residual updates summed in eq. 2. The logit lens sets those remaining updates to zero eq. 3. A variant keeps the last layer's update eq. 4 §2 and works better for some models but not all fig. 1 §2.

Step 1. Replace the zero residual with a learned translator §3 eq. 8. For each layer l the tuned lens applies an affine map A_l h + b_l to the hidden state, then the model's final norm and unembedding. The matrix is d x d, not |V| x d, so it has far fewer parameters than a new unembedding would, and is cheaper to train §3. The authors also tested a version with a learned constant in place of the zero residual eq. 7, as a step toward the final form §3.

Step 2. Train each translator by distillation §3 eq. 9. The target is the final-layer distribution f_{>l}(h_l), and the lens output is the approximation. The loss is KL(target || lens), averaged over positions. Each translator receives only its own layer's term, and the code sums those terms train_loop.py. The final-layer distribution serves as a soft label, so the probe is not trained on ground-truth labels and does not learn facts beyond what the model already encodes §3.

Step 3. Initialise and optimise §3. Translators start as the identity. The paper trains with SGD using Nesterov momentum, a linear learning-rate decay over 250 steps, gradient clipping at norm 1, and weight decay of 10^-3. Each optimiser step covers 2^18 tokens, accumulated over micro-batches. Full values are in the Hyperparameters table. Section 3 says that training with Muon instead of SGD converges much faster and reaches much lower KL, and that the authors did not redo the experiments with it §3. This was added after the first publication, so the text in section 3 is a later note.

# Tuned lens loss for one batch  [§3] [eq. 8] [eq. 9]
# Input:  frozen model, token batch (B, T), translators (A_l, b_l) per layer
# Output: scalar loss, averaged over positions

def tuned_lens_loss(model, tokens, translators):
    # Frozen forward pass that keeps every hidden state  [§3]
    outputs = model(tokens, output_hidden_states=True)
    # Final-layer log-distribution is the fixed target  [§3] [eq. 9]
    target = log_softmax(outputs.logits)             # (B, T, V)
    loss = 0.0
    # One translator per probed layer; the last hidden state is dropped  [§3]
    for layer, hidden in enumerate(outputs.hidden_states[:-1]):
        # Affine translator, then the final norm and unembedding  [eq. 8]
        translated = translators[layer](hidden)      # (B, T, d)
        lens_logp = log_softmax(unembed(final_norm(translated)))
        # KL(target || lens), summed over the vocabulary  [eq. 9]
        kl = sum(exp(target) * (target - lens_logp), dim=-1)
        # Average over positions, then add to the running total  [eq. 9]
        loss = loss + kl.mean()
    return loss

Step 4. Measure bias and perplexity §2 §3. Bias is the KL divergence between the marginal distributions over the vocabulary, averaged over the dataset, for each layer's lens against the final layer eq. 5 eq. 6. Perplexity is measured on predictions made from the lens output. The paper uses a random sample of pretraining validation tokens §3.

Step 5. Causal basis extraction §4.1 eq. 10 eq. 11 eq. 21. This finds directions that the lens relies on most. For a unit vector v, its influence is the expected KL between the lens output before and after erasing the component along v by replacing it with its mean over the dataset eq. 10 eq. 21. The directions are found one at a time. Each new direction maximises influence over the unit sphere, subject to being orthogonal to the earlier ones eq. 11. The authors call this a principal-features search, in analogy with PCA §4.1.

# Influence of a unit direction v for one layer  [§4.1] [eq. 10] [eq. 21]
# Input:  lens f, hidden states H (N, d), unit vector v (d,)
# Output: scalar KL, averaged over states

def influence(v, lens, hidden_states):
    # Mean state over the dataset, one row per position  [§4.1] [eq. 21]
    mean_state = hidden_states.mean(dim=0)           # (d,)
    # Component of the mean-minus-state gap along v  [eq. 21]
    shift = (mean_state - hidden_states) @ v         # (N,)
    # Erase v by moving each state's component along v to its mean  [eq. 21]
    erased = hidden_states + shift[:, None] * v      # (N, d)
    # KL between lens outputs before and after the erasure  [eq. 10]
    return kl(lens(hidden_states), lens(erased)).mean()

# Causal basis extraction for one layer  [§4.1] [eq. 11]
# Input:  lens f, hidden states H (N, d), basis size k
# Output: k orthonormal directions, sorted by influence

def causal_basis(lens, hidden_states, k):
    # Start from singular vectors of the lens-transformed unembedding  [§4.1]
    start = left_singular_vectors(lens_transform(unembedding.T))[:, :k]
    found = []
    for j in range(k):
        # Maximise influence on the unit sphere, orthogonal to earlier
        # vectors  [eq. 11]
        v = maximise_influence(lens, hidden_states, start[:, j], orth_to=found)
        found.append(v)
    # The search is non-convex, so influence can rise between steps;
    # sort at the end  [§4.1]
    return sort_by_influence(found, lens, hidden_states)

Step 6. Stimulus-response alignment §4.2 eq. 13 to eq. 18. This checks that a direction moves the lens and the model in the same way. The authors use a weighted Aitchison inner product on the probability simplex eq. 13, with weights from the control distribution. A stimulus is the change in the lens output caused by an intervention, and a response is the change in the final output eq. 16 eq. 17. Their alignment is the cosine similarity of these two changes under that inner product eq. 18. The intervention resamples the top 10 causal directions of the layer, one token at a time, from another random input §4.2.

Step 7. Applications §5. Four uses are reported. (a) Eliciting secret knowledge: the top 20 lens predictions are given to an auditor model, which guesses a taboo word §5.1. (b) Few-shot overthinking: the calibrated accuracy of predictions from intermediate layers under incorrect demonstrations §5.2 fig. 11. (c) Prompt injection: the log-probability of each answer at every layer is flattened into a vector, and an outlier model is fit only on normal prompts §5.3. (d) Example difficulty: the "prediction depth" of a prompt is the number of layers after which its top-1 prediction stops changing, compared with the iteration at which the model first learned the example §5.4.

Key equations, reproduced from the source:

\mathrm{TunedLens}_{\ell}(\mathbf{h}_{\ell}) = \mathrm{LogitLens}(A_{\ell}\mathbf{h}_{\ell} + \mathbf{b}_{\ell}) eq. 8

\mathop{\mathrm{argmin\:}} \mathbb{E}_{\mathbf{x}}\Big[D_{KL}(f_{>\ell}(\mathbf{h}_{\ell})\,||\,\mathrm{TunedLens}_{k}(\mathbf{h}_{\ell}))\Big] eq. 9

\sigma(\mathbf{v}; f) = \mathbb{E}_{\mathbf{h}}\Big[D_{KL}(f(\mathbf{h})\,||\,f(r(\mathbf{h}, \mathbf{v})))\Big] eq. 10

\mathbf{v}_i = \mathop{\mathrm{argmax\:}}_{||\mathbf{v}||_2=1} \sigma(\mathbf{v}; f) \quad \textrm{s.t.} \quad \langle \mathbf{v}, \mathbf{v}_j \rangle = \mathbf{0}, \; \forall j < i eq. 11

\mathbf{x}' = \mathbf{x} + P_{\mathbf{u}}(\overline{\mathbf{x}} - \mathbf{x}) eq. 21

Hyperparameters

Values as stated in the paper. "Code" means the official tuned-lens library at code_commit. "Not stated" means neither gives a value.

Results

Author-reported. Setups are given per table. Confidence intervals are in the paper; the point estimates are shown here.

Table 1. Prompt-injection detection, AUROC, Pythia 12B tab. 1 §5.3. Five tasks reach 0.98 to 1.00 with the tuned lens and an isolation forest or LOF detector. The baseline (SRM) is best on MC TACO and SciQ, and the tuned lens beats it on ARC-Challenge only. Last column: accuracy on normal prompts and on injected prompts tab. 1.

TL = tuned lens, LL = logit lens. The source gives 95% bootstrap intervals for each cell; they are omitted here and are in tab. 1.

Table 2. Correlation of prediction depth with iteration learned tab. 2 §5.4. Spearman rho across examples, Pythia 12B. The source marks the higher value in each row in bold; the tuned lens is higher on eight of the eleven tasks.

Table 3. Static interpretability scores, Pythia 125M tab. 3 app. E.1. Mean interpretability score (method in app. E.3), real parameters against randomly shuffled matrices. Higher is more interpretable. The source notes that the real-minus-random gap is consistently larger with the tuned lens.

Table 4. Model editing for toxicity, OPT-125m tab. 4 app. E.4. The caption says these are percentages of toxic generations; the values are printed as decimals in the source. Each row projects value vectors through the logit or tuned lens and keeps the top 5, 10 or 20 tokens; the most toxic vectors have their coefficient set to zero app. E.4. Perplexity is the last column.

The source reports that these gains did not carry over to pythia-125m, pythia-350m, gpt-2-medium or gpt-neo-125m, where no decrease above 2% was found app. E.4.

Text-only results.

• Logit-lens bias is about 4 to 5 bits for GPT-Neo-2.7B for most layers, and 0.0068 bits for the final layer of Pythia 160M against Pythia 12B §2 fig. 3.
• Tuned-lens perplexity is lower than the logit lens across the models tested, with or without the final layer in the probe §3 fig. 4 fig. 5 app. A.
• Transfer penalties between layers are low for nearby layers; their correlation with covariance similarity is Spearman rho = -0.78 §3 fig. 7.
• A lens trained on a base model and copied onto its fine-tune shows, at worst, a 0.3 bits-per-byte increase in KL on RedPajama compared with a lens trained on the fine-tune. On the helpful-harmless set the paper reports no significant difference §3 app. A.
• Causal basis extraction on Pythia 410M at layer 18 gives Spearman rho = 0.89 between influence on the lens and influence on the model. The least influential directions still change the model by at least 2 x 10^-3 bits §4.1 fig. 8.
• Stimulus-response alignment grows with depth, and is higher with the tuned lens than with the logit lens at each layer of Pythia 160M §4.2 fig. 9.
• Few-shot overthinking: calibrated accuracy under incorrect demonstrations peaks at an intermediate layer, around 0.4 to 0.45, in BLOOM 560M, GPT-Neo 1.3B and 2.7B, before falling to chance at the last layer §5.2 fig. 11.
• Secret elicitation: the tuned lens helps at some layers and hurts at others; the best layer cannot be chosen without the ground truth §5.1 fig. 10.

Paper vs code

The official repository is https://github.com/AlignmentResearch/tuned-lens. The paper links it in its abstract, and author Nora Belrose is its copyright holder in the LICENSE file. It is MIT-licensed. Commit abdac0c4de23d9f6d6c8459d576ad203aec15deb was cloned on 2026-11-09.

Checked: the lens definition, the training loop and its optimiser, the loss, the causal basis extraction and the ablation helpers. Not found in the release: the experiments for prompt injection, overthinking, example difficulty, secret elicitation, static analysis, layer deletion and model editing. The repository has the library, its scripts and an interactive notebook (notebooks/interactive.ipynb), and nothing that reproduces those tables and figures.

Matching the paper:

• Batch of 2^18 tokens per step §3 train_loop.py.
• 250 steps with linear decay §3 train_loop.py ingredients.py.
• Gradient clipping at norm 1 §3 train_loop.py.
• Weight decay 10^-3 §3 ingredients.py.
• Nesterov SGD §3 ingredients.py.
• KL loss from the final distribution to the lens §3 eq. 9 train_loop.py.
• Mean ablation in the CBE erasure §4.1 app. D.1 subspaces.py subspaces.py.
• L-BFGS with strong Wolfe line search §4.1 subspaces.py.

Differences and gaps, each with both citations:

• Learning rate value. Paper: base LR 1.0 §3. Code: lr = lr_scale * (1 - momentum) = 0.1 as the PyTorch argument, with lr_scale = 1.0. The comment says PyTorch scales the step by 1/(1 - momentum), so the effective LR is 1.0 ingredients.py ingredients.py.
• Momentum. Paper: not stated §3. Code: 0.9 ingredients.py.
• Translator form. Paper: A_l h + b_l with identity initialised A §3 eq. 8. Code: h + (W h + b) with W and b zero-initialised lenses.py lenses.py. The code's own comment says this makes weight decay pull toward the identity rather than the zero transform. The two are equivalent at initialisation, but the weight-decay target differs.
• Final layer. Paper: included for GPT-Neo, excluded for GPT-2, GPT-NeoX-20B, OPT and Pythia §3. Code: the training loop and the evaluation loop both drop the last hidden state for every model train_loop.py eval_loop.py. I did not find an option to include it.
• Muon. Paper: described as dramatically faster than SGD, with the earlier lenses called undertrained §3. Code: default optimiser is SGD ingredients.py; no Muon implementation appears in the repository.
• CBE initialisation. Paper: singular vectors of the probe §4.1. Code: left singular vectors of the unembedding matrix after lens.transform_hidden, with an author TODO asking whether transposes are right subspaces.py subspaces.py.
• CBE stopping. Paper: not stated beyond L-BFGS §4.1. Code: at most 100 iterations per vector, and stop when the relative energy change is below 1e-4 subspaces.py subspaces.py.
• Two ablations in the code. The subspace erasure used by CBE subspaces.py differs from ablate_layer, which replaces a layer's whole residual update with its batch mean ablation.py. The paper does not say which one each experiment used.
• Evaluation sample. Paper: 16.4M tokens §3. Code: not found.

Pitfalls

• Learning-rate convention. Copying the paper's 1.0 into torch.optim.SGD with momentum 0.9 gives an effective step ten times too large. Use 0.1 with momentum 0.9 to match the released code §3 ingredients.py.
• Momentum is not in the paper. Choose 0.9 unless you have a reason not to, and say so ingredients.py.
• Weight decay target. The paper's A h + b form decays A toward zero, which is a different regulariser from the code's residual form that decays toward the identity §3 eq. 8 lenses.py. Pick one and report it.
• Muon is not in the code. The paper's later note says Muon is much better, but the released code cannot reproduce it §3 ingredients.py. Treat the reported numbers as SGD-trained unless the paper says otherwise.
• GPT-Neo and the final layer. The paper includes the final layer for GPT-Neo, with a different LR (0.25 instead of 1.0 when keeping it) §3. The training loop used in the release always drops the last hidden state train_loop.py. Reproducing the GPT-Neo lens needs a decision that the paper does not spell out.
• Experiments are not reproducible from the release. Prompt injection, overthinking, secret elicitation, example difficulty and model editing have no code here §5 . All numbers in Results come from the paper alone.
• Example-difficulty task count. The text says "twelve multiple-choice tasks" §5.4, and Table 2 lists eleven tab. 2. The text also says the tuned lens wins on eight of eleven, which matches the table.
• Prompt-injection baselines. The detectors use scikit-learn defaults §5.3. Detection is trained on normal prompts only, and the injected string is fixed in the paper §5.3. The reported AUROCs say nothing about adaptive attacks.
• Table 1 vs the text. The text says the tuned lens beats SRM on one task (ARC-Challenge) §5.3. On ARC-Easy the tuned-lens LOF and SRM tie at 0.73 tab. 1.
• Toxicity table units. The caption says percentages, but the values are fractions as printed tab. 4 app. E.4. The reductions there come from a single model, OPT-125m; the paper says they did not generalise app. E.4.
• Equation and figure numbers come from the LaTeX source, not from a compiled PDF. Check float placement before citing them in print §3 app. A.
• Card based on arXiv v6. The source is the v6 LaTeX file (updated 2025-11-11). Its section 3 says that lenses in earlier versions were severely undertrained, and that the Muon note was added after first publication §3.
• Compute. The "under an hour on 8 x A40" claim is the authors' own §6.
• Lens checkpoints. The paper says it released checkpoints for common models §6. This was not verified here.

Lineage

• Builds on: the logit lens of nostalgebraist (2020) §2 eq. 3; its GPT-Neo extension of nostalgebraist (2021) §2 eq. 4; model stitching with affine maps (Lenc et al., 2015; Bansal et al., 2021) §3; amnesic probing (Elazar et al., 2021) §4.1; resampling ablation (Chan et al., 2022) §4.2; "Overthinking the Truth" (Halawi et al., 2023) §5.2; prediction depth (Baldock et al., 2021) §5.4; and early exit by CALM (Schuster et al., 2022) and DeeBERT (Xin et al., 2020) §5.4.
• Compared against: the logit lens and its GPT-Neo extension §2 fig. 1 fig. 4; SRM from Bai et al. (2022), the middle-layer baseline §5.3.
• Concurrent or later: Muon, cited as Jordan et al. in a note added after first publication §3.
• Superseded by: not stated in this wiki.