Transformer anatomy

Why this matters

The 2017 Transformer paper is the most-cited deep learning paper of all time. Every LLM you use — GPT-4o, Claude 4, Gemini 2.5, Llama 4, Qwen 3 — is a direct descendant. If you can sketch the block diagram of a Transformer on a whiteboard, label every shape, and explain each component, you can pass 70% of senior ML interviews. Today's lesson gives you that fluency.

Learning objectives

1. Sketch a decoder-only Transformer block from memory.
2. Distinguish encoder-only, decoder-only, and encoder-decoder variants.
3. Read shapes through every layer.
4. Explain residuals, normalisation, and the FFN.
5. Identify which parts are the same across all modern LLMs and which differ.

1. The three families

Encoder-only (BERT, RoBERTa, DeBERTa, ModernBERT)

• Bidirectional attention (every token sees every other).
• Trained with masked-language-modelling (MLM): mask 15% of tokens, predict them.
• Used for understanding tasks: classification, retrieval embeddings, NER.
• Not used for generation.

Decoder-only (GPT-2/3/4, Llama, Mistral, Qwen, DeepSeek, Claude)

• Causal attention — each token only sees the past.
• Trained with next-token prediction (autoregressive LM).
• The default LLM architecture since 2020.
• Same model can do generation, classification (read final hidden state), embeddings (with pooling).

Encoder-decoder (T5, BART, mT5, FLAN-T5, Whisper)

• Encoder reads input bidirectionally; decoder generates output causally; decoder uses cross-attention to look at encoder outputs.
• Strong on translation, summarisation, ASR.
• More parameters and complexity than decoder-only — has fallen out of fashion for general chat models.

For the rest of this course we focus on decoder-only, since that's >95% of the field today. Lessons referencing encoder-decoder will say so explicitly.

2. The Transformer block (decoder-only, pre-norm)

                     ┌──────────────────────────┐
   x  ─────────────► │ RMSNorm                  │
                     │ Multi-head self-attention │
                     │ (causal mask, RoPE)       │
                     └──────────┬────────────────┘
                                │
                       residual + (add)
                                │
                     ┌──────────▼────────────────┐
                     │ RMSNorm                  │
                     │ FFN (SwiGLU)             │
                     └──────────┬────────────────┘
                                │
                       residual + (add)
                                ▼
                              y

In code (modern flavour):

class Block(nn.Module):
    def __init__(self, d, h, mlp_mult=4):
        super().__init__()
        self.norm1 = RMSNorm(d)
        self.attn  = MultiHeadAttention(d, h, causal=True)
        self.norm2 = RMSNorm(d)
        self.mlp   = SwiGLU(d, mlp_mult * d)
    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp (self.norm2(x))
        return x

That is the block. The whole Transformer is L of these stacked.

3. The full forward pass — step by step

Inputs: tokens of shape (B, T).

1.  emb = token_embedding(tokens)        # (B, T, d)
2.  emb = emb + positional_encoding      # if absolute pos enc; for RoPE this happens inside attention
3.  for block in blocks:
        x = block(x)                     # (B, T, d) -> (B, T, d)
4.  x = final_norm(x)                    # RMSNorm
5.  logits = x @ token_embedding.weight.T  # (B, T, V) — weight tying

Note weight tying: the output projection often shares weights with the input embedding. Saves parameters and (mildly) improves perplexity. Used by GPT-2, Llama-3, Mistral.

4. The FFN — bigger than you think

For each token independently:

ffn(x) = W2 (act(W1 x))      # vanilla
ffn(x) = W2 ((SiLU(W1 x)) * (W3 x))   # SwiGLU

The hidden dim is typically 4 × d_model (vanilla) or ~2.7 × d_model (SwiGLU, tuned to keep parameter count similar).

Surprising fact: the FFN holds ~2/3 of all transformer parameters. Attention is glamorous; the FFN is where the model stores its knowledge. Recent interpretability work (Phase 7) shows that factual associations live in the FFN's down-projection.

5. Sizes you should know by heart

You will memorise GPT-2 / Llama 8B as your default mental models.

6. What every modern model has changed (vs the 2017 paper)

You will study each of these in the next lessons.

7. Encoder-decoder details (so you can read T5/Whisper papers)

ENCODER block:
  x = x + SelfAttn(LN(x), bidirectional)
  x = x + FFN(LN(x))

DECODER block:
  y = y + SelfAttn(LN(y), causal)
  y = y + CrossAttn(LN(y), encoder_output)   # NEW
  y = y + FFN(LN(y))

In cross-attention, Q comes from y, K, V come from the encoder's final output. That is the "looking at the source" mechanic Bahdanau introduced.

T5 also uses relative position biases instead of absolute encodings — old but still relevant.

8. Putting numbers to it (parameter count)

For a decoder-only transformer with vocab V, L layers, hidden d, heads h, FFN multiplier m:

• Embedding (weight-tied): V * d
• Per block: attention 4 d^2 (Q, K, V, O projections, no bias) + FFN ~3 m d^2 (with SwiGLU)
• Total: V*d + L * (4 d^2 + 3 m d^2) ≈ L * (4 + 3m) d^2 + V*d

Llama-3 8B: L=32, d=4096, m=3.5 → 32 * (4 + 10.5) * 4096^2 ≈ 7.78B plus 0.5B for embeddings ≈ 8B. Sanity-checks the model card.

This is the formula that powers napkin-math during interviews.

Hands-on lab (3 hours)

anatomy_drills.ipynb:

1. Print the structure of gpt2 from HuggingFace: for n, m in model.named_modules(): print(n). Identify embedding, blocks, FFN, attention, final norm.
2. Compute parameter count by component. Verify you get 124M for gpt2-small.
3. Inspect a single block's c_attn (combined QKV) shape. Reshape it to (d, 3, d) and verify.
4. Forward-pass a batch of 4 sequences of length 16 through gpt2. Print output last_hidden_state shape and logits shape.
5. Replace torch.nn.LayerNorm with a custom RMSNorm. Verify forward outputs are similar (not identical).
6. Bonus: compute napkin-FLOPs per token of a 32-layer, 4096-d transformer. (Hint: ~6 × params per token forward+backward.)

Common pitfalls

1. Confusing encoder-only (BERT) with decoder-only (GPT) — they share the block but differ in attention pattern (and training objective).
2. Forgetting weight tying — your model has 2× the params it should.
3. Treating attention as the whole story — most params live in the FFN.
4. Mixing pre-norm and post-norm conventions — silently miss-trains.
5. Not zero-initializing the output projection of attention/FFN to stabilise early training (a popular trick from Megatron and GPT-J).

Self-check

1. Sketch a decoder-only transformer block and label every component.
2. Where does most of the parameter count live in a transformer?
3. Why use pre-norm instead of post-norm?
4. What is weight tying?
5. Difference between MHA, GQA, and MQA (preview — we will go deep next lesson).

References

• Vaswani et al. (2017), "Attention Is All You Need."
• Radford et al. (2019), "Language Models Are Unsupervised Multitask Learners" (GPT-2).
• Touvron et al. (2024), "The Llama 3 Herd of Models" — the architecture appendix is gold.
• Karpathy, nanoGPT — the canonical clean implementation.
• HuggingFace transformers source — read modeling_llama.py.