arXiv:2505.06708cs.CL2025-05NeurIPS被引 290

给注意力加个门控,让大模型更稳、更强、更长文本

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

  • 在注意力后加局部门控,引入非线性与稀疏调制
  • 提升150亿参数模型性能,支持更大学习率与更好扩展性
  • 解决注意力汇聚问题,适合长文本任务与模型优化研究

门控机制广泛应用于从LSTM到最新状态空间模型、线性注意力及softmax注意力等架构中,但其具体作用尚未深入探讨。本文系统评估了30种门控增强型softmax注意力变体,基于3.5万亿标记数据训练的150亿参数MoE模型和17亿参数稠密模型。核心发现:在缩放点积注意力(SDPA)后添加头级sigmoid门控,能持续提升性能。该改进增强训练稳定性,容忍更大学习率并改善扩展性。通过对比不同门控位置与计算变体,我们归因于两点:(1) 在softmax注意力的低秩映射中引入非线性;(2) 采用查询相关的稀疏门控分数调节SDPA输出。值得注意的是,该稀疏门控机制缓解了‘注意力汇聚’现象,显著提升长上下文外推性能。相关代码与模型已开源。

原文摘要 · Abstract (English)

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.

注意力机制门控长文本大模型优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。