16  Self-Attention & Multi-Head Attention

TipTL;DR

Self-attention lets every position in a sequence attend to every other position in the same sequence, in parallel, with no recurrence. Each token emits a query, a key, and a value; the output at each position is a weighted sum of all values, weighted by query-key similarity: \[\text{Attention}(Q,K,V) = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V.\] Multi-head attention runs \(h\) of these in parallel on projected subspaces and concatenates them. Vaswani et al. (2017) showed this single mechanism can replace recurrence entirely — the foundation of every Transformer.

Depends on: Additive and Multiplicative Attention

16.1 Why this matters

RNN attention (Bahdanau, Luong) let a decoder query an encoder: queries came from one sequence, keys and values from another. The recurrence was still there — hidden states had to be computed left-to-right, one step at a time.

Self-attention makes two moves that change everything:

  1. Query your own sequence. Q, K, and V all come from the same input. Every token can gather context from every other token directly — no information has to travel step-by-step through a hidden state.
  2. Drop recurrence entirely. Because every pairwise interaction is computed independently, the whole sequence is processed in one parallel matrix multiply. Training is no longer bottlenecked by sequence length the way BPTT was.

The cost is quadratic — every position attends to every other, so compute and memory grow as \(O(n^2)\) in sequence length \(n\). That trade (parallelism and direct access, paid for in quadratic cost) defines the entire Transformer era, and motivates the efficiency work that follows it (Flash Attention, GQA, MLA).

16.2 The mechanism

16.2.1 Queries, keys, and values

Given an input sequence \(X \in \mathbb{R}^{n \times d_{\text{model}}}\) (n tokens, each a \(d_{\text{model}}\)-dim vector), three learned projections produce queries, keys, and values:

\[ Q = X W_Q, \quad K = X W_K, \quad V = X W_V, \qquad W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_k},\; W_V \in \mathbb{R}^{d_{\text{model}} \times d_v}. \]

The analogy: a query is what a token is looking for; a key is what each token offers as a label; a value is the content it actually contributes. The query-key dot product scores relevance; the scores weight the values.

16.2.2 Scaled dot-product attention

The core operation scores every query against every key, normalises with a softmax over the sequence, and mixes the values:

\[ \text{Attention}(Q,K,V) = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V. \]

  • \(QK^\top \in \mathbb{R}^{n \times n}\) — the full matrix of pairwise scores.
  • The softmax runs over each row, turning scores into a distribution over positions (the attention weights).
  • The result is \(n \times d_v\) — one context vector per position.
WarningPitfall: why divide by √d_k

For large \(d_k\), the dot product \(q \cdot k = \sum_{i=1}^{d_k} q_i k_i\) has variance proportional to \(d_k\). Large-magnitude scores push the softmax into saturated regions where gradients vanish — the model stops learning. Dividing by \(\sqrt{d_k}\) rescales the variance back to \(\approx 1\), keeping the softmax in a well-conditioned range. This is the one change that turns Luong’s dot-product attention into the Transformer’s.

16.2.3 Multi-head attention

A single attention function averages everything into one weighted sum, which limits what relationships it can capture. Multi-head attention runs \(h\) attention functions in parallel, each on its own learned projection:

\[ \text{head}_i = \text{Attention}(X W_Q^i,\; X W_K^i,\; X W_V^i), \] \[ \text{MultiHead}(X) = \big[\text{head}_1; \ldots; \text{head}_h\big]\, W_O, \qquad W_O \in \mathbb{R}^{h\,d_v \times d_{\text{model}}}. \]

Each head projects into a lower-dimensional subspace (\(d_k = d_v = d_{\text{model}} / h\) in the original paper), so total cost is comparable to a single full-dimensional head. Different heads learn different relationships — some track syntax, some track coreference, some attend locally.

X n × d_model h parallel heads head₁: Q₁K₁ᵀ/√d_k → softmax → ·V₁ head₂ ⋮ head_h concat · W_O n × d_model Each head attends in its own subspace; outputs are concatenated and projected back to d_model. Q, K, V all derive from the same input X — this is what makes it self-attention.

16.2.4 Masking

Two masks adapt attention to different settings, applied by adding \(-\infty\) to scores before the softmax (so masked positions get zero weight):

  • Causal (look-ahead) mask — in a decoder, position \(t\) may only attend to positions \(\le t\). Without it, the model could “see the future” and the language-modelling objective would be trivial. This is what makes GPT autoregressive.
  • Padding mask — variable-length sequences are padded to a common length; the padding positions are masked so they contribute nothing.

16.2.5 Complexity

For sequence length \(n\) and dimension \(d\), the \(QK^\top\) product and the weighting are both \(O(n^2 d)\) in time, and the attention matrix is \(O(n^2)\) in memory. Compare to an RNN’s \(O(n d^2)\) sequential time. Self-attention wins on parallelism (no sequential dependency) but loses on scaling in \(n\) — the central tension the rest of the era engineers around.

16.2.6 Minimal sketch

Single-head self-attention over a toy sequence, with an optional causal mask.

import numpy as np

rng = np.random.default_rng(0)
n, d_model, d_k = 4, 8, 8        # 4 tokens, model dim 8, head dim 8

X = rng.normal(size=(n, d_model))
W_Q = rng.normal(scale=0.3, size=(d_model, d_k))
W_K = rng.normal(scale=0.3, size=(d_model, d_k))
W_V = rng.normal(scale=0.3, size=(d_model, d_k))

def softmax(z, axis=-1):
    z = z - z.max(axis=axis, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

def self_attention(X, causal=False):
    Q, K, V = X @ W_Q, X @ W_K, X @ W_V
    scores = (Q @ K.T) / np.sqrt(d_k)          # (n, n) pairwise scores
    if causal:                                  # mask future positions
        mask = np.triu(np.ones((n, n)), k=1).astype(bool)
        scores[mask] = -np.inf
    A = softmax(scores, axis=-1)                # attention weights
    return A @ V, A

out, A = self_attention(X, causal=True)
print("attention weights (causal — lower triangular):")
print(np.round(A, 2))
print("output shape:", out.shape)
attention weights (causal — lower triangular):
[[1.   0.   0.   0.  ]
 [0.98 0.02 0.   0.  ]
 [0.32 0.34 0.35 0.  ]
 [0.23 0.27 0.22 0.28]]
output shape: (4, 8)
TipGoing deeper
  • The full encoder/decoder stack that wraps this mechanism — residuals, layer norm, position-wise FFN — is the next chapter: The Transformer.
  • Self-attention has no notion of order on its own; position must be injected separately — see Positional Encoding.
  • The \(O(n^2)\) memory cost and its mitigations (KV-cache, GQA, MLA) are covered in KV-Cache.

16.3 Application & impact

Concept here What it becomes Where
Scaled dot-product The universal attention primitive Every Transformer, every LLM
Q/K/V projections Cross-attention (Q from decoder, K/V from encoder) Enc-dec Transformers, T5
Multi-head Attention-head analysis, head pruning Interpretability, efficient inference
Causal mask Autoregressive generation GPT family, all decoder-only LLMs
\(O(n^2)\) cost Efficient-attention research Flash Attention, GQA, MLA, sparse attention
  • The keystone result: Vaswani et al. (Vaswani et al., 2017) showed a stack of self-attention + feedforward layers beats RNN/CNN sequence models on translation while training far faster — parallelism over sequence length was the unlock.
  • Cross-attention is the same operation with queries from one sequence and keys/values from another — the encoder-decoder bridge, inherited directly from RNN attention.
  • Heads as a research surface: because each head is a separate attention map, multi-head attention became the main object of interpretability work (induction heads, attention-head pruning, GQA’s head sharing).
NoteKey takeaway

Self-attention is one matrix expression — \(\operatorname{softmax}(QK^\top/\sqrt{d_k})V\) — applied to a sequence attending to itself. Multi-head runs several in parallel subspaces. Removing recurrence buys massive parallelism at the price of quadratic cost in sequence length. Everything else in the Transformer is plumbing around this core; everything after the Transformer is either scaling it up or making the quadratic term cheaper.

→ Next: The Transformer