通过减少查询头数，显著降低注意力计算量，提升长序列处理效率。
Sparse Query Attention (SQA): A Computationally Efficient Attention Mechanism with Query Heads Reduction
- 减少查询头数量，直接降低注意力机制的浮点运算开销。
- 在32k~200k token长序列上实现最高3倍的吞吐提升。
- 适合需要高效训练与推理的长文本模型开发者使用。
Transformer架构依赖多头注意力（MHA）已成为人工智能领域最先进的模型标准。然而，MHA随序列长度呈二次增长的计算复杂度，严重阻碍了其在长上下文场景中的扩展。现有方案如多查询注意力（MQA）和分组查询注意力（GQA）通过共享键值投影缓解了内存带宽瓶颈，但未减少注意力分数计算所需的浮点运算量（FLOPs），该问题仍是训练和全序列处理的关键瓶颈。本文提出稀疏查询注意力（SQA），一种新注意力架构，通过减少查询头数量来降低计算复杂度，使总FLOPs与查询头数成比例下降。论文提供了SQA的理论基础、数学公式及多种变体。在32k~200k token长序列上的实证测试表明，SQA在计算密集型场景（如预训练、微调、编码器任务）中可实现高达3倍的吞吐提升，小规模实验显示模型质量仅轻微下降。SQA是在开发下一代反应式Transformer时偶然发现，暗示其在构建更高效、可扩展模型中的潜力。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。