# Beyond Uniform Query Distribution: Key-Driven Grouped Query Attention

arXiv: https://arxiv.org/abs/2408.08454 · HTML: https://mlsanity.com/paper/2408.08454/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 9 of its 13 references that are in the map.
- First posted on arXiv 15 Aug 2024.

## Builds on (most cited first, 9 of 9)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers](https://mlsanity.com/paper/2106.10270/) (2021)
- [CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows](https://mlsanity.com/paper/2107.00652/) (2021)
- [Optimised Grouped-Query Attention Mechanism for Transformers](https://mlsanity.com/paper/2406.14963/) (2024)
- [GQKVA: Efficient Pre-training of Transformers by Grouping Queries, Keys, and Values](https://mlsanity.com/paper/2311.03426/) (2023)
- [QCQA: Quality and Capacity-aware grouped Query Attention](https://mlsanity.com/paper/2406.10247/) (2024)
- [Weighted Grouped Query Attention in Transformers](https://mlsanity.com/paper/2407.10855/) (2024)

## Built on by (most cited first, 2 of 2)

- [A Survey on Large Language Model Acceleration based on KV Cache Management](https://mlsanity.com/paper/2412.19442/) (2024)
- [Opt-GPTQ: An Optimized GPTQ Combining Sparse Attention and Quantization Techniques](https://mlsanity.com/paper/2505.02351/) (2025)

## Abstract (arXiv)

The Transformer architecture has revolutionized deep learning through its Self-Attention mechanism, which effectively captures contextual information. However, the memory footprint of Self-Attention presents significant challenges for long-sequence tasks. Grouped Query Attention (GQA) addresses this issue by grouping queries and mean-pooling the corresponding key-value heads - reducing the number of overall parameters and memory requirements in a flexible manner without adversely compromising model accuracy. In this work, we introduce enhancements to GQA, focusing on two novel approaches that deviate from the static nature of grouping: Key-Distributed GQA (KDGQA) and Dynamic Key-Distributed GQA (DGQA), which leverage information from the norms of the key heads to inform query allocation. Specifically, KDGQA looks at the ratios of the norms of the key heads during each forward pass, while DGQA examines the ratios of the norms as they evolve through training. Additionally, we present Perturbed GQA (PGQA) as a case-study, which introduces variability in (static) group formation via subtracting noise from the attention maps. Our experiments with up-trained Vision Transformers, for Image Classification on datasets such as CIFAR-10, CIFAR-100, Food101, and Tiny ImageNet, demonstrate the promise of these variants in improving upon the original GQA through more informed and adaptive grouping mechanisms: specifically ViT-L experiences accuracy gains of up to 8% when utilizing DGQA in comparison to GQA and other variants. We further analyze the impact of the number of Key-Value Heads on performance, underscoring the importance of utilizing query-key affinities. Code is available on GitHub.