论文

门控稀疏注意力:为长上下文语言模型结合计算效率与训练稳定性

Gated Sparse Attention: Combining Computational Efficiency with Training Stability for Long-Context Language Models

模型架构Transformer

摘要

长上下文语言模型中注意力的计算负担催生了大体独立的两条工作路线:通过只关注所选 token 来降低复杂度的稀疏注意力机制,以及改善训练稳定性并缓解注意力沉降现象的门控注意力变体。我们观察到这两种方法应对的是互补的弱点,并提出门控稀疏注意力(GSA),一种兼得两者收益的架构。GSA 纳入带 sigmoid 激活、产生有界且可解释选择分数的门控轻量索引器,一个依据局部不确定性调节被关注 token 数量的自适应稀疏控制器,以及值与输出阶段的双重门控。我们为该方法建立理论基础,包括复杂度分析、表达能力结果与收敛保证。在 1.7B 参数模型、400B token 训练的实验中,GSA 达到与纯稀疏基线相当的效率(128K 上下文下 12-16 倍加速),同时获得门控注意力相关的质量收益:困惑度从 6.03 降至 5.70,128K 上下文下的 RULER 分数近乎翻倍,对首个 token 的注意力——注意力沉降的代理指标——从 47% 降至 4% 以下。训练稳定性显著改善,损失尖峰减少 98%。

门控稀疏注意力:为长上下文语言模型结合计算效率与训练稳定性
图1:Gated Sparse Attention 总览。隐藏状态被投影为 query、key 与 value;在索引器(indexer)为所有位置打分并选出 top- k k 之前,value 门(G2)先对 value 进行调制;稀疏 SDPA 从被选中的上下文中聚合信息;输出门(G1)在产生 u t u_{t} 之前提供最终调制。