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 a 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 'attention sink' and enhances long-context extrapolation performance, and we also release related $\href{https://github.com/qiuzh20/gated_attention}{codes}$ and $\href{https://huggingface.co/QwQZh/gated_attention}{models}$ to facilitate future research.
核心贡献 · Key contributions
系统研究了多达 30 种、参数量高达 15B 的门控增强型 softmax 注意力变体。 Systematically investigates gating-augmented softmax attention variants over 30 models up to 15B parameters.
发现 SDPA 后的头特定 sigmoid 门控能持续提升性能、稳定性和可扩展性。 Identifies that a head-specific sigmoid gate after SDPA consistently improves performance, stability, and scaling.
将门控的有效性归因于引入非线性和查询相关的稀疏门控分数。 Attributes gating effectiveness to introducing non-linearity and query-dependent sparse gating scores.
证明稀疏门控能缓解注意力汇聚并增强长上下文外推能力。 Demonstrates that sparse gating mitigates attention sink and enhances long-context extrapolation.
发布无注意力汇聚的模型和代码以促进未来研究。 Releases attention-sink-free models and codes to facilitate future research.
局限 · Limitations
非线性对注意力动态和训练的广泛影响尚未充分探索。 Broader implications of non-linearity on attention dynamics and training remain under-explored.
未提供注意力汇聚如何影响长上下文泛化的严格理论解释。 No rigorous theoretical explanation for how attention sink affects long-context generalization.
实验限于 15B MoE 和 1.7B 密集模型，未验证更大模型的可扩展性。 Experiments limited to 15B MoE and 1.7B dense models; scalability to larger models not verified.
门控变体仅在 softmax 注意力上测试，未研究对其他注意力类型的适用性。 Gating variants tested only on softmax attention; applicability to other attention types not studied.