The Logit Lens and Tuned Lens: Mathematical Foundations, Residual Stream Projections, and Representation Dynamics in Transformers

Large language models based on decoder-only transformer architectures process information through a sequence of stacked attention and multi-layer perceptron (MLP) blocks connected by a residual stream. While the final output distribution over the vocabulary is produced by projecting the final hidden state through an unembedding matrix, understanding how intermediate representations evolve across depth remains a central problem in mechanistic interpretability.

The logit lens and the tuned lens provide two complementary frameworks for decoding the internal representations of transformers at arbitrary layer depths. By mapping intermediate residual stream activations into vocabulary space, these techniques allow researchers to track where and how candidate predictions materialize during the forward pass.

The Residual Stream and the Linear Representation Hypothesis

A standard autoregressive transformer processes an input token sequence by mapping input tokens to initial embedding vectors $h^{(0)} \in \mathbb{R}^d$. At each layer $\ell \in \{1, \dots, L\}$, the hidden state is updated via residual connections:

$$h^{(\ell)} = h^{(\ell-1)} + a^{(\ell)} + m^{(\ell)}$$

where $a^{(\ell)}$ represents the output of multi-head self-attention and $m^{(\ell)}$ represents the output of the MLP block at layer $\ell$.

The residual stream acts as a central communication bus. Rather than transforming representations entirely from one discrete feature space to another at each step, attention heads and MLP layers write additive updates to the stream.

At the final layer $L$, the model applies a final normalization layer and projects the state onto the vocabulary of size $V$ using the unembedding matrix $W_U \in \mathbb{R}^{V \times d}$:

$$z^{(L)} = W_U \cdot \text{LayerNorm}(h^{(L)})$$ $$p^{(L)} = \text{softmax}(z^{(L)})$$

Because the residual stream accumulates updates additively, intermediate hidden states $h^{(\ell)}$ share the same dimensionality $d$ as the final hidden state $h^{(L)}$. This structural property enables direct linear probing and unembedding projections throughout the network.

The Logit Lens

Introduced by nostalgebraist in 2020, the logit lens treats intermediate hidden states as if they were final representations. Instead of waiting for layer $L$, the logit lens applies the model's existing final normalization and unembedding matrix directly to layer $\ell$:

$$z^{(\ell)}_{\text{logit-lens}} = W_U \cdot \text{LayerNorm}(h^{(\ell)})$$ $$p^{(\ell)}_{\text{logit-lens}} = \text{softmax}(z^{(\ell)}_{\text{logit-lens}})$$

The logit lens requires zero additional parameters and no training. It relies on the assumption that representations in the residual stream remain aligned with the vocabulary basis throughout the entire network depth.

Representational Drift and Limitations

While the logit lens reveals how predictions crystalize in the final third of a network, it frequently yields noisy or uninterpretable distributions in early and middle layers. This failure stems from three primary factors:

1. Subspace rotation: Intermediate representations often reside in distinct subspaces that are rotated or shifted relative to the final decision boundary.
2. LayerNorm parameter mismatches: The final LayerNorm parameters (gain and bias) are optimized exclusively for the statistics of $h^{(L)}$, which differ significantly from the activation statistics of $h^{(\ell)}$ at earlier layers.
3. Feature staging: Early layers frequently compute auxiliary features (such as positional encodings, syntax trees, or retrieval indices) that are not yet formatted as next-token predictions.

As a result, applying $W_U$ prematurely can severely underestimate the amount of predictive information already present in early residual states.

The Tuned Lens

To overcome representational drift without discarding the structural vocabulary mapping, Belrose et al. (2023) introduced the tuned lens. The tuned lens introduces a lightweight, layer-specific affine transformation $T^{(\ell)}$ that aligns intermediate hidden states with the final unembedding space prior to applying $W_U$:

$$T^{(\ell)}(h^{(\ell)}) = A^{(\ell)} h^{(\ell)} + b^{(\ell)}$$

where $A^{(\ell)} \in \mathbb{R}^{d \times d}$ is a learnable weight matrix and $b^{(\ell)} \in \mathbb{R}^d$ is a bias vector.

The predicted distribution at layer $\ell$ is computed as:

$$z^{(\ell)}_{\text{tuned-lens}} = W_U \cdot \text{LayerNorm}(T^{(\ell)}(h^{(\ell)}))$$ $$p^{(\ell)}_{\text{tuned-lens}} = \text{softmax}(z^{(\ell)}_{\text{tuned-lens}})$$

Importantly, the underlying transformer weights and the unembedding matrix $W_U$ remain completely frozen during training. Only the parameters $\{A^{(\ell)}, b^{(\ell)}\}$ for each layer are optimized.

Objective Function and Training Mechanics

Each layer's affine translator is trained independently to minimize the Kullback-Leibler (KL) divergence between its predicted distribution and the frozen model's true final output distribution $p^{(L)}$:

$$\mathcal{L}(\theta^{(\ell)}) = \mathbb{E}_{x \sim \mathcal{D}} \left[ D_{\text{KL}}\left( p^{(L)}(x) \,||\, p^{(\ell)}_{\text{tuned-lens}}(x; \theta^{(\ell)}) \right) \right]$$

Expanding the KL divergence reveals the minimization objective:

$$\mathcal{L}(\theta^{(\ell)}) = \mathbb{E}_{x \sim \mathcal{D}} \left[ -\sum_{v=1}^{V} p^{(L)}_v(x) \log p^{(\ell)}_{\text{tuned-lens}, v}(x) \right] + \text{const}$$

This is equivalent to minimizing the cross-entropy between the soft target distribution of the complete model and the affine prediction from layer $\ell$. Training against the soft distribution $p^{(L)}$ rather than the one-hot ground-truth target ensures that the tuned lens learns the full uncertainty profile, token rankings, and calibration of the parent model.

In practice, $A^{(\ell)}$ is often initialized to the identity matrix $I_d$ and $b^{(\ell)}$ to zero, ensuring that training begins exactly at the standard logit lens baseline and learns minimal necessary corrections.

Key Insights and Empirical Findings

Evaluating transformer internals through both lens techniques has established several empirical principles in LLM research:

• Monotonic convergence: In most factual retrieval tasks, the top-1 predicted token stabilizes well before the final layer, typically around 60% to 75% of total depth.
• Layer specialization: Early layers (0% to 30% depth) primarily focus on local syntactic parsing and context aggregation. Middle layers (30% to 70%) perform factual recall and relational reasoning. Late layers (70% to 100%) refine formatting, grammar, and low-level token selection.
• Anomaly and injection detection: Tuned lens trajectories exhibit distinct signature divergences during adversarial prompts and prompt injection attacks. Injected instructions often cause abrupt prediction flips in middle layers, which can be flagged before output generation completes.
• Information preservation: The tuned lens demonstrates that intermediate layers contain far more predictive certainty than the bare logit lens indicates, confirming that representations undergo smooth linear shifts rather than abrupt non-linear reorganizations.

Comparison: Logit Lens vs. Tuned Lens vs. Linear Probing

Different interpretability probes trade off parameter overhead, interpretability, and expressivity:

• Logit Lens: Zero learned parameters. Evaluates raw residual stream alignment with $W_U$. High fidelity in late layers; high distortion in early layers.
• Tuned Lens: $O(d^2)$ parameters per layer via affine alignment. Preserves vocabulary structure and token semantics. High fidelity across all depths.
• Full Linear Probes: $O(d \cdot V)$ parameters per layer mapping directly from $h^{(\ell)}$ to logits without utilizing $W_U$. High risk of overfitting or acting as an independent classifier rather than explaining the existing model.

By constraining the learned transformation to an affine mapping within the hidden dimension $d$, the tuned lens prevents the probe from solving the task independently while accurately decoding the model's internal trajectory.

Sources

• Belrose, N., Furman, Z., Smith, L., Halawi, D., McKinney, L., Ostrovsky, I., Biderman, S., Steinhardt, J. (2023). Eliciting Latent Predictions from Transformers with the Tuned Lens. arXiv:2303.08112. https://arxiv.org/abs/2303.08112
• nostalgebraist. (2020). Interpreting GPT: The Logit Lens. LessWrong. https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens
• Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., et al. (2021). A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread. https://transformer-circuits.pub/2021/framework/index.html
• Geva, M., Caciularu, A., Wang, K. R., & Goldberg, Y. (2022). Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space. EMNLP 2022. https://arxiv.org/abs/2203.14680
• Dar, G., Geva, M., Gupta, A., & Berant, J. (2023). Analyzing Transformers in Embedding Space. ACL 2023. https://arxiv.org/abs/2209.02535