Multi-Head Attention

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) \, W^O$

where each head is a standard attention operation:

$\text{head}_i = \text{Attention}(X W_i^Q, \, X W_i^K, \, X W_i^V)$

Each head $i$ has its own projection matrices $W_i^Q$, $W_i^K$, $W_i^V$ that map from $d_{\text{model}}$ down to $d_k$ dimensions.

The per-head dimension is simply:

$d_k = d_v = d_{\text{model}} / h$

For $d_{\text{model}} = 512$: $4 \times 512^2 = 4 \times 262{,}144 = \textbf{1,048,576}$ parameters.

class MultiHeadAttention:
    def __init__(self, d_model, n_heads):
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # One big projection per Q, K, V (split into heads after)
        self.W_Q = np.random.randn(d_model, d_model) * 0.01
        self.W_K = np.random.randn(d_model, d_model) * 0.01
        self.W_V = np.random.randn(d_model, d_model) * 0.01
        self.W_O = np.random.randn(d_model, d_model) * 0.01

    def split_heads(self, x):
        """(seq_len, d_model) → (n_heads, seq_len, d_k)"""
        seq_len = x.shape[0]
        x = x.reshape(seq_len, self.n_heads, self.d_k)
        return x.transpose(1, 0, 2)

    def combine_heads(self, x):
        """(n_heads, seq_len, d_k) → (seq_len, d_model)"""
        x = x.transpose(1, 0, 2)  # (seq_len, n_heads, d_k)
        return x.reshape(x.shape[0], -1)

    def forward(self, X, mask=None):
        # 1. Project to Q, K, V (full d_model)
        Q = X @ self.W_Q
        K = X @ self.W_K
        V = X @ self.W_V

        # 2. Split into heads
        Q = self.split_heads(Q)  # (n_heads, seq_len, d_k)
        K = self.split_heads(K)
        V = self.split_heads(V)

        # 3. Scaled dot-product attention PER HEAD
        scores = np.einsum('hid,hjd->hij', Q, K) / np.sqrt(self.d_k)
        if mask is not None:
            scores = scores + mask
        weights = softmax(scores, axis=-1)

        # 4. Weighted values per head
        attn_output = np.einsum('hij,hjd->hid', weights, V)

        # 5. Combine heads + output projection
        combined = self.combine_heads(attn_output)
        return combined @ self.W_O

KV Cache and Efficient Variants

During text generation, the model produces one token at a time. Without optimization, it would recompute K and V for every previous token at each step. The KV cache fixes this, and MQA/GQA shrink it. Step through the animation to see how:

Key Equations Reference

Multi-Head Attention:

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) \, W^O$

Per-Head Attention:

$\text{head}_i = \text{Attention}(X W_i^Q, X W_i^K, X W_i^V)$

Per-Head Dimensions:

$d_k = d_v = d_{\text{model}} / h$

Key References

• Multi-head attention. Vaswani et al., “Attention Is All You Need” (2017) introduced the multi-head mechanism. Section 3.2.2 explains the motivation: “multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions.”

• What do heads actually learn? Voita et al., “Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned” (2019) identified the positional, syntactic, and rare-token head types described on this page, and showed that many heads can be pruned without hurting performance.

• Multi-Query Attention (MQA). Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need” (2019) proposed sharing a single set of K/V projections across all heads to shrink the KV cache, trading a small quality loss for a large inference speedup.

• Grouped-Query Attention (GQA). Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints” (2023) found the middle ground: group heads into clusters that share K/V, recovering most of MHA quality with most of MQA speed. Used in Llama 2 and many modern LLMs.