From the final vector to a score for every word

The last operation in the stack

After the final block, the vector at the last position is one list of, say, 4,096 numbers. It has to become a score for each of 128,256 possible next tokens. That happens with a single matrix multiplication against the unembedding matrix (one row per vocabulary entry) producing one number per entry.

logits = final_vector @ W_unembed.T      # (4096,) @ (4096, 128256) -> (128256,)

Each logit is a dot product: how well does the final vector align with this token’s direction? Highest alignment, highest score.

Two consequences worth noticing. This one operation costs about 525 million multiply-adds in an 8B model, comparable to a whole layer, which is why large vocabularies are not free. And because it is a linear comparison against fixed directions, the geometry of the embedding space directly determines which tokens can be strongly preferred at all.

Tied embeddings

Many models use the same matrix for input embedding and output unembedding, transposed. The argument is that a token’s input representation and the direction that predicts it should be related, and the saving is substantial: 525 million parameters in the model we counted.

Tying helps most in small models, where the vocabulary is a large share of the total. Larger models often untie, on the finding that the two roles benefit from being learned separately. Either way, config.json tells you: look for tie_word_embeddings.

Logits are not probabilities

The output at this stage is unnormalised. Logits are real numbers, positive or negative, with no ceiling. Softmax converts them to probabilities, and the differences between logits are what matter: adding a constant to every logit changes nothing after softmax.

This is worth internalising because logits are what temperature acts on, as module three covers. Temperature divides the logits before the softmax, which is why it reshapes the distribution rather than picking differently from a fixed one.

The logit lens

Now the interesting part. Nothing stops you applying the final unembedding to an intermediate layer’s vector, since the residual stream keeps the same width all the way through. Do it at every layer and you get a running readout of what the model would say if it stopped there.

for layer, h in enumerate(hidden_states):     # h: the residual stream after layer `layer`
    logits = norm_final(h[0, -1]) @ W_unembed.T
    top = logits.argmax().item()
    print(layer, tokenizer.decode([top]))

On a factual prompt you typically see noise in the lower layers, a plausible category emerging in the middle, and the specific answer sharpening in the last quarter. It is a striking thing to watch, and it runs on a small model on a free Colab or Kaggle notebook in ten minutes with Hugging Face transformers and output_hidden_states=True.

What the logit lens does not prove

Read it carefully. The technique assumes intermediate vectors are meaningfully comparable to the final unembedding directions, an assumption that holds better in some models than others, and fails outright in a few. Refinements such as the “tuned lens” fit a small transformation per layer to correct for this, and they change the picture at some layers, which tells you the raw version was partly artefact.

More generally: seeing an answer appear at layer 22 does not mean layer 22 “decided” it. The residual stream is a sum, and later layers routinely suppress or redirect what earlier ones wrote. The lens shows a projection of a running total, not a sequence of decisions.

Why this is practically useful

Three uses that survive the caveats.

Debugging your own fine-tune. If a fine-tuned model gives the right answer at layer 20 and the wrong one at layer 32, the damage is in the top layers, and that is a different problem from the model never knowing the answer.

Understanding confidence. The gap between the top two logits at the final layer is the raw material of every confidence measure you will use in module six.

Deflating mystique. A model does not retrieve an answer and then phrase it. It accumulates evidence across a stack of identical blocks until one direction wins, and you can watch that happen.