Transformers Foundations — Attention, Q/K/V, Multi-Head, Positional Encoding

Why this session matters

Part of a 24-topic learning series on engineering, ML, and LLM systems. Each session is a 90-minute deep-dive on one topic — designed so anyone can pick it up cold. Every two topics are followed by a revision session with recall prompts and hands-on drills.

Part 1: Transformers Part 1 — Attention, Q/K/V, Multi-Head

Why this session matters

It builds on the rhythm of one focused topic, paced so you have time to actually absorb it rather than rush.

Agenda

• Why transformers won — recurrence vs attention
• The residual stream — a fixed-width bus every block reads and writes
• Self-attention as soft routing — Q, K, V and the √d_k trick
• Multi-head attention — running h heads in parallel, head specialisation
• What we'll layer on next session (RoPE, MLP, LayerNorm, KV cache)

Pre-read (skim before the session)

• Attention Is All You Need (Vaswani et al., 2017)
• The Illustrated Transformer — Jay Alammar
• A Mathematical Framework for Transformer Circuits — Anthropic
• nanoGPT model.py (Karpathy)

Deep dive

1. Why transformers won

Before 2017, sequence models meant recurrence (LSTM, GRU): read tokens one at a time, carry a hidden state. Two killers:

1. Sequential dependency. You can't compute step t before step t-1. GPUs are massively parallel — RNNs leave most of the silicon idle.
2. Long-range decay. By token 500, the gradient signal from token 1 is vanishing (or exploding). Gating helped, didn't solve it.

The transformer replaced recurrence with attention: every position computes a weighted sum over every other position in one matrix multiply. The whole sequence is processed in parallel; any position can look at any other in O(1) hops. The price: O(T²) attention cost — most modern engineering is about paying less of that price (FlashAttention, GQA, sliding windows, linear attention).

2. A concrete example

"The cat sat on the mat because it was tired."

When the model processes "it", it has to decide: does it refer to cat or mat? Attention lets it look back at every previous word, score how relevant each one is, and pull information from the most relevant (here: cat) into its own representation.

That's the one trick. Q/K/V, multi-head, RoPE, MLPs — those are engineering layered around that single idea.

3. The residual stream as a bus

Think of a decoder-only transformer as a fixed-width residual stream of shape (B, T, d_model):

• B = batch size
• T = sequence length
• d_model = model width (768 for GPT-2 small, 4096 for Llama-3-8B, 8192 for Llama-3-70B, 12288 for GPT-3 175B)

Every block reads the stream, adds a small update, passes it on. Each block is a refinement, not a replacement:

   tokens  →  embed  →  block 1  →  block 2  →  …  →  block N  →  LN  →  unembed  →  logits
                          ▲                                          (stream shape: (B,T,d))
                          │ residual: x = x + sublayer(LN(x))

4. Self-attention — soft, learned routing

For each token we project the residual stream to three things:

• Query Q = x · W_Q — "what am I looking for?"
• Key K = x · W_K — "what do I represent?"
• Value V = x · W_V — "what would I contribute if attended to?"

All (B, T, d_head). We compute:

attention(Q, K, V) = softmax( Q · Kᵀ / √d_k ) · V

Step by step:

1. Q · Kᵀ → (B, T, T) raw scores. Cell i, j = how much position i should look at j.
2. Divide by √d_k — keeps logits stable so softmax doesn't saturate.
3. Causal mask (lower-triangular −∞) for decoder-only — position i only sees ≤ i.
4. Row-wise softmax → each row a probability distribution.
5. Multiply by V → weighted sum of value vectors.

Output (B, T, d_head). Same per-position shape, but each position now carries information from every allowed other position.

5. The √d_k scaling — why it's there

For random Q, K with variance 1, Q · Kᵀ has variance d_k, so logits scale as √d_k. At d_k = 128 that's ±11. Softmax of ±11 collapses into near one-hot — one position takes all the mass. Then:

• Gradient vanishes (a peaked softmax has tiny gradient w.r.t. logits).
• The model can't learn to redistribute attention.

Dividing by √d_k keeps logits ≈ unit variance, softmax stays diffuse, gradients flow. Try it: at d_model = 4096 without scaling, training collapses in the first few hundred steps.

6. Multi-head attention

A single attention computes one routing pattern. Multi-head runs h heads in parallel with d_head = d_model / h:

MHA(x) = concat[head_1, …, head_h] · W_O
   head_i = attention(x · W_Q_i, x · W_K_i, x · W_V_i)

Typical numbers:

Heads specialise. Mechanistic interpretability work has documented:

• Previous-token heads — always attend to position i-1.
• Induction heads — copy patterns: after seeing … A B … A, attend to that previous B. These are the engine of in-context learning.
• Syntactic heads — subject-verb agreement, coreference, bracket matching.

You don't tell the model to have these — they emerge during training. That parallelism is the point of multi-head: many hypotheses, jointly trained.

7. Memory & compute cost

For batch B, layers L, heads h, head-dim d_head, sequence length T:

Two takeaways:

1. MLPs dominate FLOPs for typical T < 4k (≈ 2/3 of compute).
2. Attention dominates memory at long context. B=1, L=32, h=32, T=32k in fp16 = 64 GB just for attention matrices. This is why FlashAttention exists — tiles the computation, never materialises the full matrix.

8. From a single block to a full model

A decoder-only block is:

x = x + MHA(LayerNorm(x))     # attention sub-layer
x = x + MLP(LayerNorm(x))     # MLP sub-layer 

Pre-norm style (LN inside the residual) trains stably for very deep stacks. The residual connection is essential — it lets gradients flow directly from logits back to the embedding through 96 layers.

9. Hands-on (last 30 min)

import torch, torch.nn as nn

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_head = d_model // n_heads
        self.n_heads = n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = (t.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
                   for t in (q, k, v))
        att = (q @ k.transpose(-2, -1)) / (self.d_head ** 0.5)
        mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
        att = att.masked_fill(mask == 0, float('-inf')).softmax(-1)
        out = (att @ v).transpose(1, 2).contiguous().view(B, T, D)
        return self.proj(out)

Run on (2, 16, 128) random input with d_model=128, n_heads=4. Then remove the / d_head ** 0.5 scaling and inspect att[0,0,5] — it'll collapse toward one-hot. That's the experiment.

10. What's next

• Positional encoding — sinusoidal vs learned vs RoPE
• The MLP — where most of the facts live
• LayerNorm — pre-norm vs post-norm
• KV cache — what actually grows in memory at inference

Reading material

Books:

• The Hundred-Page Machine Learning Book — Andriy Burkov (chs. on neural networks)
• Deep Learning — Goodfellow, Bengio, Courville (ch. 9–10 for sequence models; the foundation that attention superseded)
• Natural Language Processing with Transformers — Lewis Tunstall, Leandro von Werra, Thomas Wolf (HF authors)

Papers:

• Attention Is All You Need — Vaswani et al. 2017 — the original transformer paper.
• Neural Machine Translation by Jointly Learning to Align and Translate — Bahdanau et al. 2014 — the attention paper that preceded transformers.

Official docs:

• PyTorch nn.MultiheadAttention
• Hugging Face Transformers — model docs

Blog posts:

• The Illustrated Transformer — Jay Alammar — the canonical visual walkthrough.
• The Annotated Transformer — Harvard NLP — paper + code line-by-line.
• A Mathematical Framework for Transformer Circuits — Anthropic — residual stream + heads.

In-depth research material

• nanoGPT (Karpathy) — github.com/karpathy/nanoGPT — ~40k ★, the cleanest reference implementation.
• FlashAttention (Dao et al., 2022) — the IO-aware attention kernel everyone now uses.
• Multi-Query Attention (Shazeer, 2019) and Grouped-Query Attention (Llama-2, 2023) — KV-cache savings at inference.
• In-context learning and induction heads (Anthropic, 2022) — heads that do in-context learning.
• In-context learning as gradient descent — Akyürek et al. 2022.
• xformers (Meta) — github.com/facebookresearch/xformers — efficient attention primitives.
• Stanford CS25: Transformers United — guest-lecture course page.

Videos

• Attention in transformers, step-by-step — 3Blue1Brown · 26 min — the visual intuition for Q/K/V and softmax routing; the best 30-min investment you can make.
• Let's build GPT: from scratch, in code, spelled out — Andrej Karpathy · 1 h 56 min — implements a small transformer end-to-end while explaining every line. The gold standard.
• Attention is all you need — model explanation, math, training & inference — Umar Jamil · 57 min — paper walk-through with the math fully derived; pairs perfectly with the Karpathy build.
• Attention Is All You Need (paper review) — Yannic Kilcher · 27 min — sharp paper review; great after you've seen the math once.
• Transformers, the tech behind LLMs — 3Blue1Brown · 27 min — the chapter before the attention one; sets up the residual stream picture.

LeetCode — Two Sum

• Link: https://leetcode.com/problems/two-sum/
• Difficulty: Easy
• Why this problem: Hash-map for O(n) lookup; the canonical interview opener.
• Time-box: 30 minutes. Look up the editorial only after.

Assignment / Deliverables

Give yourself a clean 2-hour window and complete all of these before moving on:

1. Read the deep-dive above end-to-end — no skimming. Take notes in your own words.
2. Solve the LeetCode problem below without help first. Only look at the hint after 15 focused minutes; only look at editorial after 30. Log your time.
3. Reproduce one code snippet locally. Pick the snippet that felt least obvious and get it running in a scratch file / notebook.
4. Draw the core diagram from memory. Paper, whiteboard, or tldraw — doesn't matter. If you can't, re-read section 2 and try again.
5. Write a 3-line takeaway in your prep journal: what surprised you, what you still don't understand, what you'd read next.
6. Skim one item from the Reading material section. Bookmark the rest for the weekend.
7. Commit any code + notes to your prep repo with message session-NN: <one-line summary>.

Stretch (optional, +30 min): explain today's topic to a rubber duck / a friend / a voice note. If you can't teach it in 5 minutes, you don't own it yet — flag it and revisit next weekend.

Post-session checklist

By the end of this session you should be able to:

• Draw the residual stream from token IDs → logits with N blocks and label every arrow.
• Derive Q, K, V from the residual stream and explain what each projection learns.
• State why scaled dot-product attention divides by √d_k and what breaks without it.
• List 3 head specialisations documented by mechanistic interpretability.
• Compute attention FLOPs and memory for B=1, L=32, h=32, d_head=128, T=4096.
• Implement a causal self-attention block from scratch (≤ 30 lines).

Generated from sessions_data.py + content_part*.py. To edit a video / leetcode / title, edit the data file and re-run write_sessions.py.

Part 2: Transformers Part 2 — Positional Encoding, RoPE, MLP, LayerNorm

Why this session matters

It builds on the rhythm of one focused topic, paced so you have time to actually absorb it rather than rush.

Agenda

• Why position has to be injected at all (attention is permutation-invariant)
• Three eras of positional encoding — sinusoidal, learned, RoPE
• RoPE math — rotate Q/K pairs, relative position emerges from QᵀK
• The MLP — 4× expansion, GELU/SwiGLU, where the facts live
• KV cache anatomy — what actually grows at inference time

Pre-read (skim before the session)

• Andrej Karpathy — Let's build GPT from scratch (video)
• RoFormer: Enhanced Transformer with Rotary Position Embedding
• EleutherAI — Rotary Embeddings: A Comprehensive Guide
• FlashAttention paper

Deep dive

1. Why position has to be injected at all

Self-attention is a set operation. Permute the inputs and the outputs permute the same way — but the content of each output doesn't change. A transformer without position info would see "dog bites man" and "man bites dog" as identical bags of words.

We need to inject position. Three eras of how:

2. Era 1 — sinusoidal (original "Attention Is All You Need")

Each position gets a deterministic vector built from sines and cosines at log-spaced frequencies:

PE(pos, 2i)   = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

Added (not concatenated) to the token embedding. Properties:

• Deterministic — no parameters to train.
• Periodic — each frequency repeats; combinations give unique position vectors up to a long horizon.
• Extrapolates poorly to lengths beyond what was trained on.

3. Era 2 — learned absolute (GPT-2)

A (max_pos, d_model) embedding table. Looked up by position, added to token embeddings.

• ✅ Trains end-to-end.
• ❌ Hard cap at max_pos (GPT-2: 1024). Can't extend without retraining.
• ❌ No structural inductive bias for "tokens close in position should attend more".

4. Era 3 — Rotary (RoPE) — used by Llama, Mistral, Qwen, DeepSeek

Instead of adding a position vector to the embedding, RoPE rotates Q and K pairs by an angle proportional to position. For position m and a pair (q_{2i}, q_{2i+1}):

[ cos(mθ_i)  −sin(mθ_i) ] [ q_{2i}   ]
[ sin(mθ_i)   cos(mθ_i) ] [ q_{2i+1} ]

with frequencies θ_i = 10000^(−2i/d).

The magic: (R_m q)ᵀ (R_n k) depends only on m − n — the relative offset — even though we encoded absolute positions. So attention scores naturally know "how far apart are these two tokens?" without any extra mechanism.

Why everyone moved to RoPE:

• ✅ Relative information for free.
• ✅ Extends to longer contexts via base-frequency scaling (NTK-aware, YaRN).
• ✅ Cheap — just two multiplies per element.

5. The MLP — where most of the parameters and facts live

After attention, every block applies a position-wise MLP:

MLP(x) = W_2 · activation(W_1 · x + b_1) + b_2

with W_1: d → 4d, W_2: 4d → d. So the hidden dim is 4× the model dim.

• For Llama-3-8B: d=4096 → hidden=14336 (≈3.5×, they use SwiGLU which changes the constant). Each block has ~117M MLP params; attention has ~67M. MLPs are roughly 2/3 of the parameters.
• The activation has evolved: ReLU → GELU (BERT, GPT-2) → SwiGLU (PaLM, Llama).

Mechanistic interpretability work (Anthropic, OpenAI) shows MLPs implementing key-value lookups: certain neurons fire on "Eiffel Tower" and write "Paris" into the residual stream. Most of the model's factual knowledge is in the MLP weights. Attention chooses what to mix; MLP decides what to write.

6. LayerNorm — pre-norm vs post-norm

Every sub-layer is wrapped in either:

• Pre-norm (Llama, GPT-NeoX, most modern): x = x + sublayer(LN(x))
• Post-norm (original Transformer): x = LN(x + sublayer(x))

Pre-norm trains more stably at depth — you can stack 80 blocks without warmup tricks. Post-norm gives slightly better final perplexity but needs careful warmup. Default to pre-norm.

LN itself: per-token, per-layer, normalise to zero mean and unit variance, then apply learnable scale γ and shift β:

LN(x) = γ * (x − mean(x)) / √(var(x) + ε) + β

RMSNorm (Llama variant) drops the mean centring — same effect, slightly cheaper:

RMSNorm(x) = γ * x / √(mean(x²) + ε)

7. KV cache — what actually grows in memory at inference

When you generate token by token, for every new token you compute Q from the latest position only, but you need K and V from every previous position. Standard impl: store K/V tensors per layer, per head, per position — the KV cache.

Size in bytes:

KV_bytes = 2 · B · T · L · d_model · dtype_bytes

(The leading 2 is for K + V.)

Concrete numbers for Llama-3-70B at fp16:

• L = 80, d = 8192, dtype = 2 bytes
• Per token, per request: 2 · 80 · 8192 · 2 = 2,621,440 bytes ≈ 2.5 MB/token
• A 4096-token conversation: ~10 GB just for KV cache for one request.

Two big consequences:

1. GQA (Grouped-Query Attention, Llama-2 onward) shares K and V across groups of heads — Llama-3-70B uses 8 KV heads for 64 query heads. Drops KV cache by 8×.
2. PagedAttention (vLLM, the next session) treats KV cache like virtual memory pages. Lets you batch many requests without pre-allocating worst-case memory.

8. Encoder-only / decoder-only / encoder-decoder

Decoder-only won because:

• Same architecture handles any task with prompting (instruction tuning).
• Causal mask is conceptually simpler than encoder→decoder cross-attention.
• Scales — pretty much every frontier model since 2022 is decoder-only.

9. Concrete numbers — Llama-3-70B by the slice

MLPs are where the weights are; attention is where the routing happens.

10. Hands-on (30 min)

Extend last session's CausalSelfAttention into a full block:

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, mlp_ratio=4):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, mlp_ratio * d_model),
            nn.GELU(),
            nn.Linear(mlp_ratio * d_model, d_model),
        )

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

Stack 6 of them, add embedding + final LN + unembed projection. Train on tiny-shakespeare (~1M chars) for 5k steps. You'll get vaguely Shakespeare-like output and a working mental model of every shape that flows through.

11. What's next (the next session — RAG Part 1)

• Why RAG exists (the limits of context windows and finetuning)
• Chunking strategies
• Embeddings overview (deep treatment in the next session)
• Vector stores (Faiss, pgvector, Pinecone, Weaviate)

Reading material

Books:

• Natural Language Processing with Transformers — Lewis Tunstall, Leandro von Werra, Thomas Wolf (chs. on architecture)
• Deep Learning — Goodfellow, Bengio, Courville (ch. 6 on MLPs, ch. 8 on optimisation — context for LayerNorm + residuals)

Papers:

• Attention Is All You Need (Vaswani et al., 2017) — sinusoidal positional encoding (§3.5).
• RoFormer: Enhanced Transformer with Rotary Position Embedding (Su et al., 2021) — the RoPE paper used by Llama, GPT-NeoX, PaLM.
• Layer Normalization (Ba, Kiros, Hinton, 2016) — the original LN paper.
• Root Mean Square Layer Normalization (Zhang & Sennrich, 2019) — RMSNorm, used by Llama/Mistral.
• GLU Variants Improve Transformer (Shazeer, 2020) — SwiGLU, the modern MLP.

Official docs:

• PyTorch nn.LayerNorm
• PyTorch nn.RMSNorm
• Hugging Face — Llama model card with RoPE details

Blog posts:

• Rotary Embeddings: A Brief Overview — EleutherAI — the most-cited RoPE explainer.
• Transformer Math 101 — EleutherAI — FLOPs, memory, KV cache math.

In-depth research material

• llama (Meta) — github.com/meta-llama/llama3 — reference Llama-3 model code with RoPE + SwiGLU + RMSNorm in <500 lines.
• mistral-src — github.com/mistralai/mistral-src — same modern stack, cleanly written.
• The Transformer Family Version 2.0 — Lilian Weng — comprehensive variant tour.
• On the Inductive Bias of Pre-Norm vs Post-Norm Transformers (Xiong et al., 2020) — why pre-norm trains stably.
• ALiBi: Train Short, Test Long (Press et al., 2021) — an alternative to RoPE used in BLOOM/MPT.
• Self-extending RoPE — Yarn & NTK-aware scaling — how to push RoPE to 100k context.

Videos

• RoPE (Rotary Positional Embeddings) Explained — DeepLearning Hero · 14 min — the visual intuition: rotation matrices applied to Q/K.
• Rotary Positional Embeddings: Combining Absolute and Relative — Efficient NLP · 11 min — clean breakdown of why RoPE encodes both absolute and relative position.
• How Rotary Position Embedding Supercharges Modern LLMs — Jia-Bin Huang · 14 min — math + visualisation pairing.
• Rotary Positional Embeddings Explained Visually — Outlier · 20 min — animated derivation; great if you like to see the rotation.
• Let's reproduce GPT-2 (124M) — Andrej Karpathy · 4 h 1 min — the back half of this video implements LayerNorm + MLP + positional encoding from scratch.

LeetCode — Longest Substring Without Repeating Characters

• Link: https://leetcode.com/problems/longest-substring-without-repeating-characters/
• Difficulty: Medium
• Why this problem: Sliding window with a hash-set; shrink left when you see a repeat.
• Time-box: 30 minutes. Look up the editorial only after.

Assignment / Deliverables

Give yourself a clean 2-hour window and complete all of these before moving on:

1. Read the deep-dive above end-to-end — no skimming. Take notes in your own words.
2. Solve the LeetCode problem below without help first. Only look at the hint after 15 focused minutes; only look at editorial after 30. Log your time.
3. Reproduce one code snippet locally. Pick the snippet that felt least obvious and get it running in a scratch file / notebook.
4. Draw the core diagram from memory. Paper, whiteboard, or tldraw — doesn't matter. If you can't, re-read section 2 and try again.
5. Write a 3-line takeaway in your prep journal: what surprised you, what you still don't understand, what you'd read next.
6. Skim one item from the Reading material section. Bookmark the rest for the weekend.
7. Commit any code + notes to your prep repo with message session-NN: <one-line summary>.

Stretch (optional, +30 min): explain today's topic to a rubber duck / a friend / a voice note. If you can't teach it in 5 minutes, you don't own it yet — flag it and revisit next weekend.

Post-session checklist

By the end of this session you should be able to:

• State why attention needs position info (set vs sequence).
• Derive that RoPE makes attention depend on relative offset.
• Compute KV cache size for Llama-3-70B at T=4096.
• Explain why MLPs hold most parameters and most facts.
• Choose pre-norm vs post-norm and defend the choice for a 60-layer model.
• Solve longest-substring-without-repeating-characters with sliding window.

Generated from sessions_data.py + content_part*.py. To edit a video / leetcode / title, edit the data file and re-run write_sessions.py.