Multi-Query Attention (MQA)

Share keys and values across query heads

Multi-query attention (MQA) keeps several query heads but uses just one key head and one value head. Every query head compares its queries with the same keys, then combines the same values using its own attention weights. Sharing keys and values does not force the heads to have identical weights.

Compared with multi-head attention, fewer key/value projection parameters and cached key/value elements are needed. This is a design trade-off, not a guarantee of higher speed or a particular change in model quality. MQA is also the one-key/value-head case of grouped-query attention.

This exercise initializes a custom Linen module and computes unmasked self-attention over a complete sequence. It creates no cache and performs no training or token generation.

Follow the head shapes

For input x of shape (T, D_in), let H = num_heads, D = qkv_features and d_head = D // H.

The singleton head axis identifies shared keys and values. The reference repeats that axis to make the matching explicit. A computation with compatible broadcasting can represent the same sharing; repeating is not the only mathematical way to compute MQA.

A small example of sharing

Here two query heads read one set of two keys and values. The heads prefer different keys even though their key/value inputs are shared:

import jax
import jax.numpy as jnp

q = jnp.array([[[1., 0.], [0., 1.]]])       # (1 query, 2 heads, 2 features)
k = jnp.array([[[1., 0.]], [[0., 1.]]])     # (2 keys, 1 head, 2 features)
v = jnp.array([[[1., 0.]], [[0., 2.]]])
shared_k = jnp.repeat(k, 2, axis=1)
shared_v = jnp.repeat(v, 2, axis=1)
scores = jnp.einsum("qhd,khd->hqk", q, shared_k) / jnp.sqrt(2.0)
weights = jax.nn.softmax(scores, axis=-1)
out = jnp.einsum("hqk,khd->qhd", weights, shared_v)

assert weights[0, 0, 0] > weights[0, 0, 1]
assert weights[1, 0, 1] > weights[1, 0, 0]
assert out.shape == (1, 2, 2)
assert jnp.allclose(out[0, 0], jnp.array([0.669762, 0.660477]), atol=1e-6)
assert jnp.allclose(out[0, 1], jnp.array([0.330238, 1.339523]), atol=1e-6)

This example starts with projected arrays to isolate sharing. The exercise below also creates the input and output projections.

Keep cache savings separate from this forward pass

For one layer, one sequence and T cached positions, ordinary multi-head attention stores 2 * T * H * d_head key/value elements. MQA stores 2 * T * d_head if the cache retains the single shared head. At the same sequence length, head width and dtype, that is H times fewer elements. Multiply by bytes per element, batch size and layer count to calculate total cache storage.

For example, with 64 heads and head width 128, a single token needs 16,384 cached elements per layer in multi-head attention versus 256 in MQA. At two bytes per element that is 32,768 versus 512 bytes. Other model memory is not included in these counts.

This reference explicitly repeats keys and values and computes a dense (H, T, T) score tensor. It does not demonstrate the storage of a compact cache or a measured decoding speedup.

Specify the Linen module

Define a compact MQA module with num_heads and qkv_features fields. Create these layers in order, with their default biases and initializers:

1. nn.DenseGeneral(features=(H, d_head), axis=-1, name="q").
2. nn.DenseGeneral(features=(1, d_head), axis=-1, name="k").
3. nn.DenseGeneral(features=(1, d_head), axis=-1, name="v").
4. After attention, nn.DenseGeneral(features=D_in, axis=(-2, -1), name="out").

The first three layers project x separately. Repeat the key/value head axis H times. Move each head axis to the front, giving queries, keys and values of shape (H, T, d_head). Compute scores divided by sqrt(d_head), apply softmax over keys, and combine values. Move the result back to (T, H, d_head) before the output projection.

Preserve the layer names and initialization setup to match the seeded reference. Add no mask, dropout, normalization, position encoding or residual connection. With one key/value head, the smaller projections still have biases; the final output projection is part of the result.