Abstract
A natural way to inspect a transformer is to ask what each intermediate layer already "believes" about the next token. The classic logit lens does this by taking a hidden state from some layer and applying the model's final unembedding matrix directly. The narrow question in this memo is why that simple trick often feels insightful in some models yet brittle or misleading in others. The best answer from the literature is that intermediate residual states are not expressed in the same coordinate system as the final readout layer. As activations drift across layers, direct unembedding confuses geometry with belief. Belrose et al. (2023) formalize that problem and propose the tuned lens, which learns a small affine translator for each layer before projecting into token space. My view is that this matters not just because the tuned lens is a better interpretability tool, but because it sharpens a deeper claim: transformer layers do not simply reveal progressively clearer logits in a fixed basis. They repeatedly re-encode information before handing it to the final decoder.
Related Work
The central source is Belrose et al., Eliciting Latent Predictions from Transformers with the Tuned Lens (2023). The paper starts from the popularity of the logit lens and then shows why naive unembedding is unreliable across modern autoregressive transformers. The authors document two core failures: systematic bias in the layerwise predictions and representation drift between earlier hidden states and the final-layer space expected by the unembedding matrix.
A useful conceptual predecessor is the original logit-lens idea, which treated the tied unembedding matrix as a window into intermediate computations. That framing was appealing because it required no training and made hidden states legible in token space. But the tuned-lens paper shows that this convenience comes with a strong assumption: namely, that the same decoder basis is semantically valid at every depth. The literature now treats that assumption much more cautiously.
An auxiliary perspective comes from Hou and Tsvetkov (2023), who likewise argue that internal layers do not map cleanly back to token space without additional care. Their decoding method is aimed at saliency rather than next-token forecasting, but it reinforces the same lesson: later layers are more decision-oriented, and simple projections from arbitrary hidden states can blur the actual computation.
Method/Mechanism
The logit lens is simple. Take the hidden state at layer l, apply the final layer norm if needed, multiply by the model's unembedding matrix, and interpret the resulting token distribution as the layer's current next-token prediction. If all layers lived in nearly the same representational basis, this would give a faithful trajectory of belief updating through depth.
The tuned-lens paper argues that this assumption fails for two reasons. First, the raw logit lens is biased: averaged over a dataset, it systematically overweights some vocabulary items and underweights others relative to the model's actual final output. Second, hidden-state covariance drifts across layers. Earlier states can contain outlier directions and layer-specific geometry that make the final unembedding matrix a poor decoder for them. In effect, the model is still "thinking" in an internal dialect that the last-layer readout only partially understands.
The tuned lens inserts a learned affine translator for each layer before unembedding. Each translator is trained to make the translated hidden state predict the same token distribution as the final layer, using a KL objective. Mechanistically, this is modest: it does not claim that early layers already know the full answer, only that whatever predictive information they do contain should first be expressed in the final layer's basis before we read it out. That correction turns the tool from "apply the last decoder everywhere" into "learn the right local change of coordinates, then decode."
Key Findings
Two case studies make the distinction concrete:
- Case study 1: GPT-Neo-2.7B shows systematic logit-lens bias. Belrose et al. report that the raw logit lens is not merely noisy; it is predictably biased toward certain vocabulary items across layers. That means a layerwise probability trajectory can look like gradual evidence accumulation when part of the movement is really decoder mismatch.
- Case study 2: BLOOM and OPT-125M expose brittleness on modern models. The paper finds that on some architectures the raw logit lens often predicts the input token or other implausible continuations in many layers, while the tuned lens substantially lowers perplexity. This is the strongest practical evidence that decoder mismatch is not a small calibration issue; it can qualitatively distort what we think a layer represents.
Five crisp insights follow:
- Intermediate states are not natively in token space. They may contain predictive information without being directly readable by the final unembedding.
- Layerwise interpretability needs coordinate awareness. A bad basis can create fake narratives about when a model "decides."
- Low perplexity is not the only criterion. The tuned lens also reduces systematic vocabulary bias, which matters if we want to interpret probability trajectories as beliefs.
- Transformers repeatedly reformat information. Hidden states appear to be optimized for computation first and human-readable decoding only at the end.
- Interpretability tools can fail by projection error rather than conceptual weakness. Sometimes the model contains the signal and the probe is what is misaligned.
One alignment-relevant implication is methodological. Many attempts to diagnose latent reasoning, deception, or factual recall use layerwise token-space readouts as evidence for what the model "knew" at a given depth. The tuned-lens result says those claims should be discounted when the probe itself is mismatched to the layer geometry. A misleading readout can turn a probe artifact into a story about the model's cognition.
Limitations
The tuned lens is still a probe, not a proof of computation. Matching final-layer token distributions does not guarantee that a layer internally represents a clean next-token belief state. The translator could partially hallucinate a readable distribution from imperfect clues, especially if several latent states support similar final outputs. So the method improves faithfulness relative to the logit lens, but it does not collapse the gap between "decodable" and "causally used."
There is also a dependence on the final output as supervision. The method learns translators to agree with the model's own last-layer distribution, which is excellent for consistency but limits what it can say about representations that are useful internally yet never surface in the final logits. And while the affine form is elegant, it may underfit genuinely nonlinear layer-to-layer changes in basis.
Finally, the tool is best for autoregressive models with a relatively clean iterative residual-stream story. The paper itself is most compelling in pre-LN settings. Other architectures may require a more careful notion of where and how to insert the translator.
Future Directions
One direction is causal validation: combine tuned-lens readouts with interventions that test whether the decoded token preferences are actually used downstream. Another is richer geometry. If affine translators already help this much, then more structured probes might capture when specific subspaces, not whole hidden states, carry token-space commitments.
A second direction concerns mechanistic workflows. The tuned lens could be paired with circuit tracing or feature-level methods to separate two questions that are often conflated: what token-space prediction is decodable at a layer, and which internal features or heads make that prediction possible. That would connect layerwise readout tools to more causal interpretability traditions.
A third direction matters for alignment diagnostics. If we want to detect hidden knowledge, emerging deception, or latent goal conflict before it appears in outputs, we need readout methods whose failure modes are well characterized. The tuned lens is valuable partly because it identifies a concrete probe failure mode, representation drift, and fixes part of it.
Open question: can we build layerwise readout methods that are not only better calibrated than the tuned lens, but also demonstrably faithful to the causal computations that later layers actually consume rather than merely to the final logits they eventually produce?
Summary
The tuned lens beats the logit lens because intermediate transformer states are not written in the same basis that the final unembedding expects. Belrose et al. show that raw direct decoding is biased and brittle; a learned affine translator per layer recovers more faithful token-space readouts and lowers perplexity across several models. The broader lesson is that hidden states should be treated as computational states first, readable language objects second. If we ignore that distinction, we risk misreading projection artifacts as evidence about model cognition.
References
- Primary: Belrose et al. "Eliciting Latent Predictions from Transformers with the Tuned Lens." 2023. https://arxiv.org/abs/2303.08112
- Auxiliary: Hou and Tsvetkov. "Decoding Layer Saliency in Language Transformers." ICML 2023. https://proceedings.mlr.press/v202/hou23a.html
- Auxiliary: Logit Lens original writeup by nostalgebraist. 2020. https://www.greaterwrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens