NeurIPS 2025 最佳论文 Gated Attention：非线性、稀疏性与 attention sink

发表于： 分类于： 阅读：≈ 7分钟 浏览： 评论：

Gated Attention（Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free，2025，Qwen 团队，NeurIPS 2025 Best Paper）的核心思路：在 Scaled Dot-Product Attention 的输出后加一个 head-specific 的 elementwise sigmoid gate。在 1.7B dense 与 15B MoE（A2.54B 激活）模型上各训 3.5T tokens，PPL 下降约 0.05–0.27（视模型与设置而异），训练 loss spike 显著减少，长上下文外推能力提升，attention sink 大幅减弱（BOS token 的 attention 比例从 46.7% 降至 4.8%）。

作者系统比较了 5 个候选位置与多种 gate 形式（粒度、参数共享、乘法/加法、激活函数）共 30 个变体，上述方案在所有配置中表现最优。

论文地址： arxiv.org/abs/2505.06708

为什么是 gating

门控（gating）机制在神经网络中由来已久。LSTM 用 input/forget/output 三道门控制信息流，Highway Networks 用 gate 调制残差流向，Mamba 等 state-space model 与多种 linear attention 变体也将 gating 作为标准组件。但已有工作很少系统比较 gating 在不同位置、不同形式下的实际作用。

以 Switch Heads 为例。该方法在 attention head 之间加 sigmoid gating 做 top-K 路由，性能提升被归因为路由机制。本文作者在复现时将 expert 数减到 1（即不存在路由），gate 退化为对 value 输出的调制，性能提升仍然保留。再如 DeepSeek 的 NSA（Native Sparse Attention，ACL 2025 Best），同样使用了 gating，但论文将增益主要归于 sparse attention 的设计，未单独评估 gate 的贡献。

本文的研究目标是：将 gating 从其他混淆变量中分离出来，系统比较其加入位置与形式的影响。

五个位置 × 多种形式 = 30 个变体

五个候选 gating 位置（左），15B MoE 上各位置带来的 PPL 与 MMLU 变化（中），训练 loss 曲线对比（右）

可加入 gate 的位置共五个：

  • $G_1$：SDPA 输出之后、concat 之前
  • $G_2$：value 投影之后
  • $G_3$：key 投影之后
  • $G_4$：query 投影之后
  • $G_5$：最后的 dense 输出层之后

每个位置进一步组合多种 gate 形式：elementwise vs headwise（粒度）、head-specific vs head-shared（参数共享）、multiplicative vs additive（作用方式），以及 sigmoid 与 SiLU 等激活函数。

主实验在两组模型上进行：1.7B dense 模型与 15B MoE 模型（A2.54B 激活），均在 3.5T tokens 上从头训练。主实验共 30 个 variant。

最优方案为：在 $G_1$ 位置（SDPA 输出之后）加一个 head-specific 的 elementwise sigmoid gate，gate 由当前 query 投影得到。代码改动约十行：

1
2
3
4
5
6

# 标准 attention output
attn_out = scaled_dot_product_attention(q, k, v)  # [B, H, T, D]
# 新增：每个 head 独立的 sigmoid gate
gate = sigmoid(q @ W_gate)  # [B, H, T, D]，W_gate 是 head-specific 的
attn_out = attn_out * gate
# 后面正常接 dense 层

在 1.7B dense 模型（28 层 / 3.5T tokens）上 PPL 从 6.180 降至 6.130（约 −0.05），MMLU 在 1.7B / 400B tokens 设置下从约 50.21 升至 51.15（约 +1 点）；15B MoE 模型上 PPL 从 6.026 降至 5.761（约 −0.27），MMLU 从 58.79 升至 60.82。

为什么是这个位置

作者将 $G_1$ 的优势归因为两点：non-linearity 与 query-dependent sparsity。

Non-linearity。标准 attention 中，value 投影 $W_v$ 与最终 dense 投影 $W_O$ 是两层连续的线性变换，理论上可合并为一个秩受 head_dim 限制的低秩线性映射。也就是说，从 input 到 output，attention 引入的非线性仅来自 softmax，而 softmax 作用于 attention score 而非 value。在 $G_1$ 或 $G_2$ 处加 gate，相当于在这条低秩通路中插入一个非线性激活，提升表达能力。

这也解释了 $G_3$、$G_4$（在 query/key 上 gate）效果较差的原因：query 与 key 之后即接 softmax，本已存在非线性，gate 的边际收益较小。论文同时指出 $G_5$（在 dense 之后）效果有限，原因是它没有改善 $W_V$ 与 $W_O$ 之间缺失非线性的问题。

Sparsity。作者统计了 gate 激活值的分布：sigmoid gate 学到的分布高度稀疏，大量 token 的 gate 值接近 0，少量 token 接近 1。这种稀疏性并非显式约束，而是模型自然学到的。

进一步分析显示，该稀疏是 query-dependent 的：每个 query token 自行决定 value 信息的通过比例。换言之，gate 在 soft 的 attention 之后再增加一道 hard 的 token-level 过滤。

两点合起来：先通过非线性扩展表达能力，再通过 query-level 过滤抑制噪声 value，效果叠加。

Attention sink 的消失

Attention sink 由 Xiao et al. 在 Efficient Streaming Language Models with Attention Sinks（2023）中提出：训练后的 LLM 倾向于将大量 attention 权重分配给序列开头的若干 token（通常为 BOS），即使这些 token 在语义上不携带信息。该现象在主流 LLM 中普遍存在，被认为源自 softmax 必须输出和为 1 的概率分布，模型在无明确目标时将权重堆积在序列开头。

Attention sink 的影响有两方面：一是占用 attention 容量；二是制约长上下文外推，因为模型在训练长度内将权重集中于前几个 token，超出训练长度时该模式难以迁移。后续工作（StreamingLLM、Sink Attention 等）多从规避或利用 sink 入手，未能消除该现象。

Gated Attention 在不显式针对 sink 的前提下使该现象大幅减弱：加入 sigmoid gate 后，BOS token 上的"亮带"在 attention map 中明显弱化。

Figure 2：上排为 baseline，第一个 token 的 attention 平均占 46.7%，layer 21 上 BOS 形成显著亮线；下排为加入 SDPA 输出 gate 后，比例降至 4.8%，亮线大幅减弱

机制上的解释：在标准 attention 中，softmax 强制概率和为 1，模型在某一行不需要任何 value 时只能将权重分配给某个无害的 token。加入 sigmoid gate 之后，softmax 仍可正常分配概率，gate 可乘以 0 关闭整个 head 的输出。该额外通路为 softmax 提供了释放途径，sink 不再必要。

直接收益体现在长上下文外推。论文在 RULER 长上下文 benchmark 上对比了标准 attention 与 gated attention，后者在训练长度（如 32K）外的衰减更平缓。

训练稳定性

Gated attention 同时显著减少了训练 loss spike。

Loss spike 是大模型训练中的常见问题。15B 以上规模的模型训练中后期偶尔出现 loss 突升，可能直接发散，或需要回滚数千步重训。常见缓解方案包括 embedding scale、QK norm、加严 gradient clipping、降低学习率等。

作者观察到，加入 SDPA 输出 gate 之后，1.7B 与 15B 模型的训练曲线显著平滑，原本会出现 spike 的位置不再出现。这使得训练可使用更大的 learning rate，scaling 行为更接近线性。

论文给出的解释指向 hidden states 的 massive activation：加 gate 后 SDPA 输出的 massive activation 显著降低，对应 BF16 训练中的数值误差减小，从而抑制了 spike 的形成（参考 Budzinskiy et al. 2025）。论文进一步追踪到 massive activation 通常起源于 layer 5 的 FFN 输出。

一些细节

为什么必须是 head-specific。head-shared gate（所有 head 共用一个 $W_{gate}$）效果显著弱于 head-specific。作者将原因归为 head 之间学到的特征语义不同，统一调制会强制同步其稀疏 pattern，损失 head 多样性。head-specific 增加的参数量较小（每层多一个 $\text{head_dim} \times \text{hidden_dim}$ 的矩阵），相对收益更优。

sigmoid 优于 SiLU。论文主要比较 sigmoid 与 SiLU 两种激活函数。SiLU 表现接近但略弱（PPL、MMLU、GSM8k 等指标均小幅低于 sigmoid）。sigmoid 的优势在于输出区间为 $[0, 1]$ 且天然倾向稀疏，许多 gate 值被推向 0，对应 hard 过滤。

multiplicative 优于 additive。将 gate 改为加法（$\text{out} = \text{attn} + \text{gate}$）后效果有所下降。论文给出的解释是：加法形式虽然在 G1 位置仍然经过一道 SiLU 引入了一定非线性，但其增益小于乘法形式。

与 NSA、SwiGLU 的关系。NSA 中已隐式包含 gating，但与 sparse 设计耦合，未单独剥离评估。FFN 中的 SwiGLU 同样是 gating 的一种形式，但作用于 channel 维度；本文的 gate 位于 attention 内部，作用于 token 维度，两者位置不冲突。论文在已含 SwiGLU 的模型上加入 gated attention，收益仍然存在。

Attention sink 是减弱还是消除。论文在不同位置使用了不同表述：摘要中写为 “mitigates”，引言、第 5.2 节与结论中则写为 “eliminates”。BOS token 的 attention 比例从 46.7% 降至 4.8%（layer 21 上从 83% 降至 4%），实际效果接近消除，但仍未严格降至零。

小结

Gated Attention 在 Scaled Dot-Product Attention 输出后引入一个 head-specific 的 elementwise sigmoid gate，代码改动约十行。在 30 个变体的系统比较下，该方案 PPL 显著下降、训练 loss spike 大幅减少、长上下文外推性能提升、attention sink 现象大幅减弱。

本文的方法论贡献在于将 gating 从已有工作的多种混淆设计中分离出来，单独评估其位置与形式的影响。Switch Heads 将增益归于路由、NSA 将增益归于 sparse attention 设计，本文的实验显示 gate 本身在两种场景中均贡献了相当大的部分。