备注好友：方向-学校/公司-姓名/昵称
【AIGC 先锋科技】交流群
门控机制已被广泛应用于各种模型中，从早期的LSTM（Hochreiter & Schmidhuber, 1997）和高路网络（Srivastava et al., 2015）到近年来的状态空间模型（Gu & Dao, 2023）、线性注意力（Hua et al., 2022）以及softmax注意力（Lin et al., 2025）。
然而，现有文献很少研究门控机制的具体效果。在这项工作中，作者通过全面的实验系统地研究了门控增强的softmax注意力变体。具体而言，作者对15B专家混合（MoE）模型的30个变体和1.7B密集模型进行了全面比较，这些模型在3.5万亿个token的数据集上进行训练。
作者的核心发现是，一个简单的改进——在缩放点积注意力（SDPA）之后应用一个特定头的sigmoid门控——始终如一地提升了性能。这种改进还增强了训练稳定性，容忍更大的学习率，并改善了扩展性。通过比较不同的门控位置和计算变体，作者将这种有效性归因于两个关键因素：
（1）在softmax注意力中的低秩映射上引入非线性，以及（2）应用 Query 相关的Sparse门控分数来调节SDPA输出。
值得注意的是，作者发现这种Sparse门控机制缓解了“注意力陷阱”问题，并提升了长上下文外推性能，作者还发布了相关的代码和模型以促进未来的研究。
1 引言
门控机制在神经网络中已得到广泛应用。早期的架构，如长短期记忆网络（LSTMs，Hochreiter & Schmidhuber, 1997）、高速公路网络（Highway Networks，Srivastava et al., 2015）和门控循环单元（GRUs，Dey & Salem, 2017），开创了使用门控来控制时间步或层间信息流并改善梯度传播的先河。这一原则在现代架构中依然存在。最近的序列建模工作，包括状态空间模型（Gu & Dao, 2023；Dao & Gu, 2024）和注意力机制等，通常也应用门控，常用于调节 Token 混合器组件的输出。尽管门控机制得到了广泛采纳并取得了实证成功，但其功能和影响在超出初始直觉之外仍未得到充分探索。
对门控机制的理解不足阻碍了对其真实贡献的评估，尤其是在与其他架构因素相互混淆的情况下。例如，Switch Heads (Csordas et al., 2024a; b) 引入Sigmoid门控机制来选择top-K注意力头专家，但作者的实验揭示了一个有趣的现象（附录A.1）：即使减少到单个专家，性能提升依然显著，此时门控机制仅对输出值进行调节。这强烈表明门控机制本身提供了显著的内在价值，独立于路由机制。类似地，在Native Sparse Attention (NSA) (Yuan et al., 2025)中，虽然展示了整体性能提升，但并未将门控机制的贡献与其Sparse注意力设计的效应进行分离。这些考虑突显了将门控机制的效应与其他架构组件效应严格分离的必要性。
在本工作中，作者研究了标准softmax注意力机制中的门控机制（Vaswani，2017）（第2.2节）。具体而言，作者在不同位置引入门控（图1）：在 Query
、键
和值投影
之后；在缩放点积注意力（SDPA）输出
之后；以及最终的全连接输出层
之后。作者的探索涵盖了包括元素级和头级、特定头和共享头，以及加性和乘性形式的门控变体。作者发现：(i)应用SDPA输出特定头的门控
带来了最显著的性能提升（例如，PPL降低了0.2，MMLU提升了2分）；(ii)SDPA输出的门控也提高了训练稳定性，几乎消除了损失尖峰，使得可以使用更大的学习率，并增强了模型的扩展性。
作者识别出两个主要因素对门控机制的有效性有贡献：
（i）非线性。两个连续的线性层——值
和密集
投影——可以重写为一个低秩线性投影。因此，通过在位置
或
引入非线性，可以增强这种低秩线性变换的表达能力（第4.1节）。
（ii）Sparse性。尽管非线性的门控变体始终提升性能，但观察到它们的增益存在差异。作者的分析进一步揭示，门控分数的显著Sparse性是另一个关键因素，为SDPA输出引入了输入相关的Sparse性（第4.2节）。此外，Sparse门控消除了注意力陷阱（Xiao et al., 2023）：初始 Token 在注意力分数中占据主导地位（图2，第4.3节）。
先前工作（Xiao et al., 2023; Sun et al., 2024; Gu et al., 2024）将注意力陷阱解释为由于非负softmax归一化导致的冗余注意力累积。经验上，作者验证当在SDPA输出处应用 Query 相关的Sparse门控时，作者的密集和MoE模型（在3.5T Token 上训练）均无注意力陷阱。此外，这些模型在长度泛化方面表现出色，在RULER（Hsieh et al., 2024）上实现了超过10分的增益（第4.4节）。
综上所述，作者的研究强调了标准注意力层中门控机制对模型性能和行为的影响。通过评估门控变体，作者揭示了它们引入非线性、Sparse性以及消除注意力陷阱的能力。这些发现加深了作者对门控注意力机制的理解。作者将开源无注意力陷阱的模型，以推动未来研究的发展。
2 门控注意力层
2.1 初步：多头Softmax注意力
给定输入
，其中
是序列长度，
是模型维度，Transformer 的注意力层（Vaswani，2017）的计算可分为四个阶段。
QKV线性投影：输入
通过学习权重矩阵
线性变换为 Query
、键
和值
，其中
；
缩放点积注意力（SDPA）：计算 Query 和键之间的注意力分数，随后进行softmax归一化。输出是值的加权和：
注
意
力
其中
表示归一化的点积相似度矩阵，而softmax
确保注意力权重非负且每行权重之和为1。
最终输出层：拼接的SDPA输出通过一个输出层
2.2 使用门控机制增强注意力层
门控机制形式化为：
θ
σ
σ
θ
其中
是待调制的输入，
是用于计算门控分数的另一个输入，
指代门控的可学习参数，
是激活函数（例如sigmoid），
是门控输出。门控分数
实际上充当了一个动态过滤器，通过选择性地保留或擦除其特征来控制从
流出的信息。
在本工作中，作者全面研究了注意力层中几种门控机制的不同变体。作者的探索聚焦于五个关键方面：(1) 位置。作者研究了在不同位置应用门控的效果，如图1（左）所示：(a) 在
投影之后（等式1），对应图1（左）中的位置
；(b) 在SDPA（等式3）输出
之后
；(c) 在最终拼接的多头注意力输出
之后。
(2) 粒度。作者考虑了门控分数的两个粒度 Level ：(a) 头 Level ：单个标量门控分数调制注意力头的全部输出。(b) 元素 Level ：门控分数是维度与
相同的向量，实现细粒度、按维度调制。
(3) 头特定或共享。鉴于注意力的多头特性，作者进一步考虑：(a) 头特定：每个注意力头具有其特定的门控分数，使每个头能够独立调制。(b) 头共享：
和门控分数在头之间共享。(4) 乘法或加法。对于将门控分数应用于
，作者考虑：(a) 乘法门控：门控输出
计算为：
。(b) 加法门控：
。(5) 激活函数。
作者主要考虑两种常见的激活函数：SiLU（Shazeer, 2020）和sigmoid。由于SiLU具有无界输出范围，作者仅用于加法门控，而sigmoid仅在
范围内给出分数。此外，为了进一步剖析门控有效性的底层机制，作者还考虑了恒等映射或RMSNorm（Zhang & Sennrich, 2019）（详细见第4.1节）。
除非另有说明，作者采用基于 Head 的乘法门控，使用Sigmoid激活函数
进行计算。
3 实验
3.1 实验设置
模型架构与训练设置
作者对MoE模型（总参数量15B，激活参数量2.54B，型号15A2B）和密集模型（总参数量1.7B）进行了实验。15A2B MoE模型采用128个总专家，使用top-8 softmax门控、细粒度专家（Dai等，2024）、全局批处理LBL（Qiu等，2025）以及
-损失（Zoph等，2022）。注意力部分采用组 Query 注意力（GQA）（Ainslie等，2023）。作者在包含多语言、数学和通用知识内容的
高质量token子集上训练模型。上下文序列长度设置为4096。更详细的配置，如学习率和批大小（bsz），将在各部分中介绍。其他超参数遵循AdamW优化器的默认值。由于门控引入的参数和FLOPs较小，门控引入的墙间延迟小于
。
评估
作者在流行的基准测试上测试了少样本结果，包括用于英语的Hellaswag（Zellers等。2019）、用于通用知识的MMLU（Hendrycks等。2020）、用于数学推理的GSM8k（Cobbe等。2021）、用于编程的HumanEval（Chen等。2021）、用于中文能力的C-eval（Huang等。2024）和CMMLU（L）等。2023）。作者还报告了在多种测试集上的语言建模困惑度（PPL），包括英语、中文、代码、数学、法律和文学等领域。
3.2 主要结果
3.2.1 门控注意力机制用于MoE模型
作者首先比较了不同门控注意力层在训练高效的MoE-15A2B模型上的结果。所有模型均使用一个在1k步内将学习率预热到最大值2e-3，并使用余弦退火衰减到3e-5的调度器。作者使用全局批次大小为1024，包含100k次优化步骤。结果汇总于表1。为了进行公平比较，作者对原味MoE Baseline （第1行）补充了参数扩展方法，包括增加键值头的数量（第2行）、增加 Query 头的数量（第3行）以及增加专家的总数和激活数量（第4行）。这些方法引入的参数数量与门控机制相当或更多。
从表1中，作者观察到：(i) SDPA和值输出门控是有效的。在SDPA的输出
或值映射
处插入门控是最有效的，其PPL更低，整体基准性能优于其他变体。作者将在第4.2节进一步研究为何在这两个位置进行门控是有效的。(ii) 头特定门控很重要。在
和
处应用逐头门控引入的额外参数很少（对于MoE-15A2B模型少于2M），但仍能带来显著改进（第10行和第11行）。当跨不同注意力头共享门控分数（作者通过对 Query 头维度
进行平均，从原始的
得到
分数）时，基准改进效果小于逐头门控（第12行与第10行，第13行与第11行）。这突出了为不同注意力头应用不同门控分数的重要性。(ii) 乘法门控更受青睐。加性SDPA输出门控表现不如乘法门控，尽管其相较于 Baseline 有所改进。(iv) Sigmoid激活更好。用SiLU（第15行）替换最有效门控配置（第5行）中的激活函数，改进效果较小。
总体而言，在值层（
）和SDPA输出（
）处添加门控机制，能使PPL降低超过0.2，优于各种参数扩展 Baseline 模型。然而，在
处进行门控能够获得更好的PPL和基准测试结果。只要不同的注意力头接收不同的门控分数，门控的粒度和激活函数的选择影响相对较小。作者将进一步分析这些观察结果背后的原因（分析部分，第4.2节）。
3.2.2 基于门控注意力的密集模型
作者同样遵循(Yang et al., 2024a)的方法，对密集模型进行了实验，以验证SDPA输出Sigmoid门控的有效性。在使用门控时，作者减小了FFN的宽度以保持参数规模。大多数实验使用了 Baseline 的优化超参数。例如，对于在400B个token上训练的1.7B模型，作者使用最大学习率为4e-3，批次大小为1024。对于在3.5T个token上训练的情况，作者将最大学习率提高到4.5e-3，并将批次大小增加到2048。已有研究表明，增加网络深度、使用较大的学习率以及较大的批次大小虽然可以显著提升模型性能（McCandlish et al., 2018; Wang et al., 2022; D'Angelo et al., 2024）和分布式训练效率，但通常会导致训练不稳定（Wang et al., 2022; Zeng et al., 2022; Takase et al., 2023）。作者观察到，应用门控机制显著减少了训练过程中的损失峰值现象（Chowdhery et al. 2023; Takase et al., 2023），这表明门控在增强训练稳定性方面具有重要作用。基于这一发现，作者引入了另一个实验设置，其特点是层数增加、最大学习率提高以及批次大小增大，以进一步探究门控的稳定效果。
综上所述，作者识别出逐元素门控SDPA（Sparse Distributed Power-Averaged）是增强注意力机制的最有效方法。将该方法应用于密集型Transformer进一步证明，门控机制能够实现稳定训练，支持更大的批处理大小和学习率，从而提升性能。
4 分析：非线性、Sparse性和无注意力陷阱
在本节中，作者进行了一系列实验，以探究为何如此简单的门控机制能够显著提升性能和训练稳定性。根据作者的分析，主要结论如下：(1) 增强非线性的门控操作始终能带来性能提升（第4.1节）；(2) 最有效的SDPA逐元素
门控为SDPA输出引入了强输入依赖Sparse性（第4.2节），这有助于消除"注意力陷阱"现象。
4.1 非线性增强了注意力低秩映射的表达能力
受先前利用群体规范进行SDPA输出的研究启发（Sun等人，2023；Ye等人，2024），在3.2.1节相同设置下，作者在拼接前独立地将RMSNorm（Zhang & Sennrich，2019）应用于每个注意力头的输出。如表3第5行所示，应用几乎不引入额外参数的RMSNorm，同样能显著降低PPL。
在多头注意力机制中，对应第
个头的第
个token的输出可以表示为：
其中
是输出层
对应第
个头的参数。这里，
表示第
个头中第
个词头关注第
个词头的注意力分数，
是词头
的注意力输入，而
表示第
个头中词头
的值输出。根据公式 6，作者可以将
合并成一个对所有
应用的一组低秩线性映射，因为
。在 GQA 中，
在同一组内的头之间共享，进一步降低了表达能力。
鉴于在两个线性映射之间添加非线性可以提升其表达能力（Montufar等人，2014），作者提出了两种修改方法以缓解低秩问题：
非
线
性
映
射
非
线
性
映
射
值得注意的是，在
（表3第3行）位置添加门控对应于首次修改（公式7），而在
位置添加门控（第4行）或组归一化（第5行）则对应于第二次修改（公式8）。这也解释了为什么在
之后在
位置添加门控或归一化没有效果（表1第9行）——它并未解决
与
之间非线性缺失的问题。对于
的加性门控，门控的输出会通过SiLU（表3第4行），这也引入了一些非线性，解释了观察到的性能提升，尽管不如乘性门控实现的提升大。基于这些见解，作者进行了两个额外的实验：(i) 仅在
位置添加SiLU而不引入额外参数（表3第6行）。请注意，这个简单的修改也导致了PPL的适度降低，但大多数基准分数保持不变。(ii) 从加性门控中移除SiLU，使得门控后的
输出直接在
位置相加（表3第7行）。这进一步削弱了加性门控的增益。
综上所述，有效门控变体所关联的增强性能很可能归因于在
和
之间引入非线性。尽管在位置
和
应用门控均能引入这种非线性，但这些应用带来的性能提升存在差异。这一观察到的差异促使作者进一步分析在这些两个位置应用门控的影响。
4.2 门控引入输入相关的Sparse性
作者分析了在值
和SDPA输出
位置应用门控的模型门控分数（表1，“门控分数”列），这些模型在测试语言建模数据上进行了评估。所有层的平均门控分数展示在表4中，分数分布通过图3进行可视化（逐层分数在附录A.2中）。主要观察结果包括：
(i) 有效门控分数是Sparse的。SDPA输出的门控分数（按元素/头计算）表现出最低的平均门控分数。此外，SDPA输出门控分数的分布显示出在0附近的高度集中，这表明存在显著的Sparse性，与其优越的性能一致。(ii) 头特定Sparse性很重要。强制在注意力头之间共享门控分数会增加整体门控分数并降低性能提升。观察结果(i)和(ii)强调了头特定门控的重要性，这与先前研究一致，这些研究表明单个注意力头捕获输入的不同方面（Voita等人，2019年；Wang等人，2021年；Olsson等人，2022年；Wang等人，2023年）。
(iii) Query 依赖性至关重要。值门控（
）的得分高于SDPA输出门控（
）的得分，且性能较差。这表明门控得分Sparse性在 Query 依赖时比由 Key和Value 决定时更有效。具体而言，SDPA输出门控得分源自当前 Query 对应的隐藏状态（例如，公式8中的非线性映射依赖于
），而值门控得分源自与过去 Key和Value 相关的隐藏状态（例如，公式7中的非线性映射依赖于每个
）。这意味着门控得分Sparse性可能过滤掉了与 Query 无关的上下文信息。为进一步验证 Query 依赖性的重要性，作者引入了输入无关门控，通过将可学习参数
置零、应用Sigmoid函数并将其与SDPA输出相乘。如第(6)行所示，输入无关门控提升了 Baseline 性能，这可能是由于引入了非线性。此外，较高的门控得分表明有效的Sparse性应依赖于输入。
(iv) 较少的Sparse门控更差。为了进一步验证门控分数Sparse性的重要性，作者从门控公式中减少Sparse性。具体来说，作者用改进的非Sparse（NS）版本替换了Sigmoid函数：
该机制将门控分数约束在[0.5, 1.0]区间内。这确保了在引入非线性特性的同时消除了门控分数的Sparse性。如表4第(7)行所示，NS-sigmoid门控的增益劣于SDPA输出Sigmoid门控的增益。在附录A.2中，作者将更详细地讨论Sparse门控分数如何影响SDPA隐藏状态中的Sparse性（低于阈值的值的比例）。下一节作者将讨论不同Sparse水平对模型行为的影响，包括减少"注意力陷阱"。
4.3 SDPA输出门控减少注意力陷阱
基于观察到门控机制以输入依赖的方式为SDPA输出引入Sparse性，作者假设该机制能够过滤掉与当前 Query token不相关的上下文，从而减轻注意力陷阱（Xiao et al., 2023；Sun et al., 2024）。为验证这一假设，作者分析了注意力分数的分布（所有 Head 的平均值）以及分配给第一个token的注意力分数比例（图2，表4，“F-Attn”列）。受隐藏状态中大量激活和注意力陷阱讨论的启发（Sun et al., 2024），作者还计算了跨层的最大隐藏状态激活的平均值，如表4中的“M-Act”列所示。更详细的逐层结果在附录A.3中提供。
作者可以观察到： (i) 在SDPA输出的头 Level 和元素 Level 的 Query 依赖型Sigmoid门控
显著降低了分配给第一个token的注意力分数，并减少了大量激活。 (ii) 强制跨头共享门控分数或仅在值投影
之后应用门控，减少了大量激活，但并未降低对第一个token的注意力分数。这进一步强调了头特定门控的重要性，并表明大量激活并非注意力陷阱的先决条件。 (iii) 降低门控的输入依赖性（第6行）或使用NS-sigmoid减少Sparse性（第7行）会加剧大量激活和注意力陷阱。
综合来看，这些观察结果表明，SDPA输出的输入依赖性、头特定门控引入了显著的Sparse性，从而缓解了注意力陷阱问题。此外，SDPA输出的Sparse性减少了模型内的巨大激活值，Sparse性增加会导致激活值更小。这可以解释门控机制下训练稳定性的提升：通过减少巨大激活值，模型在BF16训练过程中（Budzinskiy等人，2025）对数值误差的敏感性降低。作者还观察到，巨大激活值主要源自早期层（例如第5层），在这些层中， FFN （FFN）输出较大值，这与（Yona等人，2025）的研究结果一致。当这些激活值被添加到残差流中时，它们通过预规范化机制在后续层中传播。这与三明治规范化（Ding等人，2021）在提升训练稳定性方面的有效性相符（表2，第7行）：将LayerNorm应用于FFN输出可以防止这些大激活值进入残差流。
4.4 SDPA输出门控有助于上下文长度扩展
基于无注意力陷阱模式，作者在长上下文环境下评估了SDPA门控的效果。具体而言，作者扩展了在3.5T token上训练的模型的上下文长度。作者将RoPE（Su等人，2024）的基从10k增加到1M，并继续在数据上进行训练。
将序列长度扩展至32k以增加80B个token。这使llm-SDPA_2505具有
的上下文长度。随后，作者使用YaRN (Peng et al., 2023) 将上下文长度扩展至
。作者在RULER基准测试 (Hsieh et al., 2024) 上评估模型，并在表5中总结结果。作者观察到以下现象：(i) 在32k设置下，带有门控的模型略微优于 Baseline 模型。这表明在训练长度内，注意力吸收现象可能不会损害模型的长期上下文性能。(ii) 当使用YaRN将上下文长度扩展至
时，在原始
范围内， Baseline 模型和门控模型均出现性能下降。这一观察结果与先前通过修改RoPE扩展上下文长度的研究工作一致 (Chen et al., 2023; Peng et al., 2023; Dong et al., 2025)。尽管门控模型的下降幅度较小。(iii) 在64k和
的上下文长度下，门控注意力模型显著优于 Baseline 模型。基于这些观察结果，作者假设添加门控有助于模型适应上下文长度的扩展。一个可能的解释是， Baseline 模型依赖注意力吸收现象来调整注意力分数的分布。Dong et al. (2025) 推导了基于注意力和隐藏状态分布的RoPE变化效应。当应用YaRN等技术修改RoPE基础时，注意力吸收模式可能难以以无训练的方式适应，导致性能明显下降。相比之下，带有门控的模型主要依赖输入相关的门控分数来控制信息流，使其更能适应此类变化。
5 相关工作
5.1 神经网络中的门控机制
门控机制在神经网络中得到了广泛应用。早期工作如LSTM（Hochreiter & Schmidhuber, 1997）和GRU（Dey & Salem, 2017）引入门控结构以调节时间步长间的信息流，通过选择性保留或丢弃信息来解决梯度消失/爆炸问题。Highway Networks（Srivastava et al., 2015）将这一概念扩展到 FFN ，使得非常深的架构得以成功训练。SwiGLU（Shazeer, 2020）将门控机制引入Transformer FFN层，增强了其表达能力，并成为许多开源大语言模型（LLM）的标准组件（Grattafiori et al., 2024; Yang et al., 2024a）。
关于状态空间模型（Gu & Dao, 2023; Dao & Gu, 2024）和线性注意力机制的研究，如FLASH（Hua et al., 2022）、RetNet（Sun et al., 2023）、Lightning Attention（Qin et al., 2024a;b; Li et al., 2025）以及Gated Delta Networks（Yang et al., 2024b），均采用了门控模块来控制token-mixer模块的信息。遗忘Transformer（Lin et al., 2025）将门控机制应用于softmax注意力的输出，观察到显著的性能提升。尽管这些研究证明了门控的有效性，但其精确机制及其有效性的原因仍需深入探索。这有助于更广泛地认识门控在RNN之外的重要性，并促进更有效地利用门控的独特优势。例如，虽然Switch Heads（Csordas et al., 2024b;a）、NSA（Yuan et al., 2025）和MoSA（Piekos et al., 2025）采用基于sigmoid的门控（Csordas et al., 2023）进行选择，但进一步研究门控的特定贡献将提供有价值的见解。与标准Transformer中包含类似门控机制的 Baseline 模型进行比较，可以更精细地评估其提出的门控选择机制的有效性。与作者最相关的研究是Quantizable Transformers（Bondarenko et al., 2023），该研究也发现将门控应用于softmax注意力可以缓解编码器模型（如BERT和ViT）中的极端注意力集中和隐藏状态中的异常值。尽管这项研究主要利用门控来消除异常值以实现模型量化，但作者详细分析了各种门控变体，揭示了通过增强非线性、Sparse性和提高训练稳定性所带来的优势。基于这些见解，作者扩展了门控注意力模型，展示了门控的广泛适用性和影响力。
5.2 注意力汇聚
小等人（2023）正式识别了“注意力陷阱”现象，即特定token获得较大的注意力分数。类似地，达尔塞特等人（2023）在视觉Transformer中发现，一些冗余token充当“寄存器”来存储注意力分数。随后，孙等人（2024）表明，与大量激活值相关的token也被分配了过高的注意力分数。
然而，作者的工作揭示，在值投影的输出端应用门控可以消除大量激活值，但注意力陷阱依然存在，表明大量激活值并非注意力陷阱的必要条件。类似地，古等人（2024）将注意力陷阱描述为存储冗余注意力分数的非信息性“键偏差”，认为softmax的内在归一化依赖性驱动了这种行为。实验尝试修改softmax注意力，如用未归一化的sigmoid注意力替代softmax（拉马普拉姆等人，2024；古等人，2024）、添加softmax注意力门控或裁剪（博达伦科等人，2023）、修改softmax计算（祖里等人，2025）和分母（米勒，2023）等方法，显示出缓解注意力陷阱的潜力。作者的工作表明，在SDPA之后的Sparse门控可以消除密集（1B参数）和MoE（15B参数）模型中的注意力陷阱，即使是在3.5T token上训练时也是如此。此外，作者揭示了消除注意力陷阱的潜力有助于扩展上下文长度。
参考
[1]. Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free