The Logit Lens: How Direct Unembedding Decodes Intermediate Representations in Large Language Models

Autoregressive large language models process input tokens through deep stacks of transformer blocks, generating a final probability distribution over the vocabulary at the final layer. For years, the intermediate computations occurring within these hidden layers were treated as opaque black boxes. Mechanistic interpretability research has demonstrated that modern transformer architectures operate through a structured, iterative refinement process across their residual stream.

The logit lens and its refined counterpart, the tuned lens, provide a mathematical window into this internal process. By applying vocabulary projection matrices directly to intermediate hidden states, these techniques allow engineers and researchers to decode the model's evolving predictions layer by layer.

The Residual Stream as a Shared Linear Bus

In standard transformer decoder architectures, input tokens are mapped into an initial embedding space $h_0 \in \mathbb{R}^d$ . As vectors propagate through $L$ successive layers, each layer modifies the state through additive residual connections:

$h_l = h_{l-1} + a_l(h_{l-1}) + m_l(h_{l-1} + a_l(h_{l-1}))$

where $a_l$ denotes the multi-head self-attention output at layer $l$ , and $m_l$ denotes the multi-layer perceptron (MLP) output.

As formalized in Anthropic's transformer circuits research, the residual stream functions as a shared linear communication bus. Individual attention heads read from the residual stream via query-key projections and write back additive updates. MLP sub-layers similarly read intermediate states, perform non-linear feature transformations, and add their updates directly back into the stream.

Because each layer adds its output vector directly to the existing stream rather than transforming the entire state destructively, the overall representation preserves a shared semantic basis across substantial portions of the network depth. At the final layer $L$ , the hidden state $h_L$ is normalized and projected onto the vocabulary matrix $W_U \in \mathbb{R}^{d \times |V|}$ to produce unnormalized log-probabilities:

$\text{logits} = \text{RMSNorm}(h_L) W_U$

$\mathcal{P}(t_{i+1}) = \text{softmax}(\text{logits})$

Mechanics of the Standard Logit Lens

The logit lens method, originally introduced by independent researcher nostalgebraist in 2020, exploits this persistent linear structure. Instead of waiting for the vector to reach layer $L$ , the logit lens applies the final normalization layer and unembedding matrix $W_U$ directly to intermediate hidden states $h_l$ :

$\text{logits}_l = \text{RMSNorm}(h_l) W_U$

$\mathcal{P}_l(t_{i+1}) = \text{softmax}(\text{logits}_l)$

Evaluating this projection across layers reveals how the model's internal hypothesis space evolves during a single forward pass.

Layer 04: [the: 0.18, a: 0.12, in: 0.08, of: 0.05, to: 0.04]      (Grammar / syntax)
Layer 12: [city: 0.24, capital: 0.19, region: 0.11, place: 0.08]   (Semantic categorization)
Layer 20: [France: 0.41, Europe: 0.22, Paris: 0.15, French: 0.09]  (Entity retrieval)
Layer 28: [Paris: 0.82, Lyon: 0.06, Marseille: 0.03, London: 0.02] (Factual convergence)
Layer 32: [Paris: 0.96,  Lyon: 0.01, Marseille: 0.01, capital: 0.01] (Final calibration)

Empirical tracking across models such as LLaMA, Mistral, and GPT-style architectures reveals three distinct operational regimes:

Early Layers (0% to ~25% depth): The network processes local token syntax, token positioning, and shallow n-gram statistics. Vocabulary projections in these layers often yield high Shannon entropy and generic function words (e.g., articles and prepositions).
Middle Layers (25% to ~75% depth): Semantic abstraction and factual retrieval dominate. Attention heads route contextual information between distant tokens, while MLP layers act as key-value associative memories to extract factual knowledge. Candidate tokens matching the correct conceptual category surge in probability.
Late Layers (75% to 100% depth): The model performs final syntactic formatting, token capitalization, subword boundary alignment, and probability calibration.

Direct Logit Attribution and Circuit Decomposition

Because the residual stream is strictly additive, the final hidden state can be expanded as a linear sum of all component operations:

$h_L = h_0 + \sum_{l=1}^L a_l + \sum_{l=1}^L m_l$

Substituting this expansion into the unembedding projection yields Direct Logit Attribution (DLA):

$\text{logits}(y) = \text{RMSNorm}\left(h_0 + \sum_{l=1}^L a_l + \sum_{l=1}^L m_l\right) W_{U, y}$

Under linear approximation of the final normalization layer, the contribution of an individual attention head $k$ in layer $l$ to the logit of a target token $y$ is computed as:

$\text{DLA}(a_l^k, y) \approx a_l^k \cdot W_{U, y}$

Similarly, the logit difference between a correct token $y_{\text{target}}$ and an incorrect competitor $y_{\text{alt}}$ can be attributed to specific components:

$\Delta \text{logit} = \left(y_{\text{target}} - y_{\text{alt}}\right) \cdot W_U^T a_l^k$

This decomposition allows mechanistic interpretability researchers to isolate specific functional sub-networks, such as induction heads that duplicate repeated patterns or factual extraction circuits in deep MLP layers.

The Tuned Lens: Addressing Basis Drift

While the standard logit lens provides immediate visibility without requiring extra training, it possesses theoretical and practical limitations:

Basis Drift and Coordinate Rotation: Intermediate layers do not share an identical coordinate alignment with the final layer. Layers often store information in temporary subspaces that rotate relative to $W_U$ .
Representation Scale Mismatches: Feature magnitudes and activation norms fluctuate across depths, leading to elevated perplexity when using the raw unembedding matrix on early layers.
Premature Convergence Artifacts: The raw logit lens can fail to detect latent predictions that are actively represented in the stream simply because they have not yet been rotated into the final output basis.

To resolve these distortions, Belrose et al. (2023) developed the Tuned Lens.

Rather than using the static unembedding projection directly on $h_l$ , the tuned lens trains a lightweight affine transformation $L_l(h_l) = A_l h_l + b_l$ for each layer $l$ . The underlying language model weights remain entirely frozen during training:

$\text{logits}_l^{\text{tuned}} = \text{RMSNorm}(A_l h_l + b_l) W_U$

Each affine probe $L_l$ is trained via cross-entropy loss against the target token distribution. Because the probe is strictly affine and operates on individual layers in isolation, it cannot perform multi-step computational reasoning on its own; it merely translates the intermediate hidden state into the vocabulary coordinates of the final layer.

Layer Depth vs. Predictive Perplexity:
Depth (%) | Raw Logit Lens Perplexity | Tuned Lens Perplexity
------------------------------------------------------------
10%       | 420.5                     | 24.8
25%       | 85.2                      | 8.1
50%       | 18.4                      | 3.9
75%       | 4.2                       | 2.4
90%       | 2.1                       | 1.9
100%      | 1.8                       | 1.8

Empirical evaluations show that the tuned lens uncovers stable token predictions between 3 and 6 layers earlier than the uncalibrated logit lens, providing a more faithful representation of when information is actually computed.

Practical Engineering and Serving Applications

Beyond mechanistic interpretability research, intermediate probing techniques are being applied across several production domains:

1. Dynamic Early Exiting and Speculative Drafting

In autoregressive inference, computing every layer for every token incurs substantial latency. When intermediate probing indicates that the top candidate token reaches high confidence ( $p_l(t^*) > 0.98$ ) at layer $L/2$ with minimal entropy, inference runtimes can trigger an early exit or use the intermediate prediction as a speculative draft token for verification by later layers.

2. Hallucination and Confabulation Auditing

By tracking the logit trajectory of factual entities across intermediate layers, automated monitoring pipelines can pinpoint where factual retrieval breaks down. If a correct factual entity dominates intermediate layers (e.g., layer 18) but is supplanted by a common frequency prior in the final layers, the system can detect late-stage suppression and flag potential hallucinations.

3. Layer-Wise Safety and Refusal Monitoring

Safety alignment and refusal behaviors (such as refusing malicious prompts) are often driven by distinct directional vectors within the residual stream. Logit lenses and directional probes enable monitoring systems to detect the emergence of refusal states mid-network before generating the full completion.

4. Entropy Signatures and Phase Transitions

Recent work on information-theoretic lenses tracks the Shannon entropy of intermediate distributions:

$H(\mathcal{P}_l) = -\sum_{t \in V} \mathcal{P}_l(t) \log \mathcal{P}_l(t)$

Sudden drops in entropy across specific layer intervals indicate discrete phase transitions, such as the resolution of an ambiguous reference, the completion of an in-context reasoning step, or the activation of an induction circuit.

Summary

The logit lens and tuned lens transform transformer hidden states from opaque numerical representations into interpretable token probability trajectories. By leveraging the additive geometry of the residual stream, these probing methods expose how language models iteratively construct, refine, and calibrate their outputs across depth.