Sparse Query Attention (SQA): A Computationally Efficient Attention Mechanism with Query Heads Reduction
Adam Filipek
TL;DR
This work addresses the compute-bound bottleneck of Transformer self-attention by proposing Sparse Query Attention (SQA), which reduces FLOPs by decreasing the number of query heads rather than KV projections. The authors provide a formal formulation, complexity analysis, and a family of variants (including sSQA and xSQA), demonstrating up to 3x throughput gains on long sequences for compute-bound tasks with only modest quality trade-offs in early experiments. SQA complements existing memory-focused techniques (MQA, GQA, MLA) and can synergize with Sliding Window Attention to handle very long contexts, while remaining drop-in compatible with current architectures. The results suggest SQA is a valuable addition to the efficiency toolbox, particularly for pre-training, fine-tuning, and encoder-based processing where full-sequence computation dominates, with open-source implementations to facilitate adoption and further research.
Abstract
The Transformer architecture, underpinned by the Multi-Head Attention (MHA) mechanism, has become the de facto standard for state-of-the-art models in artificial intelligence. However, the quadratic computational complexity of MHA with respect to sequence length presents a significant barrier to scaling, particularly for applications involving long contexts. Prevailing solutions, such as Multi-Query Attention (MQA) and Grouped-Query Attention (GQA), have effectively addressed the memory bandwidth bottleneck that dominates autoregressive inference latency by sharing Key and Value projections. While highly successful, these methods do not reduce the fundamental number of floating-point operations (FLOPs) required for the attention score computation, which remains a critical bottleneck for training and full-sequence processing. This paper introduces Sparse Query Attention (SQA), a novel attention architecture that pursues an alternative and complementary optimization path. Instead of reducing Key/Value heads, SQA reduces the number of Query heads. This architectural modification directly decreases the computational complexity of the attention mechanism by a factor proportional to the reduction in query heads, thereby lowering the overall FLOPs. This work presents the theoretical foundation of SQA, its mathematical formulation, and a family of architectural variants. Empirical benchmarks on long sequences (32k-200k tokens) demonstrate that SQA can achieve significant throughput improvements of up to 3x in computation-bound scenarios such as model pre-training, fine-tuning, and encoder-based tasks, with only a minimal impact on model quality in preliminary smallscale experiments. SQA was discovered serendipitously during the development of the upcoming Reactive Transformer architecture, suggesting its potential as a powerful tool for building more efficient and scalable models