Eliciting Latent Predictions from Transformers with the Tuned Lens

• Intermediate hidden states decode into a prediction trajectory that tends to converge smoothly to the final output distribution, with each successive layer achieving lower perplexity.The method uses early exiting to convert hidden states at each layer into distributions over the vocabulary.Source: §1 · p. 1
• The logit lens (nostalgebraist, 2020) can fail on BLOOM (Scao et al., 2022) and GPT Neo, while other early exits require changing training and cannot analyze pretrained models.Even when it works, representational drift across layers makes logit-lens outputs hard to interpret.Source: §1 · pp. 1–2
• The tuned lens trains a separate affine translator per layer with distillation to match final-layer logits, then composes each translator with the pretrained unembedding to decode vocabulary distributions.The translators map layer-specific representations into the basis expected by the final layer.Source: §1 · p. 2
• Tuned-lens predictions have substantially lower perplexity than logit-lens predictions and better represent the final-layer distribution; causal-basis ablations show influential tuned-lens features also affect model outputs.Causal basis extraction identifies influential residual-stream directions, and ablating these directions tends to disproportionately affect model output.Source: §1 · p. 2
• The tuned lens provides qualitative insight into pretrained language models by showing how their latent predictions evolve during a forward pass.The introduction points to Figure 1 and Section B for this analysis.Source: §1 · p. 2
• The tuned lens extends Halawi et al. (2023) to new models, detects prompt injection attacks (Perez and Ribeiro, 2022) often with near-perfect accuracy, and finds that examples requiring many training steps tend to be classified in later layers.Source: §1 · p. 2

• Nostalgebraist (2020) introduced the logit lens after observing roughly monotonic convergence of GPT-2 layerwise predictions, and later work used it to study few-shot processing and transformer weights.The cited applications include Halawi et al. (2023), Dar et al. (2022), Geva et al. (2022), and Millidge and Black (2022).Source: §2 · p. 2

The method

• The logit lens estimates outputs from an intermediate hidden state by setting later-layer residuals to zero and decoding the result with the unembedding matrix.The model is divided into layers up to a given layer, which produce hidden states, and later layers, which map those states to logits.Source: §2 · p. 2

Bias

• For GPT-Neo-2.7B, the logit lens is systematically biased toward some vocabulary items over others until the network’s very end, assigning them more probability mass than the final layer does.Source: Figure 3 · p. 3, §2 · p. 3

Unreliability

• The logit lens struggles to elicit predictions from BLOOM and OPT 125M, and its top-1 prediction is often the input token in more than half their layers.The information it extracts varies by model and layer, making it difficult to interpret.Source: §2 · p. 3

Bias

• Because predictable bias makes the trajectory unlike rational belief updating, an informed bettor could profit unboundedly by betting against downweighted tokens (Yudkowsky, 2007; Ramsey, 1926; De Finetti et al., 1937; Hacking, 1967; Garrabrant et al., 2016).Source: §2 · p. 3

Method

• Learned offsets and basis changes address arbitrary zero-residual assumptions and layerwise representation drift in the logit lens.The tuned lens replaces zero with a learned constant b_ℓ and uses a learned matrix A_ℓ to map each layer's output space to the final layer's input space.Source: §3 · pp. 3–4

Loss function

• Translators minimize KL divergence to the final-layer distribution, distilling the model's predictions without incentivizing probes to learn extra information.The final-layer distribution serves as a soft label.Source: §3 · p. 4

Muon optimizer

• Muon dramatically accelerates tuned-lens training and achieves much lower KL losses than the initial SGD experiments, revealing that earlier lenses were severely undertrained.Muon-trained lens matrices often have Frobenius norms several times larger than SGD-trained ones; the authors encourage practitioners to use Muon but did not rerun all experiments with it.Source: §3 · p. 5

Results

• The tuned lens has lower bias and lower perplexity than the logit lens across models, with lower perplexity variance across independently trained models.The reported perplexity comparison covers Pythia and GPT-NeoX-20B across layer indices and model sizes.Source: Figure 5 · p. 4, §3 · p. 5

Transferability across layers

• Nearby-layer transfer penalties are low and strongly negatively correlated with covariance similarity (Spearman ρ = −0.78), but are higher from layers with outlier dimensions to those without.The asymmetry is reported for transfer from Layer 5 and later versus the reverse direction.Source: §3 · p. 5, Figure 7 · p. 6

Transfer to fine-tuned models

• Base-model lenses transfer to fine-tuned models without retraining, outperforming the logit lens and comparing well with lenses trained specifically on the fine-tuned models.For Vicuna 13B, the transferred lens had at worst a 0.3 bits per byte KL increase on RedPajama and no significant difference on the Helpful Harmless corpus.Source: §3 · pp. 5–6

Relation to model stitching and probing

• The tuned lens stitches an intermediate layer to the unembedding with an affine transform, using d × d translators rather than |V| × d learned unembeddings.The authors report that training new unembeddings takes more steps and larger batches than training translators, and often converges to worse perplexity.Source: §3 · p. 6

4.1. Causal basis extraction

• Causal basis extraction finds orthogonal directions whose erasure maximally changes a function’s output, ranking them by expected KL divergence before and after erasure.The method extends amnesic probing’s search for one accuracy-degrading direction to multiple orthogonal directions.Source: §4.1 · p. 7
• Spearman ρ = 0.89 between CBE feature influence on the tuned lens and on the model at Pythia 410M layer 18, with no features influential only to the tuned lens.The model is somewhat more causally sensitive than the tuned lens.Source: §4.1 · p. 7, Figure 8 · p. 7

4.2. Stimulus-response alignment

• Stimulus-response alignment is the Aitchison cosine similarity between the tuned-lens prediction change under an intervention and the resulting final-layer prediction change.The weighted Aitchison inner product uses the final-layer control prediction distribution to set token weights.Source: §4.2 · p. 8
• Alignment is higher at later layers and somewhat higher at all layers for tuned-lens than logit-lens principal subspaces and stimuli in Pythia 160M.The interventions resample each layer’s top 10 causal basis directions from another input; Figure 9 reports average alignment.Source: §4.2 · p. 8, Figure 9 · p. 8

5.1. Eliciting Secret Knowledge from Language Models

• Cywi´nski et al. (2025) found that giving an auditor the logit lens’s top 20 latent predictions improves taboo-word guesses in clue-seeking and direct-reveal settings, but tuned-lens results vary by layer.Without peeking at ground-truth labels, it is not possible to determine in advance which layer to use.Source: §5.1 · p. 9

5.2. Extending Overthinking the Truth

• For BLOOM 560M, Neo 1.3B, and Neo 2.7B, calibrated performance under incorrect demonstrations peaks at 0.4-0.45 in an early layer, then falls to random-baseline performance at the final layer.This replicates earlier findings that intermediate predictions can resist incorrect demonstrations, including for models where the logit lens has high perplexity.Source: §5.2 · p. 9

5.3. Detecting Prompt Injections

• The tuned-lens detector achieves perfect or near-perfect AU-ROC on BoolQ, MNLI, QNLI, QQP, and SST-2, but SRM performs consistently well and wins on MC TACO and SciQ.The detector is fit only on normal-prompt trajectories and evaluated on held-out normal and injected prompts; the logit-lens version performs worse on most tasks.Source: Table 1 · p. 10, §5.3 · p. 10

5.4. Measuring Example Difficulty

• Tuned-lens prediction depth correlates positively with iteration learned on every investigated task and beats logit-lens depth in 8 of 11 tasks, sometimes dramatically.Prediction depth is the number of layers after which a prompt’s top-1 prediction stops changing; iteration learned is the earliest training step when that prediction becomes fixed.Source: §5.4 · pp. 10–11, Table 2 · p. 11

Discussion

• The tuned lens is a drop-in replacement for the logit lens that elicits interpretable prediction trajectories from pretrained language models and supports applications such as prompt-injection detection.The authors describe it as yielding qualitative and quantitative insights into model functioning.Source: §6 · p. 11
• Causal basis extraction identifies influential neural-network features but is computationally intensive, sequentially optimizing d_model causal basis vectors per layer.Future work could improve scalability by optimizing a k-dimensional subspace at each iteration.Source: §6 · p. 11
• Unlike the logit lens’s out-of-the-box use, the tuned lens requires per-layer translator training, but a full probe set trains in under an hour on one 8×A40 node.Released checkpoints for commonly used pretrained models should eliminate this training requirement for most applications.Source: §6 · p. 11
• The work focuses on language models because of space and time limits, though the authors consider applicability to other modalities likely.This is presented as a possibility, not demonstrated in the paper.Source: §6 · p. 11

A. Additional evaluation results

• Figure 13 compares a Vicuna 13B-specific lens, a lens transferred from LLaMA 13B, and Vicuna 13B’s logit lens using KL divergence from the model’s final output.The caption does not provide numerical results or identify a best-performing lens.Source: Figure 13 · p. 16
• Perplexities are reported for predictions from GPT-2, GPT-Neo, and OPT models, excluding OPT 350M because it uses a post-LN architecture.The figure caption identifies GPT-2 models (top), GPT-Neo models (middle), and OPT models (bottom).Source: Figure 14 · p. 17, A. Additional evaluation results · p. 17

B. Qualitative Results

• After “It was the best of times,” Pythia 12B’s tuned-lens prediction becomes very confident at early layers, suggesting some degree of memorization.The trajectory is shown for the first several words of Charles Dickens’ “A Tale of Two Cities.”Source: Figure 15 · p. 18
• A separate tuned-lens trajectory shows Pythia 12B prompted with the abstract of Brown et al. (2020).Source: Figure 16 · p. 18

B.1. Logit lens pathologies

• For BLOOM 560M, the logit lens assigns very high probability to the input token at many layers and token positions, complicating interpretation.Source: Figure 17 · p. 19
• Across BLOOM 560M and OPT 125M, logit-lens trajectories exhibit the input-token probability pathology that complicates interpretation.Source: Figure 17 · p. 19, Figure 18 · p. 19
• OPT 125M exhibits a similar logit-lens pathology to BLOOM 560M.Source: Figure 18 · p. 19

C. Transformers Perform Iterative Inference

• A first-order Taylor expansion shows that residual blocks are encouraged to reduce the inner product between each residual update and the loss gradient, aligning updates with the negative gradient.This gives the section’s theoretical framing of transformer residual connections as iterative inference.Source: C. Transformers Perform Iterative… · p. 20
• Jastrz˛ebski et al. (2017) found residual-update and loss-gradient cosine similarity consistently negative in ResNet image classifiers, especially in the final stage.The section reproduces their theoretical analysis for residual networks, which it says also applies to transformers.Source: C. Transformers Perform Iterative… · p. 20
• At least 95% of Pythia 6.9B layers’ residual-gradient cosine similarities were negative; observed magnitudes never exceeded 0.05, and virtually all fell below the random-vector 5th percentile, −6 × 10−4.Results cover roughly 16M Pile validation tokens; the gradient vectors have dimensionality 8,388,608.Source: Figure 19 · p. 20, C. Transformers Perform Iterative… · p. 20
• The Figure 19 comparison also reports Pythia 6.9B Pile-validation cross-entropy after replacing each of its 32 layers with the identity, against the intact-model perplexity baseline.The dashed line marks Pythia’s perplexity with all layers intact.Source: C. Transformers Perform Iterative… · p. 20

C.1. Zero-shot robustness to layer deletion

• Stochastic depth, also called LayerDrop, randomly drops layers during training and can be applied to ResNets and transformers.The passage attributes stochastic depth to Huang et al. (2016) and LayerDrop to Fan et al. (2019).Source: §C.1 · p. 20
• ResNets tolerate layer deletion even without stochastic depth, unlike CNN architectures without skip connections, supporting similarity between adjacent ResNet representations (Greff et al., 2016).Veit et al. (2016) reported the first layer's exceptional importance in ResNets, paralleling the transformer result.Source: §C.1 · p. 21
• Only the first transformer layer is crucial for performance; ablating every other layer causes a nearly imperceptible perplexity increase.The authors report this as a replication of a layer-deletion experiment in transformers and note that Veit et al. (2016) also found the first layer exceptionally important in ResNets.Source: §C.1 · p. 21

D.1. Ablation techniques

• The usual erasure strategy projects a representation onto the orthogonal complement of a subspace, zeroing its components along that subspace.Prior work has pointed out that zeroing activations may push a network out of distribution, complicating interpretation.Source: §D.1 · p. 21
• A probe may use a consistently positive direction as an implicit bias term, so zeroing it can raise loss despite its carrying no useful input information; the authors instead replace ablated directions with dataset-mean values.For a direction u, the ablation uses its projection matrix Pu and the mean representation x.Source: §D.1 · p. 21

E. Static interpretability analysis

• Static analysis with the tuned lens never decreases performance and sometimes improves it; toxicity editing outperforms the logit lens on OPT-125m and matches it on other models.The editing comparison concerns toxicity reduction.Source: E. Static interpretability analysis · p. 21
• In static analysis, the tuned lens never decreases performance and increases it for some models.Source: E. Static interpretability analysis · p. 21
• For toxicity reduction, tuned-lens model editing outperforms the logit lens on OPT-125m and performs equivalently on other models.Source: E. Static interpretability analysis · p. 21
• Because interpretability methods analyze parameters that write to intermediate hidden states, the authors test whether the tuned lens preserves their effectiveness as a less biased view.They test both static analysis and model editing.Source: E. Static interpretability analysis · p. 21

E.1. Static Analysis

• Static analysis projects parameter-derived hidden-state vectors through the unembedding matrix to identify top matching tokens; the tuned lens first maps a layer’s hidden state to the final hidden state.The parameter-derived vectors can yield interpretable semantic or syntactic token clusters.Source: §E.1 · pp. 21, 23
• For Pythia 125M, both lenses made parameters more interpretable than random matrices, and the tuned lens showed a larger random-versus-real score gap; improvement was much smaller for larger models.Many singular vectors were no more interpretable than random, with only a minority in a long right tail, so comparing against random matrices is necessary.Source: §E.1 · p. 23, Table 3 · p. 23

E.2. Interpretability score distributions

• Most singular vectors were no more interpretable than random, while a minority formed a long right tail; without random-matrix controls, tuned-lens gains would appear much larger than they were.This distribution was shown for OV-circuit singular vectors in Pythia 125M and was reported to occur for other parameters as well.Source: §E.2 · p. 23, Figure 21 · p. 24

E.3. Automated interpretability analysis method

• The automated metric averages pairwise cosine similarities between BPEmb token embeddings, and its scores usually matched human judgments but sometimes missed syntactic or set-level coherence.The pretrained encoder makes word similarities reproducible; it rarely treats human-uninterpretable lists as interpretable.Source: §E.3 · p. 24

E.4. Improved Model Editing

• On OPT-125m, tuned-lens editing outperformed logit-lens editing without significant perplexity increases, but toxicity reductions did not generalize to the other tested models.The other tested models showed no significant toxicity decrease (> 2%) compared with their unedited versions.Source: Table 4 · p. 25, §E.4 · p. 25

E.5. Implementation

• The initial implementation of many static-analysis experiments used the transformer_lens library.Source: §E.5 · p. 25