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

文章标题：门控注意力在大型语言模型中的应用：非线性、稀疏性与无注意力沉溺
作者/机构：Zihan Qiu∗1, Zekun Wang∗1, Bo Zheng∗1, Zeyu Huang∗2, Kaiyue Wen3, Songlin Yang4, Rui Men1, Le Yu1, Fei Huang1, Suozhi Huang5, Dayiheng LiuB1, Jingren Zhou1, Junyang LinB1
1 Qwen Team, Alibaba Group 2 University of Edinburgh 3 Stanford University 4 MIT 5 Tsinghua University

A1 主要贡献

本文系统性地研究了在标准 Softmax 注意力机制中引入门控（gating）机制的影响。

  • 核心问题：尽管门控机制在从早期的 LSTM 到现代的状态空间模型和注意力机制中被广泛使用，但其具体作用和影响尚未得到充分探讨。现有研究往往将其与其他架构因素（如专家路由）混为一谈，导致难以评估门控本身的真实贡献。

  • 研究目标：本文旨在通过全面的实验，系统性地研究门控增强的 Softmax 注意力变体，以厘清门控机制在大型语言模型中的具体作用和价值。

  • 核心发现与创新点：

    1. 最佳门控方案：通过对 30 多种 15B 参数的混合专家（MoE）模型和 1.7B 参数的密集模型进行训练（数据量高达 3.5 万亿 token），研究发现一个简单的修改——在缩放点积注意力（Scaled Dot-Product Attention, SDPA）之后应用一个逐头的（head-specific）Sigmoid 门控（图 1 中的 G1 位置），能够持续且显著地提升模型性能（例如，PPL 降低 0.2，MMLU 提升 2 个点）。
    2. 提升训练稳定性与扩展性：SDPA 输出门控不仅提升了性能，还增强了训练的稳定性，几乎消除了损失尖峰（loss spikes），使得模型能够容忍更高的学习率，并改善了模型的可扩展性（scaling properties）（图 1 右图）。
    3. 有效性归因：非线性与稀疏性：通过对比不同门控位置和计算变体，研究将门控的有效性归因于两个关键因素：
      • 非线性（Non-linearity）：在 Softmax 注意力中，值（value）投影（$W_v$）和输出（dense）投影（$W_O$）是两个连续的线性层，可以被视为一个低秩线性映射。在它们之间（如 G1 或 G2 位置）引入门控，增加了非线性，从而提升了该低秩变换的表达能力。
      • 稀疏性（Sparsity）：有效的门控（特别是 G1 位置的门控）会产生依赖于查询（query-dependent）的稀疏门控分数，从而对 SDPA 的输出进行调制，过滤无关信息。
    4. 消除“注意力沉溺”并提升长文本性能：研究发现，这种稀疏门控机制能够有效缓解“注意力沉溺”（attention sink）现象——即初始 token 获得过高注意力权重的问题（图 2）。经过 3.5 万亿 token 训练后，带有门控的模型（密集和 MoE 模型）均未表现出注意力沉溺。此外，这些模型在长文本外推性能上表现更优，例如在 RULER 基准测试上取得了超过 10 个点的增益。

  • 开源贡献：本文将开源其无注意力沉溺的模型及相关代码，以促进未来的研究。

图 1：左：研究中应用门控操作的位置。中：在不同位置应用门控的 15B MoE 模型的性能比较（测试 PPL 和 MMLU）。在 SDPA 之后（G1）应用门控取得了最佳的综合效果。在 Value 层之后（G2）应用门控也显示出显著的改进，尤其是在 PPL 方面。右：在相同超参数下，基线模型和 SDPA 门控的 1.7B 密集模型在 3.5T token 上的训练损失比较（平滑系数 0.9）。门控带来了更低的最终损失和显著增强的训练稳定性，减少了损失尖峰。这种稳定性使得模型可能使用更高的学习率，并有助于更好的扩展。
图 1：左：研究中应用门控操作的位置。中：在不同位置应用门控的 15B MoE 模型的性能比较（测试 PPL 和 MMLU）。在 SDPA 之后（G1）应用门控取得了最佳的综合效果。在 Value 层之后（G2）应用门控也显示出显著的改进，尤其是在 PPL 方面。右：在相同超参数下，基线模型和 SDPA 门控的 1.7B 密集模型在 3.5T token 上的训练损失比较（平滑系数 0.9）。门控带来了更低的最终损失和显著增强的训练稳定性，减少了损失尖峰。这种稳定性使得模型可能使用更高的学习率，并有助于更好的扩展。

A3 背景知识

多头 Softmax 注意力机制回顾

给定输入 $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】可分为四个阶段。

  • QKV 线性投影：输入 $X$ 通过学习到的权重矩阵 $W_Q, W_K, W_V \in \mathbb{R}^{d_{model} \times d_k}$ 被线性变换为查询（Queries）$Q$、键（Keys）$K$ 和值（Values）$V$，其中 $Q, K, V \in \mathbb{R}^{n \times d_k}$。

    公式1
    公式1

  • 缩放点积注意力 (SDPA)：该步骤计算查询和键之间的注意力分数，然后进行 softmax 归一化。输出是值的加权和。

    公式2
    公式2

    其中，$QK^T \in \mathbb{R}^{n \times n}$ 表示缩放点积相似度矩阵，softmax(·) 确保注意力权重非负且每行之和为 1。

  • 多头拼接：在多头注意力中，上述过程对 $h$ 个头并行重复，每个头都有自己的投影矩阵 $W_Q^i, W_K^i, W_V^i$。所有头的输出被拼接起来。

    MultiHead($Q, K, V$) = Concat(head$_1$, ..., head$_h$)，其中 head$_i$ = Attention($QW_Q^i, KW_K^i, VW_V^i$)。

    图 2：左：分配给初始 token 的注意力比例（在测试困惑度数据集上逐层计算）。基线模型存在严重的注意力沉溺问题，各层平均有 46.7% 的注意力分数指向第一个 token。引入门控有效缓解了这一问题，将该比例降低至 4.8%。右：每个头的平均注意力图权重。基线模型的第 21 层表现出强烈的注意力沉溺（83% 在第一个 token 上），而门控显著减少了这一现象（4%）。在最终输出层，门控增强了模型关注序列中单个 token 的现有趋势。
    图 2：左：分配给初始 token 的注意力比例（在测试困惑度数据集上逐层计算）。基线模型存在严重的注意力沉溺问题，各层平均有 46.7% 的注意力分数指向第一个 token。引入门控有效缓解了这一问题，将该比例降低至 4.8%。右：每个头的平均注意力图权重。基线模型的第 21 层表现出强烈的注意力沉溺（83% 在第一个 token 上），而门控显著减少了这一现象（4%）。在最终输出层，门控增强了模型关注序列中单个 token 的现有趋势。

  • 最终输出层：拼接后的 SDPA 输出通过一个输出层 $W_o \in \mathbb{R}^{hd_k \times d_{model}}$。

    公式4
    公式4

A2 方法细节

使用门控机制增强注意力层

  • 门控机制的形式化定义：门控机制可以被形式化地表示为：

    公式5
    公式5

    其中，$Y$ 是待调制的输入，$X$ 是用于计算门控分数的另一个输入，$W_{\theta}$ 是门控的可学习参数，$\sigma$ 是一个激活函数（例如 sigmoid），$Y'$ 是门控后的输出。门控分数 $\sigma(XW_{\theta})$ 作为一个动态滤波器，通过选择性地保留或擦除 $Y$ 的特征来控制信息流。

  • 门控变体的综合研究：本文在注意力层内全面研究了几种门控机制的变体，重点关注以下五个方面：

    1. 位置（Positions）：研究在不同位置应用门控的效果，如图 1（左）所示：
      • (a) 在 Q、K、V 投影之后（公式 1），对应图 1（左）中的 G2、G3、G4 位置。
      • (b) 在 SDPA 输出之后（公式 3），即 G1 位置。
      • (c) 在最终拼接的多头注意力输出之后（公式 4），即 G5 位置。
    2. 粒度（Granularity）：考虑两种门控分数的粒度级别：
      • (a) 逐头（Headwise）：一个单一的标量门控分数调制整个注意力头的输出。
      • (b) 逐元素（Elementwise）：门控分数是与 $Y$ 维度相同的向量，允许进行细粒度的、逐维度的调制。
    3. 逐头特定或共享（Head Specific or Shared）：考虑到注意力的多头特性，进一步考虑：
      • (a) 逐头特定（Head-Specific）：每个注意力头有其特定的门控分数，实现对每个头的独立调制。
      • (b) 头共享（Head-Shared）：$W_{\theta}$ 和门控分数在所有头之间共享。
    4. 乘法或加法（Multiplicative or additive）：对于将门控分数应用于 $Y$ 的方式，考虑：
      • (a) 乘法门控（Multiplicative Gating）：门控输出 $Y' = Y \cdot \sigma(X\theta)$。
      • (b) 加法门控（Additive Gating）：$Y' = Y + \sigma(X\theta)$。
    5. 激活函数（Activation Function）：主要考虑两种常见的激活函数：SiLU【索引 41，GLU variants improve transformer，Noam Shazeer，2020，arXiv】和 Sigmoid。SiLU 因其无界的输出范围仅用于加法门控，而 Sigmoid 的输出分数在 [0, 1] 范围内。此外，为进一步剖析门控有效性的机制，还考虑了恒等映射（Identity Mapping）或 RMSNorm【索引 60，Root mean square layer normalization，Biao Zhang and Rico Sennrich，2019，Advances in Neural Information Processing Systems】（详见 4.1 节）。

  • 默认配置：除非另有说明，本文采用的是逐头特定（head-specific）、乘法（multiplicative）的门控，并使用 Sigmoid 激活函数（$\sigma(x) = \frac{1}{1+e^{-x}}$）。

表 1：门控变体性能与结果。我们在 400B token 上训练了 15A2B MoE 模型。dk 是头维度，dmodel 是模型的隐藏维度，n 是 token 数量。q 指的是查询头的数量，k 指的是键值头的数量。‘Act Func’ 是公式 5 中的激活函数。‘Score Shape’ 是输入 X ∈ Rn,dmodel 的门控分数形状。‘added param’ 表示增加的参数（百万）。
表 1：门控变体性能与结果。我们在 400B token 上训练了 15A2B MoE 模型。dk 是头维度，dmodel 是模型的隐藏维度，n 是 token 数量。q 指的是查询头的数量，k 指的是键值头的数量。‘Act Func’ 是公式 5 中的激活函数。‘Score Shape’ 是输入 X ∈ Rn,dmodel 的门控分数形状。‘added param’ 表示增加的参数（百万）。

A4 实验环境

  • 模型架构：

    • 混合专家模型 (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】。
    • 密集模型 (Dense)：1.7B 总参数。

  • 数据集：

    • 在一个包含 3.5 万亿高质量 token 的数据集子集上进行训练。
    • 数据内容涵盖多语言、数学和通用知识。
    • 上下文序列长度设置为 4096。

  • 硬件配置：未明确说明。

  • 软件与训练配置：

    • 优化器：AdamW，超参数遵循默认值。
    • 学习率与批大小：
      • MoE-15A2B 模型：学习率从 0 预热 1k 步到最大值 2e-3，然后使用余弦衰减到 3e-5。全局批大小为 1024。
      • Dense-1.7B 模型 (400B tokens)：最大学习率为 4e-3，批大小为 1024。
      • Dense-1.7B 模型 (3.5T tokens)：最大学习率增加到 4.5e-3，批大小增加到 2048。
    • 延迟：由于门控引入的参数和浮点运算量很小，其引入的壁钟时间延迟小于 2%。

  • 评测：

    • 少样本评测 (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】 (中文能力)。
    • 语言建模困惑度 (PPL)：在包括英语、中文、代码、数学、法律和文学等领域的留出测试集上报告。

A4 实验结果

主要结果

门控注意力在 MoE 模型中的表现

  • 实验设置：在 MoE-15A2B 模型上进行对比实验，训练 400B token。为公平比较，基线模型通过增加键值头、查询头或专家数量的方式增加了相当或更多的参数。
  • 实验结论 (Tab. 1)：
    1. SDPA (G1) 和 Value (G2) 输出门控最有效：在这两个位置插入门控，相比其他变体，取得了更低的 PPL 和更好的基准性能。
    2. 逐头特定门控至关重要：在 G1 和 G2 位置应用逐头（headwise）门控，虽然只引入极少参数（<2M），但带来了显著提升。若在头之间共享门控分数，性能增益会减小，这凸显了为不同注意力头应用不同门控分数的重要性。
    3. 乘法门控优于加法门控：加法 SDPA 输出门控虽然优于基线，但效果不如乘法门控。
    4. Sigmoid 激活函数更佳：在最有效的门控配置中，将激活函数替换为 SiLU 会导致性能提升减少。
  • 总结：在 Value 层 (G2) 和 SDPA 输出 (G1) 添加门控可将 PPL 降低超过 0.2，优于各种参数扩展基线。其中，G1 位置的门控在 PPL 和基准测试上均表现最佳。

门控注意力在密集模型中的表现

  • 实验设置：在 1.7B 密集模型上验证 SDPA 输出 Sigmoid 门控的效果。为保持参数量不变，使用门控时会减小 FFN 的宽度。实验涵盖了不同训练数据量、模型层数、学习率和批大小的设置。
  • 实验结论 (Tab. 2)：
    1. 门控在各种设置下均有效：无论模型配置、训练数据量还是超参数如何变化，应用 SDPA 输出门控都能持续带来收益。
    2. 门控提升稳定性并促进扩展：在 3.5T token 的训练设置下，门控显著提升了训练稳定性，大幅减少了损失尖峰（见图 1 右图）。当提高最大学习率时，基线模型会出现收敛问题，而带有门控的模型则能稳定训练并获得性能提升。相比之下，添加 sandwich norm【索引 16，Cogview: Mastering text-to-image generation via transformers，Ming Ding et al.，2021】虽然能恢复收敛，但性能提升微乎其微。
  • 总结：SDPA 逐元素门控是增强注意力机制的最有效方法。它能使密集 Transformer 在更大的批大小和学习率下稳定训练，从而获得更好的性能。

表 2：不同方法在不同学习率、批大小和模型配置下的性能表现。'SDPA' 指的是在公式 3 的 SDPA 之后应用 sigmoid 门控，'sandwitch norm'【索引 16】表示在将 attention/ffn 输出添加到残差连接之前对其进行归一化。使用门控时，我们减小了 FFN 的宽度，以使所有方法的参数数量相同。'-' 表示模型在训练过程中发散。
表 2：不同方法在不同学习率、批大小和模型配置下的性能表现。'SDPA' 指的是在公式 3 的 SDPA 之后应用 sigmoid 门控，'sandwitch norm'【索引 16】表示在将 attention/ffn 输出添加到残差连接之前对其进行归一化。使用门控时，我们减小了 FFN 的宽度，以使所有方法的参数数量相同。'-' 表示模型在训练过程中发散。

分析：非线性、稀疏性与无注意力沉溺

非线性提升了注意力中低秩映射的表达能力

  • 动机：在多头注意力中，第 $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】。
  • 实验验证 (Tab. 3)：
    • 在 G1 位置应用 RMSNorm（引入非线性，几乎无参数增加）也能显著降低 PPL，验证了非线性的重要性。
    • 在 G2 位置添加门控对应于对值输出 $X_j W_V^k$ 应用非线性（公式 7），而在 G1 位置添加门控或归一化则对应于对加权和后的值输出应用非线性（公式 8）。这两种方式都缓解了低秩问题。
    • 在 $W_O$ 之后的 G5 位置添加门控无效，因为它没有解决 $W_V$ 和 $W_O$ 之间缺乏非线性的问题。
    • 仅添加 SiLU 非线性（无参数）也能带来 modest 的 PPL 降低。移除加法门控中的 SiLU 会进一步削弱其增益。
  • 结论：有效门控变体的性能提升，很大程度上归因于在 $W_V$ 和 $W_O$ 之间引入了非线性。

表 3：不同（非）线性增强方法的性能。
表 3：不同（非）线性增强方法的性能。

门控引入了输入依赖的稀疏性

  • 分析 (Tab. 4, Fig. 3)：对 G1 和 G2 位置的门控分数进行分析，发现：
    1. 有效门控的分数是稀疏的：SDPA 输出门控（G1）的平均门控分数最低，且其分数分布高度集中于 0 附近，表现出最强的稀疏性，这与其最佳性能相符。
    2. 逐头特定的稀疏性很重要：强制共享门控分数会提高整体门控分数并削弱性能增益，表明不同注意力头需要不同的稀疏度。
    3. 查询依赖性很重要：Value 门控（G2）的分数高于 SDPA 输出门控（G1），且性能较差。G1 的门控分数依赖于当前查询 $X_i$，而 G2 的门控分数依赖于过去的键和值 $X_j$。这表明，当稀疏性依赖于查询时，能更有效地过滤掉与查询无关的上下文信息。输入无关的门控（通过可学习参数生成）性能提升较小，也证实了这一点。
    4. 降低稀疏性会损害性能：通过修改 Sigmoid 函数（NS-sigmoid，公式 9）使其输出在 [0.5, 1.0] 之间，从而在保留非线性的同时去除稀疏性。实验表明，这种非稀疏门控的增益远不如标准的 SDPA 输出 Sigmoid 门控。

图 3：SDPA 逐元素（左）、Value 逐元素（中）以及 SDPA 逐元素头共享门控（右）的门控分数均值和分布。大多数门控分数小于 0.5，表明门控分数是稀疏的。其中，SDPA 输出门控分数表现出最强的稀疏性。
图 3：SDPA 逐元素（左）、Value 逐元素（中）以及 SDPA 逐元素头共享门控（右）的门控分数均值和分布。大多数门控分数小于 0.5，表明门控分数是稀疏的。其中，SDPA 输出门控分数表现出最强的稀疏性。

SDPA 输出门控减少了注意力沉溺

  • 分析 (Tab. 4, Fig. 2)：基于门控引入输入依赖稀疏性的观察，研究了其对“注意力沉溺”现象的影响。
    1. G1 门控消除注意力沉溺：在 SDPA 输出（G1）应用逐头特定和逐元素门控，能大幅减少分配给第一个 token 的注意力分数（从 46.7% 降至 4.8%），并降低了隐藏状态中的“巨幅激活”（massive activations）。
    2. 巨幅激活非注意力沉溺的必要条件：在 Value 投影后（G2）应用门控，虽然也减少了巨幅激活，但并未减少对第一个 token 的注意力分数。这表明巨幅激活不是注意力沉溺的先决条件。
    3. 稀疏性是关键：降低门控的输入依赖性或使用非稀疏的 NS-sigmoid 会加剧巨幅激活和注意力沉溺。
  • 结论：输入依赖、逐头特定的 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】上进行评测。
  • 实验结论 (Tab. 5)：
    1. 在 32k 上下文长度内，门控模型略优于基线模型。
    2. 当使用 YaRN 扩展到 128k 时，所有模型在原始 32k 范围内的性能都有所下降，但门控模型的下降幅度更小。
    3. 在 64k 和 128k 的上下文长度上，门控注意力模型显著优于基线模型。
  • 假设：添加门控有助于模型适应上下文长度的扩展。一个可能的解释是，基线模型依赖注意力沉溺来调整注意力分数的分布。当 YaRN 等技术修改 RoPE base 时，注意力沉溺模式可能难以在无训练的情况下适应，导致性能下降。相比之下，门控模型主要依靠输入依赖的门控分数来控制信息流，因此对这类变化更具鲁棒性。

表 5：不同方法在不同序列长度下的性能表现。‘YaRN Extended’ 表示扩展上下文长度的变体。‘(values)’ 表示扩展上下文长度后性能下降的值。
表 5：不同方法在不同序列长度下的性能表现。‘YaRN Extended’ 表示扩展上下文长度的变体。‘(values)’ 表示扩展上下文长度后性能下降的值。

A7 补充细节

相关工作

  • 神经网络中的门控机制：

    • 历史：门控机制在神经网络中被广泛采用。早期的工作如 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 的内在归一化依赖性驱动的。
    • 缓解方法与本文贡献：已有工作尝试通过修改 Softmax 注意力来缓解注意力沉溺。本文证明了在 SDPA 之后应用稀疏门控可以在大规模（1B-15B 参数）模型上消除注意力沉溺，即便在 3.5T token 的训练后也是如此。此外，本文揭示了消除注意力沉溺对上下文长度扩展的潜在好处。

A5 结论

  • 核心结论：本文系统性地研究了门控机制在标准 Softmax 注意力中的作用，揭示了其对模型性能、训练稳定性和注意力动态的显著影响。通过对 30 多种 15B MoE 和 1.7B 密集模型（在高达 3.5T token 上训练）的广泛实验比较，证明了在缩放点积注意力之后应用一个 Sigmoid 门控是最有效的改进方法。

  • 机制解释：这个简单的机制通过增强非线性、引入输入依赖的稀疏性，有效消除了如“注意力沉溺”这样的低效现象。

  • 额外收益：门控机制促进了上下文长度的扩展，使得模型能够有效地泛化到更长的序列而无需重新训练。

  • 开源贡献：本文发布了首批无注意力沉溺的模型，相信这些经验性验证将为设计下一代先进基础模型铺平道路。

局限性

  • 本文主要通过一系列消融研究来分析注意力门控的原因和影响，但承认存在一些局限性。
  • 非线性对注意力和整体训练过程动态的更广泛影响仍有待探索。
  • 尽管观察到消除注意力沉溺可以改善长上下文扩展场景下的性能，但本文并未提供关于注意力沉溺如何影响模型泛化到更长序列能力的严格理论解释。

A6 附录

A.1 Switch Head 基线实验

  • 背景：Switch Head 论文【索引 10，Switchhead: Accelerating transformers with mixture-of-experts attention，Robert Csordas et al.，2024b】通过在注意力中引入稀疏激活（每个 token 从一组键/值/输出专家中选择 top-k 个）实现了与基线相当的结果。
  • 实验发现 (Tab. 6)：本文的实验发现了一个有趣的趋势：增加激活的 kv 专家数量（在相同专家参数设置下）似乎对 PPL 有一些改善，但对整体基准性能的增益不明显。值得注意的是，基准分数和 PPL 的最佳结果是由 Switch v 1top1（表 6 第 6 行）实现的，这等同于直接在值层输出上应用 Sigmoid 门控。
  • 结论：这些发现表明，在这些实验中，性能提升的主要驱动因素可能是门控机制本身，而不仅仅是专家路由。

表 6：不同 switch head 方法在不同参数增加和配置下的性能。‘switch kv’ 和 ‘switch v’ 分别指在键值和值组件中引入选择性计算。‘Switch kv, 8top8’ 意味着有 8 个键和值映射专家，每个 token 选择 top8 专家。注意‘Switch v, 1top1’ 等同于表 1 行 (11) 中的 v Headwise Gate。
表 6：不同 switch head 方法在不同参数增加和配置下的性能。‘switch kv’ 和 ‘switch v’ 分别指在键值和值组件中引入选择性计算。‘Switch kv, 8top8’ 意味着有 8 个键和值映射专家，每个 token 选择 top8 专家。注意‘Switch v, 1top1’ 等同于表 1 行 (11) 中的 v Headwise Gate。

A.2 关于稀疏门控分数的更多讨论

  • 门控对隐藏状态的影响 (Fig. 4)：分析了在 G1 位置应用门控前后 SDPA 隐藏状态的均值。结果显示：(1) 门控后，隐藏状态的平均绝对值从 0.71 降至 0.05，与通常较小的门控分数相对应；(2) 门控后的隐藏状态与基线模型非常相似，表明门控可能起到了类似注意力沉溺的作用，即过滤掉无关信息。

  • 门控对稀疏性的影响 (Fig. 5)：分析了门控前后低于特定阈值的隐藏状态比例。结果显示：(1) 门控后，隐藏状态在不同阈值下的稀疏性显著增加。(2) 如果将门控前的隐藏状态乘以平均门控分数，稀疏性的增加幅度小于原始门控。这表明，是稀疏的门控分数（而不仅仅是缩小数值）增强了隐藏状态的稀疏性。

图 4：门控前后的平均绝对值。基线和门控后的值相似。
图 4：门控前后的平均绝对值。基线和门控后的值相似。

图 5：门控后低于阈值的 SDPA 输出值比例（左：1e-2，右：1e-3）。我们还包括了通过将平均门控分数与门控前隐藏状态相乘得到的稀疏性度量。
图 5：门控后低于阈值的 SDPA 输出值比例（左：1e-2，右：1e-3）。我们还包括了通过将平均门控分数与门控前隐藏状态相乘得到的稀疏性度量。

A.3 逐层的巨幅激活和注意力沉溺

  • 逐层分析 (Fig. 6)：对模型内的巨幅激活和注意力沉溺（第一个 token 的注意力分数）进行了逐层比较分析。
    • 基线 (第 1 行)：第 6 层 FFN 的输出包含巨幅激活，这些激活被添加到残差流中，导致后续层的残差中持续存在大激活。相应地，从第 6 层开始出现显著的注意力沉溺现象。
    • SDPA 门控 (第 2 行)：网络早期层的输出总体上保持较小，巨幅激活随着层深的增加而逐渐增长。值得注意的是，在网络的任何层都没有观察到显著的注意力沉溺现象。
    • Value 层门控 (第 3 行)：模型表现出与第 2 行相似的巨幅激活，但仍然存在一定程度的注意力沉溺现象。这表明巨幅激活不是注意力沉溺出现的必要条件。
    • 降低稀疏性的门控 (第 4-5 行)：当强制在不同头之间共享门控分数或修改激活函数以抑制稀疏性时，门控引入的稀疏性降低。在这些情况下，巨幅激活和注意力沉溺都变得与基线相当。
  • 结论：这些观察表明，在注意力机制内引入足够的稀疏性可能有助于减轻巨幅激活的发生。

图 6：不同门控配置下巨幅激活和注意力沉溺现象的比较。第 1 行（基线）：第 6 层后出现显著的巨幅激活和注意力沉溺。第 2 行（SDPA 门控）：激活减少，未观察到注意力沉溺。第 3 行（Value 层门控）：激活与第 2 行相似，但存在残余的注意力沉溺。第 4-5 行（通过跨头共享和 NS-sigmoid 减少稀疏性）：巨幅激活和注意力沉溺与基线相似。
图 6：不同门控配置下巨幅激活和注意力沉溺现象的比较。第 1 行（基线）：第 6 层后出现显著的巨幅激活和注意力沉溺。第 2 行（SDPA 门控）：激活减少，未观察到注意力沉溺。第 3 行（Value 层门控）：激活与第 2 行相似，但存在残余的注意力沉溺。第 4-5 行（通过跨头共享和 NS-sigmoid 减少稀疏性）：巨幅激活和注意力沉溺与基线相似。

A.4 更多逐层门控分数分析

  • 分析 (Fig. 7)：以 SDPA 输出门控（逐元素/逐头）为基线，分析了在两个额外约束下门控分数的分布：(1) 强制不同头之间使用相同的门控分数（左图），和 (2) 限制门控分数的最小值（右图）。
  • 发现：当强制共享门控分数时，大多数层的门控分数增加。这表明不同的头需要不同的稀疏度，凸显了逐头特定门控机制的重要性。

图 7：SDPA 输出门控变体在不同约束下门控分数的分布。
图 7：SDPA 输出门控变体在不同约束下门控分数的分布。

A.5 稳定训练的其他尝试

  • 动机：观察到 sandwich normalization【索引 16，Cogview: Mastering text-to-image generation via transformers，Ming Ding et al.，2021】和门控机制都消除了巨幅激活并提高了训练稳定性，这促使研究人员探索更简单的方法来防止残差中的大激活。
  • 实验：引入了一个裁剪（clipping）操作，在注意力层和 FFN 层的输出进入残差连接之前，将其值限制在 (-clip, clip) 范围内。
  • 结果：发现无论 clip 值设置为 300 还是 100，模型在 8e-3 的学习率下仍然遇到收敛问题。
  • 结论：这表明 pre-norm 模型训练的不稳定性不仅仅是由于残差中的大激活引起的。任何产生大输出的层都可能导致稳定性问题，这表明需要进一步研究训练不稳定性的根本原因。