Multi-query attention (MQA)

Overview

In standard multi-head attention with k heads of size h, the input at each position is projected to queries, keys and values, each of shape [k, h]. Multi-query attention leaves the query projection at [k, h] but projects keys and values to [1, h], and that single key/value head is used by all k query heads 2204.02311 §2.

The motivation is decoding cost. Autoregressive decoding reads the cached keys and values for every new token, so with many heads the cache is large and memory bandwidth limits speed 2305.13245 §1. MQA cuts the key/value load by a factor of the head count 2305.13245 §2.2. The technique is credited to Shazeer (2019), arXiv:1911.02150, which has no card here.

Six cards mention MQA. PaLM uses it from the start and reports neutral quality and training speed, but its card records no MQA-specific decoding measurement. GQA is the only card that measures MQA, as a baseline on T5 encoder-decoder models. Gemini 1.0 reports MQA as one of its attention choices. StarCoder trains a 15.5B code model with it from scratch and has no MQA ablation 2305.06161 §5.4 tab. 11. Gemma uses it in its 2B model only; the 7B model uses full multi-head attention 2403.08295 tab. 1. Mamba-2 names MQA among the head-pattern ancestors of its block design 2405.21060 §7.2; its card gives no MQA measurement, so it adds no evidence on quality or speed. Treat the PaLM quality claim as the authors' qualitative report, and see Contradictions below.

Variants

Multi-head attention (baseline)

k query heads, k key heads and k value heads, each of size h 2204.02311 §2. During decoding, every layer caches 2 \cdot k \cdot h values per past token (arithmetic from the shapes above).

Multi-query attention, trained from scratch

k query heads, one shared key head and one shared value head, each of size h 2204.02311 §2. The per-layer cache drops to 2 \cdot h values per token, a factor k smaller. The K and V projection matrices also shrink from d_\textrm{model} \times kh to d_\textrm{model} \times h (arithmetic, not stated in the paper).

PaLM's settings: head size 256 at every scale, with 16, 32 and 48 query heads for the 8B, 62B and 540B models 2204.02311 tab. 1, and one KV head 2204.02311 §2. For the 540B model this means 512 cached K+V values per token per layer, where 48-head MHA would need 24,576 (derived from 2204.02311 tab. 1). PaLM combines MQA with RoPE, parallel attention/MLP blocks and bias-free projections 2204.02311 §2. Card: PaLM: Scaling Language Modeling with Pathways.

Gemini 1.0 is described as a decoder-only Transformer with 32K context and multi-query attention 2312.11805 §2 tab. 1. The report does not say whether every model size uses MQA. Card: Gemini: A Family of Highly Capable Multimodal Models.

StarCoderBase is a 15.5B decoder-only code model with 48 query heads and one shared key and value head per layer 2305.06161 §5.4 tab. 11. The card's table gives hidden size 6144, so the head size is 128 (arithmetic). The cache is then 256 values per token per layer, against 12,288 for 48-head MHA at the same width (arithmetic, not stated in the paper). The paper gives no decoding measurement or MQA ablation. Card: StarCoder: may the source be with you!.

MQA in one size of a model family

Gemma's open 2B model has 8 query heads sharing one key/value head, with head size 256 2403.08295 tab. 1 § Model Architecture. Its 7B model has 16 query heads and 16 key/value heads, which is full multi-head attention 2403.08295 tab. 1. The two sizes also differ in parameter count and training tokens (3T vs 6T) § Training Data, so the split cannot be read as an effect of attention type alone. Card: Gemma: Open Models Based on Gemini Research and Technology.

Uptrained MQA from a multi-head checkpoint

The GQA paper converts an existing MHA checkpoint instead of training from scratch. It mean-pools all key and value heads into one head, then uptrains for \alpha = 0.05 of the original steps 2305.13245 §2.1 §3.1. The paper's MQA-XXL result on T5-XXL, with its MHA reference, is in the table below 2305.13245 tab. 1. The ablation in fig. 5 reads MQA at \alpha = 0 as 53.93 against an MHA reference of 57.54, and 56.92 at \alpha = 0.05; these are figure readings on a three-task average, not table values 2305.13245 §3.3.

Card for the converted models: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.

In-between designs

Designs with several KV heads, each shared by a group of query heads, sit between MHA and MQA. The GQA paper calls this grouped-query attention, which is covered on its own concept page.

When to use which

Implementation sketch

One attention layer during decoding with a KV cache. Batch and causal-mask details are omitted.

# x: [d_model] hidden state of the new token
# Wq: [d_model, k*h]   Wk, Wv: [d_model, h]   Wo: [k*h, d_model]
q = reshape(x @ Wq, [k, h])          # k query heads
key = x @ Wk                         # [h]  one shared key head
val = x @ Wv                         # [h]  one shared value head
apply_rope(q, key, position)         # PaLM uses RoPE; settings not stated
cache.K.append(key); cache.V.append(val)  # cache: [T, h] each, not [T, k, h]

scores = q @ cache.K.T / sqrt(h)     # [k, T]: every head reads the same K
probs  = softmax(scores, axis=-1)
out    = probs @ cache.V             # [k, h]: every head reads the same V
y      = reshape(out, [k*h]) @ Wo

The only difference from MHA is the width of Wk and Wv (h instead of k*h) and the resulting cache shape. In MHA the scores would be computed per head against cache.K[:, i, :]. PaLM has no biases in these projections 2204.02311 §2. In PaLM 540B, k \cdot h = 12{,}288 differs from d_\textrm{model} = 18{,}432, so Wq and Wo are not square (arithmetic on 2204.02311 tab. 1).

Contradictions

• Quality of MQA. PaLM's authors report MQA as neutral on quality and training speed 2204.02311 §2, without an MQA ablation. The GQA paper reports uptrained MQA-XXL averaging 0.6 points below MHA-XXL 2305.13245 tab. 1, MQA without uptraining well below the MHA reference 2305.13245 fig. 5, and frequent loss spikes and divergence when T5-Large MQA models are trained from scratch 2305.13245 app. A. The two papers differ in architecture (decoder-only PaLM, encoder-decoder T5), tasks and training budget. The cards do not resolve this.
• Gemini reports MQA without any quality or speed claim about it 2312.11805 §2, so it adds no evidence either way.
• StarCoder uses MQA in a decoder-only code model and reports benchmark scores only, with no MQA-versus-MHA comparison 2305.06161 §6.1 tab. 12. It neither confirms nor contradicts the PaLM or GQA results.

Papers

• PaLM: Scaling Language Modeling with Pathways: uses MQA (one KV head, head size 256) at 8B, 62B and 540B; reports neutral quality and training speed, with no MQA-specific decoding measurement.
• StarCoder: may the source be with you!: StarCoderBase (15.5B, 48 heads, one KV head) trained with MQA from scratch on code; no MQA ablation or decoding measurement in the card.
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints: the only card that measures MQA; uptrained MQA-XXL on T5 is 0.6 points below MHA on average and 0.24 s vs 1.51 s per sample.
• Gemini: A Family of Highly Capable Multimodal Models: reports decoder-only models with MQA and 32K context; whether every size uses it is not stated.
• Gemma: Open Models Based on Gemini Research and Technology: MQA (8 query heads, one KV head, head size 256) in the 2B model only; the 7B model uses MHA. No MQA ablation in the card.
• Mamba-2 (state space duality): names MQA as a lineage influence on its head patterns; no MQA measurement in the card.