Scaled dot-product and multi-head attention, fully derived

Queries, keys and values

Start from the soft dictionary lookup of 6.7 Attention is born. A dictionary maps a key to a value; a soft dictionary returns a weighted blend of all values, with weights given by how well the query matches each key. Attention builds the query, key and value for every token from the same input vectors with three learned projections:

where $\mathbf{X}\in\mathbb{R}^{n\times d}$ holds the $n$ token vectors, $\mathbf{W}^Q,\mathbf{W}^K\in\mathbb{R}^{d\times d_k}$ and $\mathbf{W}^V\in\mathbb{R}^{d\times d_v}$. Row $i$ of $\mathbf{Q}$ is "what token $i$ is looking for", row $j$ of $\mathbf{K}$ is "what token $j$ offers", and row $j$ of $\mathbf{V}$ is "what token $j$ will contribute if chosen".

Scaled dot-product attention

Read it in three steps. Scores: $\mathbf{S} = \mathbf{Q}\mathbf{K}^\top$ is $n\times n$; entry $s_{ij} = \mathbf{q}_i\cdot\mathbf{k}_j$ measures alignment (1.1 Vectors, dot products and norms). Weights: a softmax over each row turns scores into a probability distribution, $a_{ij} = e^{s_{ij}/\sqrt{d_k}}\big/\sum_{j'} e^{s_{ij'}/\sqrt{d_k}}$, so every output is a convex combination of values. Output: $\mathbf{Z} = \mathbf{A}\mathbf{V}$; row $i$ is $\sum_j a_{ij}\mathbf{v}_j$, a blend of other tokens' values chosen by token $i$'s query.

Why divide by $\sqrt{d_k}$

Suppose the components of $\mathbf{q}$ and $\mathbf{k}$ are independent with mean $0$ and variance $1$ (roughly true after normalization and sensible initialization, 4.6 Initialization and normalization). Then

The raw scores have standard deviation $\sqrt{d_k}$; with $d_k = 64$ that is $8$, and a softmax of inputs spread over tens of units is nearly one-hot, where its gradient (recall $\partial a_i/\partial s_j = a_i(\delta_{ij}-a_j)$) is almost zero (4.2 Activation functions). Dividing by $\sqrt{d_k}$ restores unit variance and keeps the softmax in its responsive regime.

Masking

To forbid token $i$ from seeing token $j$, set $s_{ij} = -\infty$ (in code, a large negative number) before the softmax so that $a_{ij} = 0$ exactly and the remaining weights renormalize. The causal mask zeroes all $j > i$ — this is what makes a decoder autoregressive (7.6 Decoder-only: GPT and autoregressive generation) and lets one forward pass train every next-token prediction in parallel. A padding mask hides filler tokens in batched sequences.

A complete worked example

Three tokens, $d_k = d_v = 2$, with

Scores $\mathbf{S}=\mathbf{Q}\mathbf{K}^\top$ (row $i$ = query $i$ against all keys), then scaled by $\sqrt{2}$:

Row-wise softmax (using $e^{0.707}\approx 2.028$, $e^{1.414}\approx 4.113$, $e^{0}=1$):

Output $\mathbf{Z}=\mathbf{A}\mathbf{V}$: row 1 is $0.401(1,2)+0.198(3,4)+0.401(5,6)=(3.00,\ 4.00)$; row 2 is $(3.41,\ 4.41)$; row 3, whose query matches itself best, is $(3.51,\ 4.51)$. Every output lies inside the triangle spanned by the three value vectors — attention can only mix, never invent.

With a causal mask, token 1 may only see itself ($\mathbf{z}_1=\mathbf{v}_1=(1,2)$), token 2 sees tokens 1–2 (weights $0.330$ and $0.670$, giving $(2.34,\ 3.34)$), and token 3 is unchanged. The attention-calc demo below lets you edit these matrices and watch every number update.

Math check. Each row of $\mathbf{A}$ sums to 1: $0.401+0.198+0.401 = 1.000$. ✓ And the row-1 output $(3,4)$ is the midpoint of $\mathbf{v}_1$ and $\mathbf{v}_3$ nudged toward $\mathbf{v}_2$ — exactly what weights $(0.401, 0.198, 0.401)$ should produce.

Multi-head attention

One softmax produces one weighting per query; a token cannot simultaneously attend "to the subject" and "to the previous adjective". Multi-head attention runs $h$ attentions in parallel on lower-dimensional projections and concatenates:

with $d_k = d_v = d/h$ (GPT-2: $d=768$, $h=12$, $d_k=64$). Each head sees a different 64-dimensional subspace, so heads specialize: previous-token heads, heads that track syntax, and the induction heads that copy patterns and underlie in-context learning (8.11 Interpretability: what is happening inside).

Parameters. The three projections for all heads together are $3\times h\times d\times d_k = 3d^2$, and $\mathbf{W}^O\in\mathbb{R}^{d\times d}$ adds $d^2$: about $4d^2$ per layer (biases are negligible). For $d=768$ that is $4\times 768^2\approx 2.36$M.

Cost. Computing $\mathbf{Q}\mathbf{K}^\top$ for one head is an $n\times d_k$ by $d_k\times n$ product: $O(n^2 d_k)$, so all heads cost $O(n^2 d)$ time and $O(h\,n^2)$ memory for the attention matrices, plus $O(n d^2)$ for the projections. Doubling the context quadruples attention compute — the motivation for FlashAttention and the KV cache (7.10 Efficient inference: KV cache, GQA, FlashAttention, speculative decoding, quantization) and for sub-quadratic alternatives (7.12 Beyond vanilla attention: state space models, long context and hybrids).

Minimal implementation

import torch, torch.nn.functional as F

def attention(Q, K, V, causal=False):
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / d_k ** 0.5          # (..., n, n)
    if causal:
        n = scores.size(-1)
        future = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(future, float('-inf'))
    A = F.softmax(scores, dim=-1)
    return A @ V, A

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h, self.d_k = n_heads, d_model // n_heads
        self.qkv = torch.nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = torch.nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, causal=True):                      # x: (B, n, d)
        B, n, d = x.shape
        q, k, v = self.qkv(x).split(d, dim=-1)
        heads = lambda t: t.view(B, n, self.h, self.d_k).transpose(1, 2)   # (B, h, n, d_k)
        z, _ = attention(heads(q), heads(k), heads(v), causal)
        return self.out(z.transpose(1, 2).reshape(B, n, d))

Run attention on the matrices above (as float tensors) and you will reproduce the numbers to three decimals. Note that there is nothing recurrent and nothing convolutional: every operation is a matrix product or a softmax, which is why the whole layer parallelizes across positions.

Where this shows up later

• Positions must be injected because this layer is permutation-equivariant — 7.3 Positional information: sinusoidal, learned, RoPE, ALiBi.
• The attention sub-layer sits inside a residual block with LayerNorm and an FFN — 7.4 Inside the block: residuals, LayerNorm, FFN and embeddings; its $4d^2$ parameters are a third of each layer.
• The causal mask defines decoder-only models (7.6 Decoder-only: GPT and autoregressive generation); cross-attention reuses this exact function with $\mathbf{K},\mathbf{V}$ from an encoder (7.7 Encoder–decoder: T5, BART and choosing an architecture).
• Memory and compute of the $n\times n$ matrix drive FlashAttention, GQA and the KV cache (7.10 Efficient inference: KV cache, GQA, FlashAttention, speculative decoding, quantization); attention heads are the objects interpretability studies (8.11 Interpretability: what is happening inside).