SSA: 通过对齐特征空间中的全注意力和稀疏注意力输出实现稀疏注意力
SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space
AI总结:
提出SSA训练框架,通过双向注意力输出对齐同时解决稀疏注意力的注意力差距和能力差距,实现与全注意力相当的性能。
摘要:
稀疏注意力降低了全自注意力的二次复杂度,但面临两个挑战:(1)注意力差距,即对全注意力训练模型应用稀疏注意力会因训练-推理分布不匹配导致性能下降;(2)能力差距,即纯稀疏注意力训练的模型缺乏完整梯度流,无法达到全注意力性能。我们提出SSA(稀疏注意力),一个集成稀疏和全注意力并具有双向注意力输出对齐的训练框架。我们证明近似误差与稀疏注意力下丢弃的注意力质量线性相关,并表明SSA的对齐目标相比基线大幅减少了该量。实验表明,SSA在两种推理模式下均达到最先进性能,能平滑适应不同的稀疏预算,并展现出优越的长上下文能力。
英文摘要:
Sparse attention reduces the quadratic complexity of full self-attention but faces two challenges: (1) an attention gap, where applying sparse attention to full-attention-trained models causes performance degradation due to train-inference distribution mismatch, and (2) a capability gap, where models trained purely with sparse attention lack complete gradient flow, preventing them from matching full-attention performance. We propose SSA (Sparse Sparse Attention), a training framework that integrates both sparse and full attention with bidirectional attention-output alignment. We prove that the approximation error scales linearly with the attention mass dropped under sparse attention, and show that SSA's alignment objective substantially reduces this quantity compared to baselines. Experiments demonstrate that SSA achieves state-of-the-art performance under both inference modes, adapts smoothly to varying sparsity budgets, and demonstrates superior long-context capabilities.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。