Multi-Head Attention and Representation Subspaces

Project each token into h query, key, and value spaces, then compute a separate attention pattern in each head.

A single attention head uses one set of $W_Q$, $W_K$, $W_V$ matrices to produce one attention pattern per position. What if a token needs different patterns for different relationships?

For an unmasked, bidirectional example, 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, combining these relationships into one weighted average of value vectors. Multiple heads give the model separate distributions to learn, without assigning each head a linguistic role in advance.

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. The examples also use matrix shapes and parameter counting.

Splitting the representation

Each head projects the full input vector into its own query, key, and value spaces. Attention then runs independently in those projected spaces. The split is in the projection outputs, not a partition that prevents a head from reading some input coordinates.

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})$:

The parameter count

More heads need not add projection parameters. In standard multi-head attention, keeping $h d_k=d_{model}$ fixes the total width while changing how many independent attention patterns the block computes.

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$.

Within the same parameter budget, multi-head attention trades representational capacity per head (each head produces $d_k$ projected features from the full input) for multiple independent attention patterns.

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. A fused projection and batched matmul usually use the GPU more efficiently than a Python loop over heads.

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.

The big $W_Q$ matrix is 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. 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

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 heatmaps use hand-constructed weights to illustrate these two patterns.

What heads actually learn

Voita et al. (2019) studied what individual heads specialize in and found a few recurring 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 putting more attention weight toward tokens that are harder to predict.

Michel et al. (2019) pruned heads from translation and fine-tuned BERT models and found substantial redundancy in some layers and tasks. Cross-attention in their translation model was more sensitive. A recognizable attention pattern alone does not tell us whether a head is necessary; that requires an intervention and evaluation.

Pruning heads tests whether an already-trained model needs them for a particular task. Sharing K and V projections is a separate architectural choice that reduces cache storage during generation.

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. At fixed head width and cache precision, this cuts the KV cache by a factor of $h$. The paper reports faster decoding with a small quality tradeoff in its experiments; the latency benefit depends on hardware and batch shape.

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$.

At fixed head width and precision, the cache grows with both sequence length and KV-head count. GQA reduces the latter; Post 5.4 derives the resulting memory budget.

Misconceptions

"More heads means more parameters." With standard multi-head attention and $h d_k=d_{model}$, the Q/K/V/O projection budget stays at $4d_{model}^2$ (ignoring biases). If head width stays fixed while head count grows, the parameter count grows too.

"Head width must equal $d_{model}/h$." That's a common choice, not a mathematical requirement. A projection can produce $h d_k$ features, reshape them into $h$ heads, and map the concatenated outputs back through $W_O$. The fixed parameter count above assumes $h d_k=d_{model}$.

"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.

"One wide head can reproduce several heads just by changing its weight matrices." In general it cannot. Each head computes its own softmax distribution over positions. One head has only one such distribution per query; block-diagonal projections do not supply the missing independent softmax operations.

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 in models including Llama 2 70B.
• 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.