Multi-Head Attention and Representation Subspaces

Split d_model into h smaller subspaces so each head can run its own attention pattern in parallel.

So you've got single-head attention working. One set of $W_Q$, $W_K$, $W_V$ matrices, one attention pattern per position, one weighted sum of values. It works. But there's a problem worth tracing through.

Consider the sentence: "The cat sat on the mat because it was tired."

The word "it" needs to figure out what to attend to, and the answer is kind of everything, for different reasons. It needs to attend to "cat" for coreference (that's what "it" refers to). It needs to attend to "was" for syntax ("it" is the subject, "was" is its verb). And it probably wants some sense of "tired" for semantic coherence.

A single attention head produces one softmax distribution per query position. So "it" has to smear its attention across "cat" and "was" and "tired" and hope that the weighted average of their value vectors captures all three relationships simultaneously. That's a lot to ask of a single soft lookup.

You're basically forcing the model to compute one compromise attention pattern when what it really needs is several independent ones running in parallel. One for syntax, one for coreference, one for whatever else the model discovers is useful.

That's exactly what multi-head attention does, and it's the form Vaswani et al. introduced in the original transformer paper.

Prerequisites

This post assumes you've read Self-Attention: Q, K, V from First Principles (Post 4.1). You should be comfortable with scaled dot-product attention, softmax over score vectors, and the idea that $W_Q$, $W_K$, $W_V$ project each token into query, key, and value vectors. We'll also use basic matrix shapes and parameter counting, but nothing beyond what showed up in that post.

Splitting the representation

The idea is almost disappointingly simple. Instead of computing one attention function over the full $d_{model}$-dimensional space, you split the representation into $h$ smaller subspaces and compute attention independently in each one.

If $d_{model} = 512$ and $h = 8$, each head operates on vectors of dimension $d_k = d_{model} / h = 64$. Each head gets its own projection matrices: $W_Q^i$, $W_K^i$, $W_V^i$, each of shape $(d_{model}, d_k)$. So head $i$ computes:

Then each head runs the same scaled dot-product attention from Post 4.1:

Each head's output has shape $(n, d_k)$, a sequence of $n$ tokens, each represented in that head's $d_k$-dimensional subspace. We concatenate all $h$ heads back together and project through one final matrix $W_O$ of shape $(h \cdot d_k, \; d_{model})$:

Each head gets to learn its own attention pattern. One head might specialize in positional relationships, another in syntactic dependencies, another in something the model finds useful that we might not even have a name for. And then $W_O$ mixes everything back together.

The parameter count

When I first saw multi-head attention, I assumed it must add parameters. More heads, more matrices, more stuff. It doesn't. For a fixed $d_{model}$, the total parameter count is the same whether you use 1 head or 16.

Here's the accounting. With $h$ heads, each head has three projection matrices of shape $(d_{model}, d_k)$ where $d_k = d_{model}/h$. That's $3 \times d_{model} \times d_k$ parameters per head, times $h$ heads:

Plus $W_O$ which is $(d_{model}, d_{model})$, so $d_{model}^2$ more. Total: $4 \times d_{model}^2$.

A single head with full-dimensional Q, K, V? Three matrices of $(d_{model}, d_{model})$ plus an output projection. Also $4 \times d_{model}^2$.

Same parameter budget. What changes is the structure. Multi-head attention trades representational capacity per head (each head only sees $d_k$ dimensions) for the ability to maintain multiple independent attention patterns. You're not getting a bigger model. You're getting a differently shaped one.

The reshape trick

In the math I wrote above, each head has its own separate $W_Q^i$ matrix. You could implement it that way, with a loop over heads. Nobody does, because GPUs hate loops.

In practice, you use one big $W_Q$ matrix of shape $(d_{model}, d_{model})$ that computes all the queries at once. Then you reshape the output from $(b, n, d_{model})$ to $(b, n, h, d_k)$ and transpose to $(b, h, n, d_k)$. Now the $h$ dimension is a batch dimension, and a single batched matrix multiply handles all heads simultaneously.

This is one of those things where the math and the implementation look very different, but they're computing the exact same thing. The big $W_Q$ matrix is literally the concatenation of all the per-head $W_Q^i$ matrices side by side. The reshape just splits its output into the per-head chunks.

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.h = num_heads
        self.d_k = d_model // num_heads
 
        # One big projection for all heads, then reshape
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, n, _ = x.shape
 
        # Project: (b, n, d_model) -> (b, n, d_model)
        Q = self.W_Q(x)
        K = self.W_K(x)
        V = self.W_V(x)
 
        # Reshape into heads: (b, n, d_model) -> (b, h, n, d_k)
        Q = Q.view(b, n, self.h, self.d_k).transpose(1, 2)
        K = K.view(b, n, self.h, self.d_k).transpose(1, 2)
        V = V.view(b, n, self.h, self.d_k).transpose(1, 2)
 
        # Scaled dot-product attention (all heads at once)
        # (b, h, n, d_k) @ (b, h, d_k, n) -> (b, h, n, n)
        scores = (Q @ K.transpose(-2, -1)) / (self.d_k ** 0.5)
        attn = F.softmax(scores, dim=-1)
 
        # (b, h, n, n) @ (b, h, n, d_k) -> (b, h, n, d_k)
        out = attn @ V
 
        # Concat heads: (b, h, n, d_k) -> (b, n, d_model)
        out = out.transpose(1, 2).contiguous().view(b, n, self.d_model)
 
        # Final projection
        return self.W_O(out)

That .view(b, n, self.h, self.d_k).transpose(1, 2) pattern is everywhere in transformer code. Once you recognize it, you can read pretty much any attention implementation. The view splits the last dimension into (heads, d_k), and the transpose moves the head dimension next to the batch dimension so the matmul broadcasts across both.

A worked example

Let me trace through a tiny case to make this concrete. We have $d_{model} = 8$, $h = 2$ heads, so $d_k = 4$. Our input is four tokens: "The", "cat", "sat", "down", each represented as an 8-dimensional vector.

After projecting through $W_Q$ and reshaping, each head gets its own 4-dimensional query for each token. Same for keys and values. Head 0 might end up with queries and keys that emphasize positional information, just because of how the weights happened to initialize and train. Head 1 might end up sensitive to semantic similarity.

Head 0 computes its own attention matrix, a $4 \times 4$ grid of softmax weights. Maybe "sat" attends mostly to "cat" and "down" (nearby positions). Head 1 computes a different attention matrix. Maybe "sat" attends mostly to "cat" (the subject that did the sitting).

Each head produces a $(4, 4)$ output, four tokens each with a 4-dimensional value. We concatenate along the last dimension: $(4, 4) \| (4, 4) = (4, 8)$. Back to $d_{model} = 8$. Then we multiply by $W_O$ to get the final $(4, 8)$ output.

The key insight is that those two heads computed genuinely different attention patterns. Token "sat" is looking at different things for different reasons, and the model gets both signals instead of having to compromise.

Here's what those two heads' attention matrices might look like. Head 0 concentrates on nearby tokens (positional), while Head 1 is doing something more semantic, linking subjects to their verbs.

What heads actually learn

This is where it gets interesting. Voita et al. (2019) did a careful study of what individual heads specialize in, and found a few clear patterns.

Positional heads learn to attend to tokens at a fixed relative offset. One head might always look at the previous token. Another might attend to the token two positions back. These are doing something like n-gram modeling inside the transformer.

Syntactic heads learn to attend along dependency parse edges. A head might consistently attend from a verb to its subject, or from a pronoun to its antecedent. These are the heads that would help with our "it" / "cat" example.

Rare-token heads light up disproportionately for infrequent words, possibly routing more computation toward tokens that are harder to predict.

Michel et al. (2019) showed you can prune a huge number of heads at test time with minimal quality loss. In some models, removing 60% or more of heads barely moved the needle on downstream tasks. The specialized heads above are doing the heavy lifting. Many of the other heads seem to be partially redundant or learning soft variants of the same pattern.

This doesn't mean the extra heads are useless during training. Having many heads gives the optimizer more paths to explore, and some of that redundancy probably helps training stability. But at inference time, the model is over-provisioned. That observation directly motivates the KV-cache optimizations we'll get to in a moment.

Remember, the concatenated multi-head output gets projected through $W_O$ and then added back to the residual stream. Each layer's attention block is writing an update into the same shared vector. Click through the stages below to see how these additive updates accumulate.

Sharing K and V: GQA and MQA

Multi-head attention as described above means every head has its own $W_K$ and $W_V$ projections. At inference time, when you're caching key-value pairs to avoid recomputation (the KV cache), that means storing separate K and V tensors for every head, at every layer, for every token in the context. For a model with 32 heads and a long context, this gets expensive fast.

Shazeer (2019) proposed Multi-Query Attention (MQA): keep separate Q projections for each head, but share a single K and V projection across all heads. Every head computes its own queries but looks up the same keys and values. This cuts the KV cache by a factor of $h$ and speeds up inference dramatically, at the cost of some quality.

Ainslie et al. (2023) proposed a middle ground: Grouped-Query Attention (GQA). Instead of one KV pair for all heads, or one KV pair per head, you divide the heads into $g$ groups. Each group shares one K and V projection. If you have $h = 32$ heads and $g = 8$ groups, each group of 4 heads shares one KV set. The KV cache shrinks by a factor of $h/g = 4$.

The reason this matters so much is tied to how autoregressive decoding works. During generation, you're producing one token at a time, and each new token needs to attend to all previous tokens. Without the KV cache, you'd recompute attention over the entire sequence at every step. With the cache, you store the K and V projections and only compute the new token's query. But the cache grows linearly with sequence length and with the number of KV heads. GQA directly attacks that second factor. More on this in Post 5.4.

Misconceptions

"More heads means more parameters." No. For a fixed $d_{model}$, changing the number of heads doesn't change the parameter count. It changes the structure: each head is smaller, but there are more of them. Same total budget. What you actually buy with more heads is more parallel relation types, not more capacity.

"Each head dimension can be anything." It can't, at least not in standard multi-head attention. The whole reshape trick only works because $d_k = d_{model} / h$, so concatenating $h$ heads of size $d_k$ gives you back something of size $d_{model}$. If you wanted heads with arbitrary dimension you'd have to give up the clean reshape and pay for separate matmuls. Some research variants do that. The vast majority of production attention does not.

"Each head has a clear, interpretable role." Some do. Voita et al. found clear positional and syntactic heads, and induction heads are a known phenomenon in mechanistic interpretability work. But many heads don't have a clean interpretation, and a lot of them can be pruned without meaningful degradation. The picture is messier than the clean "one head per linguistic phenomenon" story.

"Multi-head attention is fundamentally different from single-head attention." It's the same computation, restructured. You could replicate multi-head attention with a single head that has a very specific block-diagonal structure in its weight matrices. Multi-head just makes that structure explicit and easy for the optimizer to find.

What's next

The attention mechanism as described so far lets every token attend to every other token. For language generation that's a problem, since a model predicting the next word shouldn't be allowed to peek at future words during training. The next post covers causal masking, how a single upper-triangular matrix enforces that constraint.

Additional reading (and watching)

• Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017. Introduced multi-head attention as "h parallel attention functions."
• Voita, E., et al. (2019). Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, and the Rest Can Be Pruned. ACL 2019. Careful classification of head roles in NMT transformers.
• Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. The MQA paper.
• Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. GQA, the interpolation between MHA and MQA used by most modern open-weights LLMs.
• Michel, P., Levy, O., & Neubig, G. (2019). Are Sixteen Heads Really Better than One?. NeurIPS 2019. Shows large fractions of heads can be pruned at test time with minimal loss.
• Elhage, N., et al. (2021). A Mathematical Framework for Transformer Circuits. Anthropic. The framework that identified induction heads and other cross-layer compositions of specific attention heads.
• Clark, K., et al. (2019). What Does BERT Look At? An Analysis of BERT's Attention. ACL 2019 Workshop. Probed individual attention heads in BERT and found heads attending to syntactic relations, coreference, and separator tokens.
• Liu, A., et al. (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. Introduces Multi-head Latent Attention (MLA), compressing KV heads into a low-rank latent space as a next step beyond GQA.