Command Palette
Search for a command to run...
Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free
Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free
摘要
Gating mechanisms have been widely utilized, from early models like LSTMs and Highway Networks to recent state space models, linear attention, and also softmax attention. Yet, existing literature rarely examines the specific effects of gating. In this work, we conduct comprehensive experiments to systematically investigate gating-augmented softmax attention variants. Specifically, we perform a comprehensive comparison over 30 variants of 15B Mixture-of-Experts (MoE) models and 1.7B dense models trained on a 3.5 trillion token dataset. Our central finding is that a simple modification—applying an head-specific sigmoid gate after the Scaled Dot-Product Attention (SDPA)—consistently improves performance. This modification also enhances training stability, tolerates larger learning rates, and improves scaling properties. By comparing various gating positions and computational variants, we attribute this effectiveness to two key factors: (1) introducing non-linearity upon the low-rank mapping in the softmax attention, and (2) applying query-dependent sparse gating scores to modulate the SDPA output. Notably, we find this sparse gating mechanism mitigates 'massive activation', 'attention sink', and enhances long-context extrapolation performance, and we also release related codes and models to facilitate future research. Furthermore, the most effective SDPA output gating is used in the Qwen3-Next models.
一句话总结
Qwen Team(阿里巴巴集团)等在 15B MoE 模型和 1.7B dense 模型上,于 3.5 万亿 tokens 上训练,系统比较了 30 种门控增强 softmax 注意力变体,发现在 SDPA 之后应用头级 sigmoid 门控会引入非线性和查询相关稀疏性,从而提升性能、训练稳定性和长上下文外推能力,同时缓解注意力汇聚;该机制随后被 Qwen3-Next 采用。
核心贡献
-
在 1.7B dense 模型和 15B Mixture-of-Experts 模型中,对 30 种门控增强 softmax 注意力变体进行全面比较,并在 3.5 万亿 tokens 上训练。结果表明,在 Scaled Dot-Product Attention 之后采用头级 sigmoid 门控能够稳定提升性能。
-
论文将该提升归因于 softmax 注意力内部低秩映射上的非线性,以及调节 SDPA 输出的查询相关稀疏门控分数。门控模型还表现出更好的训练稳定性、对更大学习率的容忍度,以及更好的缩放行为。
-
SDPA 之后的稀疏门控能够消除 dense 模型和 Mixture-of-Experts 模型中的注意力汇聚,缓解大激活,并在无需重新训练的情况下改善长上下文外推和上下文长度扩展。相关无注意力汇聚模型和代码已发布,最有效的 SDPA 输出门控也被用于 Qwen3-Next 模型。
引言
门控机制被广泛用于神经网络和现代 Transformer 注意力中,以控制信息流动、改善训练并调节 token 混合。尽管门控出现在状态空间模型、稀疏注意力和混合专家设计等近期架构中,但由于门控效应可能与路由、稀疏性或其他架构变化混杂,先前工作往往未能明确其直接贡献。作者通过系统测试门控在标准 softmax 注意力中的不同位置来应对这一问题,发现在 scaled dot product attention 输出之后的逐元素门控尤其有效。这一简单改动增加了非线性,引入了输入相关稀疏性,改善了困惑度和 MMLU,减少了损失尖峰,支持更大学习率,消除注意力汇聚,并改善了长度泛化。
数据集
-
训练数据组成:作者在 4T-token 高质量语料库的子集上进行训练,该语料库涵盖多语言、数学和通用知识内容。
-
来源细节:论文没有说明确切原始数据源、子集大小、混合比例或详细过滤规则,除了将语料库描述为高质量且领域覆盖广泛。
-
处理细节:训练样本按固定序列长度 4096 制备;未说明额外的裁剪策略、元数据结构或分词后过滤。
-
训练使用:模型变体在该高质量 token 池中选定的一部分上训练,而不是必然使用完整 4T tokens。
-
评估数据:性能通过 held-out 困惑度测试集衡量,覆盖英语、中文、代码、数学、法律和文学,并包含 HellaSwag、MMLU、GSM8K、HumanEval、C-Eval 和 CMMLU 等 few-shot 基准数据集。
方法
作者利用一种门控注意力层,在注意力计算内部加入动态滤波器,以增强标准多头注意力。基本门控形式化如下:
Y′=g(Y,X,Wθ,σ)=Y⊙σ(XWθ),其中 Y 是待调节的特征张量,X 是用于计算门控分数的输入,Wθ 是可学习门控参数,σ 是激活函数,例如 sigmoid,Y′ 是门控输出。在该形式化中,σ(XWθ) 充当可学习且输入相关的掩码,通过保留或抑制单个特征来控制信息流动。该门控与注意力权重联合训练,不需要单独的路由或顶层选择模块。
门控机制可以插入注意力块中的几个不同位置。可参考框架示意图。
所示架构标出了查询、键和值投影之后可能的门控位置,分别记为 G4、G3 和 G2;scaled dot product attention 输出之后的门控位置,记为 G1;以及最终 dense 投影之后的门控位置,记为 G5。最有效的位置是 SDPA 输出位置 G1,其中门控在 dense 投影之前直接调节聚合后的注意力输出。门控可以是逐头的,为每个头使用一个标量分数,也可以是逐元素的,使用与输出维度相同的分数向量。最强变体采用具有头级参数的逐元素调节。乘性门控是主要形式,同时也考虑使用 SiLU 的加性门控,因为加性项需要无界激活。
这种设计并不只是归一化或正则化技巧。在标准注意力头中,token i 和头 k 的输出可以写成:
oi(k)=(j=0∑iSij(k)XjWV(k))WO(k).其中 Sij(k) 是头 k 中从 token i 到 token j 的注意力分数,Xj 是 token j 的输入表示。两个线性变换 WV(k) 和 WO(k) 可以合并为一个低秩线性映射:
oi(k)=j=0∑iSij(k)Xj(WV(k)WO(k)).这种合并揭示了一个结构性限制:从值投影到 dense 输出的路径本质上是一个低秩线性变换,而分组查询注意力会进一步降低这种表达能力。在 G2 或 G1 处门控会在不同位置引入非线性,从而打破该路径的纯线性特征。例如,值投影之后的门控可以看作:
oi(k)=(j=0∑iSij(k)Map(XjWV(k)))WO(k),而 SDPA 输出之后的门控则对应于:
oi(k)=Map(j=0∑iSij(k)XjWV(k))WO(k).在这两种情况下,门控都会在值变换和 dense 变换之间注入一个非线性操作。这改善了原本低秩的注意力通道的表达能力。加性门控可以引入一些非线性,但乘性门控提供了更强的输入相关调节。类似 RMSNorm 的操作也可以通过引入归一化和轻微的非线性稳定化来改善性能,但它们无法匹配完整门控那种查询条件化的过滤行为。
在训练期间,门控参数与模型一起端到端学习。在实用的 dense 模型训练中,加入门控时可以通过减小前馈层宽度来保持参数预算不变。主要训练配置采用具有 sigmoid 分数的头级乘性门控。该配置被观察到能够降低最终损失,并在更大学习率和更大批量下缓解损失尖峰,从而实现更稳定的优化。在随架构给出的长训练比较中,在可比设置下,门控模型沿着更平滑的损失轨迹,并达到比基线更低的最终损失。
门控分数分布揭示了一些门控位置为何比其他位置更强。如下图所示:
所示分布表明,有效门控会产生稀疏的分数分布。SDPA 输出门控具有最低的平均分数,并且在零附近具有最强的集中性,这意味着该门控经常抑制许多注意力特征。值位置门控产生更高的平均分数和较弱的稀疏模式,而头间共享门控也会提高平均分并减弱过滤行为。这支持了头级门控的必要性,因为每个注意力头可以专门化自身的过滤规则,而不必被迫与其他头共享单个门控。
另一个进一步区别是查询依赖性。在 SDPA 输出门控中,门控由与当前查询相关的隐藏状态计算得到。这使得门控能够过滤与当前处理 token 无关的上下文贡献。相比之下,值位置门控由过去值隐藏状态驱动,因此与当前查询的对齐性较弱。观察到的性能差异表明,有用的门控不仅应当稀疏,还应当是输入相关和查询条件化的。为检验这一点,作者还通过固定或零初始化的可学习参数配合 sigmoid 来评估输入无关门控。这仍然改善了基线,可能来自新增的非线性,但未能达到查询相关门控的性能。进一步消融非稀疏 sigmoid 变体,即将分数约束在较高范围并移除强抑制,会降低收益,证实稀疏门控分数是该方法的核心组件。
实验
实验在稀疏 MoE 和 dense transformer 模型上进行,这些模型在大规模多语言、编码和推理语料上训练,并在覆盖知识、推理、编码、语言能力和困惑度的广泛基准上评估,以检验替代注意力门控位置和门控设计的有效性。结果表明,在 SDPA 输出处采用逐元素或逐头 sigmoid 门控能够稳定改善困惑度和下游任务表现,优于基线和简单参数扩展方法。进一步分析表明,SDPA 输出门控通过在低秩注意力映射中引入有用的非线性,并诱导出查询相关、头级稀疏性,从而减少大激活和注意力汇聚效应,因此能够稳定训练,并实现更稳健地向更大学习率、更大批量和更长上下文扩展。
SDPA 输出处的门控和值映射处的门控都能够在性能上优于基线和参数扩展基线。SDPA 输出门控提供最强增益,而值映射门控同样有效。头级门控尤为重要,其优势与稀疏、输入相关的门控行为相关。在各种门控变体中,SDPA 输出门控带来最佳困惑度和基准测试提升。值映射门控也会降低困惑度并提升基准测试表现,尽管不如 SDPA 输出门控显著。头级门控优于跨注意力头共享的门控。稀疏、输入相关的门控与更好的性能以及更少的注意力伪影相关。
在 dense 模型配置、训练规模和超参数设置中,相对于基线,SDPA 之后的逐元素 sigmoid 门控稳定改善验证困惑度和下游基准性能,即使为保持参数量而降低 FFN 宽度也是如此。门控机制似乎也改善了优化稳定性,使得在基线变体可能变得不稳定的更大学习率和更大批量下仍能有效训练。SDPA 输出门控在不同模型和训练设置下改善了困惑度和多个评估分数。相比基线,门控支持在更大学习率和更大批量下更稳定地训练。Sandwich 归一化可以在一些不稳定设置中恢复收敛,但增益小于 SDPA 门控。
向注意力输出加入非线性增强通常会降低困惑度,并相对于基线改善下游分数。对 SDPA 输出进行 sigmoid 门控带来最清晰的增益,而 RMSNorm 也能在几乎不增加参数量的情况下改善困惑度。分析表明,输入相关门控会增加稀疏性并减少大激活,这也可能支持更稳健的上下文长度扩展。sigmoid 逐元素门控在困惑度和下游任务上均优于基线。在每个注意力头上使用 RMSNorm 可以在不增加有意义参数量的情况下降低困惑度。输入相关性较弱的门控与更强的大激活和注意力汇聚效应相关。采用 SDPA 门控的模型在扩展上下文长度下表现出比基线更强的性能。
在所比较的门控设置中,最强的定性趋势表明,SDPA 输出门控是最有效的变体,尤其是在门控为逐元素或逐头且查询相关时。这些设置展现出更稀疏的门控分数,同时具有更低的首 token 注意力分数和更低的最大隐藏状态激活,与更弱的注意力汇聚行为一致。相比之下,共享门控、仅值门控或稀疏性较弱的门控变体会产生更高的门控分数,并且对首 token 注意力的降低更弱,这表明输入相关的头级稀疏性很重要。SDPA 输出门控具有最强的稀疏门控分数。头级门控比共享门控更显著地降低首 token 注意力。值门控会降低激活,但不如 SDPA 门控那样有效降低首 token 注意力。稀疏性较弱的门控会同时增加大激活和注意力汇聚。
在已训练的序列长度范围内,门控注意力相对于基线具有轻微优势。上下文扩展后,两个模型在先前支持的长度上均下降,但门控模型下降幅度较小。在更长的评测长度上,门控模型表现明显更好,表明其对上下文扩展具有更稳健的适应性。门控注意力在原始训练长度下带来适度增益。上下文扩展会降低两种方法的性能,其中基线下降更陡。门控模型在扩展序列长度上保持更强的准确率。
实验评估了注意力门控在 dense 模型配置、训练规模、超参数设置、门控位置消融和上下文扩展中的表现,以检验门控是否能够改善语言建模、下游任务、优化稳定性和注意力行为。SDPA 之后的逐元素 sigmoid 门控成为最强变体,值映射门控也有帮助,而头级、输入相关的稀疏门控相对于共享门控或仅值门控带来最明显收益。这些结果表明,SDPA 输出门控能够改善困惑度和基准性能,同时减少大激活和注意力汇聚效应,从而支持更稳定的训练和更稳健的上下文扩展。