§ 6.1 Attention is All You Need

1. Overview

Self-attention is the operation at the heart of every modern transformer. It takes a sequence of T token vectors and produces a new sequence of T vectors where each output position is a content-addressed mixture of all (or all earlier) inputs. The mixture weights are computed on the fly from the inputs themselves — no learned routing table, no recurrence — which is what makes attention parallelisable across the entire sequence in one matrix multiply.

Vaswani et al. (2017, arXiv:1706.03762) showed that this single primitive — stacked into N blocks alongside a feed-forward layer — replaces every prior sequence-model component (recurrence, convolution, alignment networks) and outperforms them at scale. Every subsequent topic in this curriculum (positional encodings, MQA/GQA, KV-caching, FlashAttention, MoE, speculative decoding) is either a refinement of, or a system built around, scaled dot-product attention.

2. Q, K, V — Projections and Tensor Shapes

Attention starts with three learned linear projections of the same input X ∈ ℝB×T×dmodel:

Q = X W_Q     # what each position is "looking for"
K = X W_K     # what each position "advertises"
V = X W_V     # what each position contributes if attended to

The intuition is a soft, content-addressed memory: a query vector q_i is compared against every key vector k_j, the comparisons are turned into a probability distribution, and the output is the corresponding mixture of the value vectors. Because Q, K, V are all linear projections of the same X, this is self-attention; in cross-attention (encoder–decoder) Q comes from one sequence and K/V from another.

Shape conventions used on this site

After projecting and reshaping by head, each of Q, K, V lives in [B, H, T, d_k]. The two leading dims (B, H) are pure broadcast dims — every operation that follows treats them as independent — so a single fused tensor op runs B · H independent attention computations in parallel on the GPU.

3. Core Mechanism — Scaled Dot-Product Attention

Background: Given T queries and T keys (one per token, each of dim d_k), we want a similarity score between every query and every key — a T × T matrix. The dot product is the cheapest content-aware similarity available: it can be computed for all pairs simultaneously with one BLAS GEMM.

Plan:

1. Compute raw scores S = Q Kᵀ → shape [B, H, T, T].
2. Scale by 1/√dk to keep logit variance ≈ 1 (derived below).
3. Add a mask: set forbidden positions to -∞ so they get probability 0 after softmax (causal mask, padding mask).
4. Row-wise softmax → attention probabilities P (each row sums to 1).
5. Weighted sum of values: O = P V, shape [B, H, T, d_v].

Closed form

Attention(Q, K, V) = softmax( Q K^T / sqrt(d_k) + M ) V

Shapes:
  Q : [B, H, T, d_k]
  K : [B, H, T, d_k]    -> K^T : [B, H, d_k, T]
  V : [B, H, T, d_v]
  M : [T, T]   (mask, broadcasts over B and H; -inf where forbidden, 0 elsewhere)

  Q K^T          -> [B, H, T, T]   (one logit per (i, j) pair)
  / sqrt(d_k)    -> [B, H, T, T]   (same shape, just rescaled)
  + M            -> [B, H, T, T]   (broadcast mask add)
  softmax dim=-1 -> [B, H, T, T]   (row-stochastic)
  @ V            -> [B, H, T, d_v] (one output vector per query position)

Why √d_k? — variance argument

Treat the components of q and k as iid with mean 0 and variance 1 (a reasonable approximation after standard initialisation and LayerNorm). Then the dot product is a sum of d_k iid terms:

q · k = Σ_{i=1..d_k}  q_i k_i

E[q_i k_i]    = E[q_i] E[k_i] = 0         (independent, zero-mean)
Var(q_i k_i)  = E[(q_i k_i)^2] - 0
              = E[q_i^2] E[k_i^2]          (independence)
              = 1 · 1 = 1

Var(q · k)    = sum of d_k iid unit-variance terms
              = d_k

  std(q · k)  = sqrt(d_k)

For dk = 64 the logits have std ≈ 8, so typical values land in [−24, +24]. Softmax of such large logits is essentially a one-hot: the largest logit gets probability ≈ 1, the rest get ≈ 0, and the gradient of softmax saturates (∂softmax/∂x → 0 off the argmax). Training cannot proceed.

Dividing by √dk rescales the variance back to 1 independent of head dim, keeping the softmax in its soft, gradient-friendly regime. This is exactly why Vaswani et al. called it scaled dot-product attention.

Worked numerical walkthrough

Take T = 3, d_k = 4. Pick concrete Q, K, V (rounded to 2 d.p. for readability). Trace every shape transition end-to-end:

Q = [[ 1.0,  0.0,  0.0,  0.0],     # query for token 1
     [ 0.0,  1.0,  0.0,  0.0],     # query for token 2
     [ 0.5,  0.5,  0.0,  0.0]]     # query for token 3 (a blend of 1 and 2)

K = [[ 1.0,  0.0,  0.0,  0.0],     # key for token 1
     [ 0.0,  1.0,  0.0,  0.0],     # key for token 2
     [ 0.0,  0.0,  1.0,  0.0]]     # key for token 3 (unrelated direction)

V = [[10, 20],                      # value for token 1
     [30, 40],
     [50, 60]]

Step 1   S = Q K^T          (shape 3 x 3)
         S[i,j] = <q_i, k_j>
         S = [[ 1.0,  0.0,  0.0],
              [ 0.0,  1.0,  0.0],
              [ 0.5,  0.5,  0.0]]

Step 2   S / sqrt(d_k) = S / 2     (since d_k = 4)
         S' = [[ 0.50, 0.00, 0.00],
               [ 0.00, 0.50, 0.00],
               [ 0.25, 0.25, 0.00]]

Step 3   Apply CAUSAL mask: positions j > i become -inf.
         S' = [[ 0.50, -inf, -inf],
               [ 0.00,  0.50, -inf],
               [ 0.25,  0.25,  0.00]]

Step 4   Row-wise softmax (each row sums to 1):
         row 1: [1.00, 0,    0   ]
         row 2: softmax([0.0, 0.5])   = [0.378, 0.622, 0]
         row 3: softmax([0.25, 0.25, 0.0])
                = [exp(.25), exp(.25), exp(0)] / norm
                = [1.284, 1.284, 1.0] / 3.568
                = [0.360, 0.360, 0.280]

         P = [[1.000, 0.000, 0.000],
              [0.378, 0.622, 0.000],
              [0.360, 0.360, 0.280]]

Step 5   O = P V                    (shape 3 x 2)
         O[0] = 1.000*[10,20] + 0*[30,40] + 0*[50,60]
              = [10.0, 20.0]
         O[1] = 0.378*[10,20] + 0.622*[30,40] + 0
              = [22.44, 32.44]
         O[2] = 0.360*[10,20] + 0.360*[30,40] + 0.280*[50,60]
              = [28.4, 38.4]

         O = [[10.00, 20.00],       # token 1 sees only itself
              [22.44, 32.44],       # token 2 mostly attends to itself, some to token 1
              [28.40, 38.40]]       # token 3 spreads attention across all three

The output for token 1 is exactly V[1] — the causal mask forces position 0 to attend only to itself. Token 2 splits its attention between itself and token 1. Token 3 mixes all three. This is the essential picture: each output is a data-dependent convex combination of past values.

4. Multi-Head Attention

A single attention head is already useful, but its output for each token is a single weighted mean — it can only model one type of relationship at a time (e.g., "previous noun" or "matching bracket"). Multi-head attention runs H independent attention computations in parallel — each with its own small Q/K/V projections — and concatenates their outputs. Different heads learn different attention patterns: short-range, long-range, syntactic, coreferent, induction-style copy heads, etc.

Closed form

MultiHead(Q, K, V) = Concat(head_1, ..., head_H) W_O

with   head_h = Attention( Q W_Q^h, K W_K^h, V W_V^h )

Implementation trick: instead of H separate matmuls, allocate one
[d_model, d_model] matrix for each of W_Q, W_K, W_V and just RESHAPE
the result into H heads of width d_k = d_model / H.

So H heads of width d_k cost the same FLOPs as one head of width d_model,
but force the model to spend attention capacity across H parallel subspaces.

The output projection W_O is essential: without it the different heads cannot mix. W_O takes the H concatenated d_k-vectors and produces the next residual-stream update, allowing information from head 7 to be used by future computation in any other channel.

Why split into many small heads instead of one big one?

• Specialisation. Each head can lock onto a different pattern; the model gets multiple "channels" of relational reasoning per layer.
• No extra FLOPs. Because d_k = d_model/H, the QKᵀ matmul costs the same as one big head (B·H·T²·d_k = B·T²·d_model).
• Better gradient flow. Smaller per-head QKᵀ outputs are easier to normalise and produce less-saturated softmax distributions.
• Interpretability. Heads can be analysed individually — see Elhage et al. 2021 (Anthropic) on induction heads and circuits.

5. Masks — Causal, Padding, Prefix-LM

A mask is a tensor added to the pre-softmax scores. Anywhere it equals -∞ the corresponding probability after softmax becomes exactly 0. Masks are how transformers express "you may not look here": causal masking makes a decoder, padding masks ignore padding tokens, prefix-LM masks make a bidirectional prefix followed by causal continuation.

Causal mask (decoder / GPT)

For autoregressive next-token prediction, token i must not see any token j > i; otherwise the loss is trivial (the model can just copy the next token). The mask is an upper-triangular matrix of −∞:

Green = a real score is kept; red = −∞ so the softmax drives it to 0. Row i shows what query position i is allowed to attend to.

Padding mask

When sequences in a batch have different lengths they are padded to the max length with a PAD token. A padding mask zeroes out attention to padding positions (column-wise −∞) so that legitimate tokens never spend probability mass on pad. Modern packed-sequence training (FlashAttention varlen, vLLM) avoids padding entirely with a per-example cu_seqlens index.

Encoder, decoder, prefix-LM

6. FLOPs & Memory Cost

Knowing the cost of attention by heart is part of the LLM-infra interview bar. Per layer, per sequence (drop B for clarity):

Compute (FLOPs, multiplies only — multiply by 2 if you count mul+add):

  QKV projections     : 3 * T * d_model * d_model      = 3 T d^2
  Q K^T               : H * T * T * d_k = T^2 d_model  = T^2 d
  softmax             : T^2  (O(T^2) but cheap; dominated above)
  P V                 : H * T * T * d_v = T^2 d_model  = T^2 d
  Output projection   : T * d_model * d_model          = T d^2

  Total per layer ~ 4 T d^2  +  2 T^2 d

  Activation-memory peak (the part everyone trips over):
    scores tensor S   :  B * H * T * T  floats   <-- O(T^2) PER HEAD

  For B=4, H=32, T=8192, fp16:
    S = 4 * 32 * 8192 * 8192 * 2 bytes = 17.18 GB    (one tensor!)

  This single allocation is what blows up at long context and what
  FlashAttention removes by never materialising S in HBM.

Two regimes matter when reading roofline plots:

• Short sequence (T < d_model): compute is dominated by the QKV/output projections (∝ T d²). The matmul is compute-bound — tensor cores saturate.
• Long sequence (T ≫ d_model): the QKᵀ and PV matmuls (∝ T²d) dominate, and the O(T²) score tensor dominates HBM traffic. Attention becomes memory-bandwidth-bound — this is exactly the regime FlashAttention v1–v3 (T44) is designed for.

7. Minimal Demos

Demo A — Scaled dot-product attention end-to-end

Tiny B=1, H=2, T=4, D=8 example. Every shape is printed. The final assertion confirms that with a causal mask, token 0 attends only to itself — so the output at position 0 is exactly V at position 0.

Demo B — Why divide by √d_k?

Sweeps d_k ∈ {4, 16, 64, 256, 1024} and prints the entropy and max-probability of the resulting attention distribution, with and without scaling. Without scaling, entropy collapses toward zero as d_kgrows — the model degenerates into a hard argmax. With scaling, entropy stays near log₂(T).

Demo C — Multi-head attention by hand

Builds MHA from nn.Linear only — no nn.MultiheadAttention magic. Watch the shape go [B,T,d_model] → [B,H,T,d_k] → [B,H,T,d_v] → [B,T,d_model]. The final line confirms the canonical 4·d_model² parameter count.

8. Production / Source Pointers

9. References

Papers

• Vaswani et al. (2017) — Attention is All You Need. arXiv:1706.03762
• Bahdanau, Cho, Bengio (2014) — Neural Machine Translation by Jointly Learning to Align and Translate (additive attention, the predecessor). arXiv:1409.0473
• Luong, Pham, Manning (2015) — Effective Approaches to Attention-based NMT (dot-product attention). arXiv:1508.04025
• Britz et al. (2017) — Massive Exploration of NMT Architectures (empirical case for scaled dot-product). arXiv:1703.03906
• Elhage et al. / Anthropic (2021) — A Mathematical Framework for Transformer Circuits (residual stream + head decomposition). transformer-circuits.pub
• Phuong & Hutter (2022) — Formal Algorithms for Transformers. arXiv:2207.09238
• Dao et al. (2022) — FlashAttention: Fast and Memory-Efficient Exact Attention. arXiv:2205.14135 (forward reference for §6.7 cost discussion)

Lectures

• Stanford CS224n (Manning) — Lecture 8: Attention; Lecture 9: Self-Attention & Transformers
• Stanford CS25 — Transformers United (lecture 1: introduction; lecture 2: attention deep dive)
• Stanford CS336 (Liang & Yao) — Building LLMs from Scratch (2024), Attention lecture
• CMU 11-785 (Bhiksha Raj) — Deep Learning, Attention & Transformer lectures
• NYU DS-GA 1008 (Yann LeCun & Alfredo Canziani) — Deep Learning, Self-Attention
• Karpathy — Let's build GPT: from scratch, in code, spelled out (YouTube, 2 hrs)
• 3Blue1Brown — But what is a GPT? / Attention in transformers, visually explained

Textbooks

• Jurafsky & Martin — Speech and Language Processing, 3rd ed., Ch. 10–11 (Transformers and Large Language Models)
• Bishop & Bishop (2024) — Deep Learning: Foundations and Concepts, Ch. 12 (Transformers)
• Goodfellow, Bengio, Courville — Deep Learning (foundational, predates transformers but covers attention basics)

Code / Repos

• karpathy/nanoGPT — minimal causal self-attention in ~30 lines; the canonical learning reference
• karpathy/llm.c — same model in pure C/CUDA; great for kernel intuition
• huggingface/transformers — LlamaAttention and GPT2Attention reference implementations
• Dao-AILab/flash-attention — production IO-aware attention kernels (FA-1/2/3)
• harvardnlp/annotated-transformer — line-by-line annotated PyTorch implementation of Vaswani et al.

Blog Posts

• Sasha Rush — The Annotated Transformer (Harvard NLP, 2018 / 2022 update)
• Jay Alammar — The Illustrated Transformer & The Illustrated GPT-2
• Lilian Weng — The Transformer Family & Attention? Attention! (lilianweng.github.io)
• Anthropic — A Mathematical Framework for Transformer Circuits (transformer-circuits.pub)
• Karpathy — The Unreasonable Effectiveness of Recurrent Neural Networks & later GPT notes

10. Interview Prep