给定输入 $X \in \mathbb{R}^{n \times d_{model}}$,其中 $n$ 是序列长度,$d_{model}$ 是模型维度,Transformer 注意力层的计算【索引 47,Attention is all you need,A Vaswani,2017,Advances in Neural Information Processing Systems】可分为四个阶段。
混合专家模型 (MoE):15B 总参数,2.54B 激活参数(表示为 15A2B)。采用 128 个总专家,Top-8 Softmax 门控,细粒度专家【索引 11,Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models,Damai Dai et al.,2024,arXiv】,全局批次负载均衡损失(global-batch LBL)【索引 39,Demons in the detail: On implementing load balancing loss for training specialized mixture-of-expert models,Zihan Qiu et al.,2025,arXiv】,以及 z-loss【索引 61,St-moe: Designing stable and transferable sparse expert models,Barret Zoph et al.,2022,arXiv】。注意力部分使用分组查询注意力(GQA)【索引 1,Gqa: Training generalized multi-query transformer models from multi-head checkpoints,Joshua Ainslie et al.,2023,arXiv】。
少样本评测 (Few-shots):Hellaswag (英语)【索引 58,Hellaswag: Can a machine really finish your sentence?,Rowan Zellers et al.,2019,arXiv】,MMLU (通用知识)【索引 21,Measuring massive multitask language understanding,Dan Hendrycks et al.,2020,arXiv】,GSM8k (数学推理)【索引 7,Training verifiers to solve math word problems,Karl Cobbe et al.,2021,arXiv】,HumanEval (代码)【索引 4,Evaluating large language models trained on code,Mark Chen et al.,2021】,C-eval【索引 25,C-eval: A multi-level multi-discipline chinese evaluation suite for foundation models,Yuzhen Huang et al.,2024,Advances in Neural Information Processing Systems】 和 CMMLU【索引 27,Cmmlu: Measuring massive multitask language understanding in chinese,Haonan Li et al.,2023】 (中文能力)。
动机:在多头注意力中,第 $k$ 个头的输出可以表示为 $y_{ik} = (\sum_{j=1}^{n} S_{kij} X_j W_V^k) W_O^k$。由于值投影 $W_V^k$ 和输出投影 $W_O^k$ 是连续的线性变换,它们可以合并为一个低秩线性映射(因为头维度 $d_k < d_{model}$)。在两个线性映射之间添加非线性可以提升其表达能力【索引 32,On the number of linear regions of deep neural networks,Guido Montufar et al.,2014】。
结论:输入依赖、逐头特定的 SDPA 输出门控引入了显著的稀疏性,从而缓解了注意力沉溺。稀疏的 SDPA 输出减少了模型内的巨幅激活,这可能是门控提升训练稳定性的原因:减少巨幅激活使模型在 BF16 训练中不易出现数值错误【索引 3,Numerical error analysis of large language models,Stanislav Budzinskiy et al.,2025】。
SDPA 输出门控促进了上下文长度扩展
实验设置:将在 3.5T token 上训练的模型,通过修改 RoPE base【索引 43,Roformer: Enhanced transformer with rotary position embedding,Jianlin Su et al.,2024】并继续在 32k 序列长度上训练 80B token,将上下文长度扩展到 32k。随后使用 YaRN【索引 35,Yarn: Efficient context window extension of large language models,Bowen Peng et al.,2023】将上下文长度扩展到 128k,并在 RULER 基准【索引 23,Ruler: What’s the real context size of your long-context language models?,Cheng-Ping Hsieh et al.,2024】上进行评测。
历史:门控机制在神经网络中被广泛采用。早期的工作如 LSTMs【索引 22,Long short-term memory,Sepp Hochreiter and Jürgen Schmidhuber,1997】和 GRUs【索引 15,Gate-variants of gated recurrent unit (gru) neural networks,Rahul Dey and Fathi M Salem,2017】引入门控来调节时间步间的信息流,解决梯度消失/爆炸问题。Highway Networks【索引 42,Highway networks,Rupesh Kumar Srivastava et al.,2015】将此概念扩展到前馈网络。SwiGLU【索引 41,Glu variants improve transformer,Noam Shazeer,2020】将门控引入 Transformer 的 FFN 层,成为许多开源 LLM 的标准组件。
现代应用:近期的状态空间模型【索引 19,Mamba: Linear-time sequence modeling with selective state spaces,Albert Gu and Tri Dao,2023】【索引 13,Transformers are ssms: Generalized models and efficient algorithms through structured state space duality,Tri Dao and Albert Gu,2024】和线性注意力模型,如 FLASH【索引 24,Transformer quality in linear time,Weizhe Hua et al.,2022】、RetNet【索引 45,Retentive network: A successor to transformer for large language models,Yutao Sun et al.,2023】、Lightning Attention【索引 37,Various lengths, constant speed: Efficient language modeling with lightning attention,Zhen Qin et al.,2024b】和 Gated Delta Networks【索引 54,Gated delta networks: Improving mamba2 with delta rule,Songlin Yang et al.,2024b】,也集成了门控模块。Forgetting Transformer【索引 28,Forgetting transformer: Softmax attention with a forget gate,Zhixuan Lin et al.,2025】在 Softmax 注意力输出上应用门控并观察到显著性能提升。
本文贡献:尽管这些工作证明了门控的有效性,但对其精确机制的全面理解仍有待探索。本文的工作通过对各种门控变体的详细分析,揭示了其通过增强非线性和稀疏性带来的好处,以及对训练稳定性的改善。与 Quantizable Transformers【索引 2,Quantizable transformers: Removing outliers by helping attention heads do nothing,Yelysei Bondarenko et al.,2023】(其利用门控消除异常值以进行模型量化)相比,本文提供了更深入的分析,并将门控注意力模型扩展到更大规模,展示了其广泛的适用性和影响力。
注意力沉溺 (Attention Sink):
现象定义:Xiao 等人【索引 51,Efficient streaming language models with attention sinks,Guangxuan Xiao et al.,2023】正式识别了“注意力沉溺”现象。Darcet 等人【索引 14,Vision transformers need registers,Timothée Darcet et al.,2023】在视觉 Transformer 中也发现了类似现象,即一些冗余的 token 充当“寄存器”来存储注意力分数。
机制探讨:Sun 等人【索引 44,Massive activations in large language models,Mingjie Sun et al.,2024】发现过多的注意力分数也分配给了与巨幅激活值相关的 token。然而,本文的研究表明,在值投影输出处应用门控可以消除巨幅激活,但注意力沉溺依然存在,说明巨幅激活不是注意力沉溺的必要条件。Gu 等人【索引 20,When attention sink emerges in language models: An empirical view,Xiangming Gu et al.,2024】将注意力沉溺描述为存储冗余注意力的非信息性“键偏置”,并认为这是由 Softmax 的内在归一化依赖性驱动的。