ArXiv: 2303.08112

🎯 Pitch

Hidden states inside large language models can be decoded into surprisingly coherent next-token predictions as early as the middle layers, and these latent prediction trajectories can spot prompt injection attacks with near-perfect accuracy. The tuned lens achieves this by training lightweight affine probes that eliminate the representational drift plaguing earlier methods, revealing what the model 'thinks' at each step of its computation.

1. Executive Summary

This paper introduces the tuned lens, an iterative inference probing method that decodes a pretrained transformer's hidden states at each layer into interpretable vocabulary distributions by training lightweight affine translators—one per layer—to align intermediate representations with the final unembedding space. Tested on autoregressive language models up to 20B parameters (GPT-2, GPT-Neo, Pythia, BLOOM, OPT, GPT-NeoX-20B), the tuned lens substantially outperforms the earlier logit lens technique, achieving uniformly lower perplexity and virtually eliminating the systematic marginal bias (reducing KL divergence from ~4–5 bits to near zero) that made logit lens outputs unreliable. The paper demonstrates that tuned lens prediction trajectories can be used to detect prompt injection attacks with near-perfect AUROC on several benchmarks and that the features most influential on the tuned lens are also causally influential on the model's final output (Spearman ρ = 0.89), establishing that the latent predictions reflect the model's own computational process rather than spurious probe artifacts.

2. Context and Motivation

The Core Problem: We Cannot Reliably See What Transformers Think Before They Output

The central gap this paper addresses is deceptively straightforward: when a transformer language model processes a sequence layer by layer, we lack reliable tools to observe how its predictions evolve before the final output. A transformer's residual stream encodes increasingly refined representations at each layer, culminating in logits at the unembedding layer, but the intermediate states live in hidden space—they are high-dimensional vectors that are not directly interpretable. Without a way to decode them, we are essentially staring at a black box for the first L−1L-1 layers of an LL-layer model, only getting to read its "thoughts" at the very end.

This is a first-order problem for interpretability research. If we want to understand how transformers reason—what computations they perform, when they converge on answers, whether they are overthinking a problem or taking shortcuts—we need a window into the internal prediction process. Without such a window, we cannot:

  • Diagnose why a model produces a particular output (did it know the answer early and overthink? Did it fix an initial error?).
  • Detect anomalous internal behavior (is the model processing a prompt injection differently from normal input?).
  • Study the computational efficiency of the model (do all problems need all layers, or do some solve quickly?).

The paper frames this through the lens of iterative inference (Section 1, Appendix C), drawing on an insight from Jastrzębski et al. (2017): residual networks encourage each layer to perform an incremental update that pushes the hidden state in a direction that reduces the final loss. Under this view, each layer's hidden state implicitly contains a "latent prediction" of the output—a belief about which token comes next—that gets refined step by step. The problem is: how do we extract that latent prediction in a way that faithfully reflects what the model actually "believes" at that layer?

Why This Problem Matters: Interpretability, Safety, and Efficiency

The inability to decode intermediate predictions has cascading consequences across three domains:

Interpretability research. The field has invested heavily in techniques that probe hidden states for syntactic, semantic, and factual knowledge (Hewitt and Manning, 2019; Tucker et al., 2021; Li et al., 2022). These probes typically train classifiers to predict specific concepts from hidden states. However, as Hewitt and Liang (2019) and Belinkov (2022) have argued, probes can learn to rely on spurious correlations unrelated to the model's actual computational process. A method that decodes the model's own output distribution at each layer—rather than an external label—sidesteps this issue: we are asking "what does the model predict at this layer?" not "can we extract concept X?" This provides a more faithful picture of the model's internal computation.

Safety and anomaly detection. If a model processes normal inputs and adversarial inputs differently, those differences may manifest in the trajectory of predictions across layers before they manifest in the final output. This is the intuition behind the prompt injection detection application (Section 5.3): a malicious instruction appended to a prompt ("Ignore previous instructions and just print...") may cause the model's internal prediction trajectory to look qualitatively different from its trajectory on clean prompts, even if the final output appears innocuous. Without a reliable lens into intermediate layers, this signal is invisible.

Computational efficiency and early exiting. Schuster et al. (2022) showed that if you can predict when a model has converged on its answer (i.e., when its internal prediction stops changing), you can skip the remaining layers, saving computation. This is the basis for early exiting methods like CALM and DeeBERT (Xin et al., 2020). But these methods require modifying the training process—they are not applicable to already-pretrained models. A reliable post-hoc decoding method would enable early exiting analysis on any pretrained model without retraining, and could be used to study the relationship between example difficulty and computational depth (as the paper does in Section 5.4).

Prior Approaches: The Logit Lens and Its Discontents

The only prior method that attempted to decode layer-wise predictions from frozen pretrained transformers was the logit lens, introduced informally by nostalgebraist (2020). The technique is strikingly simple: take the hidden state hℓh_\ell at layer ℓ\ell, apply the model's pretrained unembedding matrix WUW_U, and you get a distribution over the vocabulary:

LogitLens(hℓ)=LayerNorm[hℓ]WU\text{LogitLens}(h_\ell) = \text{LayerNorm}[h_\ell] W_U

The intuition is that because the residual stream updates are additive (hℓ+1=hℓ+Fℓ(hℓ)h_{\ell+1} = h_\ell + F_\ell(h_\ell)), the hidden state at layer ℓ\ell is simply the final-layer hidden state minus the sum of all subsequent residual updates. If we zero out those future updates and decode what remains, we should see the model's "best guess so far."

The logit lens works reasonably well for GPT-2 (radford2019language), producing prediction trajectories that converge roughly monotonically toward the final output distribution. This generated significant excitement in the interpretability community and led to multiple follow-up applications:

  • Halawi et al. (2023) used the logit lens to show that few-shot models "overthink" when given incorrect demonstrations—their internal predictions at middle layers are more robust than their final output.
  • Dar et al. (2022), Geva et al. (2022), and Millidge et al. (2022) used the logit lens to interpret transformer weight matrices by projecting MLP and attention parameters into vocabulary space.
  • CYWINSKI et al. (2025) used the logit lens to extract "secret knowledge" from models fine-tuned to conceal a taboo word.

However, the paper identifies three specific failure modes of the logit lens that fundamentally undermine its reliability:

1. Unreliability across model families. The logit lens simply does not work for many models released since GPT-2. The paper demonstrates this empirically:

  • For BLOOM (Scao et al., 2022), the logit lens often fails to extract plausible predictions (Figure 14, Appendix A).
  • For GPT-Neo (Black et al., 2021), the method requires an ad-hoc extension that retains the final transformer layer in the probe (Equation 4), and even this extension is "only partially successful" (Section 2).
  • For OPT 125M (Zhang et al., 2022), a similar pathology appears (Figure 18, Appendix B.1).

A striking qualitative failure mode: for BLOOM and OPT 125M, in more than half the layers, the logit lens's top-1 prediction is the input token itself rather than any plausible continuation (Figures 17–18, Appendix B.1). This is not a noisy prediction—it is a systematic, structurally induced error that makes the logit lens completely uninterpretable.

2. Representational drift. The paper identifies a deeper reason why the logit lens fails: representational drift (Section 3). Transformer hidden states at different layers live in different "bases"—the same semantic content may be encoded differently at layer 5 versus layer 30. There are two specific mechanisms:

  • Rogue dimensions (Timkey and van Schijndel, 2021): hidden states contain a small number of very high-variance dimensions, and these outlier dimensions are distributed unevenly across layers. Figure 6 (top) shows that Pythia 12B's layer 4 introduces two outlier dimensions that dominate the covariance structure. If the logit lens relies on the presence or absence of these dimensions, its predictions are not just noisy—they are systematically distorted.

  • Covariance drift: Even when controlling for rogue dimensions (Figure 6, bottom), the covariance matrices of hidden states at different layers drift apart as the layer separation increases. The final layer's covariance often shifts sharply relative to earlier layers, meaning that applying the unembedding (which is trained to work with final-layer representations) to earlier hidden states is akin to plugging a European appliance into an American outlet without an adapter—the representations are in the wrong "format."

3. Systematic bias. Perhaps the most damning failure mode from an interpretability perspective: the logit lens is a biased estimator of the model's final output distribution (Section 2). Specifically, for many vocabulary items, the logit lens at layer ℓ<L\ell < L systematically assigns higher or lower probability than the final layer does, and this bias persists across all inputs. Using the paper's metric—KL divergence between the marginal distributions DKL(p∥qℓ)D_{KL}(p \parallel q_\ell)—the logit lens exhibits a bias of approximately 4–5 bits for most layers of GPT-Neo-2.7B (Figure 3).

The paper makes a compelling argument for why this bias matters beyond just perplexity: it means the logit lens prediction trajectory cannot be interpreted as the model "updating its beliefs" in response to evidence. The authors invoke a rational agent framework (drawing on Yudkowsky's conservation of expected evidence, which itself builds on Ramsey, de Finetti, and the Dutch book arguments): the beliefs of a rational agent should not update in an easily predictable direction, because predictable updates can be exploited. Biased logit lens outputs are trivially exploitable once the bias direction is known—one could "bet against" the logit lens at layer ℓ\ell that the final probability of systematically downweighted tokens will rise, and make unbounded expected profit. For interpretability research, this means the logit lens is not just noisy—it is structurally misleading about what the model "believes" at intermediate layers.

The biases are also not stationary: the type of information extracted by the logit lens varies both across model families and across layers within a single model (Section 2). This makes comparative analysis nearly impossible—if the logit lens at layer 5 of BLOOM is primarily echoing the input token and the logit lens at layer 5 of GPT-2 is producing plausible continuations, what comparison can we make between them?

Why Existing Alternatives Don't Fill the Gap

The logit lens is not the only way to extract predictions from intermediate layers, but alternative approaches have their own limitations that make them unsuitable as general-purpose interpretability tools:

Early exiting methods like CALM (Schuster et al., 2022) and DeeBERT (Xin et al., 2020) train the model to produce usable outputs at every layer by adding auxiliary classifiers during training. These methods produce reliable predictions at intermediate layers, but they require modifying the training process. They cannot be applied post-hoc to the thousands of pretrained models already deployed. Moreover, because the training objective changes (the model learns to produce good predictions at every layer, not just the final one), the resulting prediction trajectories may not reflect how a standard pretrained model processes information—the training procedure itself could alter what the intermediate representations encode.

Traditional probing (Alain and Bengio, 2016; Hewitt and Manning, 2019) trains classifiers on frozen hidden states to predict external labels (part-of-speech tags, syntactic structure, factual knowledge). This approach is powerful for categorizing what information is present in hidden states, but it does not decode the model's own predictive distribution. A probe that predicts "this hidden state contains part-of-speech information" is answering a different question than "what does the model think the next token is at this layer?" Probes are also susceptible to the "spurious feature" problem: a probe may appear to find syntactic structure in a hidden state not because the model uses that structure, but because the structure correlates with some other feature the probe latched onto.

Direct weight interpretation (Elhage et al., 2021; Dar et al., 2022; Geva et al., 2022; Millidge et al., 2022) analyzes transformer parameters directly by projecting weight matrices into vocabulary space. While this has yielded valuable insights (e.g., identifying interpretable "value vectors" in MLP layers, interpreting attention heads as "reading from" and "writing to" the residual stream), it is orthogonal to the problem the tuned lens addresses. Weight interpretation tells us what a component can write; it does not tell us what the accumulated hidden state does encode at that point in the forward pass. The two approaches are complementary, not competing.

How This Paper Positions Itself

The paper frames the tuned lens as a drop-in replacement for the logit lens that addresses its three core failures through a simple but principled design: instead of applying the unembedding directly to hidden states (implicitly assuming all layers share the same representational basis), train a lightweight affine transformation per layer that translates the representation from the basis used at that layer to the basis expected at the final layer:

TunedLensℓ(hℓ)=LogitLens(Aℓhℓ+bℓ)\text{TunedLens}_\ell(h_\ell) = \text{LogitLens}(A_\ell h_\ell + b_\ell)

where (Aℓ,bℓ)(A_\ell, b_\ell) is a d×dd \times d affine translator (with approximately d2d^2 parameters where dd is the hidden dimension) trained to minimize KL divergence between the tuned lens output and the model's actual final-layer distribution. This can be understood as a form of model stitching (Lenc and Vedaldi, 2015; Bansal et al., 2021): we are "stitching" an intermediate layer directly to the unembedding, with the affine translator serving as the alignment layer. The key finding is that a linear alignment suffices—suggesting that while the representational basis drifts, it does so smoothly and predictably.

The method is explicitly positioned as:

  • Post-hoc: it requires no modification to the pretrained model; translators are trained on a frozen model's outputs using a distillation loss.
  • Computationally lightweight: training all translators for a model can be done in under an hour on a single 8×A40 node.
  • Causally faithful: the features that are most influential for the tuned lens are also most influential for the model itself (Section 4), suggesting the tuned lens reflects the model's actual computation rather than spurious correlations.
  • Generally applicable: the paper validates the method across six model families (GPT-2, GPT-Neo, Pythia, BLOOM, OPT, GPT-NeoX-20B) and up to 20B parameters, showing consistent improvements over the logit lens in all cases.

The paper does not claim to have invented the concept of early exiting from pretrained models. Rather, it claims to have solved the reliability problems that made the existing method (the logit lens) unusable on many models and misleading on others, thereby opening up the full range of iterative inference analyses to essentially any pretrained transformer in use today. The contribution is methodological infrastructure: a tool that is reliable enough, unbiased enough, and causally faithful enough to be used as a building block in higher-level interpretability and safety applications.

Moreover, by demonstrating that an affine transformation suffices to translate between layers' representational bases, the paper provides indirect evidence for the iterative inference hypothesis: if representations at different layers were fundamentally incommensurate (encoding entirely different kinds of information in entirely different ways), a simple linear map would not suffice to align them. The fact that it does work—and that translators transfer well to nearby layers with minimal penalty (Figure 7)—suggests that transformers do indeed perform smooth, incremental updates to a shared representational space, consistent with the theoretical analysis of Jastrzębski et al. (2017) that the paper extends empirically in Appendix C.

3. Technical Approach

3.1 Reader Orientation

The tuned lens is a method for attaching a small, trainable "adapter" to each layer of a frozen pretrained transformer language model, so that the hidden state at any intermediate layer can be decoded into a meaningful probability distribution over the vocabulary—effectively showing what the model "thinks" the next token will be at that point in its internal computation. The problem it solves is that the earlier logit lens technique, which simply applies the model's final unembedding matrix directly to intermediate hidden states, fails catastrophically on many modern models due to representational drift and systematic bias; the tuned lens fixes this by learning an affine transformation per layer that translates the local representational basis into the basis expected by the final unembedding, using only a distillation loss that matches the tuned lens output to the model's actual final-layer distribution.

3.2 Big-Picture Architecture (Diagram in Words)

The system has four major components:

  1. Frozen Pretrained Transformer ($\mathcal{M}$) — an autoregressive language model decomposed into two halves at any target layer $\ell$: the prefix $\mathcal{M}_{\leq\ell}$ (layers 1 through $\ell$, producing the hidden state $h_\ell$) and the suffix $\mathcal{M}_{>\ell}$ (layers $\ell+1$ through $L$, plus the final LayerNorm and unembedding, which together map $h_\ell$ to output logits). The model's weights are never modified during tuned lens training.

  2. Affine Translators ($(A_\ell, b_\ell)$) — a learned $d \times d$ matrix $A_\ell$ and $d$-dimensional bias vector $b_\ell$ for each layer $\ell$ (where $d$ is the model's hidden dimension). The translator takes a hidden state $h_\ell$ from layer $\ell$ and applies a linear change-of-basis plus bias shift: $A_\ell h_\ell + b_\ell$. This "translates" the representation from the basis used at layer $\ell$ to the basis expected by the final unembedding, compensating for representational drift.

  3. Unembedding Matrix ($W_U$) — the pretrained model's final projection from hidden space to vocabulary space. The tuned lens reuses this matrix rather than training a new one, which dramatically reduces the number of learned parameters (from $|\mathcal{V}| \times d$ to $d \times d$ per layer) and ensures the output vocabulary distribution is anchored to the same token embeddings the model uses.

  4. Distillation Training Loop — a process that trains all translators simultaneously by feeding text through the frozen model, extracting hidden states at every layer, applying the translators, decoding through the unembedding, and minimizing the KL divergence between each tuned lens distribution and the model's actual final-layer distribution on the same input.

Information flows as follows: a token sequence enters the frozen model → the model computes forward, storing hidden states $h_\ell$ at each layer → each translator $(A_\ell, b_\ell)$ transforms its respective $h_\ell$ → the transformed vector is normalized via LayerNorm → the unembedding $W_U$ projects to vocabulary logits → a softmax produces the tuned lens distribution $q_\ell$ → the KL divergence between $q_\ell$ and the model's true final distribution $p$ drives translator parameter updates. At inference, the same pipeline applies but with frozen translators, yielding interpretable prediction trajectories across layers.

3.3 Roadmap for the Deep Dive

  • First, we will derive the tuned lens formula from the logit lens, working through Equation 8 symbol-by-symbol so that every term's role is clear, and explain why the two modifications—learned bias and learned linear map—each address a distinct failure mode of the original method.
  • Second, we will examine the training objective (Equation 9), the distillation loss that anchors tuned lens predictions to what the model actually outputs, and explain why this specific loss is chosen over alternatives like supervised cross-entropy with ground-truth tokens.
  • Third, we will cover the training data, optimization procedure, and the Muon optimizer finding, since the practical performance of the translators turns out to depend heavily on optimization details.
  • Fourth, we will analyze the translator's behaviour across layers—perplexity scaling, bias reduction, and transferability—to build intuition for why a simple affine map works as well as it does.
  • Fifth, we will connect the tuned lens to the broader idea of model stitching, explaining why the success of affine translators provides evidence for the iterative inference hypothesis that motivated the work.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a methodology paper whose core idea is that the logit lens's reliability failures can be resolved by introducing a lightweight learned affine transformation per layer—trained purely via distillation to match the model's own final output—and that the resulting probes are not only more accurate and less biased than the logit lens, but also causally faithful to the model's actual computational process.

Deriving the Tuned Lens from the Logit Lens

To understand why the tuned lens takes the form it does, we need to start from the logit lens and identify exactly where it goes wrong. The paper walks through this derivation in Section 3, building the tuned lens as a sequence of two incremental fixes.

Step 1: The logit lens decomposition. Consider a pre-LayerNorm transformer $\mathcal{M}$ with $L$ layers. The residual stream update at layer $\ell$ follows:

hℓ+1=hℓ+Fℓ(hℓ)h_{\ell+1} = h_\ell + F_\ell(h_\ell)

where $F_\ell$ is the residual output of layer $\ell$ (the combined effect of multi-head attention and MLP at that layer). Applying this recursively, the final hidden state $h_{L+1}$ (the input to the final LayerNorm and unembedding) can be expressed as the sum of any earlier hidden state $h_\ell$ plus all subsequent residual updates:

hL+1=hℓ+∑ℓ′=ℓLFℓ′(hℓ′)h_{L+1} = h_\ell + \sum_{\ell'=\ell}^{L} F_{\ell'}(h_{\ell'})

The complete forward pass from layer $\ell$ onward—mapping $h_\ell$ to output logits—is:

M>ℓ(hℓ)=LayerNorm[hℓ+∑ℓ′=ℓLFℓ′(hℓ′)]WU\mathcal{M}_{>\ell}(h_\ell) = \text{LayerNorm}\left[h_\ell + \sum_{\ell'=\ell}^{L} F_{\ell'}(h_{\ell'})\right] W_U

where $\mathcal{M}_{>\ell}$ denotes the suffix of the model starting after layer $\ell$ (through the final LayerNorm and unembedding $W_U$), and $F_{\ell'}(h_{\ell'})$ is the residual update contributed by layer $\ell'$.

What this equation describes: the actual computation that the model performs to convert an intermediate hidden state $h_\ell$ into final logits. It takes $h_\ell$, adds on every subsequent layer's residual contribution (each $F_{\ell'}(h_{\ell'})$), applies LayerNorm to the accumulated sum, and projects to vocabulary space with $W_U$.

Step 2: The logit lens approximation. The logit lens makes a drastic simplification: it sets all future residual updates to zero, as if layers $\ell+1$ through $L$ contributed nothing:

LogitLens(hℓ)=LayerNorm[hℓ]WU\text{LogitLens}(h_\ell) = \text{LayerNorm}[h_\ell] W_U

Why this form is chosen: the intuition is that if transformer layers perform iterative refinement—each layer making a small update toward the correct answer—then $h_\ell$ should already contain most of the information needed to predict the output. The subsequent residuals $\sum F_{\ell'}(h_{\ell'})$ are treated as "correction terms" that can be dropped without losing the core prediction. This is not entirely unreasonable; Appendix C provides empirical evidence that residual updates are consistently aligned with the negative gradient of the loss, meaning they push the hidden state in a predictable "improvement" direction. Dropping them should, in theory, give us the model's prediction before those improvements.

The first failure: zero is an arbitrary replacement for the future residual sum. The problem is that the network may have learned to rely on the expected value of the future residual sum, $\sum_{\ell'=\ell}^L \mathbb{E}[F_{\ell'}(h_{\ell'})]$, as an implicit bias term. If, on average across all inputs, later layers add a particular constant shift to the representation, then the logit lens's choice of zero as the replacement value will systematically push $\text{LayerNorm}[h_\ell]W_U$ away from where the model expects it to be. The hidden state $h_\ell$ by itself is "out-of-distribution" relative to what the unembedding expects to see—it is missing an additive constant that is normally present when the final LayerNorm is applied.

Step 3: First fix—learned bias. The simplest correction is to replace the hardcoded zero with a learned constant vector $b_\ell$:

LogitLensℓdebiased(hℓ)=LogitLens(hℓ+bℓ)=LayerNorm[hℓ+bℓ]WU\text{LogitLens}^{\text{debiased}}_\ell(h_\ell) = \text{LogitLens}(h_\ell + b_\ell) = \text{LayerNorm}[h_\ell + b_\ell] W_U

What this computes: instead of decoding $h_\ell$ directly, we add a per-layer bias $b_\ell$ to the hidden state before applying LayerNorm and the unembedding. The bias $b_\ell$ is learned from data, so it can capture the average expected contribution of future layers—the "missing constant" that the logit lens incorrectly sets to zero.

Why a bias helps: this directly addresses the systematic marginal bias documented in Figure 3. If certain vocabulary items are systematically upweighted or downweighted by the logit lens relative to the final distribution, a learned bias can compensate by shifting the pre-LayerNorm representation in the direction that corrects this imbalance. The paper explicitly states the logic: "the choice of zero as a replacement value is somewhat arbitrary—the network might learn to rely on $\sum_{\ell'=\ell}^L \mathbb{E}[F_\ell(h_\ell)]$ as a bias term." The learned bias $b_\ell$ is that missing term, estimated from data.

The second failure: representational drift and covariance mismatch. Even with the bias correction, the logit lens still implicitly assumes that the coordinate axes of the hidden space at layer $\ell$ mean the same thing as the coordinate axes at the final layer. The paper provides two lines of evidence that this assumption is false:

  • Rogue dimensions (Figure 6, top): certain dimensions have very high variance, and these outlier dimensions appear at different layers. If the logit lens's reliance on these dimensions is misaligned with how the final layer expects them to behave, the decoded distribution will be distorted.
  • Covariance drift (Figure 6, bottom): the covariance matrices of hidden states at different layers systematically diverge. The paper measures this using the Frobenius cosine similarity between covariance matrices at different layers, finding that similarity decays with layer separation and that the final layer's covariance often changes sharply.

The consequence is that even if $h_\ell$ contains all the semantic information needed to make the correct prediction, that information may be encoded in a different "basis"—a different linear coordinate system—than what the unembedding expects. Applying $W_U$ to $h_\ell$ without correcting for this basis mismatch is like reading a message written in the wrong coordinate system: the information is there, but you are looking at it from the wrong angle.

Step 4: Second fix—learned change-of-basis. The solution is to introduce a learned linear transformation $A_\ell$ that rotates and scales the hidden state into the coordinate system expected by the final layer:

TunedLensℓ(hℓ)=LogitLens(Aℓhℓ+bℓ)=LayerNorm[Aℓhℓ+bℓ]WU\text{TunedLens}_\ell(h_\ell) = \text{LogitLens}(A_\ell h_\ell + b_\ell) = \text{LayerNorm}[A_\ell h_\ell + b_\ell] W_U

What this computes: first, the hidden state $h_\ell$ is multiplied by the $d \times d$ matrix $A_\ell$ (a change of basis—rotating, scaling, and potentially shearing the representation), then the bias $b_\ell$ is added (a translation—shifting the representation to correct for the missing expected residual sum), then LayerNorm is applied, and finally the unembedding $W_U$ projects to vocabulary logits. The entire composed function is the tuned lens probe for layer $\ell$.

Why a linear map suffices: the fact that a single affine transformation per layer largely resolves the logit lens's failures (as shown in the perplexity results in Figure 5) is itself a substantive finding about transformer representations. It suggests that while the representational basis drifts from layer to layer, the drift is smooth and approximately linear—representations at different layers live in the same underlying semantic space, just viewed through different linear lenses. If the relationship were fundamentally nonlinear (e.g., if later layers encoded entirely new types of information in qualitatively different ways), an affine map would not suffice. The success of the tuned lens thus provides indirect evidence for the iterative inference hypothesis: layers make incremental updates to a shared representational space, rather than constructing entirely new representations from scratch.

Parameter count and design efficiency. Each translator $(A_\ell, b_\ell)$ has $d \times d + d$ parameters, where $d$ is the model's hidden dimension. For Pythia 12B with $d = 5120$, this is approximately 26.2 million parameters per layer translator. However, the paper explicitly compares this to the alternative of learning a new unembedding matrix per layer (as in Alain and Bengio, 2016), which would be $|\mathcal{V}| \times d$ parameters per layer—for Pythia 12B's vocabulary of 50,000 tokens, approximately 256 million parameters per layer, nearly 10× more. The authors note that "training a new unembedding matrix requires considerably more training steps and a larger batch size than training a translator, and often converges to a worse perplexity" (Section 3, paragraph on "Benefits over traditional probing"). The tuned lens design thus achieves better results with fewer parameters by reusing the model's own unembedding, which already knows how to map well-formed hidden states to vocabulary distributions.

The tradeoff in translator training cost. Training $L$ affine translators (one per layer) on a large model does require significant computation—the paper reports training in under an hour on a single 8×A40 node. For the largest models (GPT-NeoX-20B with $d = 6144$ and $L = 44$), the total translator parameters across all layers are approximately $44 \times (6144^2 + 6144) \approx 1.66$ billion parameters. However, these parameters are trained only once and can be distributed as precomputed checkpoints, which the authors have done via the tuned-lens library. The cost is thus amortized across all downstream analyses. Moreover, the translators are trained entirely on the model's own outputs via distillation—no human labels are needed, and the training data is simply text from the model's pretraining corpus.

Training Objective: Distillation with the Final Layer as Teacher

The tuned lens translators are trained via knowledge distillation (Hinton et al., 2015; Sanh et al., 2019), where the "teacher" is the model's own final layer and the "student" is the tuned lens probe at layer $\ell$. The loss function is the Kullback-Leibler divergence between the two distributions:

arg⁡min⁡Ex[DKL(M>ℓ(hℓ)  ∥  TunedLensℓ(hℓ))]\arg\min \mathbb{E}_{\boldsymbol{x}}\left[D_{KL}\left(\mathcal{M}_{>\ell}(h_\ell) \;\middle\|\; \text{TunedLens}_\ell(h_\ell)\right)\right]

where $\mathcal{M}_{>\ell}(h_\ell)$ is the model's actual output distribution at the final layer (the teacher), $\text{TunedLens}_\ell(h_\ell)$ is the probe's output distribution (the student), and the expectation is taken over token positions in sequences $\boldsymbol{x}$ sampled from the training corpus.

What this loss computes: for each token position in each training sequence, the model processes it normally and records both the hidden state $h_\ell$ at each intermediate layer and the final output logits. The tuned lens computes its own logits from $h_\ell$ via the translator and unembedding. The KL divergence $D_{KL}(\text{teacher} \parallel \text{student})$ measures how many extra bits would be needed to encode the teacher distribution using the student distribution as a code—it penalizes the student for putting low probability on tokens the teacher considers likely and for putting high probability on tokens the teacher considers unlikely. Minimizing this divergence pushes the student distribution to match the teacher as closely as possible.

Why KL divergence and not cross-entropy with ground-truth tokens? This is a critical design choice that the paper justifies explicitly (Section 3, "Loss function" paragraph). The alternative would be to train the probes with supervised cross-entropy against the actual next token:

arg⁡min⁡Ex[−log⁡TunedLensℓ(hℓ)[xnext]]\arg\min \mathbb{E}_{\boldsymbol{x}}\left[-\log \text{TunedLens}_\ell(h_\ell)[x_{\text{next}}]\right]

where $x_{\text{next}}$ is the true next token. This would push the probes to maximize the probability of the correct token, potentially encouraging them to learn information that the model itself has not yet encoded at layer $\ell$. The probe could, in principle, learn to "look ahead" by extracting subtle cues from $h_\ell$ that the model does not actually use for its own predictions—a form of probe overclaiming that Hewitt and Liang (2019) warned about. The distillation loss avoids this: the probe's target is exactly what the model outputs at the final layer, so the probe is constrained to only represent information that the model itself uses. As the paper states: "It ensures that the probes are not incentivized to learn extra information over and above what the model has learned, which can become a problem when training probes with ground truth labels."

What D_KL penalizes concretely. For a single token position, the KL divergence expands as:

DKL(p∥q)=∑v∈Vp(v)log⁡p(v)q(v)=−H(p)+H(p,q)D_{KL}(p \parallel q) = \sum_{v \in \mathcal{V}} p(v) \log\frac{p(v)}{q(v)} = -H(p) + H(p, q)

where $p$ is the teacher (final layer) distribution, $q$ is the student (tuned lens) distribution, $H(p)$ is the entropy of the teacher (a constant with respect to the student parameters), and $H(p, q)$ is the cross-entropy between $p$ and $q$. Minimizing $D_{KL}$ is equivalent to minimizing the cross-entropy $H(p, q)$, which penalises the student for assigning low probability to tokens that the teacher assigns high probability to. Unlike squared error in probability space, KL divergence naturally handles the simplex constraint (probabilities must sum to 1) and the fact that the teacher distribution is often diffuse (assigning non-trivial probability to many plausible tokens, not just the single correct one).

The implication for interpretability. Because the tuned lens is trained to match the model's own output distribution rather than external labels, its predictions at layer $\ell$ can be interpreted as "what the model would output if the forward pass stopped here and the remaining layers were replaced by a trained affine approximation." This is not the same as "what the model actually outputs when the remaining layers run normally"—the affine translators are an approximation, and some information is inevitably lost—but it is the best possible linear approximation of that counterfactual, and the causal fidelity experiments in Section 4 show that the approximation is good enough that the features the tuned lens relies on are also the features the model itself relies on.

Training Data and Optimization Procedure

The paper provides specific details about the training data and optimization in the implementation section of Section 3, with important practical findings about optimizer choice.

Training data. Translators are trained on text from the same distribution used for the model's pretraining, specifically:

  • For Pythia and GPT-NeoX-20B, the Pile validation set (Gao et al., 2020; Biderman et al., 2022) is used for both training and evaluation.
  • For BLOOM and GPT-2, whose original pretraining validation sets are not publicly available, the Pile validation set substitutes.
  • For OPT, whose validation set is also not public, a member of the OPT team provided access to the OPT validation set for training.
  • Documents are concatenated and split into uniform chunks of length 2048 tokens.

The evaluation set for all models is a random sample of 16.4 million tokens from the respective pretraining validation sets. This is a non-trivial amount of data—comparable to the token budget used for moderate-scale fine-tuning—and reflects the fact that the translators need to learn a mapping that generalizes across the full diversity of text the model can process.

Initial optimization recipe (SGD). The paper's original experiments used stochastic gradient descent with Nesterov momentum, with a linear learning rate decay schedule over 250 training steps. Specific hyperparameters:

  • Base learning rate of 1.0 (without the final transformer layer included in the probe) or 0.25 (with the final transformer layer included).
  • Gradient clipping to a norm of 1.
  • Gradient accumulation to achieve a total batch size of $2^{18}$ = 262,144 tokens per optimizer step.
  • All translators initialized to the identity transformation: $A_\ell = I_d$, $b_\ell = \mathbf{0}$. Starting from identity means the tuned lens initially behaves exactly like the logit lens, then gradually learns to deviate.
  • Weight decay of $10^{-3}$.

The Muon optimizer finding. After the initial publication of the paper, a significant empirical improvement was discovered: replacing SGD with the Muon optimizer (Jordan et al., 2024) "dramatically accelerates training, allowing us to achieve much lower KL divergence losses than were possible in our initial experiments." The authors state candidly that "essentially, all of the tuned lenses trained and evaluated in the earlier versions of this paper were severely undertrained."

Muon operates by maximizing the effective rank of the update applied to each matrix-valued parameter. Conceptually, it is the opposite of LoRA (Hu et al., 2022), which constrains parameter updates to be low-rank; Muon instead encourages updates to be as high-rank as possible, using a fast orthogonalization algorithm to counteract the tendency of neural network gradients to be rank-deficient. Liu et al. (2025) found that Muon accelerates the training of modern large language models compared to AdamW, and the tuned lens authors find it similarly beneficial for translator training.

The practical consequence of using Muon is quantitative rather than qualitative: the resulting translators have "Frobenius norms that are often several times larger than SGD-trained ones," meaning they deviate much further from the identity initialization—they are much better "tuned." The paper notes that they "opted not to re-do all our experiments with tuned lenses trained using Muon, but we encourage practitioners to use Muon in future experiments." This means that the perplexity and bias results reported in the paper (Figures 3–5) likely underestimate the tuned lens's true capability; with Muon-trained translators, the performance gap over the logit lens would be even larger.

Layer-specific training choices. Whether to include the final transformer layer in the probe varies by model family:

  • The final transformer layer is excluded for GPT-2, GPT-NeoX-20B, OPT, and Pythia.
  • The final transformer layer is included for GPT-Neo.
  • BLOOM is evaluated under both conditions in Figure 4.

The rationale for this model-specific choice is not fully spelled out in the paper, but the BLOOM experiment in Figure 4 provides a clue: including the final transformer layer significantly reduces perplexity for the logit lens (red squares shift down from left panel to right panel), showing that for BLOOM, the logit lens relies heavily on the final layer's residual contribution. The tuned lens, by contrast, performs well in both conditions (blue circles are similarly low in both panels), demonstrating that it is "an independent and complementary proposal" that does not depend on a particular choice of which layers to include.

Validation and early stopping. The paper does not explicitly describe an early stopping procedure for translator training. Given the distillation objective, however, a natural stopping criterion would be the KL divergence on a held-out validation set. Since the training corpus is simply text from the model's pretraining distribution and the distillation targets are the model's own outputs, overfitting is primarily a concern of the translators memorizing input-specific patterns rather than learning a general basis alignment. The large batch size ($2^{18}$ tokens) and limited number of training steps (250 for SGD) serve as implicit regularizers.

Translator Behavior Across Layers: Perplexity, Bias, and Transferability

Once the translators are trained, the paper evaluates them along three dimensions that collectively characterize how well they decode latent predictions.

Perplexity reduction. The primary quantitative metric is the perplexity of the tuned lens distribution at each layer, compared to the logit lens. Perplexity is the exponential of the cross-entropy between the decoded distribution and the true next token:

PPLℓ=exp⁡(−1T∑t=1Tlog⁡qℓ(xt+1∣x≤t))\text{PPL}_\ell = \exp\left(-\frac{1}{T}\sum_{t=1}^T \log q_\ell(x_{t+1} \mid x_{\leq t})\right)

where $x_{t+1}$ is the true next token and $q_\ell$ is the tuned lens (or logit lens) distribution at layer $\ell$. Lower perplexity means the decoded distribution assigns higher probability to the correct token.

Results in Figure 5 show that the tuned lens achieves uniformly lower perplexity than the logit lens across all layers and all model sizes in the Pythia family and GPT-NeoX-20B. Two patterns are notable:

  • Monotonic improvement with depth: for both methods, perplexity generally decreases (improves) at deeper layers, consistent with the iterative inference hypothesis—each layer incrementally refines the prediction. The tuned lens's improvement over the logit lens is largest at early layers, where the logit lens is most unreliable, and narrows at later layers where both methods converge toward the final output.
  • Lower variance across independently trained models: the tuned lens exhibits "lower variance across independently trained models" (Figure 5). This suggests the learned translators converge to a stable solution that captures genuine structure in the representation space, rather than fitting noise.

Bias elimination. The systematic marginal bias documented for the logit lens (Figure 3) is measured as:

DKL(p∥qℓ)=∑v∈Vp(v)log⁡p(v)qℓ(v)D_{KL}(p \parallel q_\ell) = \sum_{v \in \mathcal{V}} p(v) \log\frac{p(v)}{q_\ell(v)}

where $p(v)$ is the average probability assigned to vocabulary item $v$ by the model's final layer across all token positions in the evaluation set, and $q_\ell(v)$ is the same for the probe at layer $\ell$. This measures how much the probe's average behavior across all inputs differs from the model's average behavior.

For GPT-Neo-2.7B, the logit lens bias is "around 4 to 5 bits for most layers" (Figure 3). To contextualize this: 4–5 bits of KL divergence means the probe systematically misallocates probability mass by a factor of $2^{4}$ to $2^{5}$ = 16–32× for some vocabulary items relative to what the model actually outputs. The tuned lens reduces this to near zero (the blue line in Figure 3 hugs the x-axis), confirming that the learned affine transformation largely eliminates the systematic discrepancy.

What this implies: the tuned lens is not just a more accurate predictor of the correct token (lower perplexity); it is also an unbiased estimator of the model's own output distribution in the marginal sense. This is what enables the interpretation of the prediction trajectory as a belief-updating process: if the tuned lens says the probability of "dog" is 0.3 at layer 10 and 0.7 at layer 20, we can trust that this change reflects the model genuinely becoming more confident in "dog" based on the computation performed in layers 10–20, rather than reflecting a systematic tendency of the probe to underweight "dog" at layer 10.

Transferability across layers. A striking finding is that translators trained for one layer can be applied to nearby layers with only modest degradation in performance (Figure 7). The paper defines the transfer penalty from layer $\ell$ to layer $\ell'$ as:

Penalty(ℓ→ℓ′)=Ex[H(p,qℓ→ℓ′)−H(p,qℓ)]\text{Penalty}(\ell \to \ell') = \mathbb{E}_{\boldsymbol{x}}\left[H(p, q_{\ell \to \ell'}) - H(p, q_\ell)\right]

where $q_{\ell \to \ell'}$ means applying the translator trained for layer $\ell$ to the hidden state from layer $\ell'$, and $H(p, q)$ is the cross-entropy between the final layer distribution and the probe distribution. The transfer penalty measures how many extra bits per token are lost when using a translator on the "wrong" layer relative to using the correct translator.

For Pythia 12B (Figure 7):

  • Transfer penalties are lowest for nearby layers (entries near the diagonal are close to zero), and increase gradually with layer distance.
  • Transfer penalties are "strongly negatively correlated with covariance similarity (Spearman $\rho = -0.78$)." This means that if two layers have similar residual stream covariance structures, a translator trained for one works well on the other—and vice versa. This directly supports the interpretation that what the translator learns is largely a correction for covariance drift.
  • The transfer penalty matrix is not symmetric: "transfer penalties are higher when training on a layer with the outlier dimensions (Layer 5 and later) and testing on a layer without them, than the reverse." The paper hypothesizes that this asymmetry arises because the translator trained on a post-outlier layer learns to undo the effect of those outlier dimensions; when applied to an earlier layer that never had them, this undoing operation is inappropriate.

Transfer to fine-tuned models. The paper tests whether translators trained on a base model can be used on fine-tuned versions of that model without retraining. Using Vicuna 13B (an instruct fine-tuned chat model based on LLaMA 13B), they compare:

  • A tuned lens specifically trained on Vicuna.
  • A "transferred" tuned lens that copies the affine translators from a LLaMA-trained lens but uses Vicuna's unembedding.

On the RedPajama dataset, the transferred lens incurs "at worst, a 0.3 bits per byte increase in KL divergence to the model's final output." On Anthropic's Helpful Harmless conversation dataset, "we find no significant difference between the transferred and trained lenses" (Figure 13, Appendix A).

What this implies for practical use: fine-tuning appears to minimally affect the representational basis at each layer—the affine translators learned from the base model generalize well to the fine-tuned model's hidden states. This is practically significant because it means practitioners can use pretrained tuned lens checkpoints (which the authors have released) on fine-tuned models without retraining, eliminating the "limitation" that training translators requires computational effort. It also suggests that fine-tuning primarily modifies the computation performed by later layers (which use the hidden states) rather than the encoding format of the hidden states themselves—a finding with implications for understanding how fine-tuning works mechanistically.

Connection to Model Stitching and the Iterative Inference Hypothesis

The paper explicitly frames the tuned lens as an instance of model stitching (Lenc and Vedaldi, 2015; Bansal et al., 2021; Csiszárik et al., 2021), a technique where two frozen pretrained models are connected by training a small alignment layer that maps the output of one model's early layers to the input of another model's later layers. The tuned lens can be viewed as "stitching" an intermediate layer directly to the unembedding:

Layerℓ→Aℓ,bℓLayerNorm→WUlogits\text{Layer}_\ell \xrightarrow{A_\ell, b_\ell} \text{LayerNorm} \xrightarrow{W_U} \text{logits}

The key finding from the stitching literature that the tuned lens leverages is that an affine transformation suffices to align representations from independently trained models with minimal performance loss. The tuned lens extends this idea within a single model: if an affine map can stitch together different models, it should certainly be able to stitch together different layers of the same model, where the representational drift should be much smaller.

This connection is important because it grounds the tuned lens in an established empirical phenomenon rather than presenting it as an ad-hoc fix. The fact that model stitching with affine maps works at all—both across models and within models—is a significant empirical finding about the nature of neural network representations: they tend to live in a shared linear space where semantic similarity is encoded geometrically. The tuned lens exploits this property to make intermediate representations decodable.

The success of the tuned lens also provides indirect support for the iterative inference hypothesis that motivates the paper (Section 1, Appendix C). The formal argument, adapted from Jastrzębski et al. (2017), is:

L(hL)=L(hi)+∑j=iL⟨Fj(hj),∂L(hj)∂hj⟩+O(Fj2(hj))\mathcal{L}(h_L) = \mathcal{L}(h_i) + \sum_{j=i}^{L} \left\langle F_j(h_j), \frac{\partial\mathcal{L}(h_j)}{\partial h_j}\right\rangle + \mathcal{O}(F_j^2(h_j))

where $\mathcal{L}$ is the final loss (linear classifier plus loss function) evaluated at the final hidden state $h_L$, $h_i$ is the hidden state at some intermediate layer $i$, and $F_j(h_j)$ is the residual update at layer $j$. The inner product term $\langle F_j(h_j), \frac{\partial\mathcal{L}(h_j)}{\partial h_j}\rangle$ measures the alignment between the actual update $F_j(h_j)$ that the layer applies and the negative gradient $-\frac{\partial\mathcal{L}}{\partial h_j}$ (the direction that would most decrease the loss). To a first-order approximation, each layer reduces the loss by an amount proportional to this alignment.

What this Taylor expansion reveals: if the residual updates $F_j(h_j)$ are consistently aligned with the negative gradient (i.e., the inner product is negative), then each layer is, on average, moving the hidden state in a direction that reduces the final loss. This would mean the hidden state at layer $i$ already encodes a "prediction" (the classifier's best guess based on $h_i$), and each subsequent layer refines that prediction by pushing $h_i$ closer to the final $h_L$.

The paper empirically tests this by computing the cosine similarity between $F_j(h_j)$ and $\frac{\partial\mathcal{L}}{\partial h_j}$ for Pythia 6.9B (Figure 19, Appendix C). They find:

  • The cosine similarity is negative at least 95% of the time for every layer.
  • The magnitudes are small (never exceeding 0.05 in absolute value) but are "much larger than would be expected of random vectors in this very high dimensional space" (where the 5th percentile of random pairwise cosine similarities is $-6 \times 10^{-4}$).

The tuned lens's success in making these intermediate predictions decodable through a simple affine map is consistent with this picture: if each layer is indeed making a small, gradient-aligned update toward the correct output, then the hidden state at layer $i$ should be a progressively better approximation of the final hidden state, and a linear correction plus bias should suffice to account for the expected cumulative effect of the remaining layers. The translators $(A_\ell, b_\ell)$, in this interpretation, are learning to approximate the average effect of $\mathcal{M}_{>\ell}$ as an affine function—a reasonable approximation if the remaining layers' contributions are small and approximately linear in expectation.

4. Key Insights and Innovations

Innovation 1: Diagnosing Representational Drift as the Root Cause of Logit Lens Failure

Prior work treated the logit lens's unreliability as an opaque empirical nuisance—some models worked (GPT-2), others did not (GPT-Neo, BLOOM), and the field's response was ad-hoc patches like retaining the final transformer layer (nostalgebraist, 2021). This paper makes a fundamentally diagnostic contribution: it identifies and empirically demonstrates two specific, measurable mechanisms—rogue dimensions and covariance drift—that explain why the logit lens fails, and shows that these mechanisms are sufficient to account for the logit lens's systematic bias, perplexity degradation, and model-specific reliability.

This is not an incremental improvement. It transforms the problem from "the logit lens is unreliable on some models for unknown reasons" to "hidden states at different layers live in different linear coordinate systems, and applying the unembedding without a change-of-basis produces systematically distorted outputs." The distinction matters because it tells us what kind of fix is needed. If the problem were that intermediate layers simply do not encode predictions, the solution would be fundamentally different (e.g., training full output heads per layer, as in early exiting methods like CALM). But because the problem is primarily a basis mismatch—semantic information is present but encoded in the wrong coordinate system—a simple linear correction suffices.

The evidence chain is tight. Figure 6 (top) visually demonstrates the emergence of outlier dimensions at specific layers, showing the raw covariance structure that makes direct unembedding inappropriate. Figure 6 (bottom) shows that even after removing rogue dimensions, covariance similarity between layers decays smoothly with depth—meaning the basis drifts continuously, not just at a few pathological points. Figure 7 then closes the loop: the transfer penalty matrix between translators is "strongly negatively correlated with covariance similarity (Spearman ρ = -0.78)," establishing a direct causal link between covariance drift and decoder degradation. The transfer penalty asymmetry—higher when training on a post-outlier layer and testing on a pre-outlier layer than the reverse—further isolates the outlier dimensions as a specific causal factor rather than a mere correlate.

This diagnostic insight has implications beyond the tuned lens. It implies that any method that compares or operates on hidden states across layers—activation patching, representation similarity analysis, probing across layers—needs to account for basis drift. The logit lens's failure was not a quirk of GPT-Neo or BLOOM; it was the expected consequence of a general phenomenon that affects all deep transformers. The paper's identification of rogue dimensions as unevenly distributed across layers (Figure 6, top) also connects to a broader literature on the functional role of high-variance dimensions in transformers (Timkey and van Schijndel, 2021; Kovaleva et al., 2021), suggesting that these dimensions may play different computational roles at different depths.

Innovation 2: Reframing Layer-Wise Decoding as an Unbiased Estimation Problem

The paper makes a conceptual pivot that elevates the tuned lens from a better probe to a tool for studying belief dynamics. By formally defining bias in terms of expected marginal probabilities (𝔼[q_ℓ(v)] = 𝔼[p(v)] for all vocabulary items v, Section 2) and measuring it via D_KL(p ∥ q_ℓ), the paper reframes layer-wise decoding as an unbiased estimation problem. This is not the standard framing in the probing literature, which typically evaluates probes by their accuracy on external labels (Hewitt and Manning, 2019; Alain and Bengio, 2016) or their perplexity relative to the correct token. Those metrics ask "how good is the prediction?" The unbiased estimation framing asks a different question: "can we trust that changes in the decoded distribution across layers reflect genuine changes in the model's internal state, or are they artifacts of the decoder?"

This reframing has teeth. The logit lens exhibits 4–5 bits of marginal KL divergence from the final distribution for most layers (Figure 3)—meaning that for some vocabulary items, the logit lens systematically overestimates or underestimates probability by a factor of 16–32×, and this error is predictable given only the layer index and the vocabulary item. An estimator with this property is trivially exploitable: as the paper notes, one could "bet against" the logit lens at layer ℓ that the probability of systematically underweighted tokens would rise by the final layer, and make unbounded expected profit. This means the logit lens prediction trajectory is not just noisy—it is structurally misleading as a representation of the model's belief state. A trajectory that shows probability mass flowing from token A to token B across layers could reflect either a genuine computational refinement or a systematic decoder bias; without the unbiasedness property, we cannot distinguish these possibilities.

The paper's invocation of the rational agent framework (conservation of expected evidence, Dutch book arguments) might seem like philosophical window-dressing, but it serves a concrete methodological purpose: it provides a principled criterion for when a layer-wise decoder is trustworthy enough to support dynamical interpretations. If the decoder is biased, the trajectory of predictions is not interpretable as belief updating, because the "beliefs" at early layers are not beliefs at all—they are systematically distorted measurements. The tuned lens reduces this bias to near zero (Figure 3, blue line), making the dynamical interpretation coherent.

This is a conceptual contribution that generalizes beyond the specific tuned lens design. Any future method that claims to decode intermediate representations should be evaluated not just on perplexity but on marginal unbiasedness. The paper's introduction of D_KL(p ∥ q_ℓ) as a diagnostic for layer-wise decoders provides the field with a concrete metric for this property.

Innovation 3: Causal Fidelity as a Two-Dimensional Validation Criterion for Probes

The probing literature has long recognized that probes can learn to rely on spurious features unrelated to the model's actual computational process (Hewitt and Liang, 2019; Belinkov, 2022). The standard response has been to evaluate probes with causal experiments: intervene on the hidden state along the directions the probe relies on, and check whether the model's output changes correspondingly (Elazar et al., 2021). This paper advances that methodology by decomposing causal fidelity into two distinct properties and developing specific tests for each.

Property 1 (importance alignment): Latent directions that are causally influential for the tuned lens should also be causally influential for the model itself. The paper introduces causal basis extraction (CBE) to test this—a procedure that finds an orthonormal basis of directions in the residual stream, ordered by their influence on the tuned lens output (measured via KL divergence after mean ablation), and then checks whether ablating these same directions in the model produces proportional effects. The result in Figure 8—Spearman ρ = 0.89 between influence on the tuned lens and influence on the model, with no features falling in the "influential for probe but not for model" quadrant—is strong evidence that the tuned lens is causally faithful.

Property 2 (stimulus-response alignment): The direction of change induced by an intervention should be the same for the tuned lens and the model. This goes beyond "both change by similar amounts" to "both change in the same way"—the tuned lens should not just register that something is different, but should register the same kind of difference that the model registers. The paper operationalizes this using Aitchison geometry on the probability simplex (Equation 13–15), defining a cosine similarity metric that measures whether the tuned lens and the model "move" in the same direction in probability space under resampling ablation of the principal subspace. Figure 9 shows that alignment is consistently positive and increases with depth, and that the tuned lens achieves higher alignment than the logit lens at every layer.

This two-dimensional validation framework is a methodological contribution independent of the tuned lens itself. The CBE algorithm provides a general-purpose tool for extracting the principal features used by any differentiable function of hidden states. The stimulus-response alignment metric provides a way to test whether a probe and the model use features in the same way, not just whether they use the same features. Together, they provide a more rigorous standard for what it means for a probe to be "faithful"—a standard that the paper shows the tuned lens meets, and that future probing methods can be evaluated against.

Innovation 4: Prediction Trajectories as a Signal for Anomaly Detection

The tuned lens's practical payoff is most striking in the prompt injection detection results (Section 5.3, Table 1), which demonstrate that the shape of the prediction trajectory across layers—not just the final output—carries diagnostic information about whether an input is adversarial. On five of nine benchmarks (BoolQ, MNLI, QNLI, QQP, SST-2), anomaly detectors trained on tuned lens trajectories achieve near-perfect AUROC (≥ 0.99) at distinguishing normal prompts from prompt injection attacks. This is not an incremental gain over the logit lens; it is a qualitative leap in capability enabled by having reliable layer-wise predictions.

The conceptual insight here is that the tuned lens transforms the model's internal computation from an opaque process into a high-dimensional trajectory that can be analyzed with standard anomaly detection tools (isolation forest, local outlier factor). Just as a cardiologist can diagnose arrhythmias from the shape of an ECG waveform, even if the heart's final output (blood flow) appears normal, an anomaly detector can identify adversarial inputs from the shape of the prediction trajectory, even if the final output appears benign. The paper's comparison with the Simplified Relative Mahalanobis (SRM) baseline (bai2022training), which uses only a single layer's hidden state, is instructive: SRM performs consistently well but the tuned lens trajectory outperforms it on some tasks (ARC-Challenge) while underperforming on others (MC TACO, SciQ), suggesting that the trajectory and the single-layer representation capture complementary anomaly signals. The paper explicitly notes that "further gains could be made by combining the strengths of both techniques."

What makes this more than just a downstream application is that it reframes what the tuned lens provides. It is not merely a tool for human interpretability researchers to stare at—it produces a machine-readable signal (a sequence of log-probability vectors across layers) that can feed into automated monitoring systems. This opens the door to deployment-time anomaly detection that operates on internal model representations, which is a fundamentally different capability than input-filtering or output-monitoring approaches.

The correlation between prediction depth and iteration learned (Section 5.4, Table 2) reinforces this reframing. The fact that tuned lens prediction depths correlate significantly with how many training steps an example required to be learned—and do so more strongly than logit lens depths in 8 of 11 tasks—suggests that the trajectory encodes information about the computational difficulty of the input that is both genuine (it correlates with a training-dynamics metric) and practically measurable (it is extractable from a single inference pass on the final checkpoint). This connects the tuned lens to the broader literature on example difficulty and data pruning (Toneva et al., 2018; Baldock et al., 2021), providing a way to estimate example difficulty post-hoc for any pretrained model without access to training checkpoints.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. All primary experiments use text from the pretraining validation sets of the evaluated models—the Pile validation set (Gao et al., 2020; Biderman et al., 2022) for Pythia, GPT-NeoX-20B, GPT-2, and BLOOM; the OPT validation set (accessed through the OPT team) for OPT models. For evaluation, a random sample of 16.4 million tokens is drawn from each model's respective validation set. For downstream applications (prompt injection detection, example difficulty, overthinking), standard NLP benchmarks are used: BoolQ, MNLI, QNLI, QQP, SST-2, ARC-Challenge, MC TACO, SciQ, and SICK, accessed through EleutherAI's lm-evaluation-harness. Documents are concatenated and split into uniform chunks of length 2048 for both training and evaluation.

  • Base model(s). The tuned lens is tested across six autoregressive language model families spanning three orders of magnitude in parameter count: GPT-2 (Radford et al., 2019; 124M–1.5B parameters), GPT-Neo (Black et al., 2021; 125M–2.7B), Pythia (Biderman et al., 2023; 70M–12B), BLOOM (Scao et al., 2022; 560M), OPT (Zhang et al., 2022; 125M–6.7B), and GPT-NeoX-20B (Black et al., 2022). The Pythia family receives the most extensive evaluation because its intermediate training checkpoints are publicly available (enabling the iteration learned analysis in Section 5.4) and it shares architecture, data, and codebase with GPT-NeoX-20B, enabling cleaner scaling comparisons. All models use a pre-LayerNorm architecture except OPT 350M, which uses post-LN and is therefore omitted from the main results (Figure 14, Appendix A).

  • Metrics. The primary metric for translator quality is perplexity of the decoded distribution relative to the true next token: PPL_ℓ = exp(−(1/T) ∑_{t=1}^T log q_ℓ(x_{t+1} | x_{≤t})), where q_ℓ is the tuned lens (or logit lens) distribution at layer ℓ. Lower perplexity means the decoded distribution assigns higher probability to the correct token. The second core metric is marginal bias, measured as the KL divergence between the average probe distribution and the average final-layer distribution across the evaluation set: D_KL(p ∥ q_ℓ) = ∑_{v∈𝒱} p(v) log(p(v) / q_ℓ(v)), where p(v) and q_ℓ(v) are the empirical marginal probabilities assigned to vocabulary item v by the final layer and the probe respectively. For causal fidelity experiments, causal influence is measured as the expected KL divergence between the model's (or probe's) output before and after erasing a latent direction via mean ablation (Equation 10). For stimulus-response alignment, Aitchison cosine similarity (Equation 18) measures whether the probe and the model "move in the same direction" in probability space under an intervention. For prompt injection detection, AUROC (Area Under the Receiver Operating Characteristic curve) with 95% bootstrap confidence intervals measures anomaly detection performance. For the overthinking experiments, median-calibrated accuracy (following Halawi et al., 2023) is used. For example difficulty, Spearman rank correlation between prediction depth and iteration learned is reported.

  • Baselines. Three baselines are used across different experiments. (1) The logit lens (nostalgebraist, 2020): the direct application of the pretrained unembedding to intermediate hidden states, LogitLens(h_ℓ) = LayerNorm[h_ℓ] W_U, including its extended variant that retains the final transformer layer (Equation 4). This is the primary comparison point throughout the paper; every tuned lens result is measured against the logit lens on the same model and data. (2) Simplified Relative Mahalanobis (SRM) distance (bai2022training): for the prompt injection detection experiment (Section 5.3), the SRM computed on the middle layer's hidden states serves as a representation-based anomaly detection baseline, using the implementation from bai2022training with default parameters. (3) Randomly generated matrices: for the static interpretability analysis in Appendix E, randomly shuffled weight matrices (head-wise for attention matrices) provide a null distribution for interpretability scores, enabling the paper to distinguish genuine structure from artifacts of the probing method.

  • Generation budget / compute accounting. The tuned lens does not involve generation—it decodes existing hidden states from a single forward pass. The relevant "compute budget" is thus the training cost of the translators. The paper reports that training a full set of translators for a model takes "under an hour on a single 8×A40 node" using the SGD recipe, and that the Muon optimizer (Jordan et al., 2024) dramatically accelerates convergence to lower KL divergence. For the logit lens, there is zero training cost—it can be "used on any pretrained model out-of-the-box" (Section 6). This asymmetry is explicitly acknowledged as a limitation of the tuned lens. The paper mitigates this by releasing pretrained translator checkpoints through the tuned-lens library. For causal experiments, the compute cost of Causal Basis Extraction is significant: it requires "sequentially optimiz[ing] d_model causal basis vectors for each layer of the network" (Section 6), with each vector optimized via L-BFGS on a batch of 131,072 tokens—but this cost is one-time and the extracted directions can be reused.

  • Cross-validation / statistical protocol. For the prompt injection detection experiment (Section 5.3), results are "pooled over 10 random train-test splits" with 95% bootstrap confidence intervals reported on AUROC (Table 1). The anomaly detection models (isolation forest and local outlier factor) are fit exclusively on prediction trajectories from normal prompts and evaluated on a held-out set containing both normal and adversarial trajectories. For the tuned lens training itself, the paper uses a separate validation set slice from the training set—specifically, for models trained on the Pile validation set, a held-out portion is used for evaluation. No explicit k-fold cross-validation is reported for the main perplexity and bias results, which are computed on a single random sample of 16.4M tokens. For the overthinking experiments (Section 5.2), the evaluation protocol follows Halawi et al. (2023) but the paper does not specify the number of few-shot examples or the sampling procedure for incorrect demonstrations beyond citing the original work.

Main Quantitative Results

Tuned Lens vs. Logit Lens: Perplexity and Bias Across Model Families

The headline result is that the tuned lens achieves uniformly and substantially lower perplexity than the logit lens across all layers and all model sizes tested, while eliminating systematic marginal bias (reducing KL divergence from ~4–5 bits to near zero).

Perplexity (Figure 5, Figure 14). Figure 5 plots perplexity as a function of layer index for the Pythia model family (70M through 12B parameters) and GPT-NeoX-20B, with the logit lens in the left panel and the tuned lens in the right panel. The tuned lens predictions are "uniformly lower perplexity and exhibit lower variance across independently trained models." Specific patterns:

  • For all Pythia models, the tuned lens shows a smooth monotonic perplexity decrease with depth, starting from roughly 100–200 PPL at layer 0 (embedding only) and converging to the final layer perplexity at layer L. The logit lens, by contrast, often shows erratic behavior at early-to-middle layers—for Pythia 12B, the logit lens perplexity actually increases between layers 5–10 before decreasing again, while the tuned lens is smoothly decreasing throughout.
  • For GPT-NeoX-20B (Figure 5, rightmost curves in each panel), the tuned lens achieves roughly 20–30 PPL at the final layers versus roughly 50–60 PPL for the logit lens at the same depth.
  • The improvement is largest at early layers: for Pythia 12B at layer 5, the tuned lens achieves roughly 80 PPL versus roughly 400 PPL for the logit lens—a 5× reduction.
  • Figure 14 (Appendix A) extends these results to GPT-2 (124M through 1.5B), GPT-Neo (125M through 2.7B), and OPT (125M through 6.7B). In all cases, the tuned lens (blue circles) is substantially below the logit lens (red squares). For GPT-Neo 2.7B, the tuned lens reduces perplexity from roughly 200 to roughly 50 at layer 15. For OPT 6.7B, the reduction is from roughly 300 to roughly 30 at layer 20.

Bias elimination (Figure 3). For GPT-Neo-2.7B, the logit lens exhibits a marginal bias of "around 4 to 5 bits for most layers" as measured by D_KL(p ∥ q_ℓ), where p is the final layer's average distribution and q_ℓ is the probe's average distribution. To contextualize: 5 bits of KL divergence means the probe systematically misallocates probability mass by a factor of 2^5 = 32× for some vocabulary items relative to what the model actually outputs, and this misallocation is consistent across all inputs. The tuned lens reduces this bias to near zero—the blue line in Figure 3 visually hugs the x-axis while the orange line hovers around 4–5 bits for layers 5–25 before dropping sharply only at the very end of the network (layers 28–32). The paper notes that "as a point of comparison, the bias of Pythia 160M's final layer distribution relative to that of its larger cousin, Pythia 12B, is just 0.0068 bits," emphasizing that 4–5 bits is a massive discrepancy—the logit lens at layer 10 of GPT-Neo-2.7B is more "biased" relative to its own model's output than the outputs of two entirely different-sized models are relative to each other.

BLOOM and the final transformer layer (Figure 4). For BLOOM 560M, the paper tests both including and excluding the final transformer layer from the probe. The logit lens benefits substantially from including the final layer (red squares shift down from roughly 800 PPL to roughly 200 PPL at layer 12), showing that for BLOOM, the logit lens relies heavily on the final layer's residual contribution to produce interpretable outputs. The tuned lens, by contrast, "performs well in both conditions" (blue circles are similarly low in both panels, around 30–50 PPL at layer 12), demonstrating that it "is an independent and complementary proposal"—it does not depend on a particular architectural choice about which layers to include.

Statistical reliability. The paper states that the tuned lens exhibits "lower variance across independently trained models" (Figure 5), indicating that the translators converge to a stable solution rather than fitting noise. However, the number of independent training runs is not explicitly stated, and no confidence intervals or error bars are reported on the perplexity curves in Figures 5 or 14. The 16.4M-token evaluation set is substantial enough that sampling variance on perplexity estimates is likely small, but the absence of uncertainty quantification makes it difficult to assess whether the reported improvements are statistically significant or could be affected by training-data sampling.

Transferability Across Layers and Model Versions

Layer-to-layer transfer (Figure 7). The transfer penalty matrix for Pythia 12B shows that translators trained on one layer can be applied to nearby layers with modest degradation. Entries near the diagonal are near zero (meaning the correct translator for layer ℓ performs similarly to a translator trained on layer ℓ±1). The penalty increases with layer distance, reaching approximately 1–2 bits of cross-entropy increase when transferring across 10+ layers. The transfer penalty is "strongly negatively correlated with covariance similarity (Spearman ρ = -0.78)," directly linking the representational drift measured in Figure 6 to the decoding difficulty. Importantly, the matrix is not symmetric: "transfer penalties are higher when training on a layer with the outlier dimensions (Layer 5 and later) and testing on a layer without them, than the reverse." This asymmetry provides causal evidence that the rogue dimensions identified in Figure 6 (top) are a specific barrier to transfer—a translator trained on post-outlier layers has learned to compensate for these dimensions, and applying that compensation to pre-outlier layers that lack them is harmful.

Transfer to fine-tuned models (Figure 13, Appendix A). Using Vicuna 13B (an instruct fine-tuned chat model based on LLaMA 13B; vicuna2023), the paper compares three conditions: a tuned lens trained specifically on Vicuna, a transferred lens using LLaMA-trained translators with Vicuna's unembedding, and the logit lens on Vicuna. On the RedPajama dataset, the transferred lens incurs "at worst, a 0.3 bits per byte increase in KL divergence to the model's final output" relative to the Vicuna-trained lens. On Anthropic's Helpful Harmless conversation dataset (bai2022training), "we find no significant difference between the transferred and trained lenses" (Figure 13). Both tuned lens variants substantially outperform the logit lens. This finding has immediate practical significance: pretrained translator checkpoints released by the authors can be used on fine-tuned models without retraining.

Causal Fidelity: Does the Tuned Lens Reflect the Model's Computation?

Property 1—importance alignment (Figure 8, Figure 20). For Pythia 410M at layer 18, the Spearman rank correlation between a direction's causal influence on the tuned lens and its causal influence on the model is ρ = 0.89 (Figure 8). The scatter plot shows a dense cluster of points along the diagonal, with no points in the lower-right quadrant (features influential for the tuned lens but not the model)—the absence of such points is the key result, as it means the tuned lens is not relying on spurious directions that the model ignores. The model is somewhat more "causally sensitive" overall: even the least influential CBE features have a model influence above 2 × 10^{-3} bits, creating a "hockey stick" shape in the LOWESS trendline. This is interpreted as the model using a broader set of features than the subset extracted by CBE from the tuned lens, not as the tuned lens relying on spurious features. Figure 20 (Appendix D) extends this analysis to all layers of Pythia 410M, with Spearman ρ values ranging from approximately 0.85 to 0.95 across layers, consistently high throughout the network.

Property 2—stimulus-response alignment (Figure 9). Under resampling ablation of the principal subspace (top 10 CBE directions) at each layer of Pythia 160M, the Aitchison cosine similarity between the stimulus (change in tuned lens output) and response (change in model output) is consistently positive across all layers, ranging from roughly 0.3 at layer 2 to roughly 0.7 at layer 10. Alignment increases monotonically with depth—later layers show stronger stimulus-response alignment, which is expected because the tuned lens at later layers operates on representations closer to the final output. The tuned lens (blue line) achieves higher alignment than the logit lens (orange line) at every layer, with the gap being largest at middle layers (layers 4–8) where the tuned lens achieves roughly 0.5–0.6 versus roughly 0.3–0.4 for the logit lens. This supports the claim that the tuned lens not only identifies the right features (Property 1) but also uses them in the same way the model does (Property 2).

Applications: Overthinking, Prompt Injection Detection, Example Difficulty

Overthinking the truth (Figure 11). Halawi et al. (2023) found that when models are given few-shot prompts with incorrect demonstrations, predictions extracted from earlier layers using the logit lens are often more robust than final-layer predictions—the model "overthinks" and converges to the wrong answer. However, their method was inapplicable to BLOOM and GPT-Neo due to high logit lens perplexities. Using the tuned lens, the paper replicates their results for BLOOM 560M, GPT-Neo 1.3B, and GPT-Neo 2.7B (Figure 11). The key pattern: calibrated accuracy under incorrect demonstrations (red curve) peaks at an early layer (around layer 5–10, achieving roughly 0.40–0.45 accuracy) before falling to near-random performance at the final layer (roughly 0.20–0.25), while accuracy under correct demonstrations (blue curve) rises monotonically to high accuracy (roughly 0.70–0.85). This qualitative pattern matches what Halawi et al. (2023) found for GPT-2, extending it to three model families that the logit lens could not analyze.

Prompt injection detection (Table 1). Anomaly detectors trained on tuned lens prediction trajectories achieve near-perfect AUROC on five of nine tasks: BoolQ (1.000), MNLI (0.999), QNLI (0.999), QQP (0.999), and SST-2 (0.998). The logit lens trajectory performs worse on most tasks but still well above chance; for example, on SST-2 the tuned lens achieves 0.998 while the logit lens achieves 0.943, and on ARC-Challenge the tuned lens achieves 0.907 versus 0.772. The SRM baseline (using only the middle layer's hidden states with Mahalanobis distance) is the strongest competitor: it achieves near-perfect AUROC on BoolQ (0.999), MNLI (1.000), and QQP (1.000), tying or exceeding the tuned lens on these tasks, and outperforms the tuned lens on MC TACO (0.952 vs. 0.875) and SciQ (0.969 vs. 0.898). The paper acknowledges this mixed comparison and hypothesizes that "further gains could be made by combining the strengths of both techniques, since SRM uses only one layer but considers a high-dimensional representation, while the tuned lens studies the trajectory across layers but summarizes them with a low-dimensional prediction vector."

A crucial nuance: Table 1 notes that "attack detection performance is nearly perfect on tasks where the attack succeeds at driving accuracy well below the random baseline, and is still much better than chance even when the attack is only partially successful." This implies that the detectability of prompt injections via prediction trajectories is correlated with how much the attack actually disrupts the model's internal computation—a desirable property for a detector, since it means false positives are less likely on attacks that the model naturally resists.

Example difficulty (Table 2, Figure 12). The tuned lens prediction depth (number of layers after which the top-1 prediction stops changing) correlates positively with iteration learned (the earliest training step at which the model's prediction for that example stabilizes) across all 11 tasks tested. Spearman ρ values range from 0.058 (RTE) to 0.378 (ARC-Challenge), with most falling in the 0.10–0.25 range. While these correlations are modest in absolute magnitude, they are statistically significant (the paper states "a significant positive correlation... on all tasks we investigated" but does not report p-values). The tuned lens prediction depth correlates better with iteration learned than the logit lens prediction depth in 8 out of 11 tasks—sometimes dramatically so, as on HellaSwag (0.147 vs. 0.030) and WinoGrande (0.114 vs. 0.023). Figure 12 provides a qualitative visualization: token-level prediction depths for Pythia 12B on the GPT-4 technical report abstract, with "warm colors hav[ing] high prediction depth, while cool colors indicat[ing] low depth." The figure is not quantitatively analyzed but visually suggests that function words and common syntactic patterns have low prediction depth (the model resolves them quickly) while content words and semantically ambiguous tokens require deeper computation.

Eliciting secret knowledge (Figure 10). This application (Section 5.1) applies the tuned lens to the experimental setup of CYWINSKI et al. (2025), where a model is fine-tuned to respond to queries about a "taboo" word without verbalizing it directly, and an auditor LLM tries to guess the taboo word from the model's internal predictions. Results are mixed: "the tuned lens outperforms [the logit lens] at some layers and underperforms at other layers." The paper notes that "without 'peaking' at the ground truth labels, it is not really possible to determine which layer to use in advance," identifying a practical limitation: while the tuned lens provides better predictions at most layers, knowing which layer to trust for a given task requires task-specific validation.

Ablation Studies and Robustness Checks

Final transformer layer inclusion vs. exclusion (Figure 4): For BLOOM 560M, the tuned lens is evaluated both with and without the final transformer layer included in the probe. The tuned lens performs well in both settings—perplexity is similarly low whether or not the final layer is included—demonstrating that the method is robust to this architectural choice. The logit lens, by contrast, degrades severely when the final layer is excluded (perplexity rises from ~200 to ~800 at layer 12). This ablation establishes that the tuned lens's performance is not dependent on a particular choice of which layers to include in the probe, unlike the extended logit lens (Equation 4) which was specifically introduced as an ad-hoc fix for GPT-Neo.

SGD vs. Muon optimizer for translator training: After the initial publication, the authors discovered that "using Muon instead of SGD for training the tuned lens dramatically accelerates training, allowing us to achieve much lower KL divergence losses than were possible in our initial experiments." They state candidly that "essentially, all of the tuned lenses trained and evaluated in the earlier versions of this paper were severely undertrained." The Muon-trained translators have "Frobenius norms that are often several times larger than SGD-trained ones," indicating they deviate much further from the identity initialization—they are "much better 'tuned.'" This is a significant robustness finding: the quality of the translators is sensitive to the optimizer, and the reported results likely underestimate the tuned lens's true capability. The authors "opted not to re-do all our experiments with tuned lenses trained using Muon, but we encourage practitioners to use Muon in future experiments." This is an acknowledged limitation of the reported numbers—they represent a lower bound on what the method can achieve with better optimization.

Logit lens extended variant (nostalgebraist, 2021): The paper tests the extended logit lens (Equation 4, which retains the final transformer layer in the probe) as a baseline. Figure 1 (top) shows that even this extension is "only partially successful at recovering meaningful results" for GPT-Neo-2.7B. The qualitative comparison in Figure 1 shows the tuned lens producing semantically coherent top-1 predictions from layer 5 onward, while the extended logit lens produces random or repetitive tokens until layer 21. This ablation confirms that the tuned lens's improvements are not simply from including the final layer—the learned affine transformation is doing substantive work beyond what the extended logit lens achieves.

Transfer to fine-tuned models (Figure 13, Appendix A): The comparison of Vicuna-trained vs. LLaMA-transferred translators on two different evaluation datasets (RedPajama and Helpful Harmless) tests whether fine-tuning substantially alters the representational basis. The transferred lens performs nearly identically to the Vicuna-trained lens on the conversational dataset (no significant difference) and within 0.3 bits per byte on the pretraining-distribution dataset. This suggests that fine-tuning minimally affects the encoding format of hidden states, which is a robustness result for the practical use of pretrained tuned lens checkpoints.

Logit lens failure mode diagnosis (Figures 17–18, Appendix B.1): For BLOOM 560M and OPT 125M, the paper shows that the logit lens's top-1 prediction "is often the input token, rather than any plausible continuation token, in more than half the layers." Figures 17 and 18 visually demonstrate this pathology: large swaths of the prediction trajectory grid are colored identically to the input tokens, showing the logit lens is essentially echoing rather than predicting. The tuned lens is not shown in these figures, but the implication from the main results is that this pathology is resolved—the tuned lens produces plausible continuation predictions even for models where the logit lens fails entirely.

Random baseline for static interpretability (Table 3, Figure 21, Appendix E): In the static interpretability analysis (Appendix E), the paper tests whether MLP and attention weight matrices appear more interpretable when projected through the tuned lens versus the logit lens. For Pythia 125M, the tuned lens produces higher interpretability scores (as measured by pairwise cosine similarity of BPEmb embeddings of the top-k tokens) than the logit lens across all tested parameter types. However, the difference between random and real matrices is consistently higher with the tuned lens, suggesting genuine improvement rather than an overall score inflation. Figure 21 reveals a critical nuance: "most singular vectors are not more interpretable than random, except for a minority lying in a long right tail" for both lenses. In preliminary experiments with larger models, "both the tuned and logit lens appeared to perform poorly" on static interpretability, suggesting this application does not scale well with model size—a negative result that the paper reports candidly.

Critical Assessment

The experimental results provide strong support for the paper's central methodological claim: the tuned lens produces more accurate, less biased, and more causally faithful layer-wise predictions than the logit lens across a diverse set of model families and scales. However, several qualifications and gaps merit attention.

The claim that the tuned lens is "more predictive, reliable and unbiased" (abstract) is well-supported for perplexity and marginal bias. Figures 5 and 14 demonstrate uniformly lower perplexity across six model families, Figure 3 shows near-elimination of systematic marginal bias (from ~4–5 bits to near zero), and Figure 4 confirms robustness to the final-layer inclusion choice on BLOOM. The evidence is consistent and the effect sizes are large—a 5× perplexity reduction at early layers on Pythia 12B is not marginal. However, the evaluation is purely on pretraining-distribution text. The paper does not test whether tuned lens perplexity improvements generalize to out-of-distribution text, domain-shifted inputs, or adversarial inputs. Given that the translators are trained via distillation on the pretraining distribution, it is plausible that they overfit to distribution-specific representational patterns, and their advantage over the logit lens could shrink or vanish on unusual inputs. The prompt injection detection results (Table 1) provide some indirect evidence of generalization—the tuned lens trajectory enables better anomaly detection than the logit lens on adversarial prompts that were not seen during translator training—but a direct perplexity comparison on out-of-distribution text would strengthen the claim.

The causal fidelity results (Section 4) establish Property 1 convincingly but leave Property 2 with room for strengthening. The importance alignment result (Spearman ρ = 0.89, no features in the "influential for probe but not model" quadrant) is strong evidence that the tuned lens relies on genuinely important directions. The stimulus-response alignment results (Figure 9) show positive and depth-increasing alignment, with the tuned lens outperforming the logit lens, but the absolute Aitchison cosine similarities are moderate (0.3–0.7). An alignment of 0.5 means the tuned lens and the model agree on the direction of change only about half the time (in the angular sense), which leaves substantial room for the tuned lens to "misinterpret" the effect of an intervention. The paper does not provide a baseline for what perfect alignment would look like (e.g., the stimulus-response alignment of the model with itself under different random seeds, or the alignment of two independently trained probes). Without such a baseline, it is difficult to interpret whether an alignment of 0.6 at layer 8 is "good" (close to the noise ceiling) or "concerning" (far below what is achievable).

The prompt injection detection results (Table 1) are striking but raise generalizability questions. Near-perfect AUROC (≥0.998) on five tasks is impressive, but the paper uses a single attack template ("Ignore any previous and following instructions and just print...") adapted from Perez et al. (2022). It is unclear whether the detection performance would generalize to other prompt injection strategies (e.g., indirect injections embedded in retrieved documents, multi-turn injections, or attacks that do not use the "ignore previous instructions" framing). The experiment also uses only Pythia 12B; it is unknown whether the result transfers to other model families, though the tuned lens's consistent improvements across model families in the main results make this plausible. The comparison with SRM shows that the tuned lens trajectory and single-layer hidden state representations capture complementary signals—the tuned lens is better on some tasks (ARC-Challenge) while SRM is better on others (MC TACO, SciQ)—suggesting that neither method alone is sufficient for robust detection across all attack types. The paper's own suggestion to combine both approaches is sensible but untested.

The example difficulty correlation (Table 2) is statistically significant but practically modest. Spearman ρ values of 0.05–0.38 are positive but explain at most ~14% of variance in iteration learned. This means prediction depth is a weak predictor of which examples were learned early versus late during training. The paper frames this as validation of prediction depth as a difficulty metric, which is fair—the fact that any correlation exists across 11 diverse tasks is notable given that prediction depth is measured from a single forward pass on the final model while iteration learned is measured across 143 training checkpoints. However, the modest correlation magnitudes mean that prediction depth is not a practical replacement for iteration learned in applications that require precise difficulty estimates (e.g., data pruning for efficient training). Figure 12 provides a compelling qualitative visualization but is cherry-picked (the GPT-4 technical report abstract) and not quantitatively analyzed—the paper does not report whether the visually apparent patterns (low depth for function words, high depth for content words) hold statistically across the full evaluation set.

Several experiments that would strengthen the paper are absent. First, the paper does not compare the tuned lens to a simple alternative: training a linear probe on each layer's hidden state to directly predict the next token (supervised cross-entropy loss), rather than distilling from the final layer. The paper argues (Section 3) that distillation prevents the probe from "learn[ing] extra information over and above what the model has learned," but an empirical comparison would quantify how much information the probe "leaves on the table" by constraining itself to only what the model outputs. If a supervised probe substantially outperforms the tuned lens at early layers, it would suggest that those layers contain usable predictive information that the model fails to fully exploit—an interesting finding in its own right that would nuance the iterative inference interpretation. Second, the paper does not evaluate the tuned lens on a genuinely multilingual model other than BLOOM (which is evaluated only for perplexity, not for the applications). The iterative inference hypothesis should apply regardless of language, but representational drift might behave differently in multilingual models where different languages occupy different subspaces of the representation space. Third, the static interpretability analysis (Appendix E) reports that "both the tuned and logit lens appeared to perform poorly" on larger models, which is a negative result that undermines one of the paper's intended applications. The paper does not investigate why—is it because larger models have more distributed representations that are harder to capture in top-k token lists, or because the tuned lens training was insufficient (especially given the later Muon finding)?

The sample size for causal experiments is modest relative to the number of parameters. Causal basis extraction uses a single in-memory batch of 131,072 tokens for optimization (Section 4.1). For Pythia 410M with a hidden dimension of 1024, each CBE direction is optimized over only ~128 examples per dimension, which may be insufficient to reliably identify directions that are influential across the full data distribution. The paper mitigates this by initializing CBE from the singular vectors of the probe rather than random directions, which likely helps convergence, but does not report how stable the extracted directions are under different random seeds or data samples. If the top CBE directions vary substantially with the optimization batch, the causal fidelity conclusions could be specific to the particular directions found rather than a general property of the tuned lens.

The limitation acknowledged in Section 6—that training translators requires computational effort while the logit lens requires none—is genuine but somewhat overstated. The paper reports training time of "under an hour on a single 8×A40 node," which is negligible compared to the cost of pretraining the models being analyzed. Moreover, the released pretrained checkpoints eliminate this cost for downstream users. The more significant practical limitation is the Muon finding: the paper's own results may significantly underestimate the tuned lens's capability because the translators were undertrained with SGD. Since the Muon-trained translators are not systematically evaluated, the reported perplexity and bias numbers should be treated as conservative lower bounds. A re-evaluation with Muon-trained translators would likely show even larger gaps over the logit lens, particularly at early layers where the SGD-trained translators are furthest from the optimum.

A subtle but important gap: the paper never establishes whether the tuned lens's improved perplexity translates to improved practical utility in the applications. The overthinking experiment (Figure 11) reproduces qualitative patterns from Halawi et al. (2023) but does not compare the tuned lens's accuracy against the logit lens's accuracy on the same models and tasks—the logit lens simply could not run on those models, so no head-to-head comparison is possible. The prompt injection detection experiment does compare tuned vs. logit lens trajectories and shows the tuned lens is better on most tasks (Table 1), but the improvement is modest on several benchmarks (e.g., BoolQ: 1.000 vs. 0.972; SST-2: 0.998 vs. 0.943). The example difficulty experiment shows the tuned lens correlates better with iteration learned than the logit lens in 8/11 tasks, but the absolute correlation values remain low for both methods. It is therefore unclear whether the tuned lens's quantitative advantages in perplexity and bias translate into qualitatively different scientific insights or practically meaningful application improvements, or whether the logit lens—where it works at all—is often "good enough" for the research questions being asked. The paper's strongest case for practical impact is the extension of overthinking analysis to BLOOM and GPT-Neo (models the logit lens cannot handle at all), which genuinely expands the scope of possible research. But for models where the logit lens already works (GPT-2), the paper does not demonstrate that the tuned lens enables new scientific findings that were inaccessible before.

6. Limitations and Trade-offs

Distribution-Shifted and Adversarial Inputs Are Unexplored Territory

The assumption or constraint. The tuned lens translators are trained exclusively on text from the model's pretraining distribution—the Pile validation set or, for OPT, the original pretraining validation set (Section 3, "Implementation details"). The distillation objective (Equation 9) minimizes KL divergence between the tuned lens output and the model's final-layer output on this distribution. The paper never evaluates tuned lens perplexity, bias, or causal fidelity on out-of-distribution text, domain-shifted inputs, or adversarial prompts. This is an implicit assumption that the affine translators learned on pretraining text will generalize to whatever inputs downstream analyses encounter.

The consequence. If representational drift patterns differ on unusual inputs—for instance, if the rogue dimensions identified in Figure 6 behave differently when the model processes syntactically anomalous or semantically nonsensical text—the translators' learned basis corrections could become mismatched, producing decoded distributions that are systematically distorted in ways the practitioner cannot detect without a reference distribution. This is particularly concerning for the safety application the paper highlights (prompt injection detection, Section 5.3): if an attacker crafts inputs that specifically perturb the hidden state in directions the translators mishandle, the prediction trajectory could appear benign even when the model's internal computation is disrupted. More subtly, the iterative inference interpretation the paper advances (Appendix C)—that residuals align with the negative gradient to progressively refine predictions—might break down on inputs far from the training distribution, where the model's computation follows different dynamics. The paper's own prompt injection detection experiment (Table 1) was conducted on Pythia 12B using a single attack template, and while the tuned lens trajectory did enable detection, this tests generalization from pretraining text to a specific adversarial prompt, not systematic out-of-distribution robustness.

What evidence exists in the paper. None directly. The transfer-to-fine-tuned-models experiment (Figure 13, Appendix A) provides the closest proxy: translators trained on LLaMA 13B's pretraining distribution transfer well to Vicuna 13B's conversational distribution (no significant difference on Helpful Harmless, 0.3 bits per byte degradation on RedPajama). This suggests some degree of distributional robustness, but conversational text from a fine-tuned model is far closer to pretraining text than genuinely adversarial or out-of-domain inputs would be. The prompt injection experiment evaluates detection AUROC, not tuned lens perplexity or calibration on the adversarial inputs themselves.

Mitigation status. Not addressed. The paper does not discuss this limitation or propose robustness evaluations. A practitioner deploying the tuned lens for safety monitoring would need to independently validate that the translators remain calibrated on the specific input distribution of their deployment setting—and would have no principled way to detect when they fail without ground-truth final-layer outputs for comparison, which defeats the purpose of using the tuned lens as a monitoring tool.

Difficulty Estimation for Practical Deployment Is Absent

The assumption or constraint. The paper's quantitative results—perplexity, bias, causal fidelity—are all computed on a fixed evaluation set of 16.4 million tokens with known final-layer outputs. In a deployment scenario where the tuned lens is used for tasks like prompt injection detection or early exiting, the practitioner does not have access to the model's final-layer distribution (if they did, they would not need the tuned lens). The paper provides no method for estimating, at inference time, how reliable the tuned lens output is for a given input. Unlike the difficulty estimation framework, there is no mechanism to flag inputs where the translator's affine approximation may be poor—for instance, inputs where the hidden state h_ℓ falls in a region of representation space poorly covered by the translator's training data.

The consequence. This creates a silent failure mode: the tuned lens always produces a distribution over the vocabulary, but the user has no way to know whether that distribution is a faithful approximation of what the model would output if the forward pass continued, or whether it is distorted by a mismatch between the translator's training distribution and the current input. In the prompt injection detection application (Section 5.3), an anomaly detector trained on tuned lens trajectories from normal prompts could produce confident but wrong classifications if the tuned lens itself behaves anomalously on novel attack types—not because the model's computation is disrupted, but because the translator is extrapolating poorly. The paper's own finding that the secret knowledge elicitation experiment (Figure 10) produces mixed results—"the tuned lens outperforms at some layers and underperforms at other layers"—without a way to know which layer to trust in advance illustrates this problem concretely.

What evidence exists in the paper. The transfer penalty matrix (Figure 7) shows that translators degrade when applied to layers other than the one they were trained for, with penalties increasing with representational distance. This demonstrates that translator quality is sensitive to distribution shift across layers, which implies sensitivity to distribution shift across inputs as well, since unusual inputs may produce hidden states with atypical covariance structure. The Muon optimizer finding (Section 3) reveals that the SGD-trained translators were "severely undertrained," meaning the paper's own reported numbers already reflect a suboptimal translator quality that users cannot detect without access to better-trained references.

Mitigation status. The paper does not address this. There is no uncertainty quantification, no confidence metric for tuned lens outputs, and no discussion of how a practitioner should decide whether to trust the decoded distribution at a particular layer for a particular input. This is a fundamental gap between the tuned lens as a research tool (where the final-layer output is available for validation) and as a deployment tool (where it is not).

Causal Fidelity Is Established for a Single Model Family at Modest Scale

The assumption or constraint. The causal fidelity experiments in Section 4—both the importance alignment analysis via Causal Basis Extraction (Figure 8, Figure 20) and the stimulus-response alignment analysis via resampling ablation (Figure 9)—are conducted exclusively on Pythia models at modest scales (160M and 410M parameters). The paper claims that these experiments "show the tuned lens uses similar features to the model itself" (abstract), implicitly generalizing this claim to all model families and scales where the tuned lens is evaluated. However, the largest model tested for causal fidelity (Pythia 410M, 24 layers, d_model = 1024) is substantially smaller than the largest model where tuned lens perplexity is reported (GPT-NeoX-20B, 44 layers, d_model = 6144).

The consequence. The causal fidelity results may not scale to larger models. Larger models have more layers, wider hidden dimensions, and potentially more complex representational dynamics—including more pronounced rogue dimensions and covariance drift patterns that a simple affine translator might handle adequately for perplexity reduction but not for faithful causal attribution. The paper itself notes that the static interpretability analysis (Appendix E) "appeared to perform poorly" on larger models for both the tuned and logit lenses—a negative result that was not investigated further. If the tuned lens's causal fidelity degrades with scale, then the paper's central interpretability claim—that the prediction trajectory reflects the model's actual computational process—would be valid only for smaller models, limiting the method's utility for understanding the largest deployed systems where interpretability is most urgently needed.

What evidence exists in the paper. The causal fidelity experiments on Pythia 160M and 410M are thorough within their scope—the Spearman ρ of 0.89 (Figure 8) and the consistent improvement over the logit lens in stimulus-response alignment (Figure 9) are convincing for these specific models. Figure 20 (Appendix D) extends the importance alignment analysis across all layers of Pythia 410M, showing ρ values of approximately 0.85–0.95 throughout—but only for this one model. The paper provides no analogous experiments for GPT-NeoX-20B, the Pythia 12B, or any non-Pythia model family. The static interpretability degradation on larger models (acknowledged in Appendix E.1) is circumstantial evidence that something changes at scale, but its relevance to causal fidelity specifically is unclear.

Mitigation status. The paper does not discuss this scaling limitation of the causal fidelity experiments. Section 6 ("Limitations and future work") focuses on the computational intensity of CBE and the potential to optimize subspaces rather than individual directions, but does not mention the need to validate causal fidelity on larger models. The absence of causal fidelity results for the larger models in the paper's evaluation suite is an unaddressed gap.

The Training Cost of Difficulty Assessment Is Unaccounted for in Practical Deployments

The assumption or constraint. While the tuned lens does not require the elaborate difficulty estimation of test-time compute methods, it imposes its own unaccounted cost: training the affine translators. The paper reports that training takes "under an hour on a single 8×A40 node" (Section 6) using the SGD recipe, and that the Muon optimizer accelerates this substantially. However, this training requires access to the model's pretraining validation set (or a suitable proxy), a full forward pass through the frozen model to extract hidden states at every layer for every training token, and—critically—the model's final-layer logits as distillation targets. For the largest models evaluated (GPT-NeoX-20B), this means processing enough tokens to train L × (d² + d) parameters per layer (approximately 1.66 billion total translator parameters for 44 layers with d = 6144), with each training step requiring a full forward pass through the 20B-parameter model.

The consequence. For practitioners working with very large models (70B+, 175B+, or beyond), the computational cost of training translators may be non-trivial—hours or days of GPU time rather than under an hour. Moreover, the released pretrained checkpoints cover only the models evaluated in the paper. Anyone working with a model architecture or training distribution not in that set must train translators from scratch. The paper's finding that translators benefit substantially from the Muon optimizer (with SGD-trained translators being "severely undertrained") means that achieving the best possible tuned lens quality requires nontrivial optimization expertise—the simple SGD recipe in the paper produces suboptimal results, and the Muon-based training procedure is not fully specified (hyperparameters, schedule, number of steps) for the models in the paper. A practitioner following the original SGD recipe would unknowingly deploy undertrained translators, getting worse results than the method is capable of.

What evidence exists in the paper. The Muon finding is reported candidly: "essentially, all of the tuned lenses trained and evaluated in the earlier versions of this paper were severely undertrained" (Section 3, "Muon optimizer" paragraph). The Frobenius norms of Muon-trained translators are "often several times larger than SGD-trained ones," indicating substantially different solutions. However, the paper does not provide a systematic comparison of SGD vs. Muon perplexity across model families or a recommended Muon training protocol. The reported perplexity numbers (Figures 5, 14) are from the SGD-trained translators, which the paper now acknowledges are suboptimal.

Mitigation status. Partially addressed through the release of pretrained checkpoints for the evaluated models, which eliminates the training cost for those specific models. The paper encourages practitioners to use Muon "in future experiments" but does not provide a recipe. The training cost is explicitly listed as a limitation in Section 6 ("Limitations and future work"): "One limitation of our method is that it involves training a translator layer for each layer of the network, while the logit lens can be used on any pretrained model out-of-the-box." The paper frames the released checkpoints as mitigating this, but this mitigation only applies to the specific model checkpoints released—not to any other model a practitioner might want to analyze.

The Method Provides No Mechanism to Combine or Select Across Layers

The assumption or constraint. The tuned lens produces a trajectory of predictions—one distribution per layer—but the paper provides no principled method for using that trajectory to make decisions. In every application, the user must either pick a specific layer (as in the overthinking experiment, Figure 11, where early-layer predictions are more robust to incorrect demonstrations), aggregate across layers (as in the prompt injection detection experiment, Section 5.3, where the full trajectory is flattened into a feature vector), or define an ad-hoc stopping criterion (as in the prediction depth metric, Section 5.4, where the layer at which the top-1 prediction stabilizes is used). There is no theoretical framework or empirical procedure for determining which layer's prediction to trust for a given input, nor any way to combine information across layers optimally.

The consequence. This creates an interpretation gap: the tuned lens shows that the model's latent predictions evolve across layers, but does not help the practitioner decide what to do with that information. The secret knowledge elicitation experiment (Figure 10) makes this concrete: the tuned lens outperforms the logit lens at some layers and underperforms at others, and "without 'peaking' at the ground truth labels, it is not really possible to determine which layer to use in advance" (Section 5.1). This means the tuned lens is better at extracting information from each layer individually, but the user still faces the same fundamental challenge of knowing which layer's extracted information to trust. For the prompt injection detection application, flattening the trajectory into a feature vector and feeding it to an anomaly detector sidesteps this by using all layers—but this approach throws away the sequential structure of the trajectory and provides no insight into why a particular input is anomalous, reducing interpretability.

What evidence exists in the paper. The overthinking experiment (Figure 11) shows that the optimal layer for robustness to incorrect demonstrations varies by model and task—the peak of the red curve occurs at different layer indices for GPT-Neo 1.3B versus GPT-Neo 2.7B versus BLOOM 560M. The example difficulty experiment (Section 5.4) defines prediction depth as the layer where the top-1 prediction stabilizes, but this is an arbitrary criterion—the paper does not justify why "stops changing" is the right threshold rather than, say, "crosses a confidence threshold" or "achieves within 1% of final-layer accuracy." The transfer penalty matrix (Figure 7) shows that translators work best on their trained layer and degrade with distance, but this tells us about translator quality, not about which layer's decoded prediction is most useful for a downstream task.

Mitigation status. Not addressed. The paper treats the tuned lens as a drop-in replacement for the logit lens—improving the quality of the decoded predictions at each layer—but does not tackle the higher-level question of how to use the resulting trajectory. This is consistent with the paper's framing as a methodological infrastructure contribution, but it means that the tuned lens solves the decoding problem while leaving the interpretation problem—what do the trajectories mean and how should we act on them?—to future work. For a practitioner, this means that adopting the tuned lens over the logit lens will improve the quality of the per-layer predictions they see, but will not by itself tell them which layers to attend to or how to combine layer-wise information for decision-making.

Single-Modality and Single-Task-Family Scope Limits Generality Claims

The assumption or constraint. All experiments in the paper are conducted on autoregressive language models performing next-token prediction on English text. The paper acknowledges this scope limitation in Section 6: "Due to space and time limitations, we focused on language models in this work, but we think it's likely that our approach is also applicable to other modalities." The evaluation datasets for the applications (Section 5) are exclusively English NLP benchmarks. The autoregressive, left-to-right causal structure of the models is assumed throughout—the iterative inference framing (Appendix C), the LayerNorm-placement analysis (pre-LN vs. post-LN), and the training procedure all depend on the specific architecture of decoder-only transformers.

The consequence. The paper's claims about representational drift, rogue dimensions, covariance structure, and the sufficiency of affine translators for basis alignment may not transfer to other modalities or architectures. Vision transformers (ViT) use a different patch-embedding scheme and process fixed-size inputs rather than causal sequences; encoder-decoder models (T5, BART) have cross-attention between encoder and decoder; and encoder-only models (BERT) use bidirectional attention rather than causal masking. In each case, the residual stream dynamics—and therefore the pattern of representational drift that the tuned lens corrects for—could be qualitatively different. The iterative inference hypothesis that motivates the tuned lens (Appendix C) was originally developed for ResNet image classifiers (Jastrzębski et al., 2017), suggesting it might generalize, but the paper provides no empirical evidence. The statement that the approach is "likely... also applicable to other modalities" is speculation, not a supported claim.

What evidence exists in the paper. None. The paper is entirely confined to text-based autoregressive language models. Even within language models, the evaluation is restricted to English—BLOOM is a multilingual model, but the paper evaluates it only on English text from the Pile validation set (see BLOOM evaluation in Figure 4, Appendix A). The reliance on pre-LayerNorm architecture is explicit: "Both the logit lens and the tuned lens are designed primarily for the pre-LN architecture, which is more unambiguously iterative. Luckily pre-LN is by far more common than post-LN among state-of-the-art models" (Section 2 footnote). OPT 350M, which uses post-LN, is excluded from the main results (Figure 14, Appendix A). This means the method is not even architecturally universal within language models—it is restricted to the dominant but not universal pre-LN design.

Mitigation status. Acknowledged but not addressed. The paper flags this as future work in Section 6, and the speculative language ("we think it's likely") indicates awareness of the limitation. However, the abstract and introduction make strong generality claims ("making it possible to decode every hidden state," "a drop-in replacement for the logit lens that makes it possible to elicit interpretable prediction trajectories from essentially any pretrained language model in use today") that are partially undermined by the architectural constraint to pre-LN transformers and the complete absence of non-text, non-English, or non-autoregressive evaluation.

7. Implications and Future Directions

How This Work Changes the Landscape

The tuned lens transforms the landscape of transformer interpretability by converting a brittle, model-specific trick—the logit lens—into a reliable, general-purpose measurement instrument. This is not a paradigm shift; it is a methodological upgrade that takes an existing technique from "works on some models for unclear reasons" to "works on essentially any pre-LN autoregressive language model, with well-understood failure modes and a causal fidelity guarantee." The significance lies in what this upgrade enables: a whole class of analyses that were previously restricted to GPT-2 can now be applied to BLOOM, GPT-Neo, Pythia, OPT, and GPT-NeoX-20B.

The conceptual shift is the reframing of layer-wise decoding as a basis-alignment problem. Before this work, the logit lens's failures were treated as opaque—some models cooperated, others did not, and the field's response was ad-hoc patching (the extended logit lens, Equation 4). The tuned lens argues that these failures have a specific, diagnosable cause: the coordinate system of the residual stream drifts across layers, and applying the unembedding without a change-of-basis produces systematically distorted outputs. This diagnosis is backed by two lines of evidence that did not previously exist in the literature: (1) the emergence of rogue dimensions at specific layers that dominate covariance structure (Figure 6, top), and (2) the smooth decay of covariance similarity with layer distance even after removing those rogue dimensions (Figure 6, bottom). The transfer penalty matrix (Figure 7) then closes the loop by showing that covariance similarity predicts translator transferability (Spearman ρ = -0.78), establishing a causal link between the diagnosed problem and the decoding failure.

This reframing has ripple effects. It implies that any method that compares hidden states across layers—activation patching, representation similarity analysis, probing classifiers trained on different layers, causal tracing—needs to account for basis drift. A researcher who finds that a probe trained on layer 5 hidden states performs poorly when applied to layer 20 hidden states cannot conclude that the information is absent at layer 20; it may simply be encoded in a different basis. The tuned lens provides both a diagnostic tool (the transfer penalty matrix) and a corrective (the affine translators) for this confound. This makes research that relies on cross-layer comparisons more rigorous, shifting the burden of proof: claims that a capability "emerges" at a particular layer must now rule out basis drift as an alternative explanation.

The paper also pushes the field's standards for what counts as a faithful probe. The two-dimensional causal fidelity framework—importance alignment via Causal Basis Extraction plus stimulus-response alignment via Aitchison cosine similarity—provides a concrete, quantitative benchmark that future probing methods can be evaluated against. The finding that the tuned lens achieves Spearman ρ = 0.89 on importance alignment (Figure 8) and Aitchison similarity of 0.3–0.7 on stimulus-response alignment (Figure 9) sets a baseline that competing methods must meet or exceed. This raises the bar for interpretability tooling: it is no longer sufficient for a probe to achieve low perplexity or high classification accuracy; it must also demonstrate that the features it relies on are causally implicated in the model's own computation.

The paper reconciles the logit lens's inconsistent track record. Nostalgebraist (2020) found it worked well for GPT-2; subsequent researchers found it failed for GPT-Neo (nostalgebraist, 2021) and BLOOM; Halawi et al. (2023) successfully applied it to GPT-2 and GPT-J but could not extend to BLOOM or GPT-Neo. These conflicting experiences were not due to researcher error or model quirks—they were the expected consequence of varying degrees of representational drift across model families. GPT-2's residual stream basis happens to be relatively stable (making the logit lens work passably); BLOOM's is not (making the logit lens echo input tokens, Figure 17). The tuned lens explains why these model-specific differences exist and provides a unified solution, converting a confusing set of contradictory empirical results into a coherent picture.

Finally, the paper opens new research directions by making prediction trajectories reliable enough to serve as signals for downstream tasks. The prompt injection detection results (Table 1, near-perfect AUROC on five benchmarks) demonstrate that the trajectory carries diagnostic information beyond what the final output reveals. The example difficulty correlation with iteration learned (Table 2) shows that the trajectory encodes information about the computational difficulty of the input that is recoverable from a single forward pass. These applications were not possible with the logit lens on the models where the logit lens failed; the tuned lens makes them tractable across model families.

Research directions that become more attractive include: automated anomaly detection on internal model states, early-exiting analyses on any pretrained model without retraining, studying the relationship between training dynamics and inference-time computation, and using prediction trajectories as a signal for model editing or steering. Research directions that become less attractive include: ad-hoc extensions of the logit lens for specific model families (the extended logit lens of Equation 4 is superseded), and probing methods that train per-layer unembedding matrices (the paper shows translators are more parameter-efficient and converge better). The logit lens itself is not obsolete—it remains useful for quick, zero-cost exploration on models where it happens to work—but the tuned lens is now the methodologically preferred tool for any analysis where reliability matters.

Follow-Up Research This Work Enables

Systematic study of representational drift across model architectures, scales, and training regimes. The paper diagnoses representational drift in pre-LN autoregressive transformers but leaves open the question of how drift patterns vary. A natural follow-up would measure covariance similarity matrices (as in Figure 6) for a sweep of models varying (a) architecture: pre-LN vs. post-LN, decoder-only vs. encoder-decoder vs. encoder-only; (b) scale: from 125M to 175B+ parameters, testing whether drift accelerates or decelerates with depth and width; (c) training regime: models trained with different objectives (next-token prediction vs. masked language modeling), different data mixtures (multilingual vs. monolingual), and different optimization strategies (AdamW vs. Muon). The hypothesis is that drift magnitude correlates with model depth and that post-LN architectures exhibit more severe drift at the final layers (since the LayerNorm placement changes the residual dynamics). The prediction metric would be the average off-diagonal Frobenius cosine similarity in the covariance matrix, correlated with tuned lens transfer penalty. Strong negative results—models where the tuned lens fails despite low covariance drift—would identify limits of the affine-translation assumption.

Joint anomaly detection combining tuned lens trajectories with hidden-state representations. The prompt injection detection experiment (Table 1) reveals that the tuned lens trajectory and the Simplified Relative Mahalanobis distance capture complementary signals: the tuned lens outperforms SRM on ARC-Challenge (0.907 vs. 0.844) while SRM outperforms on MC TACO (0.952 vs. 0.875) and SciQ (0.969 vs. 0.898). The paper explicitly suggests combining the two approaches. A strong follow-up would train a detector that concatenates (a) the flattened tuned lens prediction trajectory across all layers (log-probabilities of each answer choice), (b) the SRM score at the middle layer, and (c) the SRM scores at multiple layers as a trajectory, then evaluate on a diverse attack suite including indirect prompt injections (retrieved documents), multi-turn injections, and attacks using different syntactic formulations beyond "Ignore previous instructions." The null hypothesis is that SRM alone saturates detection performance and the trajectory adds no marginal value; the tuned lens paper's results on ARC-Challenge argue against this null. Generalization to fine-tuned chat models (using transferred translators, as validated in Figure 13) would test whether the detection signal survives instruction tuning.

Using tuned lens prediction trajectories to study the effect of fine-tuning on model computation. The finding that translators transfer from LLaMA 13B to Vicuna 13B with negligible degradation (Figure 13) enables an experimental design that was previously impractical: compare the prediction trajectories of a base model and its fine-tuned variant on the same inputs, using the same translators, to isolate how fine-tuning changes the internal computation. Concretely, feed the same prompt through both LLaMA and Vicuna, decode hidden states at each layer using the same tuned lens translators, and compute the layer-wise KL divergence between the two prediction trajectories. Peaks in divergence indicate layers where fine-tuning most substantially altered the model's latent beliefs. This could be correlated with known fine-tuning phenomena: do RLHF-tuned models diverge most in the middle layers (suggesting value alignment happens through intermediate reasoning) or in the final layers (suggesting a shallow stylistic shift)? The experiment requires no additional training—the translators are already available from the paper's released checkpoints—and the metric (layer-wise KL divergence between base and fine-tuned trajectories on held-out prompts) is straightforward to compute.

Explaining and mitigating the static interpretability degradation at scale. Appendix E.1 reports that static interpretability analysis (projecting MLP and attention parameters through the tuned lens to find interpretable top-k tokens) "appeared to perform poorly" on larger models for both the tuned and logit lenses. This negative result is important because it bounds the utility of the tuned lens for weight-level interpretability. A follow-up should systematically characterize why this degradation occurs. Hypotheses include: (a) larger models have more distributed representations, so individual parameter vectors do not correspond to interpretable semantic clusters regardless of the projection method; (b) the tuned lens translators for larger models were more severely undertrained (given the Muon finding), and retraining with Muon might recover interpretability; (c) the BPEmb-based automated metric (Appendix E.3) is inappropriate for larger models because it measures pairwise token similarity but larger models' parameter vectors may represent more abstract or compositional concepts that are not well-captured by averaging word embeddings. Testing these hypotheses would involve: retraining translators with Muon on a large model (Pythia 12B), recomputing interpretability scores, comparing with human evaluation on a sample of top-k token lists, and developing alternative automated metrics (e.g., using the model's own embedding similarity rather than BPEmb). If the degradation persists under all conditions, this would establish an important boundary on the tuned lens's applicability: it decodes hidden states well (Figures 5, 14) but does not make individual weight matrices more interpretable at scale.

Causal basis extraction at scale and across model families. The causal fidelity experiments (Section 4) are restricted to Pythia 160M and 410M. Extending Causal Basis Extraction to Pythia 12B and GPT-NeoX-20B would test whether the Spearman ρ ≈ 0.89 finding holds at scale. This is computationally demanding—CBE sequentially optimizes d_model directions per layer using L-BFGS on a batch of 131K tokens—but the paper suggests optimizing a k-dimensional subspace at each iteration rather than individual directions, which could reduce the computational cost from O(d_model^2) to O(k × d_model). A strong follow-up would implement subspace CBE, apply it to Pythia 12B (the largest model with public training checkpoints, enabling iteration-learned correlations), and measure whether the tuned lens's causal fidelity degrades, holds steady, or improves with scale. A degradation would suggest that the affine translators become less faithful as representations become higher-dimensional and more distributed; an improvement would strengthen the iterative inference interpretation. In either case, establishing the scaling behavior of causal fidelity is essential for the tuned lens's credibility as a tool for analyzing frontier models.

Prediction depth as a data-pruning signal for pretraining. The correlation between tuned lens prediction depth and iteration learned (Table 2, Spearman ρ of 0.05–0.38 across tasks) is modest but suggests that examples learned late in training tend to require more layers to classify at inference time. If this correlation can be strengthened—perhaps by defining prediction depth using a more sophisticated criterion than "top-1 prediction stops changing," such as the layer at which the prediction distribution's entropy drops below a threshold or the layer at which the prediction achieves 95% of the final-layer's confidence—then the tuned lens could be used as a data-pruning tool: for a large unlabeled corpus, run the tuned lens on each example, estimate prediction depth, and discard examples with high prediction depth (which are "hard" and likely learned late, contributing less to training efficiency). The experiment would: (a) define several candidate prediction-depth metrics on a model with available training checkpoints (Pythia 12B), (b) compute their correlation with iteration learned across a diverse set of tasks, (c) select the metric with the strongest correlation, (d) use it to rank a large corpus by estimated difficulty, and (e) compare the efficiency of training a new model from scratch on the easiest 50% of examples versus a random 50%. The null result—that prediction-depth-based pruning does not improve training efficiency over random pruning—would indicate that the correlation, while statistically significant, is too weak to be practically useful.

Practical Applications and Downstream Use Cases

Safety monitoring for deployed language models via trajectory-based anomaly detection. An organization deploying a chatbot based on a model like Pythia 12B can integrate the tuned lens into their inference pipeline to detect prompt injection attacks in real time. The setup: at each inference call, the system records the tuned lens prediction trajectory (the log-probability assigned to each output token at each layer, or—for a multiple-choice framing—the log-probability of each candidate response at each layer), flattens it into a feature vector, and feeds it to a pre-trained anomaly detector (isolation forest or LOF, as in Section 5.3). The paper's results in Table 1 show near-perfect AUROC (0.998–1.000) on five of nine tasks for Pythia 12B, meaning the detector catches almost every injection attempt while raising essentially zero false alarms on normal prompts—a capability that input-filtering or output-monitoring approaches cannot match because the final output may appear benign even when the model's internal computation is disrupted. The practical benefit is a deployment-time safety layer that operates on internal model states, complementing existing surface-level defenses. The main limitation, unaddressed by the paper, is whether the detector generalizes to attack templates not seen during anomaly-detector training—a production deployment would require continuous retraining on newly observed attack patterns.

Post-hoc auditing of model behavior in high-stakes decisions. When a language model makes a consequential error—a medical advice system recommends a dangerous treatment, a legal AI misinterprets a clause—the tuned lens enables investigators to examine when the model converged on its wrong answer. By replaying the input through the model and inspecting the prediction trajectory, an auditor can determine whether the model "knew" the correct answer at early layers and overthought its way to an error (as in the phenomenon documented by Halawi et al., 2023 and extended in Figure 11), or whether it was confidently wrong from early layers onward. The first case suggests a problem with the model's late-layer reasoning or susceptibility to misleading context; the second suggests a fundamental lack of knowledge about the domain. This diagnostic information can guide remediation: overthinking errors might be mitigated by early-exiting or confidence-thresholding, while early-convergence errors indicate a need for better pretraining data or fine-tuning. The tuned lens makes this analysis possible on any deployed model (using the released translator checkpoints or training new ones in under an hour on modest hardware) without access to training infrastructure.

Cost-efficient evaluation of early-exiting strategies on pretrained models without retraining. Early exiting methods like CALM (Schuster et al., 2022) and DeeBERT (Xin et al., 2020) require modifying the training process to add auxiliary classifiers at intermediate layers, which is not an option for the thousands of pretrained models already deployed. The tuned lens provides a post-hoc alternative: by computing the prediction trajectory on a validation set, and measuring at which layer the prediction accuracy first exceeds a target threshold (e.g., 95% of final-layer accuracy), a practitioner can determine the optimal early-exit layer for that model-task combination without retraining the model. If the tuned lens shows that the model reaches 95% of its final accuracy by layer 18 of 32, deploying with an early-exit at layer 18 would reduce inference FLOPs by roughly 44% (assuming uniform layer cost) at a 5% accuracy cost. The paper's prediction depth analysis (Section 5.4, Table 2, Figure 12) already provides the infrastructure for this computation—the only missing piece is a cost-benefit analysis comparing the accuracy-compute tradeoff curves derived from tuned lens trajectories against the actual accuracy of the model when truncated at each layer (which could be done by manually running the model with only the first k layers). This is a direct path to inference cost savings on pretrained models without any model modification.

When to Prefer This Method

The paper positions the tuned lens as a drop-in replacement for the logit lens (Section 6, abstract) and its primary alternative is therefore the logit lens itself. The tradeoff is explicit:

  • Prefer the tuned lens when you need reliable, unbiased, causally faithful layer-wise predictions and are working with a model where the logit lens fails (BLOOM, GPT-Neo, OPT) or where systematic bias would undermine your analysis (any dynamical interpretation of the trajectory as belief updating). The tuned lens achieves 4–5× lower perplexity at early layers (Figure 5), reduces marginal bias from ~4–5 bits to near zero (Figure 3), and demonstrates causal fidelity (Spearman ρ = 0.89, Figure 8) that the logit lens has not been shown to possess. If you are doing scientific research on transformer internals—especially research that claims to interpret how beliefs evolve across layers—the tuned lens is methodologically required because the logit lens's bias makes such interpretations unreliable.

  • Prefer the logit lens when you need a quick, zero-cost exploratory look at a model's layer-wise predictions and are working with GPT-2 (where the logit lens works passably) or are in the early stages of a project where training translators would be premature. The logit lens requires no training, no data, and no GPU hours—it is a single matrix multiplication. If the goal is to glance at whether a model seems to "know" an answer early, and precision is not critical, the logit lens remains a useful reconnaissance tool.

The paper also implicitly positions the tuned lens against training new unembedding matrices per layer (the approach of Alain and Bengio, 2016, applied to language models). The argument is that training a d×d translator per layer is more parameter-efficient and converges faster than training a |𝒱|×d unembedding per layer. For a model with a 50K vocabulary and 5120-dimensional hidden states, a translator is 26M parameters while a new unembedding is 256M parameters—a 10× difference that the paper claims translates to faster training and better perplexity. If you are considering training per-layer probes for a language model, the tuned lens should be preferred over full-unembedding approaches on efficiency grounds alone, unless there is a specific reason to believe the model's unembedding is inadequate (in which case the problem is with the model, not the probing method).