---

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

Qiu Z, Wang Z, Zheng B, et al. Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free. NeurIPS 2025.

https://proceedings.neurips.cc/paper_files/paper/2025/hash/904e89bb4e632e75fb47f093b620b257-Abstract-Conference.html

https://github.com/qiuzh20/gated_attention

https://huggingface.co/QwQZh/gated_attention

NeurIPS 2025 Oralccf阿里Edinburgh

---

面向大语言模型的门控注意力：非线性、稀疏性与消除注意力汇聚

Abstract

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.

从 LSTM 和 Highway Network 等早期模型，到近期的状态空间模型、线性注意力以及 softmax 注意力，门控机制已得到广泛应用。 然而，现有文献很少考察门控的具体作用。 在本工作中，我们开展全面实验，系统研究由门控增强的 softmax 注意力变体。 具体而言，我们对 30 种变体进行全面比较，其中包括在 3.5 万亿 token 数据集上训练的 15B 混合专家（MoE）模型和 1.7B 稠密模型。

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 codes and models to facilitate future research.

我们的核心发现是，一个简单的修改，即 在缩放点积注意力（SDPA）之后应用逐注意力头的 sigmoid 门控，可以稳定提升性能。 这一修改还能增强训练稳定性、容许更大的学习率并改善扩展特性。 通过比较不同门控位置和计算变体，我们将其有效性归因于两个关键因素：（1）在 softmax 注意力的低秩映射上引入非线性；（2）使用依赖查询的稀疏门控分数调制 SDPA 输出。 值得注意的是，我们发现这种稀疏门控机制能够缓解 attention sink（注意力汇聚）并提升长上下文外推性能；我们还发布相关代码与模型，以促进后续研究。

1. Introduction

Gating mechanism is well-established in neural networks. Early architectures, such as LSTMs, Highway Networks and GRUs, pioneer the use of gating to control information flow across time steps or layers and improve gradient propagation. This principle persists in modern architectures. Recent sequence modeling works, including state-space models and attention mechanisms commonly apply gating, often to modulate the outputs of token-mixer components. Despite its widespread adoption and empirical success, the function and impact of gating mechanisms remain insufficiently explored beyond their initial intuition.

门控机制在神经网络中早已得到确立。 LSTM、Highway Network 和 GRU 等早期架构率先使用门控来控制跨时间步或跨层的信息流，并改善梯度传播。 这一原则延续到了现代架构中。 近期的序列建模工作，包括状态空间模型和注意力机制，也普遍应用门控，通常用它调制 token 混合器组件的输出。 尽管门控已被广泛采用并取得经验上的成功，但除了最初的直觉解释外，其功能与影响仍未得到充分探索。

Insufficient understanding hinders assessing gating's true contribution, especially when confounded with other architectural factors. For instance, while Switch Heads introduces a sigmoid gating to select top-K attention head experts, our experiments reveal an interesting finding (the appendix): substantial performance gains persist even when reduced to a single expert, where the gate simply modulates the value output. This strongly suggests the gating itself provides significant intrinsic value, separate from the routing mechanism. Similarly, in Native Sparse Attention (NSA), while overall performance improvements are demonstrated, they do not disentangle the contributions of its gating mechanism from the effects of the sparse attention design itself. These considerations underscore the need to rigorously disentangle the effects of gating from other architectural components.

理解不足会妨碍我们评估门控的真实贡献，特别是当它与其他架构因素相互混杂时。 例如，Switch Heads 使用 sigmoid 门控选择 top-K 注意力头专家，而我们的实验揭示了一个有趣发现（见附录）：即使将其缩减为单个专家，让门控仅调制 value 输出，显著的性能增益依然存在。 这强烈表明，门控本身具有显著的内在价值，并且独立于路由机制。 类似地，Native Sparse Attention（NSA）虽然展示了整体性能提升，却没有将其门控机制的贡献与稀疏注意力设计本身的作用区分开来。 这些考虑凸显了严格分离门控作用与其他架构组件作用的必要性。

In this work, we investigate gating mechanisms in the standard softmax attention (Section 2.2). Specifically, we introduce gating at distinct positions (Figure 1): after the query ($G_4$), key ($G_3$), and value projections ($G_2$); following the Scaled Dot Product Attention (SDPA) outputs ($G_1$); and after the final dense output layer ($G_5$). Our exploration covers gating variants including elementwise and headwise, head-specific and head-shared, as well as additive and multiplicative forms. We find that: (i) applying SDPA output head-specific gating ($G_1$) yields the most significant performance improvements (e.g., up to 0.2 PPL reduction and 2 points on MMLU); (ii) the SDPA output gating also improves training stability, nearly eliminating loss spikes, enabling larger learning rates and enhancing model scalability.

在本工作中，我们研究标准 softmax 注意力中的门控机制（第 2.2 节）。 具体而言，我们在不同位置引入门控（图1）：query（$G_4$）、key（$G_3$）和 value 投影（$G_2$）之后，缩放点积注意力（SDPA）输出（$G_1$）之后，以及最终稠密输出层（$G_5$）之后。 我们探索的门控变体包括逐元素与逐注意力头、各注意力头独立与跨注意力头共享，以及加法与乘法形式。 我们发现：（i） 对 SDPA 输出应用各注意力头独立的门控（$G_1$）可带来最显著的性能提升（例如 PPL 最多降低 0.2，MMLU 提升 2 分）；（ii） SDPA 输出门控还能提高训练稳定性，几乎消除损失尖峰，从而支持更大的学习率并增强模型的可扩展性。

We identify two primary factors contributing to the efficacy of gating: (i) Non-Linearity. The two consecutive linear layers - the value ($W_v$) and dense ($W_O$) projections - can be rewritten into one low-rank linear projection. Therefore, introducing non-linearity through gating at positions $G_1$ or $G_2$ can increase the expressiveness of this low-rank linear transformation (Section 4.1). (ii) Sparsity. Although non-linear gating variants consistently enhance performance, we observe that their gains vary. Our analysis further reveals that the pronounced sparsity of the gating scores is another crucial factor, introducing input-dependent sparsity to SDPA outputs (Section 4.2).

我们确定了门控有效性的两个主要因素：（i）非线性。 value（$W_v$）与稠密输出（$W_O$）这两个连续线性层可以改写成一个低秩线性投影。 因此，在 $G_1$ 或 $G_2$ 位置通过门控引入非线性，可以增强这一低秩线性变换的表达能力（第 4.1 节）。 （ii）稀疏性。 虽然非线性门控变体都能稳定提升性能，但我们观察到它们带来的增益有所不同。 进一步分析表明，门控分数的显著稀疏性是另一个关键因素，它为 SDPA 输出引入了依赖输入的稀疏性（第 4.2 节）。

Moreover, sparse gating eliminates the attention sink: the initial tokens disproportionately dominate attention scores (Figure 2, Section 4.3). Previous work explains attention sinks as an accumulation of redundant attention due to non-negative softmax normalization. Empirically, we verify that when query-dependent sparse gating is applied at the SDPA output, both our dense and MoE models (trained on 3.5T tokens) exhibit no attention sink. Furthermore, these models demonstrate superior performance in length generalization, achieving a gain of over 10 points on RULER (Section 4.4).

此外，稀疏门控消除了 attention sink（注意力汇聚）：初始 token 会不成比例地占据注意力分数（图2，第 4.3 节）。 先前工作将注意力汇聚解释为非负 softmax 归一化造成的冗余注意力累积。 经验上，我们验证了：当在 SDPA 输出处应用 依赖查询的稀疏门控 时，我们的稠密模型和 MoE 模型（均在 3.5T token 上训练）都不再出现注意力汇聚。 此外，这些模型在长度泛化方面表现更好，在 RULER 上取得超过 10 分的增益（第 4.4 节）。

In summary, our work highlights the impact of gating in standard attention layers on the performance and behaviors of models. By evaluating gating variants, we uncover their ability to introduce non-linearity and sparsity, and eliminate attention sinks. These findings deepen our understanding of the mechanisms of gated attention. We will open-source our attention-sink-free models to advance future research.

总之，我们的工作凸显了标准注意力层中的门控对模型性能与行为的影响。 通过评估不同门控变体，我们揭示了门控引入非线性与稀疏性并消除注意力汇聚的能力。 这些发现加深了我们对门控注意力机制的理解。 我们将开源无注意力汇聚的模型，以推动后续研究。

图1：左：所研究的门控操作位置。中：在不同位置应用门控的 15B MoE 模型性能比较；SDPA 后门控 G1 的总体结果最好，value 层后的 G2 在 PPL 上也有显著改善。右：采用相同超参数时，基线与带 SDPA 门控的 1.7B 稠密模型在 3.5T token 上的平滑训练损失比较。门控降低了最终损失并显著增强训练稳定性。

2. Gated-Attention Layer

2.1 Preliminary: Multi-Head Softmax Attention

Given an input $X\in {\mathbb{R}}^{n\times d_{\text{model}}}$, where $n$ is the sequence length and $d_{\text{model}}$ is the model dimension, the computation of transformer's attention layer could be divided into four stages. QKV Linear Projections: The input $X$ is linearly transformed into queries $Q$, keys $K$, and values $V$ using learned weight matrices $W_Q,W_K,W_V\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$ and $Q,K,V\in {\mathbb{R}}^{n\times d_k}$:

给定输入 $X\in {\mathbb{R}}^{n\times d_{\text{model}}}$，其中 $n$ 是序列长度，$d_{\text{model}}$ 是模型维度，Transformer 注意力层的计算可以划分为四个阶段。 QKV 线性投影： 使用学习得到的权重矩阵 $W_Q,W_K,W_V\in {\mathbb{R}}^{d_{\text{model}}\times d_k}$，将输入 $X$ 线性变换为 query $Q$、key $K$ 和 value $V$，其中 $Q,K,V\in {\mathbb{R}}^{n\times d_k}$：

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

Scaled Product Dot-Product Attention (SDPA): computes attention scores between queries and keys, followed by a softmax normalization. The output is a weighted sum of the values:

缩放点积注意力（SDPA）： 计算 query 与 key 之间的注意力分数，随后进行 softmax 归一化。 其输出是 value 的加权和：

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

Where $\frac{Q{K}^T}{\sqrt{d_k}}\in {\mathbb{R}}^{n\times n}$ represents the scaled dot-product similarity matrix, and $\mathrm{softmax}(\cdot )$ ensures the attention weights are no-negative and sum to 1 across each row. Multi-Head Concatenation: In multi-head attention, the above process is repeated in parallel for $h$ heads, with each head having its projection matrices $W_q^i,W_k^i,W_v^i$. All heads' outputs are concatenated:

其中，$\frac{Q{K}^T}{\sqrt{d_k}}\in {\mathbb{R}}^{n\times n}$ 表示缩放点积相似度矩阵，$\mathrm{softmax}(\cdot )$ 确保注意力权重非负且每行之和为 1。 多头拼接： 在多头注意力中，上述过程会在 $h$ 个注意力头上并行重复，每个注意力头都有自己的投影矩阵 $W_q^i,W_k^i,W_v^i$。 所有注意力头的输出会被拼接：

$$\begin{array}{rl}\mathrm{MultiHead}(Q,K,V)& =\mathrm{Concat}({\mathrm{head}}_1,\dots ,{\mathrm{head}}_h),\\ {\mathrm{head}}_i& =\mathrm{Attention}(Q{W}_Q^i,K{W}_K^i,V{W}_V^i).\end{array}$$

Final Output Layer: The concatenated SDPA output is passed through an output layer $W_o\in {\mathbb{R}}^{h{d}_k\times d_{\text{model}}}$:

最终输出层： 拼接后的 SDPA 输出经过输出层 $W_o\in {\mathbb{R}}^{h{d}_k\times d_{\text{model}}}$：

$$O=\mathrm{MultiHead}(Q,K,V)W_o.$$

图2：左：各层分配给初始 token 的注意力比例。基线模型存在显著的注意力汇聚，跨层平均有 46.7% 的注意力分数指向第一个 token；门控将该比例降至 4.8%。右：各注意力头的平均注意力图权重。基线第 21 层有 83% 的注意力落在第一个 token 上，门控将其降至 4%；在最终输出层，门控会放大模型关注序列中个别 token 的已有倾向。

2.2 Augmenting Attention Layer with Gating Mechanisms

The gating mechanism is formalized as:

门控机制形式化为：

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

Where $Y$ is the input to be modulated, $X$ is another input used to compute the gating scores (We adopt the hidden states after pre-normalization as $X$), $W_{\theta }$ refers to the learnable parameters of gate, $\sigma$ is an activation function (e.g., sigmoid), and $Y^{\prime }$ is the gated output. The gating score, $\sigma (X{W}_{\theta })$, effectively acts as a dynamic filter, controlling the information flow from $Y$ by selectively preserving or erasing its features.

其中，$Y$ 是待调制的输入，$X$ 是用于计算门控分数的另一个输入（我们采用预归一化后的隐藏状态作为 $X$），$W_{\theta }$ 表示门控的可学习参数，$\sigma$ 是激活函数（例如 sigmoid），$Y^{\prime }$ 是门控后的输出。 门控分数 $\sigma (X{W}_{\theta })$ 实际上充当动态过滤器，通过选择性保留或抹除 $Y$ 的特征来控制其中的信息流。

In this work, we comprehensively investigate several variants of gating mechanisms within the attention layers. Our exploration focuses on five key aspects:

在本工作中，我们全面研究注意力层内的多种门控机制变体。 我们的探索聚焦于五个关键方面：

1. Positions. We study the effect of applying gating at different positions, as illustrated in Figure 1 (left): (a) after the $Q,K,V$ projections, corresponding to positions $G_2,G_3,G_4$ in Figure 1 (left); (b) following the SDPA outputs ($G_1$); (c) after the final concatenated multi-head attention outputs ($G_5$).

2. Granularity. We consider two levels of granularity for the gating score: (a) Headwise: A single scalar gating score modulates the entire output of an attention head. (b) Elementwise: Gating scores are vectors with the same dimensionality as $Y$, enabling fine-grained, per-dimension modulation.

3. Head Specific or Shared. Given the multi-head nature of attention, we further consider: (a) Head-Specific: each attention head has its specific gating scores, enabling independent modulation for each head. (b) Head-Shared: $W_{\theta }$ and gating scores are shared across heads.

4. Multiplicative or additive. For applying gating score to $Y$, we consider (a) Multiplicative Gating: The gated output $Y^{\prime }$ is computed as: $Y^{\prime }=Y\cdot \sigma (X\theta )$. (b) Additive Gating: $Y^{\prime }=Y+\sigma (X\theta )$.

5. Activation Function. We mainly consider two common activation functions: SiLU and sigmoid. We only use SiLU for additive gating due to its unbounded output range, and sigmoid only gives scores in $[0,1]$. Additionally, to further dissect the mechanisms underlying gating's effectiveness, we also consider Identity Mapping or RMSNorm (detailed in Section 4.1).

1. 位置。 我们研究在不同位置应用门控的效果，如图1左侧所示：（a）在 $Q,K,V$ 投影之后，对应图1左侧的 $G_2,G_3,G_4$；（b）在 SDPA 输出之后（$G_1$）；（c）在最终拼接的多头注意力输出之后（$G_5$）。

2. 粒度。 我们考虑两种门控分数粒度：（a）逐注意力头：单个标量门控分数调制整个注意力头的输出；（b）逐元素：门控分数是与 $Y$ 维度相同的向量，从而实现细粒度的逐维调制。

3. 各注意力头独立或共享。 鉴于注意力的多头性质，我们进一步考虑：（a）各注意力头独立：每个注意力头拥有自己的门控分数，因而可以独立调制；（b）跨注意力头共享：各注意力头共享 $W_{\theta }$ 和门控分数。

4. 乘法或加法。 将门控分数应用于 $Y$ 时，我们考虑：（a）乘法门控：门控输出计算为 $Y^{\prime }=Y\cdot \sigma (X\theta )$；（b）加法门控：$Y^{\prime }=Y+\sigma (X\theta )$。

5. 激活函数。 我们主要考虑两种常用激活函数：SiLU 和 sigmoid。由于 SiLU 的输出范围无界，我们只将其用于加法门控，而 sigmoid 只会产生 $[0,1]$ 范围内的分数。此外，为了进一步剖析门控有效性的底层机制，我们还考虑恒等映射或 RMSNorm（详见第 4.1 节）。

Unless otherwise specified, we employ head-specific, multiplicative gating utilizing the sigmoid activation function ($\sigma (x)=\frac{1}{1+e^{-x}}$).

除非另有说明，我们采用各注意力头独立、使用 sigmoid 激活函数的乘法门控（$\sigma (x)=\frac{1}{1+e^{-x}}$）。

3. Experiments

3.1 Experimental Setups

Model Architecture and Training Settings. We conduct experiments on both MoE models (15B total parameters with 2.54B activated, 15A2B) and dense models (1.7B total parameters). The 15A2B MoE models utilize 128 total experts with top-8 softmax gating, fine-grained experts, global-batch LBL, and z-loss. We adopt group query attention (GQA) for the attention part. We train the models on subsets of a 3.5T high-quality tokens, encompassing multilingual, math, and general knowledge content. The context sequence length is set to 4096. More detailed configurations, such as learning rate and batch size (bsz), will be introduced in each part. Other hyperparameters follow the default values of the AdamW optimizer. Since the parameters and flops introduced by the gating are small, the wall-time latency introduced by gating is less than 2%.

模型架构与训练设置。 我们同时在 MoE 模型（总参数量 15B、激活参数量 2.54B，记为 15A2B）和稠密模型（总参数量 1.7B）上开展实验。 15A2B MoE 模型使用共 128 个专家、top-8 softmax 门控、细粒度专家、全局批次负载均衡损失（LBL）以及 z-loss。 注意力部分采用分组查询注意力（GQA）。 我们使用 3.5T 高质量 token 的子集训练模型，内容涵盖多语言、数学和通用知识。 上下文序列长度设为 4096。 学习率和批次大小（bsz）等更详细的配置将在各部分分别介绍。 其他超参数遵循 AdamW 优化器的默认值。 由于门控引入的参数量和 FLOPs 都很少，门控带来的实际运行延迟低于 2%。

Evaluation. We test the few-shots results on popular benchmarks, including Hellaswag for English, MMLU for general knowledge, GSM8k for math reasoning, HumanEval for coding, C-eval and CMMLU for Chinese proficiency. We also report the perplexity (PPL) of language modeling on diverse held-out test sets, including domains like English, Chinese, Code, Math, Law, and Literature.

评测。 我们在常用基准上测试少样本结果，包括用于英语的 Hellaswag、用于通用知识的 MMLU、用于数学推理的 GSM8k、用于代码的 HumanEval，以及用于中文能力的 C-eval 和 CMMLU。 我们还报告在多种留出测试集上的语言建模困惑度（PPL），覆盖英语、中文、代码、数学、法律和文学等领域。

表1：门控变体的性能与结果。15A2B MoE 模型在 400B token 上训练。d_k 为注意力头维度，d_model 为模型隐藏维度，n 为 token 数量，q 与 k 分别表示 query 头和 key-value 头数量；新增参数量的单位为百万。

Method	Act Func	Score Shape	Added Param	Avg PPL	Hellaswag	MMLU	GSM8k	C-eval
Reference Baselines (Baseline uses q = 32, k = 4. All methods use dk = 128.)
(1) Baseline	-	-	0	6.026	73.07	58.79	52.92	60.26
(2) k = 8	-	-	50	5.979	73.51	59.78	52.16	62.26
(3) q = 48	-	-	201	5.953	73.59	58.45	53.30	59.67
(4) Add 4 Experts	-	-	400	5.964	73.19	58.84	52.54	63.19
Gating Position Variants
(5) SDPA Elementwise G1	sigmoid	n × q × dk	201	5.761	74.64	60.82	55.27	62.20
(6) v Elementwise G2	sigmoid	n × k × dk	25	5.820	74.38	59.17	53.97	61.00
(7) k Elementwise G3	sigmoid	n × k × dk	25	6.016	72.88	59.18	50.49	61.74
(8) q Elementwise G4	sigmoid	n × q × dk	201	5.981	73.01	58.74	53.97	62.14
(9) Dense Output G5	sigmoid	n × dmodel	100	6.017	73.32	59.41	50.87	59.43
Gating Granularity Variants
(10) SDPA Headwise G1	sigmoid	n × q	1.6	5.792	74.50	60.05	54.44	62.61
(11) v Headwise G2	sigmoid	n × q	0.2	5.808	74.38	59.32	53.53	62.61
Head-Specific v.s. Head-Shared Gating
(12) SDPA Head-Shared G1	sigmoid	n × dk	201	5.801	74.34	60.06	53.15	61.01
(13) v Head-Shared G2	sigmoid	n × dk	25	5.867	74.10	59.02	53.03	60.61
Multiplicative v.s. Additive
(14) SDPA Additive G1	SiLU	n × q × dk	201	5.821	74.81	60.06	53.30	60.98
Activation Variants
(15) SDPA Elementwise G1	SiLU	n × q × dk	201	5.822	74.22	60.49	54.59	62.34

3.2 Main Results

3.2.1 Gated Attention for MoE models

We first compare the results of different gated attention layers on the training-efficient MoE-15A2B models. All models use a scheduler that warms up to a maximum LR of 2e-3 in 1k steps and decays using cosine to 3e-5. We use a global bsz of 1024, comprising 100k optimization steps. The results are summarized in Table 1. To provide a fair comparison, we supplement the vanilla MoE baseline (row 1) with parameter expansion methods, including increasing the number of key-value heads (row 2), increasing the number of query heads (row 3), and increasing both the total and activated number of experts (row 4). These methods introduce a comparable or greater number of parameters than the gating mechanisms. From Table 1, we observe:

我们首先在训练效率较高的 MoE-15A2B 模型上比较不同门控注意力层的结果。 所有模型都使用同一种调度器：在 1k 步内将学习率预热至最大值 2e-3，再通过余弦退火衰减到 3e-5。 我们采用全局批次大小 1024，共进行 100k 个优化步骤。 结果汇总于表1。 为进行公平比较，我们为原始 MoE 基线（第 1 行）补充了参数扩展方法，包括增加 key-value 头数量（第 2 行）、增加 query 头数量（第 3 行），以及同时增加专家总数和激活专家数（第 4 行）。 这些方法引入的参数量与门控机制相当或更多。 从表1中，我们观察到：

(i) SDPA and value output gating are effective. Inserting gates at the output of SDPA ($G_1$) or the value map ($G_2$) is the most effective, achieving lower PPL and better overall benchmark performance than other variants. We will further investigate why gating at these two positions is effective in Section 4.2.

（i）SDPA 与 value 输出门控有效。 在 SDPA 输出（$G_1$）或 value 映射（$G_2$）处插入门控最为有效，相比其他变体取得了更低的 PPL 和更好的整体基准性能。 我们将在第 4.2 节进一步研究这两个位置的门控为何有效。

(ii) Head-Specific Gating Matters. Applying headwise gating at $G_1$ and $G_2$ introduces very few additional parameters (less than 2M for the MoE-15A2B model) but still delivers substantial improvements (rows 10 and 11). When sharing gating scores across different attention heads (we average over the query head dimension $q$ to obtain an $n\times d_k$ score from the original $n\times q\times d_k$), the benchmark improvements are smaller than those achieved by headwise gating (row 12 v.s. 10, 13 v.s. 11). This underscores the importance of applying distinct gating scores for different attention heads.

（ii）各注意力头独立的门控很重要。 在 $G_1$ 和 $G_2$ 处应用逐注意力头门控只引入极少的额外参数（对 MoE-15A2B 模型不足 2M），但仍带来显著提升（第 10、11 行）。 当不同注意力头共享门控分数时（我们沿 query 头维度 $q$ 取平均，将原始 $n\times q\times d_k$ 分数变为 $n\times d_k$），基准性能增益小于逐注意力头门控（第 12 行对比第 10 行，第 13 行对比第 11 行）。 这凸显了为不同注意力头使用不同门控分数的重要性。

(iii) Multiplicative Gating is Preferred. Additive SDPA output gating underperforms the multiplicative one, although it shows improvements over the baselines.

（iii）乘法门控更可取。 尽管加法式 SDPA 输出门控相较基线有所提升，但其表现不如乘法门控。

(iv) Sigmoid Activation is Better. Replacing the activation function in the most effective gating configuration (row 5) with SiLU (row 15) leads to less improvement.

（iv）sigmoid 激活更好。 将最有效门控配置（第 5 行）的激活函数替换为 SiLU（第 15 行）后，提升幅度减小。

Overall, adding gating at the value layer ($G_2$) and SDPA output ($G_1$) reduces PPL by more than 0.2, outperforming various parameter-expanding baselines. However, gating at $G_1$ achieves better PPL and benchmark results. As long as different heads receive distinct gating scores, the granularity of gating and the choice of activation function have relatively minor impacts. We will further analyze the reasons behind these observations in Analysis (Section 4.2).

总体而言，在 value 层（$G_2$）和 SDPA 输出（$G_1$）处添加门控可使 PPL 降低超过 0.2，优于多种参数扩展基线。 不过，$G_1$ 处的门控取得了更好的 PPL 和基准结果。 只要不同注意力头使用不同门控分数，门控粒度和激活函数选择的影响就相对较小。 我们将在分析部分（第 4.2 节）进一步剖析这些观察背后的原因。

3.2.2 Gated Attention for Dense Models

We also conduct experiments on dense models following Qwen2 to validate SDPA output sigmoid gating. When using gating, we reduce the width of FFN to maintain the parameter size. Most experiments use optimized hyperparameters for the baseline. For instance, for the 1.7B model trained on 400B tokens, we use a maximum LR of 4e-3 and a bsz of 1024. For training on 3.5T tokens, we increase the maximum LR to 4.5e-3 and the bsz to 2048.

我们还参照 Qwen2 在稠密模型上开展实验，以验证 SDPA 输出 sigmoid 门控。 使用门控时，我们缩小 FFN 宽度，以保持参数量不变。 大多数实验都使用为基线优化过的超参数。 例如，对在 400B token 上训练的 1.7B 模型，我们采用最大学习率 4e-3 和批次大小 1024。 在 3.5T token 上训练时，我们将最大学习率提高到 4.5e-3，并将批次大小提高到 2048。

Prior work has established that while increased network depth, large learning rates, and large batch sizes can significantly improve model performance and distributed training efficiency, they often introduce training instabilities. We observe that applying gating mechanisms demonstrably reduces the occurrence of loss spikes during training, suggesting a promising role for gating in enhancing training stability. Motivated by this finding, we introduce another experimental setting characterized by an increased number of layers, a higher maximum learning rate, and a larger batch size to further probe gating's stabilizing effects.

先前工作已表明，增加网络深度、采用较大学习率和较大批次虽然能显著改善模型性能与分布式训练效率，却常常会引入训练不稳定性。 我们观察到，应用门控机制明显减少了训练期间的损失尖峰，说明门控有望增强训练稳定性。 受这一发现启发，我们又引入了一种具有更多层数、更高最大学习率和更大批次的实验设置，以进一步探究门控的稳定作用。

表2：不同学习率、批次大小和模型配置下各方法的性能。SDPA 表示在 SDPA 后应用 sigmoid 门控；Sandwich Norm 表示在将注意力或 FFN 输出加入残差前进行归一化。使用门控时会缩小 FFN 宽度，使所有方法参数量相同；“-”表示训练发散。

Method	Max LR	Avg PPL	HumanEval	MMLU	GSM8k	Hellaswag	C-eval	CMMLU
28 Layer, 1.7B Parameters, 400B Tokens, Batch Size = 1024
(1) Baseline	4.0 × 10−3	7.499	28.66	50.21	27.82	64.94	49.15	49.52
(2) SDPA Elementwise	4.0 × 10−3	7.404	29.27	51.15	28.28	65.48	50.72	50.72
28 Layer, 1.7B Parameters, 3.5T Tokens, Batch Size = 2048
(3) Baseline	4.5 × 10−3	6.180	34.15	59.10	69.07	68.02	68.19	64.95
(4) SDPA Elementwise	4.5 × 10−3	6.130	37.80	59.61	70.20	68.84	68.52	65.76
48 Layer, 1.7B Parameters, 400B Tokens, Batch Size = 1024
(5) Baseline	4.0 × 10−3	7.421	28.05	52.04	32.98	65.96	51.11	51.86
(6) Baseline	8.0 × 10−3	9.195	21.34	44.28	15.24	57.00	43.11	42.63
(7) Baseline + Sandwich Norm	8.0 × 10−3	7.407	30.49	52.07	32.90	66.00	52.04	51.72
(8) SDPA Elementwise	4.0 × 10−3	7.288	31.71	52.44	32.37	66.28	52.06	52.29
(9) SDPA Headwise	4.0 × 10−3	7.370	31.10	53.83	34.12	65.59	55.07	52.38
(10) SDPA Elementwise	8.0 × 10−3	7.325	31.10	54.47	36.62	66.40	53.91	53.80
48 Layer, 1.7B Parameters, 1T Tokens, Batch Size = 4096
(11) Baseline	5.3 × 10−3	7.363	29.88	54.44	32.22	65.43	53.72	53.37
(12) Baseline	8.0 × 10−3	-	-	-	-	-	-	-
(13) SDPA Elementwise	5.3 × 10−3	7.101	34.15	55.70	36.69	67.17	54.51	54.68
(14) SDPA Elementwise	8.0 × 10−3	7.078	31.71	56.47	39.73	67.38	55.52	55.77

Table 2 reveals that:

表2揭示了以下结论：

(i) Gating is effective across various settings. Across various model configurations (row 1 v.s. 2, 5 v.s. 8), training data (row 3 v.s. 4), and hyperparameters (row 11 v.s. 13), applying SDPA output gating consistently yields benefits.

（i）门控在各种设置下都有效。 在不同模型配置（第 1 行对比第 2 行，第 5 行对比第 8 行）、训练数据规模（第 3 行对比第 4 行）和超参数（第 11 行对比第 13 行）下，应用 SDPA 输出门控都能稳定带来收益。

(ii) Gating improves stability and facilitates scaling. Under the 3.5T token setting, gating improves training stability, largely reducing the loss spike (Figure 1, right). When increasing the maximum LR, baselines encounter convergence issues (row 6, 12). While adding sandwich norm restores convergence, the improvement is negligible. In contrast, increasing the maximum LR in models with gating results in a noticeable improvement.

（ii）门控提高稳定性并促进扩展。 在 3.5T token 设置下，门控增强了训练稳定性，大幅减少损失尖峰（图1右侧）。 提高最大学习率时，基线会遇到收敛问题（第 6、12 行）。 添加 Sandwich Norm 虽然恢复了收敛，但提升可以忽略不计。 相比之下，提高带门控模型的最大学习率会带来明显改善。

In summary, we identify SDPA element-wise gating as the most effective method to augment the attention mechanism. Applying this method to dense transformers further demonstrates that the gate enables stable training with larger batch sizes and learning rates, resulting in improved performance.

总之，我们确认 SDPA 逐元素门控是增强注意力机制最有效的方法。 将该方法应用于稠密 Transformer 进一步表明，门控能够在更大批次和更高学习率下保持稳定训练，从而改善性能。

4. Analysis: Non-Linearity, Sparsity, and Attention-Sink-Free

In this section, we conduct a series of experiments to explore why such a simple gating mechanism can yield significant improvements in performance and training stability. Here are the takeaways according to our analysis:

在本节中，我们通过一系列实验探究为何如此简单的门控机制能显著改善性能与训练稳定性。 我们的分析得到以下结论：

1. Gating operations enhancing non-linearity consistently lead to performance gains (Section 4.1).

2. The most effective SDPA elementwise $G_1$ gate introduces strong input-dependent sparsity to the SDPA outputs (Section 4.2), which then helps to eliminate the `attention sink' phenomenon.

1. 增强非线性的门控操作能够稳定带来性能增益（第 4.1 节）。

2. 最有效的 SDPA 逐元素 $G_1$ 门控为 SDPA 输出引入了很强的依赖输入的稀疏性（第 4.2 节），进而有助于消除 attention sink 现象。

4.1 Non-linearity Improves the Expressiveness of Low-Rank Mapping in Attention

表3：不同非线性与无非线性增强方法的性能。

Method	Activation Function	Avg PPL	Hellaswag	MMLU	GSM8k	C-eval
(1) Baseline	-	6.026	73.07	58.79	52.92	60.26
(2) SDPA Elementwise Gate	Sigmoid	5.761	74.64	60.82	55.27	62.20
(3) v Elementwise Gate	Sigmoid	5.820	74.38	59.17	53.97	61.00
(4) SDPA Additive Gate	SiLU	5.821	74.81	60.06	53.30	60.98
(5) SDPA GroupNorm	RMSNorm	5.847	74.10	60.15	53.75	61.14
(6) SDPA SiLU	SiLU	5.975	73.34	59.55	53.19	60.90
(7) SDPA Additive Gate	Identity	5.882	74.17	59.20	52.77	59.86

Inspired by prior works that utilize group norm for the SDPA output, with the same setting in Section 3.2.1, we apply RMSNorm independently to the output of each attention head before concatenation. As shown in Table 3 row 5, applying RMSNorm, which introduces almost no additional parameters, also leads to a significant reduction in PPL. In multi-head attention, the output of the $i$-th token, corresponding to the $k$-th head, can be expressed as:

受先前在 SDPA 输出上使用分组归一化的工作启发，我们采用与第 3.2.1 节相同的设置，在拼接前分别对每个注意力头的输出应用 RMSNorm。 如表3第 5 行所示，应用几乎不引入额外参数的 RMSNorm 同样能显著降低 PPL。 在多头注意力中，与第 $k$ 个注意力头对应的第 $i$ 个 token 的输出可表示为：

$$\begin{array}{rl}{o}_i^k& =(\sum _{j=0}^i{S}_{ij}^k\cdot X_j{W}_V^k)W_O^k\\ & =\sum _{j=0}^i{S}_{ij}^k\cdot X_j\left(W_V^k{W}_O^k\right).\end{array}$$

Where $W_O^k$ is the parameters of the output layer $W_O$ corresponding to the $k$-th head (Note that concatenating outputs from different heads and then multiplying with $W_O$ is equivalent to multiplying each head's output with its corresponding $W_O^k$ before concatenation). Here, $S_{ij}^k$ denotes the attention score of the $i$-th token attending to the $j$-th token in the $k$-th head, $X_j$ is the input to the attention for token $j$, and $X_j{W}_V^k$ represents the value output of token $j$ in the $k$-th head. From the equation, we can merge $W_V^k{W}_O^k$ into one low-rank linear mapping applied over all $X_j$ as $d_k<d_{\text{model}}$. With GQA, $W_V$ is shared among heads within the same group, further diminishing the expressiveness. Given that adding non-linearity between two linear mappings can improve their expensiveness, we have two modifications to mitigate the low-rank problem:

其中，$W_O^k$ 是输出层 $W_O$ 中与第 $k$ 个注意力头对应的参数（将不同注意力头的输出拼接后再乘以 $W_O$，等价于先让每个注意力头的输出乘以对应的 $W_O^k$，再进行拼接）。 这里，$S_{ij}^k$ 表示第 $k$ 个注意力头中第 $i$ 个 token 关注第 $j$ 个 token 的注意力分数，$X_j$ 是 token $j$ 的注意力输入，而 $X_j{W}_V^k$ 表示 token $j$ 在第 $k$ 个注意力头中的 value 输出。 由该公式可知，由于 $d_k<d_{\text{model}}$，我们可以将 $W_V^k{W}_O^k$ 合并为一个应用于所有 $X_j$ 的低秩线性映射。 采用 GQA 时，同一组内的注意力头共享 $W_V$，进一步削弱了表达能力。 由于在两个线性映射之间加入非线性可以提高其表达能力，我们提出两种修改来缓解低秩问题：

$$\begin{array}{rl}{o}_i^k& =(\sum _{j=0}^i{S}_{ij}^k\cdot \mathrm{NonLinearMap}(X_j{W}_V^k))W_O^k,\\ o_i^k& =\mathrm{NonLinearMap}(\sum _{j=0}^i{S}_{ij}^k\cdot X_j{W}_V^k)W_O^k.\end{array}$$

Notably, adding gating at the $G_2$ (Table 3 row 3) position corresponds to the first modification, while adding gating (row 4) or group normalization (row 5) at the $G_1$ position corresponds to the second. This also explains why adding gating or normalization at the $G_5$ position after $W_O$ has no effect (Table 1 row 9) — it does not address the lack of non-linearity between $W_V$ and $W_O$. For additive gating at $G_1$, the output of gating passes through SiLU (Table 3 row 4), also introducing some non-linearity, which explains the observed performance gains, albeit smaller than those achieved by multiplicative gating.

值得注意的是，在 $G_2$ 位置添加门控（表3第 3 行）对应第一种修改，而在 $G_1$ 位置添加门控（第 4 行）或分组归一化（第 5 行）对应第二种修改。 这也解释了为何在 $W_O$ 之后的 $G_5$ 位置添加门控或归一化没有效果（表1第 9 行）：它并未解决 $W_V$ 与 $W_O$ 之间缺少非线性的问题。 对于 $G_1$ 处的加法门控，门控输出会经过 SiLU（表3第 4 行），从而也引入一定的非线性；这解释了观察到的性能增益，尽管增益小于乘法门控。

Based on these insights, we conduct two additional experiments: (i) Adding SiLU only at the $G_1$ position without introducing additional parameters (Table 3 row 6). Notice this simple modification also leads to a modest reduction in PPL, but most benchmark scores remain unchanged. (ii) Removing SiLU from additive gating, such that the output of $X_j$ after gating is directly added at the $G_1$ position (Table 3 row 7). This further diminishes the gains of addictive gating.

基于这些认识，我们又开展了两个实验：（i） 仅在 $G_1$ 位置添加 SiLU，且不引入额外参数（表3第 6 行）。 这一简单修改也会小幅降低 PPL，但大多数基准分数保持不变。 （ii） 从加法门控中移除 SiLU，使 $X_j$ 经门控后的输出直接加到 $G_1$ 位置（表3第 7 行）。 这进一步削弱了加法门控的增益。

In summary, the enhanced performance associated with effective gating variants is likely attributable to the introduction of non-linearity between $W_V$ and $W_O$. Although applying gating at positions $G_1$ and $G_2$ can can both introduce this non-linearity, these applications yield differing performance gains. This observed difference motivates us to further analyze the impacts of gating at these two positions.

总之，有效门控变体带来的性能提升很可能源于在 $W_V$ 与 $W_O$ 之间引入了非线性。 虽然在 $G_1$ 和 $G_2$ 位置应用门控都能引入这种非线性，但两者带来的性能增益并不相同。 这一差异促使我们进一步分析这两个位置上门控的影响。

4.2 Gating Introduces Input-Dependent Sparsity

We analyze the gating scores (Table 1, `Gate Score' column) of models with gating applied at the value ($G_2$) and SDPA output ($G_1$) positions, evaluated on the test language modeling data. The mean gating scores for all layers are presented in Table 4, with the score distributions visualized in Figure 3 (layer-wise scores in the appendix). Key observations include:

我们在语言建模测试数据上分析在 value（$G_2$）和 SDPA 输出（$G_1$）位置应用门控的模型门控分数（表1的 Gate Score 列）。 所有层的平均门控分数见表4，分数分布可视化于图3（逐层分数见附录）。 主要观察包括：

(i) Effective Gating Scores are Sparse. SDPA output gatings (Element/head-wise) exhibit the lowest mean gating scores. Furthermore, the SDPA output gating score distribution shows a high concentration near 0, indicating substantial sparsity, consistent with its superior performance.

（i）有效的门控分数是稀疏的。 SDPA 输出门控（逐元素或逐注意力头）具有最低的平均门控分数。 此外，SDPA 输出门控的分数分布高度集中在 0 附近，说明其稀疏性很强，这与它更优的性能一致。

(ii) Head-Specific Sparsity Matters. Enforcing shared gating scores across attention heads increases the overall gating scores and diminishes performance gains. Observations (i) and (ii) underscore the importance of head-specific gating, aligning with previous research demonstrating that individual attention heads capture distinct aspects of the input.

（ii）各注意力头独立的稀疏性很重要。 强制不同注意力头共享门控分数会提高总体门控分数，并削弱性能增益。 观察（i）和（ii）凸显了 各注意力头独立门控 的重要性，这与先前关于不同注意力头捕获输入中不同方面的研究结论一致。

表4：不同激活函数下各门控方法的性能与平均门控分数。Act-Func 表示用于计算门控分数的激活函数；M-Act 表示模型各层隐藏状态最大激活值的取整均值；F-Attn 表示第一个 token 的注意力分数，值越高说明注意力汇聚越明显。

Method	Act-Func	Gate Score	M-Act	F-Attn	PPL	Hellaswag	MMLU	GSM8k
(1) Baseline	-	-	1053	0.467	6.026	73.07	58.79	52.92
(2) SDPA Elementwise Gate	Sigmoid	0.116	94	0.048	5.761	74.64	60.82	55.27
(3) SDPA Headwise Gate	Sigmoid	0.172	98	0.073	5.792	74.50	60.05	54.44
(4) SDPA Elementwise Head-shared Gate	Sigmoid	0.271	286	0.301	5.801	74.34	60.06	53.15
(5) v Elementwise Gate	Sigmoid	0.221	125	0.297	5.820	74.38	59.17	51.33
(6) SDPA Input Independent Gate	Sigmoid	0.335	471	0.364	5.917	73.64	59.02	52.40
(7) SDPA Elementwise Gate	NS-sigmoid	0.653	892	0.451	5.900	74.05	60.05	52.75

图3：SDPA 逐元素门控（左）、value 逐元素门控（中）以及跨注意力头共享的 SDPA 逐元素门控（右）的门控分数均值与分布。大多数门控分数低于 0.5，说明门控分数具有稀疏性；其中 SDPA 输出门控的稀疏性最强。

(iii) Query-Dependency Matters. The scores for value gating ($G_2$) are higher than those for SDPA output gating ($G_1$), and the performance is inferior. This suggests that gating score sparsity is more effective when query-dependent rather than determined by the key and value. Specifically, SDPA output gating scores are derived from the hidden states corresponding to the current query (e.g. the Non-Linearity-Map in the second equation depends on $X_i$), whereas value gating scores are derived from hidden states associated with past keys and values (e.g. the Non-Linearity-Map in the first equation depends on each $X_j$). This implies that gating score sparsity may filter out irrelevant contextual information for the query. To further validate the importance of query-dependency, we introduce input-independent gating by zero-initializing learnable parameters ($q\times d_k$), applying a sigmoid function, and multiplying it with the SDPA output. As shown in row (6), input-independent gating improves upon the baseline, likely due to the introduction of non-linearity. Moreover, the high gating scores reinforce that effective sparsity should be input-dependent.

（iii）查询依赖性很重要。 value 门控（$G_2$）的分数高于 SDPA 输出门控（$G_1$），且性能更差。 这说明，相比由 key 和 value 决定，门控分数的稀疏性在依赖 query 时更加有效。 具体而言，SDPA 输出门控分数来自与当前 query 对应的隐藏状态（例如第二个公式中的非线性映射依赖 $X_i$），而 value 门控分数来自与过去 key 和 value 相关的隐藏状态（例如第一个公式中的非线性映射依赖每个 $X_j$）。 这意味着，门控分数的稀疏性可能会为 query 过滤掉无关的上下文信息。 为了进一步验证查询依赖性的重要性，我们将可学习参数（$q\times d_k$）初始化为零、应用 sigmoid 函数并与 SDPA 输出相乘，从而引入与输入无关的门控。 如第（6）行所示，与输入无关的门控优于基线，这可能源于它引入了非线性。 此外，较高的门控分数进一步说明，有效的稀疏性应当依赖输入。

(iv) Less Sparse Gating is Worse. To further validate the importance of gating score sparsity, we reduce sparsity from the gating formulation. Specifically, we replace the sigmoid function with a modified Non-Sparse (NS) version:

（iv）较低的稀疏性表现更差。 为了进一步验证门控分数稀疏性的重要性，我们降低门控形式中的稀疏性。 具体而言，我们将 sigmoid 函数替换为经过修改的非稀疏（NS）版本：

$$\mathrm{NS}\text{-}\mathrm{sigmoid}(x)=0.5+0.5\cdot \mathrm{sigmoid}(x).$$

Which constrains the gating scores between [0.5, 1.0]. This ensures introducing non-linearity while removing gating score sparsity. As shown in Table 4 row (7), the gains of NS-sigmoid gating are inferior to those of SDPA output sigmoid gating. In the appendix, we provide a more detailed discussion on how sparse gating scores affect the sparsity (the proportion of values below the threshold) in SDPA hidden states. We will discuss the impact of different sparsity levels on model behavior, including reducing the `attention sink', in the next section.

这会将门控分数限制在 [0.5, 1.0] 之间。 这样既能引入非线性，又能移除门控分数的稀疏性。 如表4第（7）行所示，NS-sigmoid 门控带来的增益低于 SDPA 输出 sigmoid 门控。 附录进一步讨论了稀疏门控分数如何影响 SDPA 隐藏状态的稀疏性，即低于阈值的数值比例。 下一节将讨论不同稀疏程度对模型行为的影响，包括其对减少 attention sink 的作用。

4.3 SDPA Output Gating Reduces Attention-Sink

Based on the observation that gating introduces sparsity to the SDPA output in an input-dependent manner, we hypothesized that this mechanism can filter out context irrelevant to the current query token, thereby mitigating the attention sink. To verify this, we analyze the distribution of attention scores (averaged over all heads) and the proportion of attention scores allocated to the first token (Figure 2, Table 4, `F-Attn' column). Inspired by the discussion about massive activation in hidden states and attention sinks, we also compute the mean of the maximum hidden state activations across layers, as shown in the `M-Act' column of Table 4. More detailed layer-wise results are provided in the appendix. We can observe:

基于门控以依赖输入的方式为 SDPA 输出引入稀疏性的观察，我们推测该机制可以滤除与当前 query token 无关的上下文，从而缓解注意力汇聚。 为验证这一点，我们分析了注意力分数的分布（对所有注意力头取平均），以及分配给第一个 token 的注意力分数比例（图2、表4的 F-Attn 列）。 受隐藏状态中大规模激活与注意力汇聚相关讨论的启发，我们还计算了跨层隐藏状态最大激活值的均值，如表4的 M-Act 列所示。 更详细的逐层结果见附录。 我们可以观察到：

(i) Head-wise and element-wise query-dependent sigmoid gating at the SDPA output ($G_1$) largely reduces the attention score allocated to the first token and decreases massive activations. (ii) Enforcing shared gating scores across heads or applying gating only after the value projection ($G_2$) decreases massive activations, but does not reduce attention scores to the first token. This reinforces the importance of head-specific gating and suggests that massive activations are not a prerequisite for attention sinks. (iii) Reducing the input-dependence of gating (row 6) or using NS-sigmoid to reduce sparsity (row 7) intensifies both massive activations and attention sink. Collectively, these observations indicate that input-dependent, head-specific gating of the SDPA output introduces significant sparsity, thereby mitigating the attention sink.

（i） 在 SDPA 输出（$G_1$）处应用逐注意力头或逐元素、依赖查询的 sigmoid 门控，会大幅减少分配给第一个 token 的注意力分数，并降低大规模激活。 （ii） 强制不同注意力头共享门控分数，或仅在 value 投影后（$G_2$）应用门控，会降低大规模激活，但不会减少第一个 token 获得的注意力分数。 这进一步证明了各注意力头独立门控的重要性，并说明 大规模激活并不是出现注意力汇聚的必要条件。 （iii） 降低门控对输入的依赖（第 6 行），或使用 NS-sigmoid 降低稀疏性（第 7 行），都会加重大规模激活和注意力汇聚。 综合来看，这些观察表明，在 SDPA 输出处应用依赖输入、各注意力头独立的门控会引入显著稀疏性，从而缓解注意力汇聚。

Furthermore, sparsity in the SDPA outputs reduces massive activations within the model, with increased sparsity leading to smaller activations. This may explain the improved training stability with gating: by reducing massive activations, the model is less susceptible to numerical errors during BF16 training. We also observe that massive activations originate primarily from early layers (e.g., layer 5), where the FFN outputs large values, consistent with Yona et al. Once added to the residual stream, these activations are propagated through subsequent layers via the pre-norm mechanism. This aligns with the effectiveness of sandwich normalization in enhancing training stability (Table 2, row 7): applying LayerNorm to the FFN output prevents these large activations from entering the residual stream.

此外，SDPA 输出中的稀疏性会减少模型内部的大规模激活，稀疏性越强，激活值越小。 这或许解释了门控为何能提高训练稳定性：通过减少大规模激活，模型在 BF16 训练期间更不容易受到数值误差影响。 我们还观察到，大规模激活主要源于早期层（例如第 5 层），这些层中的 FFN 会输出很大的数值，这与 Yona 等人的观察一致。 这些激活一旦加入残差流，就会通过预归一化机制传播到后续层。 这与 Sandwich Norm 能够提高训练稳定性的结果一致（表2第 7 行）：对 FFN 输出应用 LayerNorm，可以防止这些大激活进入残差流。

4.4 SDPA Output Gating Facilitates Context Length Extension

表5：不同方法在不同序列长度下的性能。YaRN Extended 表示扩展上下文长度的变体；括号内数值表示扩展上下文长度后的性能下降。

Method	4k	8k	16k	32k	64k	128k
Baseline	88.89	85.88	83.15	79.50	-	-
SDPA-Gate	90.56	87.11	84.61	79.77	-	-
YaRN Extended
Baseline	82.90 (-6.0)	71.52 (-14.4)	61.23 (-21.9)	37.94 (-41.56)	37.51	31.65
SDPA-Gate	88.13 (-2.4)	80.01 (-7.1)	76.74 (-7.87)	72.88 (-6.89)	66.60	58.82

Based on the attention-sink-free pattern, we evaluate the SDPA gating's effect in the long-context setting. Specifically, we extend the context length for the models trained on 3.5T tokens. We increase the RoPE base from 10k to 1M and continue training on data with a sequence length of 32k for an additional 80B tokens. This gives us models with a context length of 32k. Subsequently, we use YaRN to extend the context length to 128k. We evaluate models on the RULER benchmark and summarize results in Table 5. We observe the following:

基于无注意力汇聚的模式，我们评估 SDPA 门控在长上下文设置中的作用。 具体而言，我们扩展了在 3.5T token 上训练的模型的上下文长度。 我们将 RoPE 基数从 10k 提高到 1M，并在序列长度为 32k 的数据上继续训练 80B token。 由此得到上下文长度为 32k 的模型。 随后，我们使用 YaRN 将上下文长度扩展到 128k。 我们在 RULER 基准上评估模型，并将结果汇总于表5。 我们得到以下观察：

(i) Under the 32k setting, models with gating slightly outperform the baseline. This suggests that within the training length, the attention sink phenomenon may not hurt the model's long-context performance.

（i） 在 32k 设置下，带门控的模型略优于基线。 这说明在训练长度范围内，注意力汇聚现象可能不会损害模型的长上下文性能。

(ii) When the context length is extended to 128k using YaRN, both the baseline and gated models experience a decline within the original 32k range. This observation is consistent with previous works on extending context length by modifying RoPE. Even though the decline is less pronounced for models with gating.

（ii） 使用 YaRN 将上下文长度扩展到 128k 后，基线模型和门控模型在原始 32k 范围内的性能都会下降。 这一观察与通过修改 RoPE 扩展上下文长度的先前工作一致。 不过，带门控模型的下降幅度更小。

(iii) At context lengths of 64k and 128k, the gated attention models outperform the baseline signifantly. From these observations, we hypothesize that adding gating helps the model adapt to the context-length extension. A possible explanation is that baseline models rely on attention sinks to adjust the distribution of attention scores. Dong et al. derives the effects of changing the RoPE based on the attention and hidden state distributions. When techniques like YaRN are applied to modify the RoPE base, the attention sink pattern may struggle to adapt in a training-free manner, leading to a noticeable drop in performance. In contrast, models with gating primarily rely on input-dependent gating scores to control information flow, making them more robust to such changes.

（iii） 在 64k 和 128k 上下文长度下，门控注意力模型显著优于基线。 基于这些观察，我们推测添加门控有助于模型适应上下文长度扩展。 一种可能的解释是，基线模型依赖注意力汇聚来调整注意力分数的分布。 Dong 等人根据注意力与隐藏状态分布推导了改变 RoPE 所产生的影响。 当使用 YaRN 等技术修改 RoPE 基数时，注意力汇聚模式可能难以在免训练情况下适应，从而导致明显的性能下降。 相比之下，带门控模型主要依靠依赖输入的门控分数控制信息流，因此对这类变化更加稳健。

5. Related Works

5.1 Gating in Neural Networks

Gating mechanisms have been widely adopted in neural networks. Early works such as LSTMs and GRUs introduce gates to regulate information flow across time steps, addressing gradient vanishing/exploding issues by selectively retaining or discarding information. Highway Networks extend this concept to feedforward networks, enabling the successful training of very deep architectures. SwiGLU introduce gating mechanisms into transformer FFN layers, enhancing their expressive power and becoming a standard component in many open-source LLMs. Several works on state-space models and Linear Attention, such as FLASH, RetNet, Lightning Attention, and Gated Delta Networks, also incorporate gating modules to controlinformation of token-mixer modules. Forgetting Transformer applies gating mechanisms to the output of softmax attention, observing significant performance improvements. Although these works demonstrate the effectiveness of gating, a comprehensive understanding of its precise mechanisms and the reasons behind its effectiveness still needs exploration. This could contribute to a broader appreciation of gating's importance beyond RNNs and facilitate designs that better leverage gating's unique advantages.

门控机制已在神经网络中得到广泛采用。 LSTM 和 GRU 等早期工作引入门控来调节跨时间步的信息流，通过选择性保留或丢弃信息来解决梯度消失与爆炸问题。 Highway Network 将这一概念扩展到前馈网络，使非常深的架构能够成功训练。 SwiGLU 将门控机制引入 Transformer 的 FFN 层，增强其表达能力，并成为许多开源 LLM 的标准组件。 FLASH、RetNet、Lightning Attention 和 Gated Delta Network 等状态空间模型与线性注意力工作，也引入门控模块来控制 token 混合器模块中的信息。 Forgetting Transformer 将门控机制应用于 softmax 注意力输出，并观察到显著的性能提升。 尽管这些工作证明了门控的有效性，但仍需全面理解其精确机制及其有效的原因。 这有助于更广泛地认识门控在 RNN 之外的重要性，并促进更充分利用门控独特优势的设计。

For example, while Switch Heads, NSA, and MoSA employ sigmoid-based gating for selection, further investigation into isolating gating's specific contribution could offer valuable insights. Comparisons with baselines incorporating similar gating mechanisms in standard transformers could offer a more refined perspective on the effectiveness of their proposed selection mechanisms.

例如，Switch Heads、NSA 和 MoSA 使用基于 sigmoid 的门控进行选择，但进一步分离门控自身的具体贡献仍可带来有价值的认识。 与在标准 Transformer 中加入类似门控机制的基线进行比较，可以更细致地审视这些方法所提选择机制的有效性。

The work most closely related to ours is Quantizable Transformers, which also finds that applying gating in softmax attention alleviates extreme attention concentration and outliers in hidden states in encoder models like BERT and ViT. While this work primarily leverages gating to eliminate outliers for model quantization, we provide a detailed analysis of various gating variants, uncovering their benefits through enhanced non-linearity and sparsity, as well as improved training stability. Building on these insights, we scale up gated attention models, demonstrating gating's broad applicability and impact.

与我们最相关的工作是 Quantizable Transformers；该工作同样发现，在 softmax 注意力中应用门控可以缓解 BERT 和 ViT 等编码器模型中的极端注意力集中与隐藏状态离群值。 该工作主要利用门控消除离群值以进行模型量化，而我们详细分析了多种门控变体，揭示了它们通过增强非线性与稀疏性以及改善训练稳定性所带来的收益。 基于这些认识，我们进一步扩大门控注意力模型的规模，证明了门控广泛的适用性与影响。

5.2 Attention Sink

Xiao et al. formally identifies the `attention sink' phenomenon, in which specific tokens receive large attention scores. Similarly, Darcet et al. finds in the vision transformer, some redundant tokens act as `registers' to store attention scores. Later, Sun et al. shows that excessive attention scores are also assigned to tokens associated with massive activation values. However, our work reveals that applying gating at the output of value projection eliminates massive activations, yet attention sinks persist, indicating that massive activations are not a necessary condition for attention sinks.

Xiao 等人正式识别出 attention sink 现象，即特定 token 会获得很大的注意力分数。 类似地，Darcet 等人发现，在视觉 Transformer 中，一些冗余 token 会充当存储注意力分数的“寄存器”。 随后，Sun 等人表明，过高的注意力分数也会分配给与大规模激活值相关的 token。 然而，我们的工作发现，在 value 投影输出处应用门控会消除大规模激活，但注意力汇聚仍然存在，这说明大规模激活不是出现注意力汇聚的必要条件。

Similarly, Gu et al. characterizes attention sinks as non-informative `key biases' that store redundant attention scores, arguing that softmax's inherent normalization dependency drives this behavior. Experimental attempts to modify softmax attention, such as replacing softmax with unnormalized sigmoid attention, adding softmax attention gate or clip, and modifying softmax computation and denominator, show promise in mitigating attention sinks. Our work demonstrates that sparse gating after SDPA eliminates attention sinks in both dense (1B-parameter) and MoE (15B-parameter) models, even when trained on 3.5T tokens. Furthermore, we uncover the potential of eliminating attention sinks to benefit context-length extension.

类似地，Gu 等人将注意力汇聚描述为存储冗余注意力分数的无信息“key bias”，并认为 softmax 内在的归一化依赖驱动了这一行为。 修改 softmax 注意力的实验尝试，例如用未归一化的 sigmoid 注意力替换 softmax、添加 softmax 注意力门控或裁剪，以及修改 softmax 计算和分母，都显示出缓解注意力汇聚的潜力。 我们的工作证明，即使在 3.5T token 上训练，在 SDPA 后应用稀疏门控也能消除稠密模型（1B 参数）和 MoE 模型（15B 参数）中的注意力汇聚。 此外，我们揭示了消除注意力汇聚有利于上下文长度扩展的潜力。

6. Conclusion

This work systematically investigates the role of gating mechanisms in the standard softmax attention, revealing their significant impact on performance, training stability, and attention dynamics. Through extensive experimental comparisons over 30 variants of 15B MoE and 1.7B dense models trained on up to 3.5T tokens, we demonstrate that applying a sigmoid gate after scaled dot-product attention yields the most substantial improvements. This simple mechanism enhances non-linearity, introduces input-dependent sparsity, and eliminates inefficiencies like the `attention sink' phenomenon. Additionally, gating facilitates context length extension, allowing models to generalize effectively to longer sequences without retraining. We also release the first attention-sink-free models. We believe these empirical validations will pave the way for engineering the next generation of advanced foundation models.

本工作系统研究了门控机制在标准 softmax 注意力中的作用，揭示了它们对性能、训练稳定性和注意力动态的显著影响。 通过对在最多 3.5T token 上训练的 15B MoE 与 1.7B 稠密模型的 30 种变体开展广泛实验比较，我们证明，在缩放点积注意力之后应用 sigmoid 门控可带来最显著的提升。 这一简单机制增强了非线性，引入了依赖输入的稀疏性，并消除了 attention sink 等低效现象。 此外，门控有助于扩展上下文长度，使模型无需重新训练即可有效泛化到更长序列。 我们还发布了首批无注意力汇聚模型。 我们相信，这些实证验证将为下一代先进基础模型的工程设计铺平道路。

Limitations

Our work primarily focuses on analyzing the reasons and impacts of attention gating through a series of ablation studies. However, we acknowledge several limitations. The broader implications of non-linearity on the dynamics of attention and the overall training process remain under-explored. Although we observe that eliminating attention sinks improves performance in long-context extension scenarios, we do not provide a rigorous theoretical explanation for how attention sinks influence the model's ability to generalize to longer sequences.

我们的工作主要通过一系列消融研究分析注意力门控有效的原因及其影响。 不过，我们承认本工作仍存在若干局限。 非线性对注意力动态和整体训练过程的更广泛影响仍未得到充分探索。 尽管我们观察到消除注意力汇聚可以改善长上下文扩展场景中的性能，但我们没有就注意力汇聚如何影响模型泛化到更长序列的能力给出严格的理论解释。