From vector to next token

Key idea

logits = LN(x_final) · W_U (d × vocab); probabilities = softmax(logits); W_U is often the embedding table transposed

Real world

A compass needle against a dial. The needle (the final vector) points somewhere; you read how well it lines up with each of the marked directions (each word), and the best match is your answer.

In an LLM

GPT-2 multiplies a 768-number vector by a 768 × 50,257 matrix: 50,257 dot products, one score per token. For “The capital of France is”, the top scores go to “ the”, “ now”, “ a”, “ France” and “ Paris”.

The last three steps

After the final block, every position holds a vector of $d$ numbers. To predict the next token from the last position, the model:

1. Normalises it with one final LayerNorm (or RMSNorm). In a pre-norm model the stream has grown a lot by now (see the residual stream), and this brings it back to a standard size.
2. Unembeds it: multiplies by the Unembedding matrix $W_U$, of shape $d \times \text{vocab}$. Each column of $W_U$ belongs to one token, so the result is one score per token: a logit.
3. Softmax turns the logits into probabilities that add up to 1.

Each logit is a dot product: the final vector against one token’s column. A token scores high when the final vector points the same way as its column. Try it with a toy vocabulary of eight words:

Tied embeddings

At the bottom of the model, the embedding table $E$ (vocab × $d$) turns token IDs into vectors. At the top, $W_U$ ($d$ × vocab) turns vectors back into token scores. These are opposite jobs, and many models use one matrix for both: $W_U = E^\top$. This is Tied embeddings.

It saves a lot of parameters (38.6 million in GPT-2 small, almost a third of the model), and it makes sense: a final vector that points like the embedding of “ Paris” should predict “ Paris”. GPT-2 and Gemma tie their embeddings. Llama 3 8B does not: with 8 billion parameters the extra 525 million for a separate output matrix is cheap, and an untied output layer gives the model more freedom. Untick “tied embeddings” in the widget to see a separately learned $W_U$.

The logit lens

The read-out is just “final norm, then $W_U$”. Nothing stops you applying it to the residual stream in the middle of the model, as if the model stopped after layer 5. That is the Logit lens (nostalgebraist, 2020). It works because of the residual stream: every layer adds to the same vector, and the output reads that vector, so the intermediate vectors live in roughly the same “language” as the final one.

Here it is on GPT-2 small, with real numbers. Rows are layers, columns are positions; each cell is the top guess for the next token.

What the lens shows, and what it doesn’t:

• Layer 0 is junk. It is just the token embedding (plus position) put through the final norm and $W_U$. GPT-2 decodes it to odd tokens like “ destro” and “ mathemat”. Nothing has been computed yet, and the read-out was trained for the top of the stack.
• Early layers are rough. Layers 1 to 4 give generic, high-frequency guesses (“ now”, “ not”, “ also” after “ is”) that would fit almost any sentence.
• The answer forms in the late middle. At the last position of “The capital of France is”, “ France” becomes the top guess at layer 10 (62%), with “ Paris” second (18%). In the Mary and John sentence, layers 9 and 10 guess “ them”, and layer 11 switches to “ Mary” (61%), the correct name.
• The last layer often hedges. GPT-2’s final layer spreads the probability out again: after “France is” the output puts “ the” first. The lens shows that the model had the facts in mind earlier than its final answer suggests.

The logit lens is a rough tool: it assumes every layer speaks the output’s language, which is only roughly true. The tuned lens (Belrose et al., 2023) learns a small correction for each layer and gives cleaner readings. More in interpretability.

logits = model.lm_head(model.transformer.ln_f(x))   # x: (T, 768) → (T, 50257)
probs = logits.softmax(-1)
 
# logit lens: the same read-out on every layer's hidden state
out = model(ids, output_hidden_states=True)
for layer, h in enumerate(out.hidden_states[:-1]):
    p = model.lm_head(model.transformer.ln_f(h)).softmax(-1)