Grouped Query Attention › Mathematical Formulation

Let's formalize GQA mathematically. The core idea is straightforward: instead of each query head having its own key-value projections (as in MHA) or all query heads sharing one key-value projection (as in MQA), GQA divides query heads into groups where each group shares a single key-value projection.

Given an input sequence $X \in \mathbb{R}^{T \times d}$, GQA is parameterized by:

• $h$: number of query heads (determines query diversity)
• $g$: number of key-value groups, where $g$ divides $h$ evenly (controls memory vs. capacity trade-off)
• $d_k = d / h$: dimension per head (same as standard multi-head attention)

The ratio $h/g$ determines how many query heads share each key-value head. For example, with $h = 32$ query heads and $g = 8$ KV groups, each group of 4 query heads shares one key-value representation.

Query projections work exactly as in standard multi-head attention. Each query head $i$ has its own learned projection matrix that turns the input into a query representation:

$$Q_i = X W^Q_i \quad \text{for } i \in \{1, \ldots, h\}$$

where:

• $Q_i \in \mathbb{R}^{T \times d_k}$: the query matrix for head $i$, containing a $d_k$-dimensional query vector for each of the $T$ positions
• $X \in \mathbb{R}^{T \times d}$: the input sequence with $T$ tokens, each represented as a $d$-dimensional vector
• $W^Q_i \in \mathbb{R}^{d \times d_k}$: the learned query projection matrix for head $i$, transforming from model dimension to head dimension

Having $h$ separate query projections preserves the model's ability to ask $h$ different "questions" of the input. Each query head can still specialize in detecting different patterns, maintaining expressive power.

Here is where GQA diverges from MHA. Instead of $h$ key-value projections, we have only $g$ projections that are shared among groups of query heads:

$$K_j = X W^K_j, \quad V_j = X W^V_j \quad \text{for } j \in \{1, \ldots, g\}$$

where:

• $K_j \in \mathbb{R}^{T \times d_k}$: the key matrix for group $j$, containing key vectors for all $T$ positions
• $V_j \in \mathbb{R}^{T \times d_k}$: the value matrix for group $j$, containing value vectors for all $T$ positions
• $W^K_j \in \mathbb{R}^{d \times d_k}$: the learned key projection matrix for group $j$
• $W^V_j \in \mathbb{R}^{d \times d_k}$: the learned value projection matrix for group $j$

With only $g$ key-value projections instead of $h$, the KV cache during inference stores $g$ key-value pairs rather than $h$, reducing memory by a factor of $h/g$.

To connect query heads to their shared key-value representations, we define a grouping function that maps each query head index to its corresponding key-value group:

$$\text{group}(i) = \left\lfloor \frac{(i - 1) \cdot g}{h} \right\rfloor + 1$$

where:

• $i$: the query head index, ranging from 1 to $h$
• $g$: the total number of key-value groups
• $h$: the total number of query heads
• $\lfloor \cdot \rfloor$: the floor function (rounds down to nearest integer)

This function partitions the $h$ query heads into $g$ groups of size $h/g$ each. For example, with $h = 8$ query heads and $g = 2$ groups:

• Query heads 1, 2, 3, 4 map to group 1 (using $K_1$, $V_1$)
• Query heads 5, 6, 7, 8 map to group 2 (using $K_2$, $V_2$)

Each query head computes attention using its own query but the shared key-value from its group:

$$\text{head}_i = \text{softmax}\left(\frac{Q_i K_{\text{group}(i)}^T}{\sqrt{d_k}}\right) V_{\text{group}(i)}$$

where:

• $\text{head}_i \in \mathbb{R}^{T \times d_k}$: the output of attention head $i$
• $Q_i \in \mathbb{R}^{T \times d_k}$: the query matrix for head $i$ (unique to this head)
• $K_{\text{group}(i)} \in \mathbb{R}^{T \times d_k}$: the key matrix shared by all heads in group $\text{group}(i)$
• $V_{\text{group}(i)} \in \mathbb{R}^{T \times d_k}$: the value matrix shared by all heads in group $\text{group}(i)$
• $Q_i K_{\text{group}(i)}^T \in \mathbb{R}^{T \times T}$: the attention score matrix
• $\sqrt{d_k}$: scaling factor to prevent dot products from growing too large

The key insight is that different query heads in the same group can still learn different attention patterns. Even though heads 1-4 all use $K_1$ and $V_1$, their different query projections $Q_1, Q_2, Q_3, Q_4$ produce different attention weights over the same key-value pairs.

Finally, all head outputs are concatenated and projected back to the model dimension:

$$\text{GQA}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O$$

where:

• $\text{Concat}(\text{head}_1, \ldots, \text{head}_h) \in \mathbb{R}^{T \times (h \cdot d_k)}$: concatenation of all $h$ head outputs along the feature dimension
• $W^O \in \mathbb{R}^{(h \cdot d_k) \times d}$: the output projection matrix
• $\text{GQA}(X) \in \mathbb{R}^{T \times d}$: the final output, matching the input dimension

Since $h \cdot d_k = d$, the concatenated heads have dimension $d$, and $W^O$ projects back to dimension $d$, preserving the input shape for residual connections.