What is multi-head attention and why use multiple heads?

Multi-head attention is the mechanism that lets a transformer model attend to different representation subspaces simultaneously, computing multiple self-attention operations in parallel and concatenating their results. Each head learns its own query, key, and value projections, enabling the model to capture distinct relational patterns — syntactic, semantic, positional — within the same layer. This parallel design is what gives transformers their expressive power over single-head attention or recurrent architectures.

How multi-head attention works

The computation follows a straightforward pattern: split the model dimension into h heads, project queries, keys, and values independently per head, run scaled dot-product attention on each, then concatenate and project back to the model dimension.

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, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        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)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)

        # Project and split into heads: (batch, seq, d_model) -> (batch, num_heads, seq, d_k)
        q = self.w_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = self.w_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.w_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # Scaled dot-product attention per head
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)

        # Apply attention to values
        out = torch.matmul(attn, v)  # (batch, num_heads, seq, d_k)

        # Concatenate heads: (batch, num_heads, seq, d_k) -> (batch, seq, d_model)
        out = out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        return self.w_o(out)

The key insight is in the view and transpose operations. By reshaping (batch, seq, d_model) into (batch, num_heads, seq, d_k), we create independent attention computations that share no parameters across heads. Each head sees a different linear projection of the same input, so head 0 might learn to track subject-verb agreement while head 3 tracks coreference chains — all in the same forward pass.

The output projection w_o then mixes information across heads. Without it, heads would remain isolated; with it, the model can compose multi-head representations into richer features for the next layer.

Why multiple heads matter

A single attention head computes one weighted sum per position. That weighted sum is a convex combination of value vectors, which means it can only express a single “perspective” on the context. Multiple heads give you multiple perspectives simultaneously.

Consider what a single head can represent: it produces attention weights that sum to 1, so the output is always an average of values. If position i needs to attend to position j for syntactic reasons and position k for semantic reasons, a single head must compromise — it can’t fully attend to both without diluting each. Multi-head attention removes this bottleneck.

Empirically, different heads specialize. Research using attention visualization and probing classifiers has shown consistent patterns: some heads attend to nearby tokens (local syntax), others to distant tokens (long-range dependencies), some to specific positional offsets (relative position), and some broadly (global context). This specialization emerges from training, not architectural constraints.

The parameter count stays roughly constant. A single head with dimension d_model has three projection matrices of size d_model × d_model. h heads with dimension d_model/h have three projection matrices of size d_model × d_model each — same total parameters, but factorized into independent subspaces. The output projection adds another d_model × d_model, which is negligible relative to feed-forward layers.

Concrete example: disambiguating “bank”

Take the sentence: “The bank approved the loan after the river bank flooded.”

Token positions 1 and 8 both contain “bank” but mean different things. A single head must choose one attention pattern. With multiple heads, the model can route:

Head 0: attends “bank” (position 1) → “approved”, “loan” (financial sense)
Head 2: attends “bank” (position 8) → “river”, “flooded” (geographical sense)
Head 5: attends both “bank” tokens → each other (coreference or contrast)

The feed-forward network after attention receives the concatenated result and can cleanly separate these senses because they occupy different subspaces in the residual stream. This is why ablation studies show performance drops sharply when you reduce head count below a threshold — you lose the capacity to maintain simultaneous, distinct relational representations.

Head count and dimension trade-offs

The standard configuration uses d_k = d_v = d_model / num_heads. Common settings:

Model	d_model	num_heads	d_k
BERT-base	768	12	64
GPT-2 small	768	12	64
GPT-3 175B	12288	96	128
Llama-2 7B	4096	32	128

Smaller d_k means each head has less capacity but you get more heads. Larger d_k means richer per-head representations but fewer perspectives. The 64-dimension-per-head heuristic (from the original Transformer paper) works well across scales, but newer models like Llama increase it to 128. There’s no universal optimum — it interacts with depth, dataset, and training compute.

Grouped-query attention (GQA) and multi-query attention (MQA) modify this structure for inference efficiency. GQA uses fewer key/value heads than query heads (e.g., 32 query heads, 8 KV heads), reducing KV cache memory during generation while retaining most quality. This is now standard in production models like Llama-2 70B and PaLM-2.

Common misconceptions

Misconception: “More heads always help.” Past a point, heads become redundant. Ablation studies on BERT show that removing 30-40% of heads at inference time often has minimal impact on downstream tasks. Many heads learn similar patterns or attend to [CLS] / padding tokens. The marginal utility of each additional head diminishes.

Misconception: “Heads are interpretable by design.” While some heads show clear patterns (e.g., “attends to next token”, “attends to verb”), many are polysemantic or distributed. Don’t assume head 7 is the “coreference head” — it may participate in coreference alongside heads 3 and 11. Interpretability requires careful probing, not inspection.

Misconception: “Multi-head attention is just ensemble averaging.” The output projection w_o learns to combine heads non-linearly (via the subsequent feed-forward network and residual connections). It’s not a simple average — the model learns which heads to trust for which contexts. The residual stream carries the combined signal forward, and later layers can route around unhelpful heads.

Misconception: “Attention weights explain model decisions.” Attention weights show where information flows, not necessarily what the model decides. A head might attend heavily to a token but the value projection could zero out that dimension. Conversely, a low-weight token might carry a critical feature in its value vector. Use attention weights as hints, not explanations.

Implementation notes for production

When implementing multi-head attention for serving, two optimizations matter:

Fused kernels: Use FlashAttention or xFormers memory-efficient attention. The naive implementation materializes the full (batch, num_heads, seq, seq) attention matrix, which is O(seq²) memory. FlashAttention tiles the computation to stay in SRAM, enabling longer contexts.
KV cache layout: For autoregressive generation, store keys and values as (batch, num_kv_heads, seq_len, d_k) — not interleaved with queries. This matches the access pattern during incremental decoding and avoids transposes. With GQA, num_kv_heads < num_query_heads, so you repeat KV heads across query heads during the attention computation.

# KV cache update during generation (simplified)
def update_kv_cache(kv_cache, new_k, new_v, layer_idx):
    # kv_cache: list of (k_cache, v_cache) per layer
    # new_k, new_v: (batch, num_kv_heads, 1, d_k)
    k_cache, v_cache = kv_cache[layer_idx]
    k_cache = torch.cat([k_cache, new_k], dim=2)
    v_cache = torch.cat([v_cache, new_v], dim=2)
    kv_cache[layer_idx] = (k_cache, v_cache)
    return kv_cache

If you’re routing requests across multiple model variants or providers, the attention implementation details are abstracted behind the inference API. The gateway handles model-specific KV cache formats and attention backends so your application code stays clean.

Summary

Multi-head attention factorizes a single large attention operation into multiple smaller, independent operations that run in parallel and combine through a learned output projection. This gives the model simultaneous access to multiple relational perspectives — syntactic, semantic, structural — without increasing parameter count proportionally. The design is simple: project, split, attend, concatenate, project. But the emergent specialization across heads is what makes deep transformers effective at modeling language structure at scale.

When you’re debugging attention patterns or designing model architectures, remember: heads are not modules with fixed roles. They’re learned subspaces that the rest of the network composes. Treat them as a distributed representation, not a collection of interpretable experts.