Multi-head attention: from equations to implementation

Start with one weighted sum. Two positions supply values [2, 0] and [0, 4]. With attention weights 0.75 and 0.25, the output is [1.5, 1]. The rest of attention explains where those weights come from, which positions are masked, and how several such computations are combined.

The core computation: match and aggregate information

Each position carries its own question (query) and asks every position's index (key). The better they match, the more it takes from that position's content (value). What comes back is a weighted average.

Each head uses its own projections to compute a weighted result, and the results are concatenated. Some heads may learn patterns such as proximity or coreference, but those roles are neither assigned in advance nor guaranteed to be individually interpretable.

What the attention matrix represents

Start with self-attention: $Q$, $K$, and $V$ are three learned linear projections of the same input $X$. Cross-attention can instead obtain K/V from another sequence; see the vanilla Transformer.

$$ \begin{gathered} Q=XW_Q\\ K=XW_K\\ V=XW_V \end{gathered} $$

They are not three individual dimensions with assigned meanings. This example sets $d_v=d_k$: with $T$ tokens, $Q,K,V$ each have shape $(T,d_k)$, although V can have a different width in general. Their computational roles differ: Q/K compute the weights; V supplies the vectors to aggregate.

The full computation has only the following five steps.

1. Compute one scalar score for every pair of tokens

$$ S=\frac{QK^\top}{\sqrt{d_k}},\qquad S_{ij}=\frac{\mathbf q_i^\top\mathbf k_j}{\sqrt{d_k}} $$

$S$ has shape $(T,T)$; this is the attention score matrix. Row $i$ means “where the current position $i$ wants to fetch information from”; column $j$ means “candidate source position $j$”. Each cell $S_{ij}$ is just one scalar; $V$ has not been used yet.

2. A decoder-only model first blocks future positions

When GPT predicts the next token at position $i$, it may use only position $i$ and everything to its left. The causal mask is written

$$ M_{ij}= \begin{cases} 0, & j\le i\\ -\infty, & j>i \end{cases} $$

For three tokens the score matrix becomes:

$$ S+M= \begin{bmatrix} s_{11} & -\infty & -\infty\\ s_{21} & s_{22} & -\infty\\ s_{31} & s_{32} & s_{33} \end{bmatrix} $$

The mask does not restrict the model to the single previous token. It lets a position see itself and every earlier token, while hiding the future. Training computes all positions in parallel; without the upper-right triangle masked, earlier positions could read the correct answers that come later, and the training objective would leak information. A bidirectional encoder normally has no causal mask, though it may still use a padding mask.

3. Apply softmax to every row

$$ \begin{gathered} A=\operatorname{softmax}_{j}(S+M) \end{gathered} $$

Softmax runs over the column index $j$, so every row satisfies

$$ \sum_j A_{ij}=1 $$

Because $e^{-\infty}=0$, masked positions receive zero weight. This assumes at least one allowed position per row, with finite scores at allowed positions. An entirely masked row has no valid normalization denominator; ordinary softmax cannot be expected to return zeros automatically. $A$ is the attention weight matrix: row $i$ specifies how much to take from each allowed position.

4. Use that row of weights to aggregate all values

$$ O=AV,\qquad \mathbf o_i=\sum_j A_{ij}\mathbf v_j $$

So $QK^\top$ only decides “what the weights are”; what is actually fetched and mixed is $V$. One cell of the attention matrix is a scalar, while the output $\mathbf o_i$ is a $d_k$-dimensional vector.

5. Every head computes its own set, then they merge

Each head has its own projections, score matrix, and attention weights, so heads can learn different ways of matching. The outputs of all heads are concatenated first, then projected through $W_O$ back to the model dimension.

The whole data flow compresses into one line:

$$ \begin{gathered} S=QK^\top/\sqrt{d_k},\\ A=\operatorname{softmax}_{\rm row}(S+M),\\ O=AV. \end{gathered} $$

That is: three projections produce Q/K/V; Q and K produce pairwise weights between tokens; the mask removes forbidden information paths; row-wise softmax normalizes; and those weights finally take a weighted sum of V.

The softmax in MHA does provide nonlinearity, but it mainly decides “whom to communicate with” along the token dimension. The FFN activation instead performs a nonlinear transformation along each token's feature dimension; the two play different roles.

What a single head computes

$$ \begin{gathered} \text{Attention}(Q,K,V)\\ = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V \end{gathered} $$

Broken down to a single query $\mathbf{q}_i$:

$$ \begin{gathered} s_{ij}=\mathbf q_i^\top\mathbf k_j/\sqrt{d_k},\\ \alpha_{ij}=\frac{\exp(s_{ij})}{\sum_{j\prime}\exp(s_{ij\prime})},\\ \mathbf o_i=\sum_j\alpha_{ij}\mathbf v_j. \end{gathered} $$

Without attention dropout, the weights are nonnegative and sum to one, so one head's output before output projection is a convex combination of its values. The opening [1.5, 1] example illustrates this. That does not describe the entire Transformer block: dropout, output projection, and residual connections change the result. The weights also depend on the input; this is not a fixed average.

Why divide by $\sqrt{d_k}$

Start with what scaling changes. Suppose a head has dimension 64 and one query gives dot-product scores [8, −8] for two keys:

Without scaling, almost all weight goes to the first position. After scaling, the second still contributes. Uniform weights are not the goal. We want to avoid making attention extremely sharp simply by increasing the head dimension.

Dividing by a positive constant changes the gaps between scores and therefore the weights. Subtracting a common constant preserves the gaps; it makes exponentiation safer without changing the mathematical result. Even [1000, 1000] should give [0.5, 0.5]. Large absolute values alone do not imply saturation.

Use the head dimension $d_k$, not the full model width $d_\text{model}$. If a model also uses QK normalization or a learned temperature, follow its actual formula. Scaling addresses one source of trouble; it does not guarantee stable training.

Writing backward: the other position matters too

Take two equal attention weights, $\alpha=(0.5,0.5)$, and let the loss be the first weight. Raising the first score increases the loss. Raising the second steals weight from the first and lowers the loss. The score gradients are therefore $(0.25,-0.25)$, not $(0.25,0)$.

import numpy as np

def softmax_backward(probabilities, upstream):
    weighted_mean = np.sum(probabilities * upstream, axis=-1, keepdims=True)
    return probabilities * (upstream - weighted_mean)

probabilities = np.array([0.5, 0.5])
assert np.allclose(softmax_backward(probabilities, np.array([1.0, 0.0])),
                   [0.25, -0.25])
assert np.allclose(softmax_backward(probabilities, np.ones(2)), [0.0, 0.0])

Why use multiple heads: the goal is not more dimensions

$$ \begin{gathered} \text{head}_i=\\ \operatorname{Attention}(XW_i^Q,XW_i^K,XW_i^V),\\ \operatorname{MultiHead}=\\ \operatorname{Concat}(\text{head}_1,\ldots,\text{head}_h)W^O. \end{gathered} $$

One head uses one set of attention weights for all value channels. Splitting a fixed width into $h$ heads gives several such relations at similar projection and matmul cost, although storing separate attention matrices and executing kernels can add overhead. The three points below explain the tradeoff.

1. The core of multi-head: multiple sets of attention relations

Suppose $d_{\text{model}}=512$. One full-width attention head computes

$$ \begin{gathered} A=\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{512}}\right)\\ O=AV. \end{gathered} $$

The key limitation is not that “512 dimensions are not enough”. It is that every value channel shares the same attention matrix $A$. To resolve “she” in “Xiao Ming gave the book to Xiao Hong because she loves reading”, the model may need to track coreference, syntactic dependency, semantic role, and local adjacency at the same time; a single head must compress all of these relations into one distribution.

Multi-head lets head $i$ learn its own projections and weights:

$$ \begin{gathered} Q_i=XW_i^Q\\ K_i=XW_i^K\\ V_i=XW_i^V, \end{gathered} $$

$$ A_i=\operatorname{softmax}\!\left(\frac{Q_iK_i^\top}{\sqrt{d_k}}\right). $$

The model therefore obtains $A_1,\ldots,A_h$: several ways to read the sequence. Some heads may lean toward coreference, others toward locality or syntax, but those jobs are not assigned by hand and can overlap. The more precise conclusion is that different features can use different attention weights instead of all sharing one distribution.

2. Splitting dimensions controls the budget

The original model uses $d_{\text{model}}=512$ and $h=8$, usually with

$$ d_k=d_v=\frac{512}{8}=64, $$

so $8\times64=512$. If all eight heads kept the full 512 dimensions, parameters and compute would grow substantially; splitting the total width is what yields eight sets of relations at roughly the budget of a single head.

A full-width single-head projection has

$$ W_Q,W_K,W_V\in\mathbb{R}^{512\times512}. $$

Each multi-head projection is $512\times64$, and eight of them still total

$$ 8\times(512\times64)=512\times512. $$

Standard MHA therefore has about $4d_{\text{model}}^2$ parameters across Q, K, V, and the output projection, independent of the head count itself. Code also usually performs one large projection and then reshapes to (B, H, T, d_head) rather than running eight small models in sequence — identical math, a single GEMM.

3. Expressivity is not a generalization guarantee

What multi-head attention increases first is expressivity: it lets several token relations, matching functions, and contextual summaries coexist. Better representations sometimes improve performance on unseen data, but “it uses multiple heads” does not by itself imply better generalization.

With too many heads, each head may become too narrow, several heads may duplicate one another, parameters may be used inefficiently, and the model may even overfit. In practice some heads can often be pruned with almost no loss in quality. Therefore:

Multi-head attention obtains multiple relations at similar cost; it does not enlarge width for its own sake.

“Different representation subspaces” should not be over-interpreted either. Each head does have independent parameters and can therefore learn a different matching function; but for an ambiguous word like “apple,” “one head must handle the company sense and another must handle the fruit sense” is neither designed in advance nor guaranteed to be interpretable. A claim that a head performs a function needs a causal test, such as patching specific heads on the IOI task; see circuits.

The six places a from-scratch implementation goes wrong

Check these six details before trying to diagnose a mistake from the final loss:

Why #1 must be that way. The projection's output is (B, T, d_model), where the d_model dimension holds the $h$ heads laid end to end. So split the last dimension into (H, dh) first, then move H forward. A direct view(B,H,T,dh) slices across the time dimension instead, so every "head" is a pile of fragments from unrelated positions — right shape, wrong numbers, and no error raised.

Check the mask with two scores. Softmax of [0, 0] is [0.5, 0.5]. Allowing only the first position gives [1, 0] when masked before softmax. Zeroing the second weight afterwards gives [0.5, 0], shrinking the output. Dividing by the remaining weight sum recovers [1, 0] mathematically, but that sum can underflow in finite precision. Masking before softmax avoids this issue. The teaching code rejects entirely masked rows instead of silently producing NaNs.

Experiment: check the handwritten version against the library

import numpy as np

def softmax(logits, axis=-1):
    values = np.asarray(logits)
    if values.dtype.kind not in "biuf":
        raise ValueError("Expected real logits")
    converted = values.astype(np.result_type(values.dtype, np.float32), copy=False)
    if values.dtype.kind in "iu" and any(
            int(original) != int(floating) for original, floating in zip(values.flat, converted.flat)):
        raise ValueError("Integer logits must be exactly representable in the computation dtype")
    values = converted
    if np.any(np.isnan(values) | np.isposinf(values)):
        raise ValueError("Logits must be finite or -inf for masked positions")
    if np.any(np.all(np.isneginf(values), axis=axis)):
        raise ValueError("Each query must have at least one allowed key")
    shifted = values - np.max(values, axis=axis, keepdims=True)
    weights = np.exp(shifted)
    return weights / np.sum(weights, axis=axis, keepdims=True)


def attention(q, k, v, mask=None):
    """q,k,v: (B, H, T, d_head)   mask: True = blocked"""
    d = q.shape[-1]
    scores = q @ k.swapaxes(-2, -1) / np.sqrt(d)      # (B, H, Tq, Tk)
    if mask is not None:
        scores = np.where(mask, -np.inf, scores)      # before the softmax
    w = softmax(scores, axis=-1)
    return w @ v, w

class MultiHeadAttention:
    def __init__(self, d_model, n_head, seed=0):
        assert d_model % n_head == 0
        rng = np.random.default_rng(seed)
        s = 1.0 / np.sqrt(d_model)
        self.n_head, self.d_head = n_head, d_model // n_head
        self.Wq, self.Wk, self.Wv, self.Wo = (
            rng.uniform(-s, s, (d_model, d_model)) for _ in range(4))

    def split(self, x):                               # (B,T,C) -> (B,H,T,dh)
        b, t, _ = x.shape
        return x.reshape(b, t, self.n_head, self.d_head).transpose(0, 2, 1, 3)

    def __call__(self, x, mask=None):
        q, k, v = self.split(x @ self.Wq), self.split(x @ self.Wk), self.split(x @ self.Wv)
        out, w = attention(q, k, v, mask)
        b, h, t, dh = out.shape
        out = out.transpose(0, 2, 1, 3).reshape(b, t, h * dh)   # merge heads
        return out @ self.Wo, w

def causal_mask(t):
    return np.triu(np.ones((t, t), dtype=bool), k=1)  # strict upper triangle = the future

import math, torch, torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_head, bias=False):
        super().__init__()
        assert d_model % n_head == 0
        self.n_head, self.d_head = n_head, d_model // n_head
        self.wq = nn.Linear(d_model, d_model, bias=bias)
        self.wk = nn.Linear(d_model, d_model, bias=bias)
        self.wv = nn.Linear(d_model, d_model, bias=bias)
        self.wo = nn.Linear(d_model, d_model, bias=bias)

    def forward(self, x, mask=None):
        B, T, C = x.shape
        split = lambda p: p(x).view(B, T, self.n_head, self.d_head).transpose(1, 2)
        q, k, v = split(self.wq), split(self.wk), split(self.wv)

        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head)
        if mask is not None:
            att = att.masked_fill(mask, float("-inf"))
        if torch.isneginf(att).all(dim=-1).any():
            raise ValueError("Each query must have at least one allowed key")
        att = att.softmax(dim=-1)

        y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)   # merge heads
        return self.wo(y), att

def causal_mask(t, device=None):
    return torch.triu(torch.ones(t, t, dtype=torch.bool, device=device), diagonal=1)

import torch.nn.functional as F

output = F.scaled_dot_product_attention(
    q, k, v,
    attn_mask=None if mask is None else ~mask,
    dropout_p=0.0,
)

../code/attention_torch.py runs four implementations on the same set of weights and checks them against one another:

  from-scratch vs F.scaled_dot_product_attention : 1.19e-07
  from-scratch vs nn.MultiheadAttention          : 1.19e-07
  from-scratch torch vs pure NumPy               : 1.19e-07
  every attention row sums to 1                  : 1.19e-07
  weight on any future position                  : 0.00e+00

Getting the from-scratch version to match nn.MultiheadAttention matters more than merely "it runs". When they disagree, locating the difference is itself an effective debugging exercise. nn.MultiheadAttention stores $W_Q, W_K, W_V$ as one concatenated in_proj_weight, while nn.Linear stores $(out, in)$, so moving weights to NumPy needs a transpose.

One matching output is a useful start, not a complete migration test. Check these next:

The script's library_version creates a separate module and copies the parameters for comparison. Its parameter gradients do not flow back to the original mha. Tests cover bias, float32 / float64, per-head masks, and gradients across the three PyTorch paths; the NumPy comparison remains a small bias-free fixture. Dropout is off throughout. These are correctness checks, not speed or memory benchmarks. See the PyTorch 2.14 MultiheadAttention API for its conventions.

Where does the causal mask start with a KV cache?

The earlier mask was square: five queries against five keys. Now suppose positions 0, 1, and 2 are cached, and we append positions 3 and 4 together. Only two queries are new, but they need access to five keys.

The rule is still “do not read the future.” The difference is that query row 0 now represents position 3, not position 0 in the full sequence.

To make the error visible, set all Q/K scores to zero and use values [10, 20, 30, 40, 50]. Attention becomes a plain average of the allowed values.

As of PyTorch 2.14 (documentation checked October 2026), F.scaled_dot_product_attention(..., is_causal=True) uses upper-left alignment for a non-square mask. It does not infer that the queries belong at the end of the KV sequence. Here, construct attn_mask from actual positions and leave is_causal off. This is an API-specific convention, not a claim that every attention kernel uses the same alignment. SDPA documentation

import torch
import torch.nn.functional as F

def cached_attention_example():
    query = torch.zeros(1, 1, 2, 2, dtype=torch.float64)
    key = torch.zeros(1, 1, 5, 2, dtype=torch.float64)
    value = torch.tensor([10., 20., 30., 40., 50.], dtype=torch.float64)
    value = value.view(1, 1, 5, 1)
    query_positions = torch.tensor([3, 4])
    key_positions = torch.arange(5)
    allowed = key_positions[None, :] <= query_positions[:, None]
    aligned = F.scaled_dot_product_attention(
        query, key, value, attn_mask=allowed, dropout_p=0.0,
    )
    upper_left = F.scaled_dot_product_attention(
        query, key, value, is_causal=True, dropout_p=0.0,
    )
    return aligned, upper_left, allowed

aligned, upper_left, allowed = cached_attention_example()
torch.testing.assert_close(aligned.flatten(), torch.tensor([25., 30.], dtype=torch.float64))
torch.testing.assert_close(upper_left.flatten(), torch.tensor([10., 15.], dtype=torch.float64))
print(aligned.flatten().tolist(), upper_left.flatten().tolist())

If only position 4 is new and all five supplied keys are valid, there is no future key to block, so no causal mask is needed. Padding, sliding windows, or different sequence lengths in a batch can invalidate that shortcut. Also, functional SDPA takes dropout from dropout_p; calling eval() on an enclosing module does not automatically change that argument to zero.

This example isolates the mask, not an entire decoder. A full model also needs correct RoPE positions, cache identity, and matching logits with and without caching. Continue in KV cache.

Common interview questions

Self-check

Next

Without positional information or an order-dependent mask, self-attention is permutation-equivariant: reorder the inputs and the outputs follow that order, rather than staying unchanged. A causal mask already distinguishes earlier from later positions; positional encodings provide more direct position information. Next, see how residual connections and normalization help training, then assemble the block in the vanilla Transformer.

The equations start from Attention Is All You Need §3.2; the tiled memory optimization is explained in FlashAttention.

Quick learning: multi-head attention in one sentence and its boundary conditions