Understanding the Logit Lens in AI Models

This article explores the mechanics of the logit lens, an interpretability technique used to understand how large language models form predictions. By applying the model's final decoding matrix to intermediate hidden layer activations, researchers can project internal numerical states directly into readable vocabulary space. This overview covers the transformer architecture's residual stream, the mathematical projection process, and how this method reveals the step-by-step evolution of an artificial intelligence model's final output.

The Residual Stream Architecture

To understand the logit lens, you must first understand how transformer models process information. A transformer consists of multiple layers stacked on top of one another. As input text passes through the model, it is converted into numerical vectors called hidden layer activations.

These activations travel through a central pathway known as the residual stream. Rather than completely replacing the data at each step, every layer reads from the residual stream, performs computations using attention mechanisms and neural networks, and adds its updated results back into the stream. Because the data is modified incrementally rather than rewritten entirely, the vector space remains consistent from the first layer to the last.

Standard Output Projection

In a standard forward pass, the model only translates these numerical activations into human-readable words at the very end of the process. Once the data passes through the final hidden layer, the model applies a final normalization step and multiplies the resulting vector by the unembedding matrix.

This unembedding matrix projects the hidden state into the vocabulary space, generating a score—known as a logit—for every possible token in the model's vocabulary. A softmax function then converts these logits into probabilities, allowing the model to select the most likely next word.

How the Logit Lens Projects Intermediate Layers

The logit lens bypasses the upper layers of the model to see what the artificial intelligence "believes" the answer is at any given midpoint. It achieves this by taking the hidden activation vector from an intermediate layer and feeding it directly into the final layer's normalization and unembedding matrix.

Because the residual stream maintains a shared representation space across all layers, the final unembedding matrix can still decode earlier activations. The process follows three straightforward steps at any chosen layer:

  1. Extract the hidden state vector from the residual stream at the end of a specific intermediate layer.
  2. Apply the model's final layer normalization to this intermediate vector to ensure the numbers are scaled correctly.
  3. Multiply the normalized vector by the final unembedding matrix to generate intermediate logits, and convert them into token probabilities.

Observing Prediction Formation

By applying this projection across every layer sequentially, researchers can read the top vocabulary predictions at each stage of computation. This creates a transparent timeline of how the model forms its final output.

In the earliest layers, the logit lens usually reveals predictions that are nonsensical or closely tied to the exact input words. Through the middle layers, the model begins to surface broader concepts, related themes, and grammatically appropriate word types. In the final layers, the model refines these broad concepts, suppresses incorrect alternatives, and locks onto the precise token required to complete the prompt. Through this direct projection into vocabulary space, the logit lens turns an opaque neural network into an observable sequence of narrowing hypotheses.