Logit Lens

The Logit Lens is a technique from mechanistic interpretability — the field that tries to reverse-engineer the internal computations of neural networks. It answers a deceptively simple question: if we stopped a language model halfway through its layers and forced it to guess the next token right now, what would it say?

Introduced by the researcher nostalgebraist in a 2020 LessWrong post titled “Interpreting GPT: the Logit Lens,” the method has become one of the most widely used and intuitive tools for peering inside transformer-based language models. Its enduring value is that it requires no training, no extra parameters, and only a few lines of code — yet it exposes the layer-by-layer “thought process” of a model.

The Core Idea: Reading the Residual Stream

To understand the logit lens, you need one concept from the Transformer architecture: the residual stream. As a token passes through a transformer, each layer doesn’t overwrite the representation — it adds its contribution to a running sum. This accumulating vector is the residual stream, and by the final layer it holds everything the model needs to predict the next token.

Normally, only the final residual state is decoded into a prediction. The model applies a final layer normalization and then multiplies by the unembedding matrix (W_U) to turn the hidden vector into logits — one score per vocabulary token. The logit lens makes a bold, elegant assumption: apply that exact same decoding step to the intermediate layers too. The result is a “pseudo-logit” distribution at every depth, letting us watch the prediction form.

The Formula

For a hidden state h_ℓ at layer ℓ, the logit lens computes:

logits_ℓ = LayerNorm_final(h_ℓ) · W_U

Applying a softmax turns these into a probability distribution over the vocabulary. Plotting the top token across all layers reveals a characteristic trajectory: early layers are noisy and generic, the answer emerges in the middle layers, and the final layers sharpen it into a confident prediction. This directly supports the “iterative refinement” view of how transformers compute.

Limitations — and the Tuned Lens

The original logit lens is powerful but imperfect. Because it forcibly reuses the final layer’s decoder on earlier representations, it can produce biased or brittle results — a problem known as basis drift, where intermediate layers operate in a different representational subspace than the output layer. On some model families the logit lens produces misleading or nonsensical intermediate predictions.

To fix this, Belrose et al. (2023) introduced the Tuned Lens. Instead of decoding intermediate states directly, it learns a lightweight, per-layer affine transformation that first “translates” each layer’s hidden state into the output basis:

logits_ℓ = LogitLens(A_ℓ · h_ℓ + b_ℓ)

Here A_ℓ and b_ℓ are small learned parameters (a “translator”) trained per layer. The tuned lens substantially improves the alignment between intermediate predictions and the model’s true final output, is more robust across architectures, and produces cleaner, less biased trajectories — at the cost of a short training step.

2025–2026: The Lens Family Expands

Logit-lens thinking has become a foundation for a growing family of interpretability tools, and the technique remains highly active in current research:

• Logit Prisms (2024–2025) decompose the final logit into contributions from individual layers, attention heads, and MLP sublayers — extending the lens from “what does each layer predict” to “which component is responsible.”
• Diffusion Steering Lens (arXiv:2504.13763, 2025) adapts lens-style decoding to vision transformers, using diffusion models to visualize what intermediate image representations encode.
• Multimodal interpretability — a 2025 survey (arXiv:2502.17516) catalogs how the logit lens and its variants are now applied across text, vision, and audio foundation models.
• Standardized tooling — libraries like TransformerLens and nnsight make the logit lens a one-line operation on modern open-weight models (LLaMA, Gemma, Qwen), and lens-style probing is now a routine step in debugging model behavior, tracing factual recall, and studying safety-relevant phenomena.

For teams building on large language models, the logit lens is one of the fastest ways to move from treating a model as a black box to actually seeing where and how a specific prediction — or a specific failure — is being formed inside the network.

import torch
from transformer_lens import HookedTransformer

# Load a model with cached activations exposed
model = HookedTransformer.from_pretrained("gpt2")

prompt = "The Eiffel Tower is located in the city of"
tokens = model.to_tokens(prompt)

# Run the model and cache every layer's residual-stream state
logits, cache = model.run_with_cache(tokens)

# Grab the residual stream after each layer (the "accumulated" stream)
# shape: [n_layers + 1, batch, seq, d_model]
resid = cache.accumulated_resid(layer=-1, incl_mid=False, return_labels=False)

# THE LOGIT LENS: apply the final layer-norm, then the unembedding (W_U)
# to every layer's hidden state to get pseudo-logits over the vocabulary
resid = model.ln_final(resid)               # final LayerNorm
layer_logits = model.unembed(resid)         # multiply by W_U

# For each layer, read off the top predicted token at the last position
for layer, lg in enumerate(layer_logits):
    top_id = lg[0, -1].argmax(dim=-1)
    token = model.to_string(top_id)
    prob = torch.softmax(lg[0, -1], dim=-1)[top_id].item()
    print(f"Layer {layer:2d}: {token!r:12s} (p={prob:.2f})")

# Typical output: early layers are diffuse/nonsensical, and the correct
# answer ("Paris") emerges and sharpens in the later layers.