一、GatedAttention [2025]

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

1. 门控机制（gating mechanisms ）已被广泛应用，从早期的 LSTMs、Highway Networks，到近期的 state space models、linear attention 、以及 softmax attention 。然而，现有文献很少深入探讨 gating 的具体作用。本文通过全面实验，系统研究了 gating-augmented 的 softmax attention 变体。

   具体而言，我们在 3.5 trillion token dataset 上训练了 15B Mixture-of-Experts (MoE) models and 1.7B dense models 的 30 种变体，并进行了综合对比。核心发现是：在 Scaled Dot-Product Attention: SDPA 后添加一个 head-specific sigmoid gate 这一简单修改，能够一致地提升模型性能。该 modification 还增强了 training 的稳定性，支持更大的学习率，并改善了模型的 scaling 特性。通过对比不同 gating positions and computational 的变体，我们将其有效性归因于两个关键因素：

   • (1)：在 softmax attention 的低秩映射（low-rank mapping）中引入非线性（non-linearity）。

   • (2)：采用 query-dependent sparse gating scores 来调制 SDPA ouput。

   值得注意的是，这种 sparse gating 机制缓解了 "注意力汇聚" （'attention sink'）问题，提升了长上下文外推（long-context extrapolation）性能。本文还公开了相关代码和模型，以促进未来研究。

2. 门控机制（gating mechanisms ）在神经网络中已十分成熟。早期架构如 LSTMs、Highway Networks、和 GRUs 率先采用 gating 来控制跨 time steps or layers 的 information flow，并改善梯度传播。这一原则在现代架构中依然延续。近期的 sequence modeling 工作，包括 state-space models 和 attention 机制普遍应用 gating，通常用于调制 outputs of token-mixer components。尽管 gating 机制被广泛采用且取得了实证成功，但除了最初的直觉外，其功能和影响尚未得到充分探索。

   对 gating 机制的理解不足，阻碍了对其真正贡献的评估，尤其是当它与其他 architectural factors 相互干扰时。例如，Switch Heads 引入 a sigmoid gating 来选择 top-K attention head experts ，但我们的实验揭示了一个有趣的发现（附录A.1）：即使简化为 a single expert （此时 gate 仅调制 value output），仍能获得显著的性能提升。这有力地表明，gating 本身具有重要的内在价值，与 routing 机制相互独立。类似地，在 Native Sparse Attention: NSA （《Native sparse attention: Hardware-aligned and natively trainable sparse attention》）中，尽管展示了整体性能提升，但并未将其 gating 机制的贡献与 sparse attention design 本身的效果区分开来。这些情况凸显了有必要严格区分 gating 与其他 architectural components 的作用。

   本文研究了 standard softmax attention（《Attention is all you need》）中的 gating 机制。具体而言，我们在不同位置引入 gating（Figure 1）：

   • query projection 之后（$G_4$$G_4$）、key projection 之后（$G_3$$G_3$）、 value projection 之后（$G_2$$G_2$）。

   • Scaled Dot Product Attention: SDPA 的输出之后（$G_1$$G_1$）。

   • 以及 final dense output layer 之后（$G_5$$G_5$）。

   我们的探索涵盖了多种 gating 变体，包括 elementwise and headwise、head-specific and head-shared ，以及 additive and multiplicative forms。研究发现：

   • (i)：应用 SDPA output head-specific gating （$G_1$$G_1$）能带来最显著的性能提升（例如，困惑度 PPL 降低高达 0.2，MMLU 数据集提升 2 个百分点）。

   • (ii)：SDPA output gating 还改善了 training 的稳定性，几乎消除了 loss spikes，支持更大的学习率并增强了 model scalability。

   

   我们确定了 gating 有效性的两个主要因素：

   • (i)：非线性（Non-Linearity）。两个连续的线性层——value projection （${\mathbf{W}}_V$$\mathbf W_V$）和 dense projection （${\mathbf{W}}_O$$\mathbf W_O$ ）——可以重写为一个低秩线性投影（low-rank linear projection）。因此，在 $G_1$$G_1$ 或 $G_2$$G_2$ 位置通过 gating 来引入非线性，能够提高这种 low-rank linear transformation 的表达能力（第 1.4.1节）。

   • (ii)：稀疏性（Sparsity）。尽管 non-linear gating 变体总能提升性能，但增益程度有所不同。进一步分析表明，gating scores 的显著稀疏性是另一个关键因素，它为 SDPA outputs 引入了 input-dependent sparsity （第 1.4.2 节）。此外，sparse gating 消除了 attention sink（《Efficient streaming language models with attention sinks》）——即initial tokens 不成比例地主导 attention scores （ Figure 2，第 1.4.3节）。先前的研究将注意力汇聚（ attention sink ）解释为由于 non-negative softmax normalization 导致的 accumulation of redundant attention。实证结果表明，当在SDPA output 处应用 query-dependent sparse gating 时，我们的 dense models 和 MoE models（均在 3.5T tokens 上训练）均未出现 attention sink 现象。此外，这些模型在长度泛化（length generalization）方面表现更优，在 RULER 数据集上的得分提升超过 10 个百分点（第 1.4.4 节）。

     

   总之，本文强调了 standard attention layers 中的 gating 机制对模型性能和行为的影响。通过评估 gating 变体，我们发现其能够引入 non-linearity 和 sparsity，并消除 attention sinks。这些发现加深了我们对 gated attention 的理解。我们将开源 attention-sink-free models，以推动未来研究。

1.1 相关工作

1.1.1 神经网络中的 Gating

1. gating 机制已被广泛应用于神经网络。

   • 早期工作如 LSTMs 和 GRUs 引入 gating 来调节跨 time steps 的信息流（information flow），通过选择性地保留或丢弃信息来解决梯度消失/爆炸（gradient vanishing/exploding）问题。

   • Highway Networks 将这一概念扩展到前馈网络，使极深架构（very deep architectures）的成功训练成为可能。

   • SwiGLU （《Glu variants improve transformer》）将 gating 机制引入 Transformer 的 FFN layers，增强了其表达能力，并成为许多开源大型语言模型的标准组件。

     $$\begin{array}{c}\text{SwiGLU}(\mathbf{X})=\text{Swish}({\mathbf{W}}_1\mathbf{X}+{\mathbf{b}}_1)\odot ({\mathbf{W}}_2\mathbf{X}+{\mathbf{b}}_2)\\ \text{Swish}(z)=z\times \text{sigmoid}(\beta z)\end{array}$$

     $\beta =1$$\beta = 1$ 或者从数据中学习；当 $\beta =1$$\beta=1$ 时，Swish 激活函数退化为 SiLU 激活函数。

   • 多项关于 state-space models 和 Linear Attention 的研究（如 FLASH 、RetNet、Lightning Attention 和 Gated Delta Networks ），也整合了 gating 模块来控制 information of token-mixer modules 。

   • Forgetting Transformer 将 gating 机制应用于 output of softmax attention ，观察到显著的性能提升。

   尽管这些工作证明了 gating 的有效性，但对其精确机制和有效性原因的全面理解仍有待探索。这可能有助于更广泛地认识 gating 在 RNNs 之外的重要性，并促进更好的设计从而利用 gating 的独特优势。例如，虽然 Switch Heads、NSA 和 MoSA 采用 sigmoid-based gating 用于 selection ，但进一步研究分离 gating 的具体贡献可能会提供有价值的见解。通过对 standard transformers 中纳入类似 gating 机制的 baselines 模型进行对比，可能会对这些模型所提出的 selection 机制的有效性提供更精确的视角。

2. 与我们的工作最相关的是 Quantizable Transformers；该研究还发现，在 softmax attention 中应用 gating 可以缓解 encoder models（如 BERT 和 ViT ）中的极端注意力集中（extreme attention concentration）和隐状态异常值（outliers in hidden states）问题。

   但是，该工作主要利用 gating 来消除异常值以实现 model quantization。而我们详细分析了各种 gating 变体，通过 enhanced non-linearity and sparsity 以及改善 training 稳定性，揭示了它们的优势。基于这些见解，我们扩展了 gated attention models 的规模，证明了 gating 的广泛适用性和影响。

1.1.2 Attention Sink

1. 《Efficient streaming language models with attention sinks》正式提出了 'attention sink' 现象，即 specific tokens 获得 large attention scores。类似地，《Vision transformers need registers》 在 vision transformer 中发现，一些冗余的 tokens 充当 “寄存器” （'registers'）来存储 attention scores。

   后来，《Massive activations in large language models》 表明，与 massive activation values 相关的 tokens 也被分配了 excessive attention scores。然而，我们的工作揭示，在 output of value projection 处应用 gating 消除了 massive activations ，但 attention sink 仍然存在，这表明 massive activations 并非 attention sink 的必要条件。

2. 同样，《When attention sink emerges in language models: An empirical view》 将 attention sink 描述为 non-informative 'key biases'，这个 'key biases' 是冗余的注意力分数（redundant attention scores）；该论文认为 softmax 固有的归一化依赖性（normalization dependency）驱动了这种行为。

   修改 softmax attention 的实验尝试，例如用 unnormalized sigmoid attention 替换 softmax attention、 adding softmax attention gate or clip 、modifying softmax computation、以及修改 softmax 的分母，在缓解 attention sink 方面显示出潜力。

   我们的工作表明，在 SDPA 之后进行 sparse gating，能够在 dense models（1B 参数量）和 MoE models（15B 参数量）中消除 attention sink，即使在 3.5T tokens 上训练也是如此。此外，我们发现消除 attention sink 有助于 context-length extension 。

1.2 Gated-Attention Layer

1.2.1 预备知识：Multi-Head Softmax Attention

1. 给定一个 input $\mathbf{X}\in {\mathbb{R}}^{n\times d_{\text{model}}}$$\mathbf X \in \mathbb R^{n\times d_\text{model}}$ （其中 $n$$n$为序列长度，$d_{\text{model}}$$d_\text{model}$ 为模型维度），transformer 的 attention layer 的计算可分为四个阶段：

   • QKV Linear Projections：input $\mathbf{X}$$\mathbf X$ 通过学到的权重矩阵 ${\mathbf{W}}_Q,{\mathbf{W}}_K,{\mathbf{W}}_V\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$$\mathbf W_Q, \mathbf W_K, \mathbf W_V \in \mathbb R^{d_\text{model}\times d_k}$ 线性变换为 query $\mathbf{Q}$$\mathbf Q$、key $\mathbf{K}$$\mathbf K$ 和value $\mathbf{V}$$\mathbf V$，其中 $\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{n\times d_k}$$\mathbf Q, \mathbf K, \mathbf V\in \mathbb R^{n\times d_k}$：

     $$\mathbf{Q}=\mathbf{X}{\mathbf{W}}_Q,\mathbf{K}=\mathbf{X}{\mathbf{W}}_K,\mathbf{V}=\mathbf{X}{\mathbf{W}}_V$$

   • Scaled Product Dot-Product Attention (SDPA)：计算 queries 和 keys 之间的 attention scores，随后进行 softmax 归一化。output 是 values 的加权和：

     $$\text{Attention}(\mathbf{Q},\mathbf{K},\mathbf{V})=\text{softmax}\left(\frac{\mathbf{Q}{\mathbf{K}}^{\mathrm{\top }}}{\sqrt{d_k}}\right)\mathbf{V}$$

     其中：$\frac{\mathbf{Q}{\mathbf{K}}^{\mathrm{\top }}}{\sqrt{d_k}}\in {\mathbb{R}}^{n\times n}$$\frac{\mathbf Q\mathbf K^\top}{\sqrt {d_k}} \in \mathbb{R}^{n\times n}$ 表示 scaled dot-product similarity matrix ；$\text{softmax}(\cdot )$$\text{softmax}(\cdot)$ 确保 attention 权重非负且每行求和为1.0 。

   • Multi-Head Concatenation：在 multi-head attention 中，上述过程通过 $h$$h$ 个 heads 来并行地执行，每个 attention head $i$$i$ 具有独立的投影矩阵 ${\mathbf{W}}_Q^i,{\mathbf{W}}_K^i,{\mathbf{W}}_V^i$$\mathbf W_{Q}^{i}, \mathbf W_{K}^{i}, \mathbf W_{V}^{i}$。所有 heads 的输出进行拼接：

     $$\text{MultiHead}(\mathbf{Q},\mathbf{K},\mathbf{V})=\text{Concat}({\text{head}}_1,\cdots ,{\text{head}}_h)$$

     其中：${\text{head}}_i=\text{Attention}({\mathbf{X}}^i{\mathbf{W}}_Q^i,{\mathbf{X}}^i{\mathbf{W}}_K^i,{\mathbf{X}}^i{\mathbf{W}}_V^i)$$\text{head}_i= \text{Attention}\left(\mathbf X^i\mathbf W_Q^i, \mathbf X^i\mathbf W_K^i, \mathbf X^i\mathbf W_V^i\right)$ 。

     其中：${\mathbf{X}}^i$$\mathbf X^i$ 为 $\mathbf{X}$$\mathbf X$ 的第 $i$$i$ 个 head 的 split 。

   • Final Output Layer：concatenated SDPA output 被馈入 output layer ：

     $$\mathbf{O}=\text{MultiHead}(\mathbf{Q},\mathbf{K},\mathbf{V}){\mathbf{W}}_O$$

     其中：${\mathbf{W}}_O\in {\mathbb{R}}^{(h{d}_k)\times d_{\text{model}}}$$\mathbf W_O\in \mathbb R^{(hd_k)\times d_\text{model}}$ 。

1.2.2 Augmenting Attention Layer with Gating Mechanisms

1. 门控机制（gating mechanism） 的形式化定义为：

   $${\mathbf{Y}}^{\prime }=g(\mathbf{Y},\mathbf{X},{\mathbf{W}}_{\theta },\sigma )=\mathbf{Y}\odot \sigma (\mathbf{X}{\mathbf{W}}_{\theta })$$

   其中：

   • $\mathbf{Y}$$\mathbf Y$ 表示等待被调制的 input ，${\mathbf{Y}}^{\prime }$$\mathbf Y^\prime$ 为 gated output。

   • $\mathbf{X}$$\mathbf X$ 为 gating 的 input ，${\mathbf{W}}_{\theta }$$\mathbf W_\theta$ 为 gate 的可学习的参数。注意，我们采用 hidden states after pre-normalization 作为 $\mathbf{X}$$\mathbf X$ 。

     在 Transformer 的标准架构中，每个子层（如注意力层或前馈网络层）通常有两种归一化方式：

     • Post-Norm（后置归一化）：先计算子层输出，再执行归一化，然后与残差连接相加。

       SubLayerOutput = Norm( SubLayer(Input) + Input )
       

     • Pre-Norm（前置归一化）：先对输入进行归一化，再输入到子层中进行计算，最后与残差连接相加。

       xxxxxxxxxx
       SubLayerOutput = SubLayer( Norm(Input) ) + Input
       

       hidden states after pre-normalization 就是这里的 Norm(Input) 。

     二者区别：

     特性	Post-Norm	Pre-Norm
     归一化位置	残差连接之后	残差连接之前
     梯度传播	梯度必须通过归一化层，容易导致梯度消失/爆炸，尤其在深层网络中。	梯度有一条“捷径”（残差路径）直接回传，不经过归一化层，梯度更稳定。
     训练稳定性	不稳定，需要精细的学习率预热（Warm-up）和调参。	非常稳定，易于训练，通常不需要复杂的学习率调度。
     收敛速度	较慢。	更快，因为稳定的梯度允许使用更大的学习率。
     模型表达能力	理论上更强，因为每个子层的输出都经过归一化，表征更一致。	可能稍弱，因为每个子层直接输出到下一层，表征尺度可能不一致。
     扩展到极深网络	困难，梯度问题会随着层数增加而加剧。	容易，是构建数百层乃至上千层 Transformer 的基础（如 GPT、LLaMA 系列）。
     结构简洁性	每个子层后都需做归一化，结构略繁。	结构更简洁，残差连接更“干净”。

   • $\sigma (\cdot )$$\sigma(\cdot)$ 为激活函数，比如 sigmoid ；$\odot$$\odot$ 为逐元素乘法。

   gating score $\sigma (\mathbf{X}{\mathbf{W}}_{\theta })$$\sigma(\mathbf X \mathbf W_\theta)$ 有效地充当一个动态的过滤器，通过选择性地保留或删除 $\mathbf{Y}$$\mathbf Y$ 的特征来控制信息流（information flow）。

2. 本文全面研究了 attention layers 内 gating mechanisms 的多种变体，重点关注五个关键方面：

   • (1) 位置：研究在不同位置应用 gating 的效果，如 Figure 1(left) 所示。

     • a)：在 Q, K, V projections 之后，对应 Figure 1(left) 中的 $G_2,G_3,G_4$$G_2, G_3, G_4$ 的位置。

     • b)：在 SDPA output 之后，对应 Figure 1(left) 中的 $G_1$$G_1$ 的位置。

     • c)：在 final concatenated multi-head attention outputs 之后，对应 Figure 1(left) 中的 $G_5$$G_5$ 的位置。

   • (2) 粒度：考虑 gating score 的两种粒度：

     • a) headwise：单个标量的 gating score 调制 an attention head 的整个输出。

     • b) elementwise：gating scores 是与 $\mathbf{Y}$$\mathbf Y$ 维度相同的向量，支持细粒度的、逐维度的调制。

   • (3) Head Specific or Shared：针对 attention 的 multi-head 特性，进一步考虑：

     • a) Head-Specific：每个 attention head 具有专属的 gating scores，支持对每个 head 进行独立地调制。

     • b) Head-Shared：${\mathbf{W}}_{\theta }$$\mathbf W_\theta$ 和 gating scores 在所有 attention heads 之间共享。

   • (4) Multiplicative or additive：当将 gating score 应用于 $\mathbf{Y}$$\mathbf Y$ 时，考虑：

     • a) Multiplicative Gating：gated output ${\mathbf{Y}}^{\prime }$$\mathbf Y^\prime$ 为：${\mathbf{Y}}^{\prime }=\mathbf{Y}\odot \sigma (\mathbf{X}{\mathbf{W}}_{\theta })$$\mathbf Y^\prime = \mathbf Y\odot \sigma(\mathbf X \mathbf W_\theta)$。

     • b) Additive Gating：gated output ${\mathbf{Y}}^{\prime }$$\mathbf Y^\prime$ 为：${\mathbf{Y}}^{\prime }=\mathbf{Y}+\sigma (\mathbf{X}{\mathbf{W}}_{\theta })$$\mathbf Y^\prime = \mathbf Y+ \sigma(\mathbf X \mathbf W_\theta)$。

   • (5) 激活函数：主要考虑两种常用激活函数：SiLU 和 sigmoid 。

     • 由于 SiLU 的输出范围无界，所以仅用于 additive gating 。

       $\text{SiLU}(x)=x\times \text{sigmoid}(x)$$\text{SiLU}(x) = x\times \text{sigmoid}(x)$

     • 由于 sigmoid 的输出范围为 [0,1]，所以仅用于 multiplicative gating。

     此外，为了进一步剖析 gating 有效性的底层机制，我们还考虑了恒等映射（Identity Mapping ）或 RMSNorm（详见第 1.4.1节）。

   除非另有说明，否则我们采用 head-specific, multiplicative gating，使用 sigmoid 激活函数 $\sigma (x)=1/(1+e^{-x})$$\sigma(x)=1/\left(1+e^{-x}\right)$。

   

1.3 实验

1.3.1 实验设置

1. 模型架构与 Training Settings：我们在 MoE 模型（总参数量 15B、激活参数量 2.54B ，记为 15A2B ）和 dense 模型（总参数量 1.7B ）上进行实验。

   • 15A2B MoE models 采用 128 total experts with Top-8 softmax gating，fine-grained experts （《Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models》），global-batch LBL（《Demons in the detail: On implementing load balancing loss for training specialized mixture-of-expert models》）和 z-loss（《St-moe: Designing stable and transferable sparse expert models》）。

     • 无论训练还是推理，每个 Token 都只激活 Top-8 个专家，即使在训练期间也是如此。

     • Fine-grained Experts（细粒度专家）是 MoE（混合专家）模型架构中的一种设计范式，指的是将专家网络划分得更小、更精细，并在更细的粒度（例如神经元级别或子模块级别）上执行路由选择，而不是传统的 “一个专家就是一个完整的 FFN 层”。

     • Global-Batch LBL（Global-Batch Load Balancing Loss） 是一种用于训练 MoE 模型的、在分布式训练环境下跨所有计算设备计算和优化的负载均衡损失函数。它的核心目标是确保所有专家在全局训练数据上得到均衡使用，从而稳定训练并充分发挥 MoE 模型的参数容量优势。

     • Z-Loss（Zero-Mean Logit Loss，零均值逻辑损失）是一种专门用于 MoE 和 稀疏门控模型 训练的辅助损失函数，旨在通过稳定路由器的数值，来间接提升整个模型的训练稳定性。它是一种 正则化项，而非直接影响模型输出的主损失。

       $${\mathcal{L}}_z=\frac{1}{n}\sum _{i=1}^n(\log \sum _{j=1}^E\exp (z_{i,j}))$$

       其中：$n$$n$ 为 tokens 数量$E$$E$ 为专家数量；$z_{i,j}$$z_{i,j}$ 为路由器为第 $i$$i$ 个 token、第 $j$$j$ 个专家输出的原始 logits。

       Log-Sum-Exp 可以近似看作 max(z) + 平滑项。因此，Z-Loss 实质上是在惩罚路由器输出的最大值不要过大，从而将整个 logits 向量温和地“压缩”到一个更合理的数值范围内。

   • attention 部分采用分 group query attention: GQA（《Gqa: Training generalized multi-query transformer models from multi-head checkpoints》）。

     多头注意力（Multi-Head Attention: MHA）：这是原始 Transformer 的标准配置。

     • 原理：对于给定的序列，模型会并行计算 h 个（例如 32 个）独立的“注意力头”。每个头都有自己的 $\mathbf{Q}$$\mathbf Q$（query ）、$\mathbf{K}$$\mathbf K$（key）、$\mathbf{V}$$\mathbf V$（value） 线性投影矩阵。

     • 问题：在自回归生成（如 GPT 逐词生成）时，为了计算下一个词的注意力，需要将所有历史时间步的 $\mathbf{K}$$\mathbf K$ 和 $\mathbf{V}$$\mathbf V$（每个头都有一份）缓存下来，称为 KV 缓存。MHA 的 KV 缓存非常大，成为内存带宽的主要瓶颈，限制了推理速度和处理的长上下文长度。

       KV 缓存大小公式（MHA）：2 * 序列长度 * 层数 * 头数 * 每个头的维度。

     多查询注意力（Multi-Query Attention: MQA）：为了极致优化推理速度而提出的方案。

     • 原理：所有 h 个 query head 共享同一套 $\mathbf{K}$$\mathbf K$ 和 $\mathbf{V}$$\mathbf V$ 投影。也就是说，只有一套 K 和 V 矩阵，但 Q 矩阵仍然有 h 套。

     • 优势：KV 缓存大小瞬间缩小为 MHA 的 1/h，极大减轻了内存压力，显著提升了推理速度。

     • 缺点：由于所有头共享相同的“视角”（K, V），模型的表达能力可能下降，有时会导致生成质量（尤其是在需要多维度、细粒度理解的任务上）相比 MHA 有所损失。

     分组查询注意力（Group-Query attention: GQA）：平衡的艺术。GQA 的设计思路是在 MHA 的表达能力和 MQA 的效率之间找到一个折中点。

     • 原理：将 h 个 query head 分组。假设有 h 个查询头，将其分成 g 个组（g < h）。每个组内的所有查询头，共享同一套 $\mathbf{K}$$\mathbf K$ 和 $\mathbf{V}$$\mathbf V$ 投影。

     • 关键参数：

       • h：查询头的总数。

       • g：组的数量（即 KV head 的数量）。当 g = h 时，GQA 退化为 MHA；当 g = 1 时，GQA 退化为 MQA。

       • 分组大小：h / g。

   • 模型在 3.5T high-quality tokens 的子集上训练，涵盖多语言、数学和通用知识内容。

   • 上下文序列长度设为 4096 。

   • 学习率和 batch size: bsz 等更详细的配置将在各部分介绍。

   • 其他超参数遵循 AdamW optimizer 的默认值。

   • 由于 gating 引入的参数和计算量较小，其导致的 wall-time latency 增加不到 2%。

     1. Wall-time（墙上时钟时间）

        • 定义：程序从开始到结束所经过的真实世界时间。

        • 包含：所有计算时间、I/O 等待时间、内存访问延迟、并行任务中的同步等待时间、CPU 空闲时间等。

        • 比喻：就像你按下秒表开始计时，直到程序运行结束再按停，所记录的时间。

     2. CPU Time（CPU 时间）

        • 定义：程序在 CPU 上实际执行指令所花费的时间。

        • 不包含：I/O 等待、网络延迟、其他进程占用 CPU 的时间等。

2. Evaluation：我们在多个主流 benchmarks, 上测 few-shots results，包括英语基准Hellaswag、通用知识基准 MMLU、数学推理基准 GSM8k、代码基准 HumanEval、中文能力基准 C-eval and CMMLU。我们还报告了语言模型在不同 held-out test sets 上的困惑度（perplexity: PPL），包括英语、中文、代码、数学、法律和文学等领域。

1.3.2 主要结果

a. Gated Attention for MoE models

1. 首先对比不同 gated attention layers 在 training-efficient MoE-15A2B models 上的结果。所有模型采用的 scheduler 为：在 1k steps 内预热至最大学习率 2e-3，然后通过余弦衰减至 3e-5 。global bsz = 1024，共进行 100k optimization steps 。结果汇总于 Table 1。为了公平对比，我们为 vanilla MoE baseline （第 1行）补充了 parameter expansion 方法，包括增加 key-value heads 数量（第 2 行）、增加 query heads 数量（第 3 行）、以及同时增加 total and activated number of experts （第 4 行）。这些方法引入的参数数量与 gating 机制相当或更多。

   从 Table 1 中可以观察到：

   • (i)：SDPA and value output gating 都是有效的。在 SDPA 的输出位置（$G_1$$G_1$）或 value map 位置（$G_2$$G_2$）处插入 gates 最为有效，相比其他变体，能够实现更低的困惑度和更好的整体 benchmark 性能。我们将在第 1.4.2 节进一步探讨为何在这两个位置应用 gating 是有效的。

   • (ii)：Head-Specific Gating 至关重要。

     • 在 $G_1$$G_1$ 和 $G_2$$G_2$ 位置应用 headwise gating 仅引入极少的额外参数（对于 MoE-15A2B 模型而言不到 2M），但仍能带来显著提升（第 10 行和第 11 行）。

     • 当在不同 attention heads 之间共享 gating scores 时（我们对 query head dimension $q$$q$ 取平均，从原始的 $n\times q\times d_k$$n\times q \times d_k$ score 得到 $n\times d_k$$n\times d_k$ score ），benchmark 性能提升小于 headwise gating （第 12 行 vs 第 10 行，第 13 行 vs 第 11 行）。

     这凸显了为不同 attention heads 应用 distinct gating scores 的重要性。

   • (iii)：Multiplicative Gating 更优。Additive SDPA output gating 的性能不如 multiplicative SDPA output gating ，尽管 Additive SDPA output gating 相比 baselines 仍有提升。

   • (iv)：sigmoid 激活函数更优。将最有效的 gating configuration（第 5 行）中的激活函数替换为 SiLU（第 15 行），提升效果减弱。

   总体而言，在 value layer（$G_2$$G_2$）和 SDPA output（$G_1$$G_1$）处添加 gating 可使困惑度降低超过 0.2，优于各种 parameter-expanding baselines 。然而，$G_1$$G_1$ 位置（相比 $G_2$$G_2$ ）的 gating 实现了更好的困惑度和 benchmark 结果。只要不同 heads 获得distinct gating scores ，gating 的粒度、以及激活函数选择的影响相对较小。我们将在分析部分（第 1.4.2 节）进一步解释这些观察结果背后的原因。

   

b. Gated Attention for Dense Models

1. 我们还在 dense models 上进行了实验（遵循 《Qwen2. 5 technical report》），以验证 SDPA output sigmoid gating 的效果。使用 gating 时，我们减小了前馈网络（FFN）的宽度，以保持所有方法的参数量相同。大多数实验采用 baseline 的 optimized hyperparameters 。例如，在 400B tokens 上训练好的 1.7B model，使用最大学习率 4e-3 和 batch size = 1024。在 3.5T tokens 上训练时，将最大学习率提高到 4.5e-3 ，batch size 增加到 2048。

   先前的研究表明，增加 network depth、使用大学习率、以及大 batch size 可以显著提升模型性能和分布式训练效率，但通常会引入 training 不稳定性。我们观察到，应用 gating 机制显著减少了训练过程中的 loss spikes，表明 gating 在增强 training 稳定性方面具有潜力。基于这一发现，我们引入了另一种实验设置，通过增加 layer num、提高最大学习率、以及扩大 batch size，进一步探究 gating 的 stabilizing 效果。

   Table 2 显示：

   • (i)：gating 在各种 settings 中均有效。在不同的模型配置（第 1 行 vs 第 2 行，第 5 行 vs 第 8 行）、训练数据（第 3 行 vs 第 4 行）和超参数（第 11 行 vs 第 13行）下，应用 SDPA output gating 始终能带来收益。

   • (ii)：gating 提高了稳定性并促进了 scaling 。在 3.5T token setting 下，gating 改善了 training 的稳定性，大幅减少了 loss spike（Figure 1 right）。当提高最大学习率时，baselines 出现收敛问题（第 6 行、第 12行 ）。虽然添加sandwich norm （《Cogview: Mastering text-to-image generation via transformers》）可以恢复收敛性，但提升效果微乎其微。相比之下，models with gating 在提高最大学习率后，性能有明显提升。

   总之，我们确定 SDPA element-wise gating 是增强 attention 机制的最有效方法。将该方法应用于 dense transformer 进一步表明，gates 支持更大 batch size 和更大学习率的 stable training ，从而提升性能。

   

1.4 Analysis: Non-Linearity, Sparsity, and Attention-Sink-Free

1. 本节通过一系列实验，探究为何这种简单的 gating 机制能显著提升性能和training 的稳定性。分析得出以下结论：

   • (1)：增强非线性的 gating 操作始终带来性能提升（第 1.4.1节）。

   • (2)：最有效的 SDPA elementwis $G_1$$G_1$ gat 为 SDPA outputs 引入了 strong input-dependent sparsity （第 1.4.2 节），进而有助于消除 'attention sink' 现象。

1.4.1 Non-linearity 提升注意力中 Low-Rank Mapping 的表达能力

1. 受先前在 SDPA output 上使用 group norm 的研究（《Retentive network: A successor to transformer for large language models》、《Differential transformer》）的启发，在第 1.3.2.a 节的相同设置下，我们在 concatenation 之前对output of each attention head 独立地应用 RMSNorm。如 Table 3 第 5 行所示，应用了几乎不引入额外参数的 RMSNorm ，也能显著降低困惑度。

   

2. 在 multi-head attention 中，第 $k$$k$ 个 head 对应第 $i$$i$ 个 token 的输出可表示为：

   $${\overrightarrow{\mathbf{o}}}_i^k=(\sum _{j=0}^i{S}_{ij}^k\times {\overrightarrow{X}}_j{\mathbf{W}}_V^k){\mathbf{W}}_O^k=\sum _{j=0}^k{S}_{ij}^k\times {\overrightarrow{X}}_j\left({\mathbf{W}}_V^k{\mathbf{W}}_O^k\right)$$

   其中：

   • $S_{ij}^k$$S_{ij}^k$ 表示第 $k$$k$ 个 attention head 中第 $i$$i$ 个 token 对第 $j$$j$ 个 token 的 attention score 。

   • ${\overrightarrow{X}}_j$$\vec X_j$ 为第 $j$$j$ 个 token 的 attention 的输入；${\overrightarrow{X}}_j{\mathbf{W}}_V^k$$\vec X_{j} \mathbf W_{V}^{k}$ 表示第 $k$$k$ 个 attention head 中第 $j$$j$ 个 token 的 value output 。

   • ${\mathbf{W}}_O^k$$\mathbf W_O^k$ 为 output layer ${\mathbf{W}}_O$$\mathbf W_O$ 对应第 $k$$k$ 个 attention head 的参数。

     注意：将不同 attention head 的输出拼接之后再与${\mathbf{W}}_O$$\mathbf W_O$相乘，等同于每个 attention head 的输出先与其对应的${\mathbf{W}}_O^k$$\mathbf W_{O}^{k}$相乘再进行拼接。

   根据上述公式，由于 $d_k<d_{\text{model}}$$d_k < d_\text{model}$ ，我们可以将 ${\mathbf{W}}_V^k{\mathbf{W}}_O^k$$\mathbf W_{V}^{k} \mathbf W_{O}^{k}$ 合并为一个应用于所有 ${\overrightarrow{X}}_j$$\vec X_{j}$ 的低秩线性映射（low-rank linear mapping ）。在 GQA 中，${\mathbf{W}}_V$$\mathbf W_{V}$ 在相同 group 内的 attention head 之间共享，进一步降低了表达能力。

3. 鉴于在两个线性映射之间添加非线性可以提高其表达能力（《On the number of linear regions of deep neural networks》），我们提出两种修改来缓解低秩（low-rank）问题：

   $$\begin{array}{c}{\overrightarrow{\mathbf{o}}}_i^k=(\sum _{j=0}^i{S}_{ij}^k\times \text{NonLinearityMap}\left({\overrightarrow{X}}_j{\mathbf{W}}_V^k\right)){\mathbf{W}}_O^k\\ {\overrightarrow{\mathbf{o}}}_i^k=\text{NonLinearityMap}(\sum _{j=0}^i{S}_{ij}^k\times {\overrightarrow{X}}_j{\mathbf{W}}_V^k){\mathbf{W}}_O^k\end{array}$$

   值得注意的是：

   • 在 $G_2$$G_2$ 位置添加 gating（ Table 3 第 3 行）对应第一种修改。

   • 而在 $G_1$$G_1$ 位置添加 gating（Table 3 第 4 行）或 group normalization （ Table 3 第 5 行）对应第二种修改。

   这也解释了为何在 ${\mathbf{W}}_O$$\mathbf W_{O}$ 之后的 $G_5$$G_5$ 位置添加 gating 或 normalization 无效（Table 1 第 9 行）——它无法解决 ${\mathbf{W}}_V$$\mathbf W_{V}$ 和${\mathbf{W}}_O$$\mathbf W_{O}$ 之间缺乏非线性的问题。

   对于 $G_1$$G_1$ 位置的 additive gating ，output of gating 经过 SiLU（ Table 3 第 4 行），也引入了一定的非线性，这解释了观察到的性能提升，尽管小于 multiplicative gating 。基于这些见解，我们进行了另外两个实验：

   • (i)：仅在 $G_1$$G_1$ 位置添加 SiLU，不引入额外参数（ Table 3 第 6 行）。注意到这种简单修改也能适度降低困惑度，但大多数 benchmark 分数保持不变。

   • (ii)：从 additive gating 中移除 SiLU，使得 ${\overrightarrow{X}}_j$$\vec X_{j}$ 的输出经过 gating 后直接添加到 $G_1$$G_1$ 位置（ Table 3 第 7 行）。这进一步削弱了 addictive gating 的增益。

   总之，有效的 gating 变体的性能提升可能归因于在 ${\mathbf{W}}_V$$\mathbf W_{V}$ 和 ${\mathbf{W}}_O$$\mathbf W_{O}$ 之间引入了非线性。尽管在 $G_1$$G_1$ 和 $G_2$$G_2$ 位置应用 gating 都能引入这种非线性，但它们带来的性能提升有所不同。这一观察到的差异促使我们进一步分析在这两个位置应用 gating 的影响。

   

1.4.2 Gating 引入 Input-Dependent Sparsity

1. 我们分析了在 value（$G_2$$G_2$）和 SDPA output（$G_1$$G_1$）位置应用 gating 的模型的 gating scores（ Table 1 中的 Gate Score 列），评估数据为 test language modeling data 。所有层的 mean gating scores 如 Table 4 所示，score distributions 如 Figure 3 所示（layer-wise scores 见附录 A.2）。

   主要观察结果包括：

   • (i)：Effective Gating Scores 具有稀疏性。SDPA output gatings （Element-wise 和 head-wise）的 mean gating scores 最低。此外，SDPA output gating score distribution 高度集中在 0 附近，表明具有显著的稀疏性，这与其优异性能一致。

   • (ii)：Head-Specific Sparsity 至关重要。在 attention heads 之间强制共享 gating scores 会提高 overall gating scores，并降低性能增益。观察结果 (i) 和 (ii) 凸显了 head-specific gating 的重要性，与先前的研究一致，即 individual attention heads 捕获 input 的不同方面。

   • (iii) ：Query-Dependency 至关重要。value gating（$G_2$$G_2$）的分数高于 SDPA output gating（$G_1$$G_1$），且性能较差。这表明，当 gating score 依赖于 query 而非 key and value 时，gating score sparsity 更为有效。

     具体而言，SDPA output gating scores 源自 current query 对应的 hidden states ，而 value gating scores 源自与past keys and values 相关的 hidden states 。这意味着 gating score sparsity 可能过滤掉与 query 无关的上下文信息。为了进一步验证 query-dependency 的重要性，我们引入 input-independent gating ：将自由的、可学习的参数（$q\times d_k$$q\times d_k$）通过零初始化（然后通过训练来学习这个参数），然后应用 sigmoid 函数，并将其与 SDPA output 相乘。如第 (6) 行所示，nput-independent gating 相比 baseline 有所提升，这可能归因于引入了非线性。此外，high gating scores 进一步证明， effective sparsity 应该是 input-dependent 的。

   • (iv)：稀疏性较低的 Gating 更差。为了进一步验证 gating score sparsity 的重要性，我们通过 gating 公式来降低稀疏性。具体而言，我们将 sigmoid 函数替换为修改后的 Non-Sparse: NS 版本：

     $$\text{NSsigmoid}(x)=0.5+0.5\times \text{sigmoid}(x)$$

     该函数将 gating scores 限制在 [0.5, 1.0] 之间。这确保了在引入非线性的同时，消除了 gating score sparsity 。如 Table 4 第 (7) 行所示，NS-sigmoid gating 的增益低于 SDPA output sigmoid gating 。在附录 A.2 中，我们详细讨论了在 SDPA hidden states 中 sparse gating scores 如何影响 sparsity（低于阈值的值的比例）。下一节将讨论不同 sparsity levels 对模型行为的影响，包括减少 'attention sink'。

   

   

1.4.3 SDPA Output Gating 减少 Attention-Sink

1. 基于观察 “ gating 以 input-dependent 的方式为 SDPA output 引入稀疏性“，我们假设该机制可以过滤掉与 current query token 无关的上下文，从而缓解 attention sink。为了验证这一点，我们分析了 attention scores 的分布（所有 attention head 的平均值）和分配给 first token 的 attention scores 的比例（Figure 2，Table 4 中的 'F-Attn'列）。受关于 massive activation in hidden states 与 attention sinks 的讨论（《Massive activations in large language models》）所启发，我们还计算了各层的 mean of the maximum hidden state activations，如 Table 4 中的 'M-Act' 列所示。更详细的 layer-wise 结果见附录 A.3。

   我们可以观察到：

   • (i)：在 SDPA output（$G_1$$G_1$）处应用 head-wise and element-wise query-dependent sigmoid gating，大幅降低了分配给 first token 的 attention score ，并减少了 massive activations。

   • (ii)：在 attention heads 之间强制 gating scores 、或仅在 value projection 后（$G_2$$G_2$）应用 gating，减少了 massive activations，但并未降低 first token 的 attention scores。这进一步强调了 head-specific gating 的重要性，并表明 massive activations 并非 attention sinks 的先决条件。

   • (iii)：降低 input-dependence of gating（第 6 行）、或使用 NS-sigmoid 来降低稀疏性（第 7 行），会加剧 massive activations 和 attention sink。

   

   

2. 综上，这些观察结果表明：

   • input-dependent, head-specific gating of the SDPA output 引入了显著的稀疏性，从而缓解了attention sink。

   • 此外，SDPA outputs 的稀疏性减少了模型内部的 massive activations，稀疏性越高，activations 越小。这可能解释了gating 带来的 training 稳定性提升：通过减少 massive activations，模型在 BF16 training 过程中不易受到数值误差（numerical errors）的影响（《Numerical error analysis of large language models》）。

   • 我们还观察到，massive activations 主要源自 early layers（例如 layer 5），其中 FFN 输出较大的值，这与 《Interpreting the repeated token phenomenon in large language models》 的研究一致。一旦添加到 residual stream 中，这些 activations 通过 pre-norm 机制在后续层中传播。这与 sandwich normalization（《Cogview: Mastering text-to-image generation via transformers》）在增强 training 稳定性方面的有效性一致（Table 2 第 7 行）：对 FFN output 应用 LayerNorm，可防止这些 large activations 进入 residual stream。

1.4.4 SDPA Output Gating 促进 Context Length Extension

1. 基于 attention-sink-free pattern ，我们评估了 SDPA gating 在 long-context setting 中的效果。具体而言，我们扩展了在 3.5T tokens 上训练的模型的上下文长度。将旋转位置编码（RoPE）（《Roformer: Enhanced transformer with rotary position embedding》）的 base 从 10k 增加到 1M，并在序列长度为 32k 的数据上继续训练 80B tokens，得到上下文长度为 32k 的模型。随后，我们使用 YaRN（《Yarn: Efficient context window extension of large language models》）将上下文长度扩展到 128k。我们在 RULER benchmark 上评估模型，结果汇总于 Table 5。观察结果如下：

   • (i)：在 32k setting 下，models with gating 略优于 baseline。这表明在 training length 范围内，attention sink 现象可能不会损害模型的 long-context 性能。

   • (ii)：当使用 YaRN 将上下文长度扩展到 128k 时，baseline 和 gated models 在 original 32k 范围内均出现性能下降。这一观察结果与先前通过修改RoPE 来扩展上下文长度的研究一致。尽管如此，models with gating 的性能下降幅度较小。

   • (iii)：在上下文长度为 64k 和 128k 时，gated attention models 显著优于 baseline 。

   

2. 基于这些观察结果，我们推测：添加 gating 有助于模型适应 context-length extension。一个可能的解释是，baseline models 依赖 attention sinks 来调整 attention scores 的分布。《Longred: Mitigating short-text degradation of long-context large language models via restoration distillation》基于 attention and hidden state distributions，推导了 changing the RoPE 的影响。当应用 YaRN 等技术修改 RoPE base 时，attention sink pattern 可能难以通过 raining-free 的方式适应，导致性能显著下降。相比之下，models with gating 主要依赖于 input-dependent gating scores 来控制 information flow，因此对这类变化更具鲁棒性。

1.5 结论

1. 本文系统研究了 standard softmax attention 中 gating 机制的作用，揭示了其对性能、training 稳定性、以及 attention dynamics 的显著影响。通过在 3.5T tokens 上训练的 15B 参数 MoE 模型和 1.7B 参数 dense 模型的 30 种变体进行广泛的实验对比，我们证明了在 scaled dot-product attention 之后应用 sigmoid gate 能带来最显著的提升。这种简单机制增强了非线性，引入了 input-dependent sparsity，并消除了低效的 'attention sink' 现象。此外，gating 促进了 context length extension，使模型能够在不重新训练的情况下有效泛化到更长的序列。我们还发布了首个 attention-sink-free 模型。我们相信这些实证验证将为下一代先进的 foundation models 的设计铺平道路。

1.6 局限性

1. 本文主要通过一系列消融研究分析 attention gating 的原因和影响。然而，我们承认存在一些局限性。

   • 非线性对 dynamics of attention 和 overall training process 的更广泛影响仍有待探索。

   • 尽管我们观察到消除 attention sinks 改善了长上下文扩展场景下的性能，但我们并未为 attention sinks 如何影响模型泛化到更长序列的能力提供严格的理论解释。

1.7 附录

略。参考原始论文。