Lesson 21 · Slide 1
Transformer Architecture
A deep dive into the architecture that powers every modern large language model — from self-attention to production optimisations like Flash Attention and KV caching.
Learning Objectives
- Understand the self-attention mechanism and the role of Q, K, V matrices
- Explain multi-head attention and why multiple heads help
- Compare positional encoding strategies: sinusoidal, RoPE, ALiBi
- Trace data flow through feed-forward layers and layer normalisation
- Distinguish encoder-only, decoder-only, and encoder-decoder architectures
- Reason about O(n²) attention complexity and how Flash Attention addresses it
- Explain how KV cache accelerates autoregressive inference
Lesson 21 · Slide 2
Self-Attention: Q, K, V
Every token in the sequence is projected into three vectors — Query, Key, and Value — via learned weight matrices W_Q, W_K, W_V.
Data Flow Diagram (text)
Input X [n x d_model] → multiply by W_Q, W_K, W_V → Q [n x d_k], K [n x d_k], V [n x d_v]
Attention(Q,K,V) = softmax( Q · K^T / sqrt(d_k) ) · V
Intuition
- Query: "What am I looking for?" — each token asks a question
- Key: "What do I contain?" — each token advertises its content
- Value: "What information do I pass on?" — the actual payload
- The dot product Q · K^T measures relevance; softmax normalises into attention weights
- Division by sqrt(d_k) prevents dot products from growing too large, which would push softmax into saturated regions with vanishing gradients
Notes: The scaling factor sqrt(d_k) is critical. Without it, for large d_k the dot products become very large, softmax concentrates on a single token, and gradients vanish.
Lesson 21 · Slide 3
Multi-Head Attention
Instead of running one attention function with d_model-dimensional keys, queries, and values, we run h parallel heads, each with dimension d_k = d_model / h.
Architecture Diagram (text)
Input X → split into h heads → each head computes Attention(Q_i, K_i, V_i) independently → Concat(head_1, ..., head_h) → multiply by W_O → Output [n x d_model]
Why Multiple Heads?
- Each head can learn different relationship types: syntactic dependencies, coreference, positional patterns
- Empirically, models with multi-head attention outperform single-head attention even at the same total parameter count
- Typical configurations: GPT-2 uses 12 heads, GPT-3 uses 96 heads, Llama 2 70B uses 64 heads with 8 KV heads (GQA)
Grouped Query Attention (GQA)
Modern models like Llama 2/3 and Mistral use GQA: multiple query heads share a single key-value head. This reduces KV cache memory by up to 8x with minimal quality loss. It sits between multi-head attention (MHA) and multi-query attention (MQA).
Lesson 21 · Slide 4
Positional Encoding
Self-attention is permutation-invariant — it has no inherent notion of token order. Positional encodings inject sequence position information.
Sinusoidal (Original)
Fixed function: PE(pos, 2i) = sin(pos / 10000^(2i/d)) and PE(pos, 2i+1) = cos(pos / 10000^(2i/d)). Added to token embeddings. No learned parameters. Can theoretically extrapolate to unseen lengths.
RoPE (Rotary)
Used by Llama, Mistral, Qwen. Encodes position by rotating Q and K vectors in 2D subspaces. Relative position information is naturally captured in the dot product. Supports length extrapolation with NTK-aware scaling.
ALiBi (Attention with Linear Biases)
Used by BLOOM, MPT. Adds a linear penalty proportional to distance directly to attention scores: score_ij = q_i · k_j - m · |i - j|. No modification to embeddings. Zero extra parameters. Excellent length extrapolation — trains on 1K tokens, infers on 8K+.
Notes: RoPE has become the dominant approach. Most production models (Llama 3, Mistral, Claude) use RoPE or variants of it. ALiBi is simpler but less common in the latest generation.
Lesson 21 · Slide 5
Feed-Forward & Layer Norm
Each transformer block contains two sub-layers: multi-head attention followed by a position-wise feed-forward network (FFN). Both are wrapped with residual connections and layer normalisation.
Transformer Block Diagram (text)
x → LayerNorm(x) → MultiHeadAttention → + x (residual) → LayerNorm → FFN → + x (residual) → output
Feed-Forward Network
- Two linear transformations with a non-linearity:
FFN(x) = W_2 · GELU(W_1 · x + b_1) + b_2
- Inner dimension is typically 4x the model dimension (e.g., d_model=4096 → d_ff=16384)
- Modern variants use SwiGLU (Llama, PaLM):
FFN(x) = W_2 · (SiLU(W_1 · x) * W_3 · x) — a gated linear unit that improves performance
Layer Normalisation
- Pre-norm (modern default): normalise before each sub-layer. More stable training, used by GPT-2+, Llama, Mistral
- Post-norm (original paper): normalise after each sub-layer. Requires careful learning rate warmup
- RMSNorm (Llama, Gemma): simplified variant that only normalises by root mean square, omitting the mean-centering step. Faster and equally effective
Lesson 21 · Slide 6
Encoder vs Decoder Architectures
The original transformer has both an encoder and a decoder. Modern models typically use only one half.
| Architecture | Attention | Examples | Best For |
|---|
| Encoder-only | Bidirectional (full) | BERT, RoBERTa, DeBERTa | Classification, NER, embeddings, retrieval |
| Decoder-only | Causal (masked future) | GPT, Claude, Llama, Mistral | Text generation, instruction following, code |
| Encoder-decoder | Bidirectional encoder + causal decoder with cross-attention | T5, BART, Flan-T5 | Translation, summarisation, seq2seq tasks |
Why Decoder-Only Won
- Simpler architecture — one unified model for all tasks
- Scales more predictably with compute
- In-context learning and instruction following emerge naturally
- No need for task-specific heads or fine-tuning formats
When Encoders Still Shine
- Embedding/retrieval: BERT-family models produce superior embeddings for search
- Classification: bidirectional context is better for understanding
- Efficiency: encoder models are much smaller and faster for specific tasks
Lesson 21 · Slide 7
Attention Complexity & Flash Attention
Standard self-attention has O(n²) time and memory complexity, where n is the sequence length. This is the primary bottleneck for long sequences.
Why O(n²)?
Computing Q · K^T produces an [n x n] attention matrix. For a 128K-token context, this is a 128K x 128K matrix — over 16 billion entries. Storing this in float16 requires ~32 GB of memory just for one layer and one head.
Flash Attention (Dao et al., 2022)
- Key insight: the bottleneck is not compute but memory bandwidth — reading/writing the n x n attention matrix to GPU HBM
- Solution: compute attention in tiles that fit in SRAM (on-chip fast memory), never materialising the full attention matrix
- Result: same exact output, but 2-4x faster and uses O(n) memory instead of O(n²)
- Flash Attention 2: further optimisations achieve ~2x speedup over Flash Attention 1
- Now the default in PyTorch 2.0+ (
torch.nn.functional.scaled_dot_product_attention)
Other Approaches to Long Contexts
Sparse attention (BigBird, Longformer): attend to local windows + global tokens — O(n). Linear attention (Performers): approximate softmax with kernels — O(n). Ring Attention: distribute sequence across GPUs for near-infinite context. In practice, Flash Attention + KV cache has proven most effective.
Lesson 21 · Slide 8
KV Cache & Lab Preview
During autoregressive generation, the model produces one token at a time. Without caching, it would recompute attention over the entire sequence for each new token.
KV Cache Diagram (text)
Step 1: Process prompt tokens [t1, t2, ..., tn] → store K and V for all positions in cache
Step 2: Generate token t_{n+1} → compute Q for new token only, reuse cached K, V → append new K, V to cache
Step 3: Repeat — each step processes only 1 token instead of the full sequence
KV Cache Memory
- Memory per token:
2 x num_layers x num_kv_heads x d_head x bytes_per_param
- Llama 2 70B at 4K context: ~2.5 GB of KV cache per request
- This is why batch size in serving is often limited by KV cache memory, not compute
- GQA and MQA reduce KV cache size by sharing KV heads across query heads
Lab Preview
- Exercise 1: Implement scaled dot-product attention from scratch in NumPy
- Exercise 2: Visualise attention patterns with heatmaps
- Exercise 3: Generate and compare sinusoidal vs. RoPE positional encodings
Notes: The lab uses NumPy only (no PyTorch required) to keep the focus on understanding the math. Students will implement the exact equations from the slides.