Logit Lens › The Tuned Lens

The standard logit lens has a known limitation: it uses the final-layer unembedding matrix to decode intermediate representations that were never trained to be decoded that way. Intermediate hidden states are in the "processing domain" rather than the "output domain". They encode information in geometric arrangements optimized for subsequent transformer layers to process, not for direct vocabulary projection. Applying the output decoder to these states may therefore produce noisy or miscalibrated probability distributions, especially in early layers.

The tuned lens, introduced by Belrose et al. (2023), addresses this by training a small affine transformation per layer that bridges the gap between intermediate representations and the output domain. Conceptually, the tuned lens learns a "translation function" for each layer that converts the intermediate representation into a form the unembedding matrix was designed to decode.

For each layer $\ell$, the tuned lens learns a linear transformation parameterized by a matrix $A^{(\ell)}$ and a bias $\mathbf{b}^{(\ell)}$:

$$\text{TunedLens}^{(\ell)}(\mathbf{h}^{(\ell)}) = W_U \cdot \text{LayerNorm}\left(A^{(\ell)} \mathbf{h}^{(\ell)} + \mathbf{b}^{(\ell)}\right)$$

where:

• $A^{(\ell)} \in \mathbb{R}^{d \times d}$ is a learned linear transformation (the "translator") for layer $\ell$
• $\mathbf{b}^{(\ell)} \in \mathbb{R}^d$ is a learned bias for layer $\ell$
• $W_U$ is the fixed unembedding matrix, shared across all layers and not updated during tuned lens training

These per-layer affine transformations are trained to minimize the Kullback-Leibler divergence between the tuned lens prediction at layer $\ell$ and the final model prediction:

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

where:

• $p^{(L)}(x)$ is the final layer's probability distribution over the vocabulary given input $x$ (the model's actual output)
• $\hat{p}^{(\ell)}(x)$ is the probability distribution produced by the tuned lens at layer $\ell$
• $D_{\text{KL}}(p \| q) = \sum_{v=1}^{V} p_v \log(p_v / q_v)$ is the KL divergence from $q$ to $p$
• The expectation is taken over a corpus of training text $\mathcal{D}$

The training objective asks: given the intermediate hidden state at layer $\ell$, can we linearly transform it into a form from which the final model output is recoverable? If the tuned lens achieves low KL divergence at layer $\ell$, it means the information needed to make the final prediction was already present at layer $\ell$ in a linearly accessible form.

In practice, the $A^{(\ell)}$ matrices are initialized to the identity and the $\mathbf{b}^{(\ell)}$ biases are initialized to zero. This means the tuned lens starts from the logit lens (identity transformation) and learns corrections. The training uses the same text data used to pretrain the language model, running forward passes with frozen model weights and updating only the translator parameters.

An important constraint is that $A^{(\ell)}$ is often regularized to stay close to the identity matrix, preventing it from learning arbitrary transformations that would overfit to the training distribution. The tuned lens is intended to be a minimal correction that brings intermediate representations into the output domain, not a full re-learning of the unembedding.

The parameter overhead is modest. For a model with $L$ layers and hidden dimension $d$, the tuned lens adds $L \cdot (d^2 + d)$ parameters. For GPT-2 Small with $L = 12$ and $d = 768$, this is approximately $12 \times (768^2 + 768) \approx 7.1$ million parameters, a small fraction of the model's 117 million parameters.

The tuned lens tends to show smoother probability trajectories than the raw logit lens, with less noise in early layers. Critically, it separates two questions that the logit lens conflates:

1. When is information first present? At which layer does the model first have enough information to make the correct prediction?
2. When is that information used? At which layer does the model commit to using that information for the final output?

These can differ significantly. A fact might be recoverable from middle-layer representations (the information is "present") but only become the dominant prediction in later layers (the information is "used"). The tuned lens helps distinguish these cases by providing calibrated probabilities throughout the network, independent of the specific geometric orientation of intermediate representations.

The tuned lens has been applied to study factual recall, arithmetic reasoning, and code generation in large models. In each domain, it has revealed that the information required for correct outputs is often established in middle layers, with final layers primarily performing refinement and confidence calibration rather than introducing new information.