9. Stacking it up and producing output

Full forward pass for a sequence of tokens:

flowchart TB
    A["Step 1: Tokenize text<br/>→ t₁, t₂, …, tₙ"] --> B["Step 2: Embedding lookup (+ positions)<br/>→ n vectors of size d_model"]
    B --> C["Step 3: Repeat L times:<br/>attention block (+ residual)<br/>MLP block (+ residual)"]
    C --> D["Step 4: Final norm"]
    D --> E["Step 5: Unembedding<br/>multiply by a d_model × vocab matrix<br/>→ a logit for every token, per position"]
    E --> F["Step 6: Softmax over logits<br/>→ probability distribution over next token"]

The unembedding and softmax steps in math:

$\text{logits}=\mathbf{h}{W}_U{p}_i=\frac{e^{z_i}}{\sum _{j=1}^V{e}^{z_j}}$

Here $\mathbf{h}$ is the final vector for a position, $W_U$ is the $d_{\text{model}}\times V$ unembedding matrix, $z_i$ is the logit (raw score) for token $i$, and $p_i$ is the resulting probability. See Sampling for what that softmax is doing.

Key detail

During training, the model outputs a prediction at every position simultaneously. Position 1 predicts token 2, position 2 predicts token 3, and so on. A single sequence of 4,000 tokens gives 4,000 training signals in one forward pass. That’s a huge part of why the training is so efficient. The causal mask guarantees no position can cheat by peeking ahead.

Many models tie the embedding and unembedding matrices (same weights used for both), which saves parameters.

Counting parameters

Rough rule: each block has about $12d_{\text{model}}^2$ parameters (4 $d^2$ in attention, 8 $d^2$ in the MLP). With $L$ layers:

$\text{params}\approx 12\cdot L\cdot d_{\text{model}}^2$

For GPT-3: $d_{\text{model}}=12,288$ and $L=96$ layers.

$12\times 96\times (12,288{)}^2\approx 174\text{ billion}(\text{reported: }175\text{B})$

The math checks out, which is a nice sanity test that you understand the architecture.

---

← 8. The transformer block · Contents · 10. Pretraining at scale →