A variant of multi-head attention in which groups of heads share a single key and value, to shrink the KV cache and speed up generation.
In multi-head attention, each head has its own query, key and value. During generation,
the keys and values of all earlier tokens are stored in the KV cache, and the
bottleneck is less arithmetic than bandwidth: at every step those tensors must be loaded
again from memory (Shazeer, 2019, sec. 1).
Multi-query (MQA). Shazeer (2019) proposed that heads keep their own queries but share a
single key and a single value. With H heads, the cache shrinks H times
(Ainslie et al., 2023, sec. 2.2); decoding becomes much faster, with what the author describes as
minor quality degradation (2019).
Grouped (GQA). Ainslie et al. (2023) proposed a middle ground: the H query heads are split into G
groups, and each group shares one key–value pair. GQA-1 is MQA; GQA-H is ordinary multi-head
attention (2023, sec. 2.2). They also showed that an already trained model can be
converted by averaging each group's projections and training it further with 5% of the
original compute. On T5-XXL, the 8-group version took 0.28 s at inference against the
original's 1.51 s (MQA: 0.24 s), with an almost identical average score (47.1 versus 47.2;
MQA: 46.6) (2023, table 1).
Where it is used. Llama 2 adopted it only in its 34- and 70-billion-parameter versions
(Touvron, Martin, et al., 2023, sec. 2.2), and Llama 3 uses it too (Alammar & Grootendorst, 2024, p. 98) (LLaMA).
Mixtral, Mistral's mixture of experts, has 32 query heads and 8 key–value heads
(Jiang et al., 2024, sec. 2), and Jamba uses it in its few attention layers (SSM hybrid (Jamba, Qwen3.5))
(Lieber et al., 2024, sec. 2).
What it shows. GQA divides the cache by a fixed factor; linear attention and
SSMs instead change how fast it grows. That sharing keys and values
costs so little suggests (our reading, not the authors') that many heads consult the same
index of the context with different questions: what varies among them is more what they look
for than how what was read is filed.
Sources
- Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training generalized multi-query transformer models from multi-head checkpoints. In H. Bouamor, J. Pino, & K. Bali (Eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (pp. 4895–4901). Association for Computational Linguistics. https://doi.org/10.18653/v1/2023.emnlp-main.298 §
- Alammar, J., & Grootendorst, M. (2024). Hands-on large language models: Language understanding and generation. O’Reilly Media. §
- Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., de Las Casas, D., Bou Hanna, E., Bressand, F., Lengyel, G., Bour, G., Lample, G., Renard Lavaud, L., Saulnier, L., Lachaux, M.-A., Stock, P., Subramanian, S., Yang, S., . . . El Sayed, W. (2024). Mixtral of experts (arXiv:2401.04088). arXiv. https://doi.org/10.48550/arXiv.2401.04088 §
- Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalev-Shwartz, S., Abend, O., Alon, R., Asida, T., Bergman, A., Glozman, R., Gokhman, M., Manevich, A., Ratner, N., Rozen, N., . . . Shoham, Y. (2024). Jamba: A hybrid Transformer-Mamba language model (arXiv:2403.19887). arXiv. https://doi.org/10.48550/arXiv.2403.19887 §
- Shazeer, N. (2019). Fast transformer decoding: One write-head is all you need (arXiv:1911.02150). arXiv. https://doi.org/10.48550/arXiv.1911.02150 §
- Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Canton Ferrer, C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., . . . Scialom, T. (2023). Llama 2: Open foundation and fine-tuned chat models (arXiv:2307.09288). arXiv. https://doi.org/10.48550/arXiv.2307.09288 §
