Logit Lens › How the Logit Lens Works

The core idea is elegantly simple: transformers end with a specific structure that converts hidden states to vocabulary probabilities, and we can apply that same structure to any intermediate hidden state.

Before diving into the logit lens, it helps to understand exactly how the final layer of a transformer produces a probability distribution over the vocabulary.

A standard transformer language model maintains a hidden state $\mathbf{h} \in \mathbb{R}^d$ for each token position, where $d$ is the hidden dimension, also called the model dimension or $d_{\text{model}}$. After all transformer layers have processed the sequence, the final hidden state passes through two components to produce a prediction.

The first component is layer normalization. The final hidden state $\mathbf{h}^{(L)}$ is normalized to have zero mean and unit variance across its $d$ dimensions, then rescaled by learned parameters:

$$\hat{\mathbf{h}} = \text{LayerNorm}(\mathbf{h}^{(L)}) = \frac{\mathbf{h}^{(L)} - \mu}{\sigma} \cdot \boldsymbol{\gamma} + \boldsymbol{\beta}$$

where:

• $\mu = \frac{1}{d} \sum_{i=1}^{d} h_i^{(L)}$ is the mean of the hidden state components
• $\sigma = \sqrt{\frac{1}{d} \sum_{i=1}^{d} (h_i^{(L)} - \mu)^2 + \epsilon}$ is the standard deviation with a small stability constant $\epsilon$
• $\boldsymbol{\gamma} \in \mathbb{R}^d$ and $\boldsymbol{\beta} \in \mathbb{R}^d$ are learned scale and shift parameters

The second component is the unembedding matrix multiplication:

$$\mathbf{logits} = W_U \hat{\mathbf{h}}$$

where $W_U \in \mathbb{R}^{V \times d}$ is the unembedding matrix and $V$ is the vocabulary size.

The resulting logits vector $\mathbf{logits} \in \mathbb{R}^V$ assigns a raw score to each vocabulary token. Applying softmax converts these to probabilities:

$$P(\text{token} = v \mid \mathbf{h}^{(L)}) = \frac{\exp(\mathbf{logits}_v)}{\sum_{v'=1}^{V} \exp(\mathbf{logits}_{v'})}$$

where:

• $\mathbf{logits}_v$ is the score for vocabulary token $v$
• $V$ is the total vocabulary size
• The denominator normalizes the scores to form a valid probability distribution over all tokens

This pipeline, LayerNorm followed by unembedding followed by softmax, is the entire machinery that converts a hidden state to a next-token prediction. The logit lens borrows this machinery and applies it to every layer, including intermediate ones.

The logit lens applies this same final-layer decoder, consisting of LayerNorm followed by the unembedding matrix, to the hidden state at each intermediate layer. For a model with $L$ layers, the forward pass produces a sequence of hidden states at each token position:

$$\mathbf{h}^{(0)}, \mathbf{h}^{(1)}, \mathbf{h}^{(2)}, \ldots, \mathbf{h}^{(L)}$$

where $\mathbf{h}^{(0)}$ is the embedding layer output (token embedding plus positional encoding, before any transformer layers) and $\mathbf{h}^{(L)}$ is the final hidden state that produces the model's output.

The logit lens at layer $\ell$ is defined as:

$$\text{LogitLens}^{(\ell)}(\mathbf{h}^{(\ell)}) = W_U \cdot \text{LayerNorm}(\mathbf{h}^{(\ell)})$$

where:

• $\mathbf{h}^{(\ell)} \in \mathbb{R}^d$ is the hidden state after the $\ell$-th transformer layer
• $\text{LayerNorm}$ uses the same learned parameters ($\boldsymbol{\gamma}$, $\boldsymbol{\beta}$) as the model's actual final layer normalization
• $W_U \in \mathbb{R}^{V \times d}$ is the fixed unembedding matrix, shared across all layers
• The result is a vector of logits $\in \mathbb{R}^V$ over the vocabulary, which can be converted to probabilities via softmax

This gives us a probability distribution over the vocabulary at every single layer, which we can inspect and visualize across layers. We are not re-running the model from each intermediate state; we are simply applying the output machinery at different points in the processing pipeline.

The logit lens connects naturally to the residual stream view of transformers. In a transformer with pre-norm architecture (like GPT-2 and most modern language models), each layer adds a delta to the residual stream:

$$\mathbf{h}^{(\ell)} = \mathbf{h}^{(\ell-1)} + \Delta_{\text{attn}}^{(\ell)} + \Delta_{\text{ffn}}^{(\ell)}$$

where $\Delta_{\text{attn}}^{(\ell)}$ is the contribution from the multi-head attention sublayer and $\Delta_{\text{ffn}}^{(\ell)}$ is the contribution from the feed-forward network sublayer at layer $\ell$. Because each layer adds to the same residual stream rather than replacing it, the final hidden state $\mathbf{h}^{(L)}$ is the sum of contributions from all layers plus the initial embedding.

The logit lens at any intermediate layer $\ell$ shows what prediction is implied by the partial sum of contributions accumulated up to that layer. In this sense, the logit lens acts as a running tally of how the cumulative residual stream converges toward the final prediction. This makes the technique especially natural for residual-stream architectures: intermediate states live in the same geometric space as the final state, and the unembedding matrix was trained to decode from this space.

In many transformer language models, including GPT-2, GPT-Neo, and others, the token embedding matrix $W_E \in \mathbb{R}^{V \times d}$ and the unembedding matrix $W_U \in \mathbb{R}^{V \times d}$ are the same matrix. This is called weight tying, and it was originally introduced to reduce parameter count and improve sample efficiency. A convenient side effect is that it gives the logit lens a clean geometric interpretation.

When $W_U = W_E$, the logit for token $v$ at layer $\ell$ becomes:

$$\text{logit}_{v}^{(\ell)} = \mathbf{e}_v \cdot \text{LayerNorm}(\mathbf{h}^{(\ell)})$$

where:

• $\mathbf{e}_v \in \mathbb{R}^d$ is the embedding vector for vocabulary token $v$
• The dot product measures geometric alignment between the normalized hidden state and each token's embedding vector

A high logit for token $v$ means the hidden state at that layer is geometrically close to $v$'s embedding in the $d$-dimensional representation space. This is an intuitive way to think about what "the model is thinking" at each layer: the model implicitly compares its current partial computation to every token it knows and reports which ones are most similar. Early layers produce hidden states that are close to many tokens' embeddings (high uncertainty). Later layers produce hidden states that concentrate geometric similarity onto the correct answer token (low uncertainty).

For models without weight tying, such as some larger-scale transformers where weight tying is removed for training stability, this geometric interpretation does not apply directly, but the operational definition of the logit lens (applying LayerNorm then $W_U$) remains valid.

The logit lens shows you an implicit prediction at each layer, but it does not show you the hidden state itself, the attention patterns at that layer, or what computation that layer performed. It is a read-out mechanism, not a causal analysis. The token that ranks highest at layer 7 was not necessarily computed at layer 7; it may have been present in the residual stream since layer 3, with layer 7 simply leaving it undisturbed.

This distinction matters for interpretation. When you see "Paris" emerge as a top prediction in a late layer, you should read that as "by this point, the accumulated computation is pointing toward Paris." You should not read it as "this layer is the one that retrieved the Paris fact." Establishing causal responsibility requires different tools, which we cover in the Activation Patching chapter.