让注意力机制按内容动态选关键词，提升效率与性能。
Mixture of Sparse Attention: Content-Based Learnable Sparse Attention via Expert-Choice Routing
- 用专家选择路由机制动态挑选每头关注的关键词
- 相同算力下困惑度比密集模型低27%，且速度更快
- 适合追求高效大模型推理与训练的开发者
大型语言模型中的自注意力机制存在二次计算开销。尽管已有诸多改进，但现有次二次注意力方法在实际中仍表现不佳。我们提出一种受混合专家（MoE）启发的内容感知可学习稀疏注意力方法——稀疏注意力混合（MoSA）。MoSA通过专家选择路由动态为每个注意力头选择重要标记，实现任意稀疏模式。从长度为T的序列中选取k个标记后，单个注意力头的复杂度从O(T²)降至O(k² + T)，使相同算力下可部署更多注意力头以增强专业化能力。实验表明，在所测试的稀疏注意力变体中，只有MoSA能超越密集基线，在相同算力下困惑度最高降低27%。此外，即使未使用优化内核，匹配困惑度的MoSA模型在实际运行时间更短、训练内存占用更低，且KV缓存大小显著减少。
原文摘要 · Abstract (English)
Recent advances in large language models highlighted the excessive quadratic cost of self-attention. Despite the significant research efforts, subquadratic attention methods still suffer from inferior performance in practice. We hypothesize that dynamic, learned content-based sparsity can lead to more efficient attention mechanisms. We present Mixture of Sparse Attention (MoSA), a novel approach inspired by Mixture of Experts (MoE) with expert choice routing. MoSA dynamically selects tokens for each attention head, allowing arbitrary sparse attention patterns. By selecting $k$ tokens from a sequence of length $T$, MoSA reduces the computational complexity of each attention head from $O(T^2)$ to $O(k^2 + T)$. This enables using more heads within the same computational budget, allowing higher specialization. We show that among the tested sparse attention variants, MoSA is the only one that can outperform the dense baseline, sometimes with up to 27% better perplexity for an identical compute budget. MoSA can also reduce the resource usage compared to dense self-attention. Despite using torch implementation without an optimized kernel, perplexity-matched MoSA models are simultaneously faster in wall-clock time, require less memory for training, and drastically reduce the size of the KV-cache compared to the dense transformer baselines.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。