动态切换注意力机制,兼顾效果与效率。
Mixture of Attention Schemes (MoAS): Learning to Route Between MHA, GQA, and MQA
- 通过学习路由机制,动态选择每令牌的MHA/GQA/MQA方案
- 在WikiText-2上验证,动态路由比静态混合更低损失(2.3074 vs 2.3093)
- 适合追求推理效率又不想牺牲性能的研究者与工程师
Transformer模型中注意力机制的选择涉及建模质量与推理效率的权衡。多头注意力(MHA)性能最佳,但推理时需大量键值缓存。多查询注意力(MQA)和分组查询注意力(GQA)虽降低内存开销,却常以性能为代价。本文提出混合注意力方案(MoAS),通过学习路由机制,动态为每个令牌选择最优注意力形式(MHA、GQA或MQA)。实验表明,动态路由优于静态平均方案,在WikiText-2数据集上实现更低验证损失(2.3074),性能接近MHA基线,同时具备条件化计算效率潜力。代码已开源。
原文摘要 · Abstract (English)
The choice of attention mechanism in Transformer models involves a critical trade-off between modeling quality and inference efficiency. Multi-Head Attention (MHA) offers the best quality but suffers from large Key-Value (KV) cache memory requirements during inference. Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) reduce memory usage but often at the cost of model performance. In this work, we propose Mixture of Attention Schemes (MoAS), a novel architecture that dynamically selects the optimal attention scheme (MHA, GQA, or MQA) for each token via a learned router. We demonstrate that dynamic routing performs better than static averaging of schemes and achieves performance competitive with the MHA baseline while offering potential for conditional compute efficiency. Experimental results on WikiText-2 show that dynamic routing (val loss 2.3074) outperforms a static mixture (2.3093), validating the effectiveness of the proposed method. Our code is available at https://github.com/Esmail-ibraheem/Mixture-of-Attention-Schemes-MoAS.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。