Logit lens

A diagnostic trick: project the residual stream into vocabulary space at any intermediate layer to read "what does the model predict right now?"

The logit lens applies the same linear projection that normally runs only at the very end of the network — the unembedding matrix WU — to the residual stream mid-network. Formally: at layer l, compute logitsl = WU · RMSNorm(xl), run softmax, and read the top token. This is not a layer the model was trained with; it is a post-hoc inspection tool. Because the residual stream is the only state the model maintains between layers, this projection is a complete readout of everything the model "knows" at that stage.

The result is a layer-by-layer sequence of guesses. In early layers the top token is often wrong or only loosely related to the correct answer — the stream mostly holds positional and shallow lexical information at that point. As the stream accumulates attention and feed-forward edits, the top token shifts toward the right answer, and its probability rises. By the last layer the guess is the token that will actually be generated. The sharpening of the logit-lens readout, plotted across layers, is a concrete trace of how the model is solving the prediction problem — syntactic features arriving in shallow layers, semantic specificity arriving in deeper ones.

The playground shows this readout as the large token label displayed above the feature-rail strip. The confidence bar beneath it is the softmax probability of that top token. Watch the label and bar change as you step through the layers: vague or incorrect early on, converging to the real next token by the final block.

The large label is the top-1 token from the logit-lens projection at this layer; the bar below is its softmax probability. The guess sharpens layer by layer as the residual stream accumulates information — early layers are uncertain, the final layer commits.

Related: The residual vector · Residual add