Logit Lens 和中间层解码

A transformer language model processes text as a stack of layers. Each token position carries a vector of numbers, called the hidden state or residual stream, that every layer reads and then adds an update to. Only after the last layer does the model convert this vector into a prediction: it applies a final normalization and multiplies by the unembedding matrix, a learned table that turns the vector into one score, or logit, per vocabulary token. A softmax turns those logits into probabilities for the next token. Intermediate layer decoding asks what would happen if that final readout were applied early, at layer 5 or layer 20, instead of at the end.

The logit lens is the simplest answer. Take the hidden state after some layer, apply the model's own final layer normalization and unembedding, and read off the top tokens. No training is needed, because it reuses the model's own weights. Doing this at every layer and every position produces a grid showing the most likely token at each depth, which researchers call a prediction trajectory. If the trajectory settles on the final answer early, the later layers are mostly refining or sharpening that answer; if it changes late, the decision is made near the end.

The technique was introduced in a LessWrong post by the pseudonymous author nostalgebraist on August 31, 2020. Running GPT-2's 1.5 billion parameter version on the abstract of the GPT-3 paper, the post reported that intermediate layers already produce plausible guesses, moving from nonsense or shallow guesses with the right part of speech in early layers to refined predictions later. It also reported a surprise: the input token itself was not preserved in early layers, as the hidden state moved almost at once toward something closer to the output. The author added caveats in the post, including that top tokens show only part of what the activations contain.

The logit lens did not generalize well. In 2023, Nora Belrose and colleagues at EleutherAI, FAR AI, the University of Toronto, Boston University and UC Berkeley reported that it fails to produce plausible predictions for model families such as GPT-Neo, BLOOM and OPT 125M. For BLOOM and OPT 125M, its top prediction was often just the input token in more than half the layers. Even where it worked, it was biased, systematically giving more probability to some vocabulary items than the final layer does. One cause is representational drift: a feature can be encoded differently at different depths, so the final layer's readout is the wrong decoder for earlier layers.

Their fix, the tuned lens, trains one small affine transformation, a matrix plus a bias called a translator, for each layer of a frozen model. Each translator maps that layer's hidden state into the coordinate system the final readout expects, and is trained so that its decoded distribution matches the model's real final output as closely as possible, measured by Kullback-Leibler (KL) divergence, a standard measure of how different two probability distributions are. Tested on models with up to 20 billion parameters, the paper reports that tuned lens predictions have substantially lower perplexity, meaning they predict the next token better, and are less biased than the logit lens. On GPT-Neo-2.7B, the logit lens produced uninterpretable output before layer 21, while the tuned lens gave sensible predictions throughout.

These lenses have been used for several research questions. Halawi, Denain and Steinhardt (2023) found that when a few-shot prompt contains wrong labels, middle-layer predictions are often more accurate than final outputs, and traced the late switch to attention heads that copy the false labels. The tuned lens paper reproduced this on models where the logit lens fails. Wendler and colleagues (2024) used the logit lens to argue that Llama-2 passes through a concept space closer to English when processing other languages. Patchscopes, from Google researchers in 2024, generalizes the idea by patching a hidden state into a separate prompt and letting the model describe it in words, and its authors argue that vocabulary-projection methods are special cases of their framework.

The main limitation is interpretive. A lens is a probe, a separate readout attached to the model, and what it shows depends on how it is built. A token that a lens ranks highly at layer 12 does not prove the model represents that fact at layer 12, and a token that does not appear may still be encoded in a form the readout cannot decode. The tuned lens is trained to imitate the final output, so it may reveal what later layers will do with the state rather than what the current layer computes. Belrose and colleagues partly addressed this with causal experiments, removing the directions the lens relies on and showing that this also changes the model's own output, but such checks are evidence, not a complete causal account. Both lenses were designed for pre-normalization transformers, the common modern layout, and work less cleanly on other designs.

In practice the logit lens costs almost nothing: a few lines of code that reuse the model's own weights, run on any open-weights model whose hidden states can be accessed. It cannot be applied to models available only through an API that hides internal activations. The tuned lens requires training one translator per layer on text similar to the model's pretraining data; the authors trained theirs on a slice of the Pile validation set, initialized each translator to the identity and used stochastic gradient descent with Nesterov momentum for 250 steps, a small cost compared with training the model itself.