Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free

门控机制已被广泛应用于各种模型中，从早期的 LSTM 和高速公路网络，到近期的状态空间模型、线性注意力机制以及 softmax 注意力机制。然而，现有文献很少深入探讨门控的具体作用。在本研究中，我们通过全面的实验系统地分析了带有门控增强的 softmax 注意力机制变体。具体来说，我们在一个包含 3.5 万亿 token 的数据集上训练了 30 种不同的模型变体，其中包括 150 亿参数的专家混合（MoE）模型和 17 亿参数的密集模型，并对它们进行了全面对比。 我们的核心发现是：一个简单的修改——在缩放点积注意力（SDPA）之后引入与头部相关的 sigmoid 门控——能够持续提升模型性能。此外，这一修改还增强了训练稳定性，允许使用更大的学习率，并改善了模型的扩展特性。通过比较不同位置的门控机制及其计算变体，我们将这种有效性归因于两个关键因素：(1) 在 softmax 注意力的低秩映射上引入非线性；(2) 使用查询依赖的稀疏门控分数来调节 SDPA 的输出。 值得注意的是，我们发现这种稀疏门控机制能够缓解“注意力汇”问题，并显著提升长上下文外推性能。为了促进未来的研究，我们还开源了相关的 $\href{https://github.com/qiuzh20/gated_attention}{代码}$ 和 $\href{https://huggingface.co/QwQZh/gated_attention}{模型}$。

该论文试图通过系统实验研究注意力机制中引入门控机制的具体效果，特别是针对softmax注意力的改进。这是一个相对较少被深入探讨的问题，尽管门控机制已在多个模型中广泛使用。

论文提出了一种简单而有效的修改方法：在缩放点积注意力（SDPA）之后应用头特定的sigmoid门控。这种设计通过引入非线性及查询依赖的稀疏门控分数来增强注意力机制的表现。相比现有研究，这种方法不仅提升了性能，还改善了训练稳定性、学习率容忍度和扩展特性。

作者对15B参数的MoE模型和1.7B参数的密集模型进行了全面对比实验，数据集规模达到3.5万亿token，体现了实验的规模与严谨性。研究发现，稀疏门控机制可以缓解'注意力陷阱'问题并提高长上下文外推能力。此外，代码和模型均已开源，方便后续研究者复现结果或进一步探索相关方向。

相关的研究包括早期利用门控机制的LSTM和Highway Networks，以及最近的State Space Models和线性注意力机制等。例如，论文《Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context》探讨了长序列建模问题；《Linformer: Self-Attention with Linear Complexity》提出了降低注意力计算复杂度的方法；还有《GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding》讨论了大规模MoE模型的设计与优化。