Multi-Query Attention › Multi-Query Attention: Extreme Sharing

Multi-Query Attention, introduced by Noam Shazeer in 2019, addresses the KV cache problem through aggressive parameter sharing. The core insight is simple: instead of having $h$ separate key and value projections (one per head), use just one shared key and one shared value projection across all heads.

To truly understand Multi-Query Attention, we need to first build intuition about what multi-head attention does and why each component exists. Think of attention as a retrieval system: queries ask questions, keys advertise what information is available, and values contain the actual content to retrieve. Multi-head attention runs multiple such retrieval systems in parallel, each learning to find different kinds of relationships.

Step 1: The Standard Multi-Head Attention BaselineLink Copied

In standard multi-head attention, each head $i$ transforms the input through its own learned projection matrices. Given an input sequence $\mathbf{X}$, we compute the three head-specific projection outputs:

where:

• $\mathbf{X} \in \mathbb{R}^{n \times d_{\text{model}}}$: the input sequence with $n$ tokens, each represented as a $d_{\text{model}}$-dimensional vector
• $\mathbf{W}^Q_i \in \mathbb{R}^{d_{\text{model}} \times d_k}$: the query projection matrix for head $i$
• $\mathbf{W}^K_i \in \mathbb{R}^{d_{\text{model}} \times d_k}$: the key projection matrix for head $i$
• $\mathbf{W}^V_i \in \mathbb{R}^{d_{\text{model}} \times d_v}$: the value projection matrix for head $i$
• $d_k$: the dimension of queries and keys per head (typically $d_{\text{model}} / h$)
• $d_v$: the dimension of values per head (often equal to $d_k$)

The subscript $i$ on each projection matrix is the key detail here. It tells us that every head learns its own transformation. Head 1 might learn to project tokens in a way that emphasizes syntactic features, while head 2 might emphasize semantic similarity. This specialization is what makes multi-head attention so powerful: the model can simultaneously attend to different types of relationships.

But this flexibility comes at a cost. During inference, we must cache the keys and values for each head separately. With 32 heads, we store 32 different $\mathbf{K}_i$ matrices and 32 different $\mathbf{V}_i$ matrices. This is where the memory bottleneck arises.

Step 2: The MQA Insight: Sharing Keys and ValuesLink Copied

MQA asks a provocative question: do we really need different keys and values for each head? The hypothesis is that while different heads benefit from asking different questions (different queries), they might be able to share a common "index" (keys) and "content store" (values).

Mathematically, MQA removes the subscript from the key and value projections:

Notice the key change:

• $\mathbf{W}^Q_i$: still has subscript $i$, meaning each head retains its own query projection
• $\mathbf{W}^K$ and $\mathbf{W}^V$: no subscript, meaning a single projection shared across all heads

The resulting matrices have dimensions:

• $\mathbf{Q}_i \in \mathbb{R}^{n \times d_k}$: unique query matrix for head $i$
• $\mathbf{K} \in \mathbb{R}^{n \times d_k}$: one shared key matrix for all heads
• $\mathbf{V} \in \mathbb{R}^{n \times d_v}$: one shared value matrix for all heads

This is a simple but effective change. Different heads can still "ask different questions" because each projects the input into its own query space. But they all look up answers from the same key-value store. It's like having 32 researchers who each have their own research questions but share access to a single library catalog and book collection.

Step 3: Computing Attention with Shared KVLink Copied

Each head now computes attention using its unique queries against the shared keys and values:

Let's trace through what happens in this formula:

1. Compute attention scores: $\mathbf{Q}_i \mathbf{K}^T$ produces an $n \times n$ matrix where entry $(p, q)$ measures how much token at position $p$ should attend to position $q$. Because $\mathbf{Q}_i$ is unique to head $i$, each head produces different attention patterns.

2. Scale the scores: Dividing by $\sqrt{d_k}$ prevents the dot products from growing too large as dimension increases. Without scaling, the softmax would produce very peaked distributions, making gradients vanish during training.

3. Normalize to weights: The softmax converts each row of scores into a probability distribution that sums to 1. These are the attention weights.

4. Aggregate values: Multiplying by $\mathbf{V}$ computes a weighted sum of value vectors. All heads aggregate from the same $\mathbf{V}$, but with different weights.

The key insight is that expressiveness comes from two sources: the diversity of attention patterns (which heads preserve through unique queries) and the diversity of retrieved content (which MQA sacrifices through shared values). Empirically, the diversity of attention patterns matters more than the diversity of values for most tasks.

To illustrate this, let's visualize how different query heads can produce distinct attention patterns even when sharing the same keys and values.

Each head learns a different way to "ask questions" of the shared key-value store. Head 1 might focus on recent context for local coherence, Head 2 might anchor to the beginning of the sequence, Head 3 might aggregate broadly, and Head 4 might pick up periodic patterns. The shared keys and values don't limit this diversity; they simply mean all heads retrieve from the same "library."

Step 4: Combining Head OutputsLink Copied

Finally, we concatenate all head outputs and project back to the model dimension:

where:

• $\text{head}_i \in \mathbb{R}^{n \times d_v}$: the output from attention head $i$, containing each head's weighted aggregation
• $\text{Concat}(\cdot) \in \mathbb{R}^{n \times (h \cdot d_v)}$: concatenation along the feature dimension
• $\mathbf{W}^O \in \mathbb{R}^{(h \cdot d_v) \times d_{\text{model}}}$: the output projection that mixes information across heads
• $h$: the total number of attention heads

This final projection is unchanged from standard MHA. Each head contributes its perspective, and the output projection learns to combine them into a coherent representation. The key benefit of MQA is that we achieve this with much less memory overhead during inference.

The architectural difference between MHA and MQA becomes clear when we visualize the projection structure.