Logit Lens

The problem, from first principles

A transformer's layers write into a shared residual stream, and only the final layer's stream is turned into next-token probabilities by the unembedding matrix. But nothing stops us from applying the same unembedding to the stream after every layer. The result, the logit lens, shows what the model would predict if it stopped early. Watching the prediction sharpen layer by layer reveals where information about the answer first appears. The final normalization must be applied before unembedding because the real model always applies it, so a lens that skips it reads garbage on streams whose scale differs from the last layer's.

From theory to code

Implement logit_lens and target_rank_by_layer.

Constraints

• resid has shape (L, d): the residual stream at one position after each of L layers. W_U has shape (d, V). norm_weight is a length-d RMSNorm gain.
• logit_lens(resid, W_U, norm_weight, eps=1e-6) applies RMSNorm to each row, x / sqrt(mean(x ** 2) + eps) * norm_weight, then multiplies by W_U. It returns logits of shape (L, V).
• target_rank_by_layer(logits, target) returns an integer array of length L giving the rank of token target in each row, where rank 0 means it is the top prediction (the number of tokens with a strictly larger logit).
• Do not modify the inputs.

Hints