9. The full forward pass
Assemble everything: text in at the top, a probability distribution out at the bottom, with every shape accounted for.
By the end of this lesson you can
- Trace a tensor's shape through an entire model
- Explain the unembedding step and logits
- Describe what changes as depth increases
- Explain why only the last position matters at generation time
End to end, with shapes
Take a 7-token prompt through a model with d_model = 4096, 32 layers, vocabulary 128,000.
| Step | Operation | Shape after |
|---|---|---|
| 1 | Tokenize the text | [7] integer ids |
| 2 | Embedding lookup | [7, 4096] |
| 3 | Block 1 … block 32 (RoPE applied inside each attention) | [7, 4096] throughout |
| 4 | Final normalisation | [7, 4096] |
| 5 | Unembedding projection [4096, 128000] | [7, 128000] logits |
| 6 | Take the last row, softmax it | [128000] probabilities |
The single most important thing on that table: the shape never changes through the entire stack. Every block maps [T, d] to [T, d]. The blocks are shape-compatible with each other, which is why you can stack 32 or 96 of them without altering anything else.
Watch a tensor flow
Step through the pass and see the shape and a sample of values at each stage.
Logits and the unembedding
The final projection produces one logit per vocabulary entry per position — an unnormalised score. Softmax converts a row of logits into probabilities:
P(token i) = exp(logit_i) / Σⱼ exp(logit_j)
If the output matrix is tied to the input embedding table (common), each logit is literally the dot product of the final residual vector with a token's embedding. The prediction becomes: which token's embedding has the largest dot product with my final vector? Lesson 3's geometry, closing the loop — and note it is the dot product, not the cosine, so a token whose embedding is long is easier to predict than a short one pointing the same way.
Note that logits themselves are what get manipulated at generation time — temperature, top-k, penalties all operate on logits before the softmax. That is Lesson 10.
What the layers do, roughly
Probing studies across many models find a consistent rough progression. Treat it as a tendency, not a law — the boundaries are soft and features are distributed.
- Early layers resolve surface form: detokenization (stitching word-pieces back into a concept), part of speech, immediate local context.
- Middle layers carry the heaviest lifting: syntax, entity tracking, factual recall, coreference. Interpretability work locates much factual knowledge in mid-layer feed-forward networks.
- Late layers move from "what is being said" to "what token comes next", sharpening toward the vocabulary — increasingly task- and output-specific.
This progression is why techniques like early-exit and layer-skipping work at all for easy tokens, and why probing a middle layer is often the best place to read off a model's belief about a fact.
Prediction sharpening with depth
A simulated view of how the top-token distribution concentrates as you read out from deeper layers.
Only the last row matters (when generating)
The model computes logits at every position, but when generating the next token you only need the last row. Why compute the rest?
During training, you need all of them. One pass over a 4096-token document yields 4096 separate predictions and 4096 loss terms — every position predicts its own successor. That parallelism over positions is precisely why transformers train efficiently and why the causal mask is essential.
During generation, all earlier positions produce logits you discard. Worse, without caching you would recompute the entire prefix for every single new token. The fix is the KV cache, which is Lesson 11 — and it splits inference into two phases with completely different performance characteristics.
Lesson in one breath
Tokenize, embed to [T, d], pass through N blocks keeping [T, d] throughout, apply a final norm, project to [T, vocab] logits, and softmax the last row into next-token probabilities.
Practice
Answers are checked in your browser and saved to this device. Get one wrong and you can retry as many times as you like.
A 12-token prompt enters a model with d_model = 768. What is the shape of the activations between block 5 and block 6?
What is the name for the raw, unnormalised scores the model produces for each vocabulary entry before softmax is applied?
With d_model = 4096 and a vocabulary of 128,000, how many parameters are in the unembedding projection, in millions? (Ignore bias.)
During training, why does the model compute logits at every position rather than only the last one?
Order the stages of a full forward pass.
Done with this lesson?
A lesson counts as complete once it is marked read and every exercise is solved.
Tip: press ← and → to move between lessons.