Attention, intuitively

Why this matters

Every modern LLM is "attention all the way down." Most engineers can recite "Q, K, V" without having a real picture of what attention is computing. This lesson installs that picture before we hit the formal Transformer architecture in Phase 2. Read this lesson once, do the lab, and you will never be confused by attention again.

Learning objectives

Describe attention as a soft, differentiable dictionary lookup.
Implement scaled dot-product attention from scratch.
Reason about masking, padding, and the role of the sqrt(d_k) scaling.
Distinguish self-attention from cross-attention from causal attention.
Implement multi-head attention and explain why we need multiple heads.

1. The big idea — soft lookup

Imagine a Python dictionary:

python

db = {"apple": 5, "banana": 3, "cherry": 9}
db["banana"]  # 3 — exact key match

Attention is the same idea, but soft and differentiable:

The query "kind-of matches" each key.
We blend values according to how well the query matches each key.

Mathematically:

attention(Q, K, V) = softmax(Q K^T / sqrt(d_k)) V

Q: the queries (what I am looking for).
K: the keys (what each item advertises).
V: the values (the actual content to fetch).
softmax(Q K^T / sqrt(d_k)): a row of attention weights summing to 1 — the soft "lookup."

That is all of attention. Everything else is plumbing.

2. Why `sqrt(d_k)` scaling?

If Q and K are random vectors with unit-variance components, Q . K has variance d_k. As d_k grows, the dot product magnitudes grow → the softmax becomes very sharp (mostly 0s and 1s) → gradients through softmax vanish.

Dividing by sqrt(d_k) keeps the variance ~1 regardless of d_k, keeping the softmax in a "useful temperature" zone. Forget this and your transformer trains terribly.

3. Implementation in 12 lines

python

import torch, math
import torch.nn.functional as F

def attention(Q, K, V, mask=None):
    # Q, K, V: (B, T, d_k)
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)   # (B, T_q, T_k)
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))
    weights = F.softmax(scores, dim=-1)                # (B, T_q, T_k)
    return weights @ V, weights                        # (B, T_q, d_v), weights for inspection

Try it:

python

B, T, d = 2, 5, 16
Q = K = V = torch.randn(B, T, d)        # self-attention: same source
out, w = attention(Q, K, V)
print(out.shape, w.sum(-1))             # (2, 5, 16), all ones

4. Self-attention vs cross-attention

Self-attention: Q, K, V are all derived from the same sequence. Used in the encoder of T5/BERT and in every layer of GPT-style models.
Cross-attention: Q comes from one sequence (e.g., decoder so far), K/V from another (e.g., encoder output). Used in seq2seq (T5, Whisper, NMT models) and in some adapters / multimodal models.

Same equation, different inputs.

5. Causal masking — the GPT secret

For autoregressive models, position t must not see positions t+1...T-1 (it is generating them, after all). Implement with a lower-triangular mask:

python

T = 5
causal_mask = torch.tril(torch.ones(T, T, dtype=torch.bool))     # True where allowed
# rows are queries, columns are keys
# row 0: only key 0
# row 1: keys 0,1
# row 2: keys 0,1,2 ...

Apply via masked_fill(~causal_mask, -inf). After softmax, future positions get weight 0.

This is the only algorithmic difference between BERT-style (bidirectional) and GPT-style (causal) transformers.

Padding mask

In a batch, sequences have different lengths; we pad the short ones. The padding tokens must also get attention weight 0:

mask = causal_mask & (key_positions != PAD)

6. Multi-head attention

Empirical observation: a single attention head learns to look at one kind of relationship (e.g., "the next noun"). Multiple heads attend to different kinds of relationships in parallel.

Mechanically, multi-head attention is:

1. Project x to Q, K, V each of shape (B, T, d). [Linear projections.]
2. Reshape to (B, T, h, d/h) -> (B, h, T, d/h).
3. Run scaled-dot-product per head -> (B, h, T, d/h).
4. Concatenate heads -> (B, T, d).
5. Output projection W_O.

In PyTorch:

python

import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, causal=True):
        super().__init__()
        assert d_model % n_heads == 0
        self.h = n_heads
        self.dk = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)
        self.causal = causal

    def forward(self, x, mask=None):
        B, T, d = x.shape
        qkv = self.qkv(x).reshape(B, T, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)
        Q, K, V = qkv[0], qkv[1], qkv[2]                  # (B, h, T, dk)
        scores = (Q @ K.transpose(-2, -1)) / (self.dk ** 0.5)
        if self.causal:
            cm = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))
            scores = scores.masked_fill(~cm, float("-inf"))
        if mask is not None:
            scores = scores.masked_fill(~mask[:, None, None, :], float("-inf"))
        w = scores.softmax(dim=-1)
        out = w @ V                                       # (B, h, T, dk)
        out = out.transpose(1, 2).reshape(B, T, d)        # concat heads
        return self.out(out)

That is ~30 lines. You will rebuild this from memory in interviews.

Why multiple heads?

Different heads capture different patterns (subject-verb agreement, coreference, syntactic role, semantic similarity).
Mathematically, multi-head attention with h heads of size d/h has the same parameter count as one head of size d but strictly more expressive (the projection split lets different subspaces specialise).
Modern wrinkle: at scale, many heads are redundant; GQA (Grouped-Query Attention) and MQA (Multi-Query Attention) share K/V across heads to save memory. Phase 2 covers it.

7. Attention is `O(T^2)` — the bottleneck

For sequence length T, the score matrix is (T, T) — quadratic in time and memory. That is why long contexts are expensive and why FlashAttention (re-orders the computation to use SRAM and avoid materialising the full T×T matrix) is so important. We will return to this in Phase 5.

Variants that try to break the quadratic wall:

Sparse attention (Longformer, BigBird) — only compute attention to a subset of positions.
Linear attention (Performer, RWKV) — replace softmax with a kernel that factorises.
State-space models (Mamba) — drop attention entirely; Phase 6.

For 2026 production: flash + sliding window + GQA is the standard cocktail; full long-context now goes to 1M+ tokens (Gemini 2.x, Claude 4, Llama 4).

Hands-on lab (3 hours)

attention_lab.ipynb:

Implement attention(Q, K, V, mask) from scratch as above.
Generate Q, K, V with T=8, d=4. Pretty-print the attention weight matrix.
Add a causal mask and verify each row's leftmost N entries sum to 1.
Multiply every component of Q by 10. Show that the softmax becomes near-one-hot. Now apply scaling by sqrt(d) and discuss.
Implement MultiHeadAttention with d=64, h=8. Sanity-check by feeding random input and printing output shape (B, T, 64).
Visualise an attention head's weights as a heatmap on a real sentence using a small pretrained transformer (e.g., DistilBERT). Look for syntactic patterns.
Bonus: replace softmax with linear attention (Performer-style feature map) and compare on a tiny task.

Common pitfalls

Forgetting the sqrt(d_k) scaling — your loss diverges or never moves.
Mask sign confusion — pass mask=True for allowed positions in our convention; check carefully when calling library helpers (torch.nn.functional.scaled_dot_product_attention uses additive masks).
Forgetting to apply causal mask in decoder-style attention → leakage from future tokens. Looks like training works but model "memorises" instead of generalising.
Wrong head reshape: (B, T, h*dk).reshape(B, T, h, dk) is fine; .reshape(B, h, T, dk) is not equivalent and silently breaks.
Computing Q @ K^T instead of Q @ K.transpose(-2, -1) — same mistake, different bug.

Self-check

Why divide by sqrt(d_k)?
Difference between causal attention and self-attention.
Why do we use multiple heads instead of one big head?
Memory cost of attention as a function of sequence length?
What does GQA share across heads?

References

Vaswani et al. (2017), "Attention Is All You Need" — read this carefully; you can read it in 30 minutes.
Bahdanau et al. (2014), "Neural Machine Translation by Jointly Learning to Align and Translate."
Karpathy, "Let's build GPT: from scratch, in code, spelled out."
Jay Alammar, "The Illustrated Transformer."
3Blue1Brown — "Attention in transformers, visually explained."

Why this matters

Learning objectives

1. The big idea — soft lookup

2. Why sqrt(d_k) scaling?

3. Implementation in 12 lines

4. Self-attention vs cross-attention

5. Causal masking — the GPT secret

Padding mask

6. Multi-head attention

Why multiple heads?

7. Attention is O(T^2) — the bottleneck

Hands-on lab (3 hours)

Common pitfalls

Self-check

References

2. Why `sqrt(d_k)` scaling?

7. Attention is `O(T^2)` — the bottleneck