Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free
Authors: Zihan Qiu
Zekun Wang
Bo Zheng
Zeyu Huang
Kaiyue Wen
Songlin Yang
Rui Men
Le Yu
Fei Huang
Suozhi Huang
Dayiheng Liu
Jingren Zhou
Junyang Lin
Venue: NeurIPS 2025 oral ()
Keywords: Attention
Large Language Model
Gating
Primary area: deep_learning
TL;DR: We find applying a query-dependent head-specific sigmoid gate after the Scaled Dot-Product Attention (SDPA) consistently improves performance, improves scaling properties and mitigates the `massive activation' and `attention sink'.

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.
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 `massive activation`, `attention sink` and enhances long-context extrapolation performance. 
We also release related codes (https://github.com/qiuzh20/gated_attention}) and models (https://huggingface.co/QwQZh/gated_attention) to facilitate future research.
Furthermore, the most effective SDPA output gating is used in the Qwen3-Next models (https://huggingface.co/collections/Qwen/qwen3-next).

Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free
Zihan Qiu"!, Zekun Wang" Bo Zheng" Zeyu Huang
Kaiyue Wen Songlin Yang', Rui Men Le Yu', Fei Baeangg,Suozie Huang?,
Dayiheng Liu Jingren Zhou', Junyang Lin
'Qwen Team, Alibaba Group University of Edinburgh *Stanford University *MIT TTsinghua University
Abstract
Gating mechanisms have been widely utilized, from early models like LSTMs [I] and Highway Networks (2 to recent state space models [3]. linear attention [4] and also softmax attention 6). Yet, existing literature rarely examines the specific effects of gating. In this work, we conduct comprchensive experiments to system- atically investigate gating-augmented softmax attention variants. Specifically, we perform a comprehensive comparison over 30 variants of 15B Mixture-of -f-exes (MoE) models and 1.7B dense models trained on a 3.5 trillion token dataset. Our central finding is that a simple modificationn-aplyiin an head-specific sigmoid gate after the Scaled Dot-Product Attention (SDPA) consiitenyy improves perfor- mance. 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 the softmax attention, and (2) applying query-deeendent sparse gating scores to modulate the SDPA output. Notably, we find this sparse gating mechanism mitigates 'massive activation' (7). 'attention sink' [8], and enhances long-context extrapolation performance, and we also release related codes and models to facilitate future research. Furthermore, the most effective SDPA output gating is used in the Qwen3-Next models
Introduction
Gating mechanism is well-established in neural networks. Early architectures, such as LSTMs [
Highway Networks 22 and GRUs [9], pioneer the of gating control information flow across time steps or layers and improve gradient propagation. This principle persists modern architectures. Recent sequence modeling works, including state-space models 3 [T] and attention mechanisms [1T]
12 13
6 commonly apply gating, often to modulate the outputs of token- mixer components. Despite its widespread adoption and empirical success, most recent works do look into the gating mechanisms like the gating scores and their effect on the model's hidden states.
14
[5] 16,
17 18,
Insufficient understanding hinders assessing gating's true contribution. especially when confounded with other architectural factors. For instance, while Switch Heads [99 [0] introduces a sigmoid gating to select top-K attention head experts, experiments reveal interesting finding (Appendix [A. 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) (11, while overall performance improvements are demonstrated, they do disentangle the contributions its gating mechanism from the effects of the sparse attention design itself. These considerations underscore the need rigorously disentangle the effects of gating from other architectural components.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

6.05
40,015
Wo
Dense Layer
3.55-
-2.105
Saseline
5.55-
Most Eftecivv!
G)
Concat
5.75-
apary
Scaled Dot Product Attention (SDPA)
60.5
+2.03
W
W
W V
-0.30 2M
+0.62
Daveline
outpul gata
Query Layer
Kay
Value Layer
58 52YA
1.5
(otal b00k
Figure 1: Left: Investigated positions for applying gating operations within self-attention layer;; Middle:
Performance comparison (Test PPL and MMLU) of 15B MoE models with gating applied at various positions. Gating after SDPA (G1) yields the best overall results. Gating after the Value layer (G2) also demonstrates notable improvements, particularly in PPL. Right: Training loss comparison (smoothed, 0.9 coeft.) over 3T tokens between baseline and gated 1.7B dense models under identical hyperparameters. Gating results lower final loss and substantially enhanced training stability, mitigating loss spikes. This stability allows for potentially higher learning rates and facilitates better scaling.
In this work, we investigate gating mechanisms in the standard softmax attention (22 (Sec 2.2 Specifically, we introduce gating distinct positions (Fig. [ after the query (G4), key (Ga), and value projections (G2); following the Scaled Dot Product Attention (SDPA) outputs (G1); and after the final dense output layer (G5). Our exploration covers gating variants including elementwise and headwise, head- specifcc and head-shared, as well as additive and multiplicative forms. We find that: (i) applying SDPA output head-specific gating (G1) 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.
We identify two factors contributing to the efficacy of gating: Non-Linearity. The consecutive linear layers the value (Ww) and dense (Wo) projections can be rewritten into one low-rank linear projection. Therefore, introducing non-linearity through gating at positions G G2 can increase the expressiveness of this low-rank linear transformation (Sec. 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 SDPA outputs (Sec. 4.2). Sparse alse gating eliminates the massive activation 77] and attention sink [ the initial tokens have large activation values in the corresponding hidden states (Tab. 4 and disproportionately dominate attention scores (Fig. 2 Sec. (.33) Previous work B 7 23 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 gain of over 10 points RULER 24 (Sec 4.4).
Practical Recommendation. For best results, apply elementwise SDPA gating G₁ (i.e., gating after the serention-weighted value projection) and train with a moderately increased learning rate.
2
Gated-Attention Layer
2.1 Preliminary: Multi-Head Softmax Attention
Given an input E Raxdaa where the sequence length and dmwstel the model dimension, the computation of transformer's 's attention layer [22 could be divided into four stages.

---

baseline
baseline, 21
baseline.
Iist token attuntion score
fist tolon score: 0.88
Bist tokcn score: 0.41
meen 6666)
Biee tokenn attertion NOOEe
20-
Quary
05
0.4
64
0.2
60
60
10
20
20
30
50
10
50
layer
attent scare atten outout gate
atten cutput gate, layer
atten output cate, layer
1.0
first token scome: 04
First token
Neee tolen ettenton SCOUM
Query
30-
20-
0.
10
20
50
10
50
byye
Kay
Figure 2: Left: Proportion of attention allocated to the initial token per layer (test perplexity dataset). The baseline model suffers from a significant attention sink, with average of 46.7% of attention layers directed towards the first token. Introducing gate effectively alleviates this, reducing the proportion to 4.8%. Right: Average attention map weights for each head. Layer 21 the baseline model demonstrates a strong attention sink (83% on the first token), which is substantially reduced by the gate (4%). In the final oulput layer, the gate amplifics the existing tendency for the model to attend to individual tokens within the sequenee.
across
QKV Linear Projections: The input X is linearly transformed into queries Q, keys K, and values V using learned weight matrices Wo. WK. Wv E deda xdy and Q, K, E Raxd;
XWQ,
K XWR, V==WX
Scaled Product Dot-Product Attention (SDPA): computes attention scores between queries and keys, followed by a softmax normalization. The output is weighted sum of the values:
QKT
Attention(Q.. K, V) softmax
V.
(2)
VdL
where OKT
<dz E 1R represents the scaled dot-product similarity matrix, and softmax(-) ensures the attention weights are no-negative and sum to across each row.
Multi-Head Concatenation: In multi-head attention, the above process repeated for h heads, with each head having its projection matrices W W Wi. All heads' outputs are concatenated:
MultiHead K, V) Concat(head,, ..... head,,
where head, Attention(QW3 KWi, VW().
Final Output Layer: The concatenated SDPA output is passed the output layer W, Rhdd Xdemnd
0 MultiHead( KVV..
(4)
2.2 Augmenting Attention Layer with Gating Mechanisms
The gating mechanism formalized as:
Y' = gYY We. YYor))YYooo)i
(5)
where Y is the input be modulated, X is another input used compute the gating scores refers to the learnable parameters of gate, a is an activation function (e.g., sigmoid), and Y is
W
We adopt the hidden states after pre- pre-normalization as X.

---

the gated output. The gating score, (XW)). effectively acts as a dynamic filter, controlling the information flow from Y by selectively preserving or crasing its features.
This work comprehensively investigates variants of gating mechanisms within the attention layer. Our exploration focuses on five key aspects: (1) Positions. We study the effect of applying gating at different positions, as illustrated in Fig. [eft): (a) after the Q, K, V projections (Equ. D). corre- sponding to positions G2, G3. G4 in Fig. [ (b) following the SDPA (Equ. 3 outputs (G1). (c) after the final concatenated multi-head attention outputs (Equ. 4 G5). (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- graiiedd per-dimension modulation. (3) Head Specific or Shared. Given the multi-head nature of attention, we further consider: (a) Head-Specifie: each attention head has its specific gating scores, enabling independent modulation for each head. (b) Head-Sharee 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' computed as: Y Y o (X0). (b) Additive Gating: Y = Y + (X0). (5) Activation Function. We mainly consider two common activation functions: SiLU (25) and sigmoid. We only use SiLU for additive gating due to its unbounded output range, and sigmoid only gives scores in [0, 1J. Additionally, to further dissect the mechanisms underlying gating's ss effectiveness, we also consider Identity Mapping or RMSNorm [26) (detailed in See 4.1). Unless otherwise specified, employ head-specific, multiplicative gating utilizing the sigmoid activation function (azz Ite 2).
Wo
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- greingriied experts (27], global-batch
LBL [ and z-loss (29] We adopt group query attention (GQA) [0 for the attention part. More detailed model architecture configurations are discussed in Appendix A.2 We train the models on subsets of a 4T high-quality tokens, encompassing multilingual, math, and general knowledge content. sequence length of 4096 is used. 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 relatively small, the wall-time latency introduced by gating is less than 2%. Evaluation We test the few-shots results popular benchmarks, including, Hellaswag [31] for English, MMLU (32 for general knowledge, GSM8k [33] for math reasoning, HumanEval (4] for coding, C-eval 35 and CMMLU (66 for Chinese proficiency. We also test the perplexity (PPL) on diverse held-out test sets, including domains like English, Chinese, Code, Math, Law and Literature.
3.2
Main Results
3.2.1
Gated Attention for MoE models
We first compare different gatings on the training-efffcient MoE-15A2B models. All models use scheduler that warms up to a maximum LR of 2e-3 in 1K steps and decays using cosine to 3e-5. We use global bsz of 1024, comprising 100k optimization steps. The results are summarized in Tab. 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 Tab. I we observe: (i) SDPA and value output gating are effective. Inserting gates at the output of SDPA (G1) or the value map (G2) 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 effective in See 4.2 (ii) Head-Speeific Gating Matters. Applying headwise gating at G and G2 introduces very few additional parameters (less than 2M for the MoF- 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 x dy score from

---

Table 1:
Gating variant performance and results. We train the 15A2B MoE models on 400B tokens, de is the head dim, dmudk is the model's 's hidden dim, and n is the number of tokens. refers to the number of query heads, k refers to the number of key-value heads. 'Act Func' is the activation function in Eq 5
'Score Shape' is the
gating score shape for an input X E IR" RR.ddoou "added param' indicates added parameters.
Method
Act F Fune
Score Shape Added Param
Avg PPL
Hellaswag
MMLU
GSM8K
C-eval
Reference Baselines (Baseline uses 32,
4. All methods use d,
128.)
(1) Bascline
0
6.026
73.07
58.79
52,92
60.26
(2) l =
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 G2
sigmaid
x g x de
201
5.761
74.64
611.82
62.20
(6) v Elementwise Ga
sigmoid
x & x de
25
5.820
74.38
59.17
53,97
61.00
(7) k Elementwise Go
sigmoid
1 x & x de
25
6.016
72.88
59.18
50.49
61.74
(8) Elementwise G1
sigmoid
x y
201
5.981
73.01
58.74
53,97
62.14
(9) Dense Output Gs
sigmoid
100
6.017
73.32
59.41
50.87
59.43
Gating Granularity Variants
(10) SDPA Headwise Gj
sigmoid sigmaid
1.6
5.792
74.50
60:05
54.44
62.61
(11) v Headwise G2
n Head Specific
0.2
5.808
74.38
59.32
53.53
62.61
Head Shared Gating
(12) SDPA Seau-- G1
sigmoid sigmoid
x de a x de
201
5.801
74.34
60.06
53.15
61.01
(13) v Head Shared G2
25
5.867
74.10
59.02
53.03
60.61
Multiplicative v.s. Additive
(14) SDPA Additive G
SiLU x y x da
201
5.821
74.81
60.06
$3.30
60.98
Activation Variants
(15) SDPA Flementwise G2
SILU
a x q x dg
201
5.822
74.22
60.49
54.59
62.34
the original n x q x dg), 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. (iii) Multiplicative Gating is Preferred. Additive SDPA output gating underperforms the multiplicative one, although it shows improvements over the baselines. (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.
Overall, adding gating at the value layer (G2) and SDPA output (G1) reduces PPL by more than 0.2, outperforming various porameter-expanding and benchmark results. As long as different heads receive distinet 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 (Sec 4.2).
baselines. However, gating at G1 achieves better PPL
3.2.2 Gated Attention for Dense Models.
We also conduct experiments on dense models following (37 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 2048. Prior work has established that while increased network depth, large learning rates, and large batch sizes can significantly improve
model performance (38
39
44 and distributed training efficiency, they often introduce training
instabilities [39, 41
42
We observe that applying gating largely reduces the loss spikes 43] 42
during training (Fig. 2 right), suggesting a promising role for gating in enhancing training stability. Therefore, we introduce another setting characterized by an increased number of layers, higher maximum learning rate, and a larger batch size to further probe gating's 's stabilizing effects.
Tab. 2 reveals that: (i) Gating is effective across various settings Across various model configu- rations (row 1 v.s. 2, v.s. 8), training data (row v.s. 4), and hyperparameters (row v.s. 13), SDPA output gating consistently yields benefits. (ii) Gating improves stability and facilitates scaling. Under the 3.5T token setting, gating improves training stability, largely reducing the loss spike (Fig. [ right). When increasing the maximum LR, baselines encounter convergence issues (row 6, 12). While adding sandwich norm 44 restores convergence, the improvement is negligible. In contrast, increasing the maximum LR in models with gating results in noticeable improvement.
5

---

Table Performance of different methods with varying learning rates, batch sizes model configurations.
'SDPA' refers to the sigmoid gating after SDPA in Eq [] and 'sandwitch norm' + indicates normalizing attention/fffn outputs before adding them to the residual. When using gating, we reduce FFN's width so that all
methods have the same number of parameters. means the model diverges during training.
Method
Max LR
Arg PPL HumanEval
MMLU GSM8K
Hellaswag
C-val
CMMLU
Layer, 1.7B Purameters, 400B Tokens, Batch Size=1024
(1) Baseline
4.0 x 10-3
7.499
28.66
50.21
27.82
64.94
49.15
49.52
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 x 10-3
6,180
34.15
59.10
69.07
68.02
68. 19
64.95
(4) SDPA Elementwise
4.5 x 10-3
6.130
37.80
59.61
70.20
68.84
68.52
65.76
Layer, 1.7B Purameters, 400B Tokens, Batch Size=1024
Baseline
4.0%
7.421
28.05
52.04
32.98
65.96
51.11
$1.86
(6) Baseline
8.0 x 10-3
9.195
21.34
44.28
15.24
57.00
43.11
42.63
(7) SSndwwich Norm
8.0 x 10-3
7.407
30.49
52.07
32,90
66.00
52.04
51.72
(8) SDPA Elementwise
10-2
7,288
52.44
32.37
52.06
(9) SDPA Headwise
4.0% 10"
7.370
31.10
53.83
34.12
65.59
55.07
52.38
(10) SDPA Elementwise
8.0 x 10-3
7.325
31.10
54.47
36.62
66.40
53.91
53.80
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 x 10
7.101
34.15
55.70
36.69
67.17
54,51
$4.68
(14) SDPA Elementwise
8.0 10-3
7.078
31.71
56.47
39.73
67.38
55.52
55.77
In summary, we identify SDPA clement-wise gating as the most effective method to augment the attention mechanism. Incorporating the SDPA output gate enables stable training under larger learning rates and batch sizes-regimes where the baseline often becomes unstable. This suggests that the optimal hyperparameter configuration shifts when using gating. In practice, one effective way to leverage the gate is to start from the baseline's 's optimal batch size and moderately increase the learning rate. Further jointly tuning batch size and learning rate may yield additional gains.
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 yield significant improvements in performance and training stability. Here are the takeaways according to our analysis: (1) Gatings enhancing non-linearity consistently lead to performance gains (Sec 4. b (2) The most effective SDPA elementwise gate introduces strong input-dependent sparsity (Sec which then helps to eliminate the 'massive activation' and "attention sink' phenomenon.
4.1 Non-linearity Improves the Expressiveness of Low-Rank Mapping in Attention
Inspired by prior
Table 3: Performance different (non)-lincanie augmentations.
works that utilize
Method
Activation Function
Asg PPL
Hellaswag
MMLJ
GSMSK
C-eval
group norm for the
Baseline
73.07
52.92
SDPA output [14 (5),
(3) Flementwise Gate
(2) SDPA Elementwise Gate
Sigmoid Sigmoid
5,820
5,761
74.64
74.38
59.17
60.82
53.97
55.27
61:00
62.20
with the same setting
SDPA Additive Gate
SILU
5.821
74.81
60.06
53.30
60.88
in Sec. 3.2.1] we apply RMSNorm (26) independently to the
SDPA GreupNorm
RMSNom
5,847
74.30
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
71.17
59.20
52.77
39.86
output of each attention head before concatenation. As shown in Tab. row 5, applying RMSNorm, which introduces almost no additional parameters, also leads to sigmificant reduction in PPL.
In multi- head- attention, the output of the i-th token, corresponding to the k-th head, can be expressed:
of = 1 ==
st x,Wt)wt = E st K,WWWW))
(6)

---

where W is the parameters of the output layer Wo corresponding to the k-th head Here, S is the attention score of the i-th token attending to the j-th token in the k-th head, X, the imput to the attention for token j. and X, W is the value output of token j in the k-th head. From Equ. 6 we can merge A W into one low-rank linear mapping applied over all X, as d2 < dmodet- With GQA, Wy 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 46 we have two modifications to mitigate the low-rank problem:
of
Σ
Non-LLncariiy-MappXX W
( == St X,W W
of = Non-Linearity-May
(8)
Notably, adding gating at the G2 (Tab. row 3) position corresponds to the first modification (Equ. 7 while adding gating (row 4) group normalization (row 5) at the G1 position corresponds to the second (Equ. [8. This also explains why adding gating or normalization at the G5 position after Wo has no effect (Tab. rrow 9) - it does not address the lack of non- non-linearity between Wv and Wo.
(Tab.
For additive gating at G₁, the output of gating passes through SiLU (Tab. row 4), also introducing some non-linearity, which explains the observed performance gains, albeit smaller than those achieved by multiplicative gating. Based on these insights, we conduct two additional experiments: (i) Adding SiLU only at the G1 position without introducing additional parameters (Tab, 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, after gating is directly added at the 11 position (Tab. row 7). This further diminishes the gains of addictive gating.
In summary, the enhanced performance associated with effective gating variants is likely attributable to the introduction of non-linearity between Wv and Wo. Although applying gating at positions G1 and G2 can can both introduce this non- non-linearity, these applications yield differing performance gains. This observed difference motivates us to further analyze the impacts of gating at these two positions.
4.2 Gating Introduces Input-Dependent Sparsity
We analyze the gating scores (Tab. [ 'Gate Score' column) of models with gating applied at the value (G2) and SDPA output (G1) ) positions, evaluated on the test language modeling data. The mean gating scores for all layers are presented in Table g with the score distributions visualized in Fig. 3 (layer-wise scores in Appendix A.3). Key observations include:
(i) Effective Gating Scores are Sparse. SDPA output gatings (Elemeetwise//teadwise) exhibit the lowest mean gating scores. Furthermore, the SDPA output gating score distribution shows high concentration near 0, indicating substantial sparsity, consistent with its superior performance. (ii) Head-Specifie 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-specife gating, aligning with previous research demonstrating that individual attention heads capture distinct aspects of the input [4
50
(iii) Duery-Dependency Matters. The scores for value gating (G2) are higher than those for SDPA output gating (G1), and the performance is inferior. This suggests that gating score sparsity is more effective when query- 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-Man in Eq 8 depends on X2). whereas value gating scores are derived from hidden states associated with past keys and values (e.g. the Lon--iicarity-aaa in Eq 7 depends on each X₃). This implies that gating score sparsity may filter out irrelevant contextual information for the query. To further validate the importance of query-dependencyy we introduce impu-inndeenndntt gating by zero-initializing learnable parameters (g dg). applying 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.
Note that concatenating outputs from different heads and then multiplying with Wo is equivalent to multiplying each head's output with its correspouding W3 before concatenation
7

---

Table 4: Performance of different gating methods with varying activation functions and average gate scores.
'Act-Func' refers to the activation function used for computing the gating scores, while 'M-Act' denotes the rounded maximum activation values of the hidden states output by each layer of the model. Additionally, 'F-Attn'
represents the attention score of the first token, with higher values indicating more pronounced 'attention sink'.
Method
Act-Fune
Gate Score
M-Act
F-Atin
PPL
Hellaswag
MMLU
GSM8L
(1) Baseline
1053
0.467
6.026
73.07
58.79
$2.92
(2) SDPA Elementwise Gate
Sigmoid Sigmoid
0.116
94
0.048
5.761
74.64
60.82
55.27
(3) SDPA Headwise Gate
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
SDPA Inpot Independent Gate
0.335
471
0.364
5.917
73.64
59.02
52.40
(7) SDPA Elementwise Gate
NS -S-iigmoid
0,653
892
0.451
5,900
74.05
60.05
52.75
SDWA Blementwiie Case Sconss Distr Dsttiibution
Elemensaise Gate Scores Detribution
50W Bemeenoeie HS Gate Scores Disatibtioon
0.36
2.36
0:30
Wean 0.116
Mean 0.221
Mean 00223
0.34
Datrbbtiio
0.34
Datrtbution
0:34-
Ditrr bution
2.0.2
2:12
0:12
0.30
5.10
6:30
9.56
0.00
6.00
0.30
0.2
0.4
3.6
0.8
3.0
2.00
62
0.4
0.6
3.6
1.0
0.00
0.0
0.2
0.4
0.6
1.0
Figure 3: Gating score means and distributions for SDPA elementwise (Left), value Elementwise (Middle), and SDPA elementwise with head-shared gating (Right). Most gating scores are less than 0.5, indicating that the gating scores are sparse. Among them, the SDPA output gating score exhibits the strongest sparsity.
(iv) Less Sparse Gating is Worse To further validate the importance of gating score sparsity, reduce sparsity from the gating formulation. Specifically, we replace the sigmoid function with a modified Non-Sparse (NS) version:
(9) which constrains the gating scores between [0.5, 1.0]. This ensures introducing non-linearity while removing gating score sparsity. As shown in Tab. Frow (7), the gains of NS-sigmoid gating are inferior to those of SDPA output sigmoid gating. In Appendix [A.3] we provide 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.
NS-sigmoid (z) 0.5 + 0.5 sigmoid(z).
4.3 SDPA Output Gating Reduces Massive Activation and 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 (8, 7). Correspondingly, early work [51] finds that in explicit top-k sparse attention- where only the most relevant tokens are attended to to-the attention sink phenomenon does not occur. 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 (Fig. [ Tab, 4 'F-Attn' column). Massive activations 77] (large values in hidden states) are belied to lead to the attention sink to their corresponding tokens. Inspired by this, we also compute the mean of the maximum hidden state activations across layers, as shown in the 'M-Act' column of Tab. 4 More detailed wise results are provided in the Appendix A..4
We can observe: (i) Hcad-wise and clement- query- query-dependent sigmoid gating the SDPA output (G1) largely reduces the attention score allocated to the first token and decreases massive activations. (ii) Enforcing shared gating scores across heads applying gating only after the value projection (G2) decreases massive activations, but does not reduce attention scores to the first token. This reinforces the importance of spec-fp gating and suggests that massive activations
not necessary condition for attention sinks. (iii) Reducing the -epedddeeeneence 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- dependee,,,,,,,, read-specific gating of the SDPA output introduces significant sparsity, thereby mitigating the attention sink. Furthermore, sparsity in
8

---

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 152). We also observe that massive activations originate primarily from early layers (e.g., layer 5), where the FFN outputs large values, consistent with [53] 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 [44] in enhancing training stability (Table 2 row 7): applying LayerNorm to the FFN output prevents these large activations from entering the residual stream.
4.4 SDPA Output Gating Facilitates Context Length Extension
Based the attention- sink-free pattern, we evaluate the SDPA gat- ing's 's effect in the long- context setting. Specif- ically, we extend the context length for the models trained on 3.5T tokens. We increase the RoPE [54] base
Table 5: Performance of different methods across varying sequence lengths. "YaRN Extended' indicates the expanded context length variant. "(values)" indicate the performance declines after extending the context length.
Method
4k
8k
16k
32k
64k
128k
Baseline SDPA-Gate
88.89 90.56
85,88 87.11
83.15 84.61
79.50
79.77
YaRN Extended
Baseline SDPA-Gaa 88. 18-14()).4) (-.01)
82, 9.000)))))
71.52(-14.4)
61.23(-21.9
37.94(441.56)
37,51
31,65
76.74(-7.87)
72.88(-6.89)
66.60
58.82
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 context length of 32k. Subsequently, we use YaRN [55] to extend the context length to 128k. We evaluate models on the RULER benchmark [24] and summarize results in Tab. [ We observe the following: (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. (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 extending context length by modifying RoPE [66 55 57]. Even though the decline is less pronounced for models with gating. (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 attention sinks to adjust the distribution of attention scores. [57] 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.
5 Related Works
5.1 Gating in Neural Networks
Gating mechanisms have been widely adopted in neural networks. Early works such as LSTMs N and GRUs (9 introduce gates to regulate information flow across time steps, addressing gradient vanishing/exploding issues by selectively retaining or discarding information. Highway Networks 2 extend this concept to feedforward networks, enabling the successful training of very deep archi- tectures. SwiGL.U 25 introduce gating mechanisms into transformer FFN layers, enhancing their expressive power and becoming a standard component in many open-source LLMs
37].
Several works on state-space models 10 [9] and Linear Attention, such as FLASH 4), RetNet [1). Lightnng Attention [1 60 6T], and Gated Delta Networks [18. also incorporate gating modules to controlinformation of token-mixer modules. AlphaFold2 [5] and Forgetting Transformer [6) introduce
gating mechanisms the output of softmax attention. GaAN [T] uses gating to control each attention
head's 's importance for learning on graphs. Some works 62 13
63
[
16
64] 65 66 also apply
operations similar to gating to augment softmax attention. Attention on Attention (AoA) [12]
also modulates the attention output with a sigmoid gating, depending on the query. 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 s importance beyond RNNs and facilitate designs that better leverage gating's 's unique advantages. For example, while Switch Heads [9] NSA [2]]. and MoSA [67] employ sigmoid-based gating [68] for selection, further investigation into isolating gating's 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. The work most closely related to ours is Quantizable Transformers (69), 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, analysis of various gating variants, uncovering their benefits through enhanced non-line linearity and sparsity, well improved training stability. Building these insights, we scale up gated attention models, demonstrating gating's s broad applicability and impact.
provide detailed
5.2 Attention Sink
StreamingLLM formally identifies "attention sink', in which specific tokens receive large attention scores. Similarly, in VIT, some redundant tokens act 'registers' to store attention scores [70]. Later, Massive Activation [7] shows that excessive attention scores are assigned to tokens associated with massive activation values. However, work reveals that value output (G2) gating eliminates massive activations, yet attention sinks persist, indicating that massive activations are necessary for attention sinks. Similarly, attention sinks are characterized as con-informative "key biases' that store redundant attention scores, arguing that softmax's 's inherent normalization dependency drives this behavior 233 Experimental attempts modify softmax attention, such as replacing softmax with unnormalized sigmoid attention 77 [3], using explicit top-k sparse attention [51]. adding softmax attention gate clip [69], calibrating attention scores [72). and modifying softmax computation (73 and denominator [74). show promise in mitigating attention sinks. Another stream of works try to move the sink tokens from input tokens manually added components, like 'registers' [70], 'meta tokens' (75) and learnable 'sink' (76). Our work shows that sparse gating after SDPA eliminates attention sinks in both dense (1 (IB-parameter) and MoE (5B-paareette) models, even when trained on 3.5T tokens. Furthermore, we uncover the potential of eliminating attention sinks to benefit context-length extension.
6 Conclusion and Limitations
This work systematically investigates gating mechanisms in softmax-attention, revealing their sig- nificant impact on performance, training stability, and attention dynamics. This simple mechanism enhances non-linearity, introduces input-dependent sparsity, and eliminates "attention sink'. Ad- ditionally, gating facilitates context length extension, allowing models to generalize effectively to longer sequences without retraining. We will release the ciiientioonsinnk-fe models, providing foundation for future research into attention mechanisms.
The broader implications of non- non-riiey on the dynamics of attention and the overall training process remain under-explored. We don't provide a theoretical explanation for how attention sinks influence the model's 's ability to generalize to longer sequences.
References
[1] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 98):1735-1188) 1997.
[2] Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber. Highway networks. arXiv preprint arXiv:1505.000 2015.
[3] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprini arXiv:2312.000) 2023.
[4] Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoe V. Le. Transformer quality in linear time.
In International Conference Machine Learning, ICML 2022, 17-23 July 2022, Baltimore,
10

---

Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 9099-9117. PMLR, 2022.
[5] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger. Kathryn Tunyasuvunakool, Russ Bates, Augustin Zídek. Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 59(7787)5)833-55888 2021.
[6] Zhixuan Lin, Evgenii Nikishin, Xu Owen He, and Aaron Courville. Forgetting transformer: Softmax attention with forget gate, arXiv preprint arXiv:2503.021300 2025.
[7] Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu. Massive activations in large language models. arXiv preprint ar arXiv:2402. 17762, 2024.
[8] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint (tti::2309.145533 2023.
[9] Rahul Dey and Fathi M Salem. Gate-variants of gated recurrent unit (gru) neural networks. In 2017 IEEE 60th international midwest symposium on circuits and systems (MWSCAS), pages 1597-1600. IEEE, 2017.
[10] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024.
[11] Jiani Zhang, Xingjian Shi, Junyuan Xie, Hao Ma, Irwin King, and Dit- Yan Yeung. Gaan: Gated attention networks for learning large and spatiotemporal graphs. arXiv preprint arXiv: 2018.
[12] Lun Huang, Wenmin Wang, Jie Chen, and Xiao Yong Wei. Attention on attention for image captioning. In International Conference on Computer Vision, 2019.
[13] Lanqing Xue, Xiaopeng Li, and Nevin L Zhang, Not all attention is needed: Gated attention network for sequence data. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 6550-6557, 2020.
[14] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models, 2023.
[15] Da Xiao, Qingye Meng, Shengping Li, and Xingyuan Yuan. Improving transformers with dynamically composable multi-head attention. arXiv preprint arXiv:2405.0855: 2024.
[16] Xuechen Zhang, Xiangyu Chang, Mingchen Li, Amit Koy-Chowdhury, Jiasi Chen, and Samet Oymak. Selective attention: Enhancing transformer through principled context control. Ad- vances in Neural Information Processing Systems, 37:1 1061-11086, 2024.
[17] Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Lightning attention-2: A free lunch for handling unlimited sequence lengths in large language models. arXiv preprint arXiv:2401 044558, 2024.
[18) Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06440 2024.
[19] Robert Csordas, Kazuki Irie, Jurgen Schmidhuber, Christopher Potts, and Christopher D Manning. Moeut: Mixture-of-experts universal transformers. arXiv preprint arXiv:2405.16039, 2024.
[20] Robert Csordas, Piotr Pickos, Kazuki Irie, and Jurgen Schmidhuber. Switchhead: Accelerating transformers with mixture-of-expee ooeeeeerts attention. Advances in Neural Information Processing Systems, 17:74411-77443 2024.
[21] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprint arXiv:2502.1108) 2025.
11

---

[22] A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
[23] Xiangming Gu, Tianyu Pang. Chao Du, Qian Liu, Fengzhuo Zhang, Cunxiao Du, Ye Wang, and Min Lin. When attention sink emerges in language models: An empirical view. arXiv preprint arXiv:2410.1078) 2024.
[24] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint wrXiv:2404006644 2024.
[25] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020.
[26] Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
[27] Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization 240100666666 2024.
in mixture-of-eeperts language models. arXiv preprint
[28] Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu. Jingren Zhou, and Junyang Lin. Demons in the detail: On implementing load balancing loss for training specialized mixture-of-expee -x-pert models, 2025.
[29] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08900 28002009020 2022.
[30] Joshua Ainslie, James Lee-Thorp. Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprini arXiv:2305. 13245, 2023.
[31] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.077 077000 2019.
[32] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint MXXi:20009003000 2020.
[33] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 14168, 2021.
[34] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan. Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei. Sam McCandlish, llya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained code, 2021.
[35] Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangjun Su, Junteng Liu, Chuancheng Lv, Yikai Zhang, Yao Fu, et al. C-eval: multi-level multi-discipline chinese evaluation suite for foundation models. Advances in Neural Information Processing Systems, 36, 2024.
[36] Haonan Li, Yixuan Zhang, Fajri Koto, Yifei Yang, Hai Zhao, Yeyun Gong, Nan Duan, and Timothy Baldwin. Cmmlu: Measuring massive multitask language understanding in chinese, 2023.
12

---

[37] An Yang, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2. 5 technical report. arXiv preprint arXiv:24/2. 15115, 2024.
[38] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. arXiv preprint /812v:/8 2018.
[39] Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers to 1,000 layers, 2022.
[40] Francesco D' Angelo, Maksym Andriushchenko, Aditya Vardhan Varre, and Nicolas Flammar- ion. Why do need weight decay in modern deep learning? Advances in Neural Information Processing Systems, 222231991222233 2024.
[41] Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang. Yifan Xu, Wendi Zheng, Xiao Xia, et al. Glm-130b: An open bilingual pre-trained model. arXiv preprint arXiv:2210.024144 2022.
[42] Sho Takase, Shun Kiyono, Sosuke Kobayashi, and Jun Suzuki. Spike no more: Stabilizing the pre-training of large language models. arXiv preprint arXiv:2312.16903 2023.
[43] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1- 113, 2023.
[44] Ming Ding, Zhuoyi Yang, Wenyi Hong, Wendi Zheng, Chang Zhou, Da Yin, Junyang Lin, Xu Zou, Zhou Shao, Hongxia Yang, and Jie Tang. Cogview: Mastering text-to-image generation via transformers, 2021.
[45] Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer. arXiv preprint arXiv:2410.05258, 2024.
[46] Guido Montufar, Razvan Pascanu, Kyunghyun Cho, and Yoshua Bengio. On the number of linear regions of deep neural networks, 2014.
[47] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:1905.094115 2019.
[48] Hanrui Wang, Zhekai Zhang, and Song Han. Spatten: Efficient sparse attention architecture with cascade token and head pruning. In 202/ IEEE International Symposium on High-Perffrmance Computer Architecture (HPCA), pages 97-110. IEEE, 2021.
[49] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 1895, 2022.
[50] Zekun Wang, Jingchang Chen, Wangchunshu Zhou, Haichao Zhu, Jiafeng Liang, Liping Shan, Ming Liu, Dongliang Xu, Qing Yang, and Bing Qin. Smarttrim: Adaptive tokens and attention pruning for efficient vision-language models. arXiv preprint arXiv2205..1033. 2023.
[51] Guangxiang Zhao, Junyang Lin, Zhiyuan Zhang, Xuancheng Ren, Qi Su, and Xu Sun. Ex- plicit sparse transformer: Concentrated attention through explicit selection. arXiv preprint 192. 2019.
[52] Stanislav Budzinskiy, Wenyi Fang, Longbin Zeng, and Philipp Petersen. Numerical error analysis of large language models. arXiv preprint arXiv:2503. 10251, 2025.
[53] Itay Yona, Ilia Shumailov, Jamie Hayes, Federico Barbero, and Yossi Gandelsman. Interpreting the repeated token phenomenon in large language models. arXiv preprint arXiv:2503.088: 08908, 2025.
[54] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, Nenrocomputing. 127063, 2024.
13

---

[55] Bowen Peng. Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint urXIv2230.00071..22 2023.
[56] Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation, 2023.
[57] Zican Dong, Junyi Li, Jinhao Jiang, Mingyu Xu, Wayne Xin Zhao, Bingning Wang, and Weipeng Chen. Longred: Mitigating short-text degradation of long-context large language models via restoration distillation. ArXiv, abs/2502.073655, 2025.
[58] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 2024.
3 herd of models. arXiv preprini
[59] Yingfa Chen, Xinrong Zhang, Shengding Hu, Xu Han, Zhiyuan Liu, and Maosong Sun. Stuffed mamba: Oversized states lead to the inability to forget. arXiv preprint arXiv:2410.07145, 2024.
[60] Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Various lengths, constant speed: Efficient language modeling with lightning attention. arXiv preprint arXiv:2405. 2024.
[61] Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. Minimax-01: Scaling foundation models with lightning attention. arXiv preprint arXiv:2501.08313 2025.
[62] Noam Shazeer, Zhenzhong Lan, Youlong Cheng, Nan Ding, and Le Hou. Talking- heads attention. arXiv preprint arXiv:2003.024330 2020.
[63] Qiming Bao, Alex Yuxuan Peng, Tim Hartill, Neset Tan, Zhenyun Deng, Michael Witbrock, and Jiamou Liu. Multi-step deductive reasoning over natural language: An empirical study on out-of-distribution generalisation. arXiv preprint urXiv:2207.14000 2022.
[64] Giorgio Roffo, Carlo Biffi, Pietro Salvagnini, and Andrea Cherubini. Hard-attention gates with gradient routing for endoscopic image computing. arXiv preprint arXiv:2407.044000 2024.
[65] Olga Golovneva, Tianlu Wang, Jason Weston, and Sainbayar Sukhbaatar. Multi-token attention. arXiv preprint arXiv:2504.00927, 2025.
[66] Zelin Peng. Yu Huang, Zhengqin Xu, Feilong Tang, Ming Hu, Xiaokang Yang, and Wei Shen. Star with bilinear mapping. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 25292-25302, June 2025.
[67] Piotr Piękos, Róbert Csordás, and Jürgen Schmidhuber. Mixture of sparse attention: Content- based learnable sparse attention via expert-choice routing, 2025.
[68] Robert Csordas, Kazuki Irie, and Jurgen Schmidhuber. Approximating two-layer feedforward networks for efficient transformers. arXiv preprint arXiv:2310. 2023.
[69] Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Quantizable transformers: Removing outliers by helping attention heads do nothing. Advances in Neural Information Processing $ystems,,6:750077 36:75067-755966 2023.
[70] Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
[71] Jason Ramapuram, Federico Danieli, Eeshan Dhekane, Floris Weers, Dan Busbridge, Pierre Ablin, Tatiana Likhomanenko, Jagrit Digani, Zijin Gu, Amitis Shidani, et al. Theory, analysis, and best practices for sigmoid self-attention. arXiv preprint arXiv:2409.0444 2024.
[72] Zhongzhi Yu, Zheng Wang, Yonggan Fu, Huihong Shi, Khalid Shaikh, and Yingyan Celine Lin. Unveiling and harnessing hidden attention sinks: Enhaneing large language models without 2024.
training through attention calibration. arXiv preprint
14

---

[73]
Zayd M. K. Zuhri, Erland Hilman Fuadi, and Alham Fikri Aji. Softpick: No attention sink, no massive activations with rectified softmax, 2025.
[74] Evan Miller. Attention off by one, 2023.
[75] Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabalosh- warkar, Shih- Shih-Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, al. Hymba: hybrid-head architecture for small language models. arXiv preprint 1.X.x:24111113771 2024.
[76] OpenAl. gpt-oss-1 100 10b & gpt-oss- model card, 2025.
15

---

A Supplement Material
A.1 Switch Head Baselines
In this section, we present detailed experiments related to Switch Heads. The Switch Head paper demonstrates that introducing sparse activation in attention-where each token selects the top-k experts from a pool key/value/output experts via learnable sigmoid routing-enable the model to achieve comparable results to the baseline. This suggests that, within the Switch Head framework, both expert parameters and activated parameters are beneficial, with more being better under the same total parameter budget.
Table Performance of different switch head methods with varying parameter additions configurations. 'switch kv' and 'switch v' refer to introducing selective computing in key-value and value components, respec- tively. 'Switch kv, Stop8' means there are 8 key and value map experts, and each token select top8 experts. Notice 'Switch Itopl' is equivalent to Hcadwise Gate Tab.
row (11).
Method
Added Param (M)
PPL
MMLU
GSM8K
Hellaswag
C-eval
(1) Baseline (q32, kv4)
6.026
58.79
52.92
73.07
60.26
(2) Switch kv, Stop8
38
5.847
59.17
52.54
73.32
61.01
Switch 4top4
5,935
58.14
53.27
73.75
59.67
(4) Switch v, 4top4
13
5,820
59.02
52.77
73.34
61.74
(5) Switch v, Stop2
5,870
59.10
53.53
74.17
62.34
Switch v. Itop1
5.808
59.32
53.53
74.38
62.61
Looking the results in Tab. 6 observe an interesting trend: while increasing the number of activated kv experts (with the same expert parameter settings) appears to offer some improvement in PPL (row 4 vs. 5). the gains in overall benchmark performance are less pronounced. Notably, the best results for both benchmark scores and PPL were achieved by 'Switch v Itop1' (row 6), which, as mentioned earlier, is analogous applying sigmoid gating directly to the output of the value layer. These findings raise an intriguing question about the primary driver of the performance improvements observed in these experiments. It suggests that the introduction of gating itself plays a significant role in the effectiveness of this approach.
A.2 More Discussion on Sparse Gating Score
In our experiments, we evaluate three distinct LLM architectures: two dense variants with 1.7 billion parameters and one sparse MoE model with an effective size of approximately 15 billion parameters (denoted as 15A2B). The two dense models differ primarily in depth and hidden dimension: the 1.7B-28 layer model uses 28 transformer layers with hidden size of 2048, while the 1.7B-48 layer variant employs 48 layers but a reduced hidden size of 1536 to maintain comparable parameter count. Both dense models tie input and output embeddings and apply query-key normalization for training stability. Full architectural hyperparameters are summarized Table 7 To support reproducibility and future research, we plan to open-source both 1.7B dense models.
A.3 More Discussion on Sparse Gating Score
In this section, we analyze the impact of gating score sparsity on attention output. First, we examine the mean values of SDPA output before and after applying gating to the hidden states. Specifically, we calculated the mean absolute values of Y and y" before and after G at each layer, as shown in Fig. 4 We also included results from baseline without gating for comparison. The results indicate that: (1) after gating, the mean value of hidden states decreased from 0.71 to 0.05, corresponding to the generally small gating scores; (2) the gated hidden states closely resemble the bascline, suggesting that gating might serve a similar function as attention sink in filtering out irrelevant information.
We further analyze the proportion of hidden states below certain thresholds before and after gating, as shown in Fig [5] The results reveal that: (1) after gating, the sparsity in hidden states significantly increases across different thresholds. Since the mean gating scores are already small, multiplying hidden states by a small number naturally pushes some values below the threshold. Therefore, (2) we further multiply the pre-gating hidden states by the average gating score and observed that the increase sparsity is smaller than with original gating. This suggests that sparse gating scores enhance sparsity in hidden states.
16

---

Table Architectural specifications of the target LLMs used in our experiments. The 15A2B model is a Mixture-oun centuree-ofEaaee of-Experts (MoE) architecture with 128 experts and top-8 routing. All models use head dimension of 128 and apply query-key normalization. Embedding weights are tied the dense models but not in the MoE model.
Model
1.7B-28 layer
1.7B-48 layer
15A2B MoE
Layers
28
48
24
Query Heads
32
Key / Value Heads
8
8
4
Head Dimension
128
128
128
Tie Embedding
Yes
Yes
No
QK Normalization
Yes
Yes
Yes
Hidden Size
2048
1536
2048
FFN Size
6144
4608
768
Number of Experts
128
Top-k
Hidden State Mean
hidden gate, avg
hidden state mean after gate, avg 0.05
hidden avg
Abs Mean
10
20
Layer
Figure 4: Mean absolute values before and after gating. The baseline and post-gating values are similar.
A.4 Layerwise Massive Activations and Attention Sinks
In this section, we compare and analyze the presence of massive activations and attention sinks (the attention score the first token) within the model. From the results, we observe the following:
For the baseline (row 1), the output of the 6th layer's FFN contains massive activations, which are subsequently added to the residual stream, causing large activations persist in the residuals of subsequent layers. Correspondingly. significant attention sink phenomena emerge starting from the 6th layer. After applying gating to the SDPA output (row 2), the outputs of the earlier layers in the network remain relatively small overall, with massive activations growing gradually as the layer depth increases. Notably, no significant attention sink phenomenon is observed in any layer of the network.
When gating is applied only at the value layer (row 3), the model exhibits massive activations similar to row 2. However, certain degree of attention sink phenomenon persists. This indicates that massive activations are not a necessary condition for the emergence of attention sinks. When enforcing shared gating scores across different heads (row 4) modifying the activation function of gating suppress sparsity (row 5), the sparsity introduced by gating is reduced. In these cases, both massive activations and attention sinks become comparable to those observed in the baseline.
These observations suggest that introducing sufficient sparsity within attention mechanism may help mitigate the occurrence of massive activations. However, further investigation is needed to fully understand the interplay between sparsity, massive activations, and attention sinks, particularly the context of scaling to deeper and larger models.
17

---

Sparsity Ratlo SDPS Output. Threshold 1e-2
Sparsity Ratio SDMM Output, Threshold 1e-3
1.0
spertby before gate. avg
sparsity ratio gate, avg 0.003
spentty netio atter gate. avg 0.44
sparsity ratio after gate, avg 6.126
sporaty ado after average gate, avg 0.33
sparsity ratio after average gate, avg 0.080
Sparsity Ratio
Sparsity Ratin
0.6-
0.4
0.2
0.2
15
20
Layer
Layer
Figure 5: Proportion of SDPA output values below threshold after gating (Left: 1c-2, Right: 1c-3). We also include sparsity measurements obtained by multiplying the average gating score with pre-gating hidden states.
A.5 More Layerwise Gating Scores
In this section, we analyze the distribution of gating scores under two additional constraints while using SDPA output gating as the baseline (row 1, elementwisehheadwis)) (1) enforcing the same gating score across different heads (row 2, left), and (2) restricting the minimum value of the gating scores (row 2, right). When enforcing shared gating scores across different heads, the gating scores for most layers increase. This indicates that different heads require different sparsity, highlighting the importance of head-speche -pecifiec gating mechanisms.
A.6 Other Attempt to Stabilize Training
We observe that both the addition of sandwich normalization massive activations while improving training stability. This prompts us to explore whether simpler methods could prevent large activations within residuals. Specifically, we introduce a clipping operation to constrain the outputs of attention and FFN layers before they enter the residual connection, limiting their values to the range (-clip, clip). However, we find that regardless of whether the clip value set 300 or 100, the model still encounters convergence issues at learning rate of 8c-3. This suggests that the instability in pre-norm model training is not solely due to large activations within residuals. It is likely that any layer producing large outputs can lead to stability problems, indicating the need for further investigation into the root causes of training instability.
46 and gating mechanisms eliminate
A.7 Adding Gating in Continue Training
We also experiment with incorporating an attention output gating mechanism during the continued training phase. However, find that this approach neither mitigates the massive activations and attention sinks already present in the model nor significantly affects the final performance. We believe this is largely because the effectiveness of gating mechanisms stems primarily from their affects the model's training dynamics, and thus their impact is limited when applied to models that are not trained from scratch.
A.8 Broader Impacts
This work focuses on improving the efficiency, stability, and -handdig capabilities of LLMs. The potential positive socictal impacts include: Improved Accessibility: More efficient and stable training methods can lower the computational cost of developing and deploying LLMs, potentially making them more accessible to researchers and organizations with limited resources. Enhanced Performance in Long-Core Context Applications: The ability to handle longer contexts more effectively can lead to improvements in applications such as document summarization, question answering, and code generation, where understanding relationships across extended text is crucial. Advancements in Model Scalability: Our findings offer insights and practical guidance for designing advanced models, contributing to ongoing efforts to scale L.MMs while maintaining or improving their performance and stability.
18

---

Mardive Activation baseline
rest toker attention scores of base Ine
1000
10
1400
abn sutpu
0.6
800
alonn residha
suttur
9.4-
Seat.
25
0.0
15
20
Layer
layer
Massive Activation
400
First toker witention scoew attier sulput gate
10-
resiusl
100
250
200
0.-
150
0.4-
100
Rest
10
15
20
0.0
20
Layer
layar
Massive Activ Altin allmmnnimiue
first taken allantion of slammentiiie
oftn culgs
Tr. Estpp:
Bist iohtn wtters acore
8thh dual
Rh, ree daa
meen: 5,247
15st tolen attenston score
3.a
Nox Actovation Vallu
253-
450
Activation output, gate share
first tolon attontion score aitton putp goze hoad
abornion
mean 6.360
Max Actraation Value
400
thrt talon attore on scor
1302
atty, surpat altcrradie
my, suppr tr
100
3.0
Layer
Tayer
Massive Activation of atatioon output gate wo sparsity
1 sttention ottee__orpptttooeeo wwtpptttttttttooooooooo www.tttwowsoorrtttt
3402
4or save smiecian cowe
mean: 481
Pla Adtivatinn Salle
3230
032-
thss Irrres
dudl
002-
235
Figure Comparison massive activations and attention sink phenomena across different gating configurations. Row (Baseline): Significant massive activations and attention sinks emerge after the 6th layer. Row 2 (SDPA Gating): Reduced activations and no attention sinks ohserved. Row 3 (Value Layer Gating): Similar activations to Row 2 but with residual attention sinks. Rows 4-5 (Reduced Sparsity via cross-head share and NS-sigmoid): Massive activations and attention sinks resemble the baseline.
A.9 Licenses for Existing Assets
We conduct experiments based on Megatron-LMM, We acknowledge the original authors and contribu- tors of Megatron-LM, LLM and respect their licensing terms. The specific license terms are available at
19

---

atten @iten_eutput.oom gale gate values
attention, wwtppptppttt..ooo gale, scaler values
1.0
std
0.8
min
03te value
0.6
min
0.4-
that
0.2
0.2
0.0
layer
layer
output, head share gate values
atten, output, sparsity values
1.0
meen
min
gate value
0.6
0.4
0.44
max
0.2
layer
Figure 7: Distribution of gating scores under different constraints for SDPA output gating variants.
hutps://githu..com/N www..///DD..AAAaaaeooooooooaaann our proposed gated attention mechanism will be released under the MIT License, allowing for open research and further development.
Any released code implementing
20

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
We elucidate our contributions in the Abstract and Intrudctions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations in Sec
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-texx system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess..
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
21

---

Justification:
Our study is mainly empirical exploration.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We report our detailed training settings in Sec 3.2. 1 and Sec 3.2.2
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
22

---

Answer: [Yes]
We provide a Pytorch version implementation for the main method in the paper.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nipp) cc/
publ /c www.gddid/eeeeeeeeeeeeeee
ss Loonolics cy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (ht
/ ctc.ps.ccpuubbic
CooeeoodeSSSaa Submisbbbsss ss cessoooPllcyy for more details. The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We report our detailed training settings including hyperparameters in Sec 3.2. and Sec 3.2.2
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: LLMs pre-training consumes a significant amount of computational resources, making impractical conduct multiple experiments to obtain error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
23

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative emrr rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We discussion computation cost in Sec 3.1
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http://....ss............ooooo/a
Answer: [Yes]
Justification:
We have read this code.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:
[Yes]
Justification:
We discuss broader impacts in Appendix A.8
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [No]
Justification: Our study is an empirical exploration and the codes and models we release are intended solely for further research and are not meant for direct industrial application.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We discuss this in Appendix A.9
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, paperssww.. www.osocittcceeeeeoo com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
25

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Related codes will be released with suitable license.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/cood/mmodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
26

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration is not required.
Answer: [No]
Justification: LLM is used only for writing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (ht /ttp:::/neurppp httperrrrpppseeeoooooeenneeee//e 2025/ LLM) for what should or should not be described.
27


Reviews:
[{"review_id": "FWaHfxBjin", "reviewer": "Reviewer_cPo5", "summary": "The paper studies the effect of introducing gating at different stages of multi-head attention. The strongest point of the paper is the comprehensive study and comparison of various options. The authors find that applying gating after the value matrix or after head concatenation provides most of the benefits, with gains of up to 2% on MMLU. Additionally, the paper demonstrates that such gating reduces the effect of attention sink and massive activations, leading to easier long-context finetuning.", "questions": "It will be great if authors can comment on weaknesses above.", "limitations": "Mentioned", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n1) Proposed analysis is novel and makes a lot of sense.\n2) The depth of ablations is impressive: position, granularity, head-specific or shared, multiplicative or additive, activation functions — all are studied.\n3) MoEs and dense models are considered, training scale is reasonable.\n4) Long context, sparsity, and attention sink are studied.\n5) Each experiment is followed by a heat summary; it is quite enjoyable to read and learn as you go.\n\n**Weaknesses**\n1) No major weaknesses, only a couple of suggestions.\n2) One suggestion is to summarize a single learning and a single recommendation on the best way to apply gating. Currently, multiple options are highlighted, but narrowing down to one can be beneficial for readers.\n3) The attention sink can be suppressed by using \"Quiet Attention\" or Meta tokens from [R1]. It would be great to see if such techniques are complementary.\n4) The reduction of massive activations should be beneficial for quantization. Adding post-training quantization results would be beneficial.\n\n[R1] Dong, Xin, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih-Yang Liu et al. \"Hymba: A hybrid-head architecture for small language models.\" arXiv preprint arXiv:2411.13676 (2024).", "was_revised": false, "final_justification": "", "rebuttal": "> One suggestion is to summarize a single learning and a single recommendation on the best way to apply gating. Currently, multiple options are highlighted, but narrowing down to one can be beneficial for readers.\n> \n- Thank you for the suggestion! Our top recommendation is to apply **elementwise gating at the SDPA output** (i.e., after the attention-weighted value projection) in combination with a moderately increased learning rate.\n- Furthermore, our follow-up experiments reveal that increasing the head dimension (e.g., changing from q=32, kv=4, head_dim=128 to q=16, kv=2, head_dim=256) while keeping the total model size constant can further improve performance. However, such a configuration leads to more frequent training instabilities (e.g., loss spikes) in the baseline. In contrast, with **elementwise SDPA gating**, the model remains stable during training.\n- Therefore, more generally, we recommend using **elementwise SDPA output gating** as a foundational improvement, upon which more aggressive training strategies (e.g., larger learning rates, larger head dimensions) can be safely combined.\n\n---\n\n> The attention sink can be suppressed by using \"Quiet Attention\" or Meta tokens from [R1]. It would be great to see if such techniques are complementary.\n> \n- Thank you for the insightful suggestion. We will include a discussion on **Meta tokens** in the appendix. In our view, Meta tokens do not truly *resolve* the attention sink issue but rather *redirect* it to a designated special token. Once the model learns to suppress the sink phenomenon via our gating mechanism, the role of the Meta token may be reduced to that of a soft prompt, providing only marginal additional gains.\n\n---\n\n> The reduction of massive activations should be beneficial for quantization. Adding post-training quantization results would be beneficial.\n> \n- Thank you for the comment. While we currently do not include quantization results in the paper, we plan to release **quantized versions** of our models alongside the full-precision checkpoints to facilitate community research in this direction."}, {"review_id": "LVRM4i6YbZ", "reviewer": "Reviewer_bMKL", "summary": "The paper explores the effects of gating at various positions of self-attention, i.e., after query/key/value projections, after the self-attention output, and after the final dense layer.  Two different model types (MoE, Dense) with architectural/training recipe variations are trained with the aforementioned gating placements.  The paper reports qualitative results for both MoE and Dense models, showing SDPA output gating effectively improves performance across standard natural tasks.  The source of this improvement is thus explored, showing that gating location can greatly affect the sparsity of attention scores and, subsequently, this can mitigate massive activations and attention sinks during training.", "questions": "-Why were the same placement experiments in Table 1 not repeated for the dense model in Table 2?  The missing \"Max LR\" experiments for different configurations (e.g., (9) SDPA Headwise trained with max LR 8.0 × 10−3) make it difficult to draw definitive conclusions.\n\n## Incorrect citatios\n> Prior work has established that while increased network depth, large learning rates, and large batch sizes can significantly improve\nmodel performance [33, 34, 35] and distributed training efficiency, they often introduce training instabilities [34, 36, 37].\n\nCitations 34 and 36 do not support the above claim.  34 attributed instability to pre vs post residual placement of the norm, while 36 did not attribute instability to any particular factor (but rather explored layernorm placement's effect on stability).", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "# Strengths\n## Clarity, Significance, and Originality\nThe paper is well written and easy to understand.  The topic is interesting; a nuanced, controlled study of the role gating plays in Transformer models can have a large impact.  In terms of originality, the specific study has the potential to separate itself from previous works in the area.  \n\n## Quality\nThe experiments make sense and are in line with the motivations of the paper.  The explainability results in Section 4 were also good contributions to explain the effects of gating-placement has on attention scores.  Although it is not surprising that SDPA Elementwise Gate induces the most sparsity, it is nice to see this verified.\n\n# Weaknesses\n## Quality\nAlthough the paper seeks to assess gating-location's contribution without confounding factors, key takeaways are made without sufficient evidence.  In particular,\n> 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.\n\nIt is difficult to make this claim definitively, since the batch size is varied with the total training tokens.  Thus, a carefully isolated experiment would compare: a) 28 Layer, 1.7B Parameters, 400B Tokens, Batch Size=1024, (b) 28 Layer, 1.7B Parameters, 3.5T Tokens, Batch Size=1024, (c) 28 Layer, 1.7B Parameters, 400B Tokens, Batch Size=2048, (d) 28 Layer, 1.7B Parameters, 3.5T Tokens, Batch Size=2048.  Instead, the only comparison we have is between (a) and (d).  This is the same for the 48 layers experiments.\n\n## Significance\nSome of the performance improvements are not as substantial as reported.  For example, overall, results in Table 1 are not significant performance improvements.  It is difficult to say \"significant reduction in PPL\" in Table 3.  Similarly, it is difficult to call a 0.2 PPL reduction a significant performance improvement.  For dense models, while gating-placement was much more impactful on 48 layer 1T pretraining token models, gains are modest for the 28 layer model.  It is possible that substantial improvements may have occurred for the missing experiments above.\n\n## Originality\n> Despite its widespread adoption and empirical success, the function and impact of gating mechanisms remain insufficiently explored beyond their initial intuition. \n\nThis claim is too broad, the impact of gating has been explored in linear attention [1] and standard attention [2,3] networks.\n\n# References\n[1] Chen, Yingfa, et al. \"Stuffed Mamba: Oversized States Lead to the Inability to Forget.\" arXiv preprint arXiv:2410.07145 (2024).\n\n[2] Lin, Zhixuan, et al. \"Forgetting transformer: Softmax attention with a forget gate.\" arXiv preprint arXiv:2503.02130 (2025).\n\n[3] Xue, Lanqing, Xiaopeng Li, and Nevin L. Zhang. \"Not all attention is needed: Gated attention network for sequence data.\" Proceedings of the AAAI conference on artificial intelligence. Vol. 34. No. 04. 2020.", "was_revised": true, "final_justification": "# Post rebuttal update\n\nThe authors have added additional experiments and discussion during the rebuttal to address my primary concerns regarding disconnects in the original experiments and the claims from the paper.  Indeed, the main takeaway regarding stability under gating looks sound and I've increased my score accordingly.", "rebuttal": "## Weaknesses\n### Quality\n> 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.` It is difficult to make this claim definitively...\n  - Thank you for this suggestion. We increased the batch size when increasing the number of training tokens because we aimed to compare under the best possible baseline performance. Our hyperparameter search for the baseline revealed that both batch size and learning rate need to be increased as the number of training tokens grows, which is consistent with existing scaling laws (`Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining`). We understand that this setup may introduce additional confounding factors (i.e., token count) when comparing the effect of gating under different batch sizes. Therefore, we have further added experiments under setting (c):\n    \n    \n    | Method | LR | Tokens | Bsz | Avg PPL | MMLU | HumanEval | GSM8k | Hellaswag | C-eval | CMMLU |\n    | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n    | Baseline | 4e-3 | 400B | 1024 | 7.499 | 50.21 | 28.66 | 27.82 | 64.94 | 49.15 | 49.52 |\n    | SDPA Elementwise | 4e-3 | 400B | 1024 | **7.404** | **51.15** | **29.27** | **28.28** | **65.48** | **50.72** | **50.72** |\n    |  |  |  |  |  |  |  |  |  |  |  |\n    | Baseline | 4e-3 | 400B | 2048 | 7.538 | 49.02 | 24.37 | 25.72 | 63.54 | 48.02 | 48.82 |\n    | SDPA Elementwise | 4e-3 | 400B | 2048 | 7.487 | 50.41 | 27.56 | 27.13 | 64.11 | 49.64 | 50.03 |\n  - We observe that increasing the batch size larger than the optimal one leads to a performance drop. However, adding the SDPA output gate still results in better performance than the baseline, with more pronounced improvements on benchmarks such as MMLU, HumanEval, and GSM8k.\n  - Nevertheless, we acknowledge that the phrase `resulting in improved performance` could be misleading and that the comparison remains challenging. In the next version, we will revise this statement to `also resulting in performance improvement relative to the baseline.`\n\n---\n\n### Significance\n\n> For example, overall, results in Table 1 are not significant performance improvements. It is difficult to say \"significant reduction in PPL\" in Table 3. Similarly, it is difficult to call a 0.2 PPL reduction a significant performance improvement...\n\n- We would like to provide additional context to demonstrate that a 0.2 PPL reduction is a meaningful performance gain:\n    - In our Table 2, for the 48-layer model, increasing training tokens from 400B to 1T leads to only a 0.06 reduction in PPL for the baseline (from 7.4 to 7.34), while the SDPA Elementwise gate achieves a 0.187 reduction (from 7.28 to 7.093). In contrast, in Table 1, the gate reduces PPL from ~6.0 to ~5.8—a 0.2 drop—even though architectures differ, this remains a substantial improvement.\n    - Additionally, we supplement Table 1’s setting with baseline models trained for 500B and 600B tokens (with full cosine schedulers), as shown below:\n        \n        \n        | Method | Tokens | Avg PPL | MMLU | GSM8k |\n        | --- | --- | --- | --- | --- |\n        | Baseline | 400B | 6.026 | 58.79 | 52.92 |\n        | Baseline | 500B | 5.851 | 60.17 | 54.85 |\n        | Baseline | 600B | **5.704** | **61.03** | 55.15 |\n        |  |  |  |  |  |\n        | SDPA Elementwise | 400B | 5.761 | 60.82 | **55.27** |\n    - Notably, the SDPA Elementwise model trained on 400B tokens outperforms the baseline trained on 500B tokens and is only slightly behind the 600B baseline.\n    - For further reference, in the Fox paper (mentioned by the reviewer), Table 3 shows that adding an output gate in a 360M model trained on 7.5B tokens reduces PPL from 7.08 to 6.88—a 0.2 drop. While test PPLs are not directly comparable due to different datasets, this supports our claim that a 0.2 PPL reduction is meaningful.\n    - Moreover, we will remove the word `significant` to avoid mis-understandings.\n\n---\n\n> For dense models, while gating-placement was much more impactful on 48-layer 1T pretraining token models, gains are modest for the 28-layer model. It is possible that substantial improvements may have occurred for the missing experiments above.\n> \n- Thank you for raising this point. We have further supplemented experiments on the 28-layer 1.7B model trained on 400B tokens with increased learning rates. The results are shown in the table below:\n    \n    | Method | LR | Avg PPL | HumanEval | MMLU | GSM8k | Hellaswag | C-eval | CMMLU |\n    | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n    | Baseline | 4e-3 | 7.499 | 28.66 | 50.21 | 27.82 | 64.94 | 49.15 | 49.52 |\n    | SDPA Elementwise | 4e-3 | **7.404** | 29.27 | 51.15 | 28.28 | **65.48** | 50.72 | 50.72 |\n    | Baseline | 8e-3 | 8.941 | 20.13 | 41.17 | 16.15 | 58.10 | 44.36 | 45.13 |\n    | SDPA Elementwise | 8e-3 | 7.425 | **30.37** | **52.51** | **29.24** | 65.30 | **51.76** | **51.88** |\n    \n    As shown, when the learning rate is further increased, the baseline model encounters convergence issues and suffers significant performance degradation. In contrast, the model with SDPA Elementwise gating maintains stable training and achieves further improvements across all benchmarks except PPL and Hellaswag. This supports our earlier claim that the gating mechanism enables stable and effective training under larger learning rates.\n---\n### Originality\n\n> \"Despite its widespread adoption and empirical success, the function and impact of gating mechanisms remain insufficiently explored beyond their initial intuition.\"\n> This claim is too broad, the impact of gating has been explored in linear attention [1] and standard attention [2,3] networks.\n> \n- Thank you for the feedback. We have discussed the relation to [2] in line 300. In the revised version, we will add the following clarifications regarding [1] and [3]:\n    - [1] discusses the difficulty of Mamba’s decay mechanism in forgetting historical hidden states and also uses an output gate at the SSM layer. However, neither [1] nor the original Mamba2 paper analyzes the behavior or impact of the output gate in detail.\n    - [3] proposes a Gumbel-Softmax-based gate in LSTM attention to mask irrelevant tokens for the current query. In contrast to [3], which applies a hard mask *before* SDPA to explicitly remove irrelevant token information, our attention output gate operates *after* SDPA and uses gating scores to suppress irrelevant information, effectively mitigating the attention sink phenomenon.\n    - We will modify the expression to `Despite its widespread adoption and empirical success, most recent works do not look into the gating mechanisms like the gating scores and their effect on the model’s hidden states.`\n\n---\n## Questions:\n\n> Why were the same placement experiments in Table 1 not repeated for the dense model in Table 2? The missing \"Max LR\" experiments for different configurations (e.g., (9) SDPA Headwise trained with max LR 8.0 × 10⁻³) make it difficult to draw definitive conclusions.\n> \n- Initially, we focused on evaluating gating positions in MoE models due to their lower training cost and strong performance, expecting similar trends in dense models. We now supplement the results for different gating placements in a 28-layer, 1.7B dense model trained on 400B tokens. These will be added to the appendix:\n    \n    \n    | Method | Avg PPL | MMLU | GSM8k | Hellaswag | C-eval | CMMLU |\n    | --- | --- | --- | --- | --- | --- | --- |\n    | Baseline | 7.499 | 50.21 | 27.82 | 64.94 | 49.15 | 49.52 |\n    | SDPA Elementwise G1 | **7.404** | **51.15** | **28.28** | **65.48** | **50.72** | **50.72** |\n    | v Elementwise G2 | 7.453 | 50.53 | 27.97 | 65.12 | 49.82 | 49.91 |\n    | k Elementwise G3 | 7.485 | 50.25 | 28.01 | 64.99 | 49.03 | 49.17 |\n    | q Elementwise G4 | 7.466 | 50.33 | 27.89 | 65.04 | 49.15 | 49.33 |\n    | Dense Output G5 | 7.489 | 50.34 | 28.08 | 65.01 | 49.38 | 49.25 |\n- We further supplement results for `SDPA Headwise` at LR = 8e-3:\n    \n    \n    | Method | Avg PPL | MMLU | GSM8k | Hellaswag | C-eval | CMMLU |\n    | --- | --- | --- | --- | --- | --- | --- |\n    | SDPA Elementwise | **7.325** | **54.47** | **36.62** | **66.40** | 53.91 | **53.80** |\n    | SDPA Headwise | 7.384 | 54.13 | 34.56 | 65.97 | **54.10** | 52.38 |\n- We observe that under more aggressive training setting, the elementwise gate outperforms the headwise variant, and the training curve is smoother. Hence, we recommend the elementwise gating design.\n\n---\n### Incorrect citations\n\n> \"Prior work has established that while increased network depth, large learning rates, and large batch sizes can significantly improve model performance [33, 34, 35] and distributed training efficiency, they often introduce training instabilities [34, 36, 37].\"\n> \n> \n> Citations 34 and 36 do not support the above claim. 34 attributed instability to pre vs post residual placement of the norm, while 36 did not attribute instability to any particular factor (but rather explored layernorm placement's effect on stability).\n> \n- Thank you for pointing this out. The confusion likely arose from combining multiple factors across citations.\n    - For [34], we intended to use its Table 1 to support the claim that increased depth can cause instability (many models diverge when depth increases), and its Table 2 to support that depth can improve performance (both baseline and DeepNet show gains with more layers).\n    - For [36], a more appropriate citation would be `A Theory on Adam Instability in Large-Scale Machine Learning`, which thoroughly discusses how larger learning rates and batch sizes can lead to loss spikes and other instabilities. [36] primarily proposes techniques like Embedding Layer Gradient Shrink (EGS) to address such spikes .\n    - We will modify the expression accordingly.\n\nThanks again for your suggestions!"}, {"review_id": "URvPmnhgv5", "reviewer": "Reviewer_Wjax", "summary": "This paper conducts an extensive empirical investigation into incorporating gating mechanisms into the softmax attention module and provides a detailed analysis of the resulting gains and learned patterns. The authors explore over 30 gating variants, including different gating positions (post-q, k, v, Wo. output), granularity (token-wise, head-wise, or head-shared), gating types (additive vs. multiplicative), and activation functions. The study spans both dense models (e.g., 1.7B models trained on 400B or 3.5T tokens) and mixture-of-experts models (e.g., 15A2B trained on 400B tokens), all under a well-optimized training pipeline—covering training data quality, architectural tuning, global batch size, label smoothing, z-loss, and more.\n\nBased on these comprehensive experiments, the paper delivers credible takeaway messages: adding a gating mechanism before the weighted output (Wo) projection in multi-head attention can significantly improve perplexity and performance on a range of downstream benchmarks, including MMLU, GSM8K, and C-Eval.\n\nThe paper further attributes these improvements to two key effects: (1) introducing non-linearity in the value-dense low-rank mappings, and (2) enabling sparse, input-dependent modulation that mitigates excessive activations and reduces attention sinks. Detailed ablations support these claims, and the method is shown to improve training stability and generalization to long-context settings up to 128k tokens.", "questions": "Table 1 could be further improved by including two additional experiments to strengthen the paper’s findings and conclusions:\n\n1). One of the sub-questions explored in the paper is how the placement of additional parameters affects performance. Currently, the baselines include only GQA settings and expert counts. It would be valuable to add a “more-layer” baseline, especially for the 2.54B activation model. Adding ~200M parameters via extra layers is feasible and would help compare against the current approaches under a similar parameter budget.\n\n2) We also noticed that setting (6) achieves fairly strong performance, despite having only 1/8 the additional parameters of setting (5). This raises the question of whether the performance difference is primarily due to parameter count. Could you consider adding an experiment comparing v-elementwise G2 with multi-head (i.e., n × q × d_k) to control for parameter count and isolate the architectural impact?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The topic addressed in this paper is highly practical and valuable, with strong applicability to structural improvements in large language models (LLMs).\n2. The experimental scale and setup are at a production level, lending high credibility and reference value to the conclusions.\n3. The takeaway messages are well-reasoned, with thorough analysis and comprehensive experimental support, especially in the ablation and insight sections.\n\nWeaknesses:\nN/A", "was_revised": true, "final_justification": "After the rebuttal, all my concerns have been addressed, and I have no further questions. I will maintain my strong accept score.", "rebuttal": "> It would be valuable to add a “more-layer” baseline, especially for the 2.54B activation model. Adding ~200M parameters via extra layers is feasible and would help compare against the current approaches under a similar parameter budget.\n> \n- Thank you for this insightful suggestion. We fully agree that a “more-layer” baseline is a strong and valuable architectural comparison, as increasing depth often brings significant performance gains. Following your recommendation, we have conducted additional experiments and will include them in the revised version of the paper:\n    \n    \n    | Method | Added Total Param (M) | Added Activate Param (M) | Avg PPL | Hellaswag | MMLU | GSM8k | C-eval |\n    | --- | --- | --- | --- | --- | --- | --- | --- |\n    | Baseline | 0 | 0 | 6.026 | 73.07 | 58.79 | 52.92 | 60.26 |\n    | k = 8 | 50 | 50 | 5.979 | 73.51 | 59.78 | 52.16 | 62.26 |\n    | q = 48 | 201 | 201 | 5.953 | 73.59 | 58.45 | 53.30 | 59.67 |\n    | Add 4 Experts | 400 | 0 | 5.964 | 73.19 | 58.84 | 52.54 | **63.19** |\n    | Add More Layer | 625 | 100 | 5.871 | 73.96 | 59.74 | 53.53 | 61.45 |\n    |  |  |  |  |  |  |  |  |\n    | SDPA Elementwise | 201 | 201 | **5.761** | **74.64** | **60.82** | **55.27** | 62.20 |\n    \n    For the \"Add More Layer\" MoE configuration, we separately report both the increase in total parameters and activated parameters, as adding a layer increases both. As shown, adding an extra layer—increasing total parameters by approximately 625M (from 15B)—leads to a noticeable improvement compared to other parameter-increasing methods. However, its performance still falls short of the **SDPA Elementwise** gating approach, which achieves superior results with only 201M additional (and activated) parameters.\n    \n\n---\n\n> We also noticed that setting (6) achieves fairly strong performance, despite having only 1/8 the additional parameters of setting (5). This raises the question of whether the performance difference is primarily due to parameter count. Could you consider adding an experiment comparing v-elementwise G2 with multi-head (i.e., n × q × d_k) to control for parameter count and isolate the architectural impact?\n> \n- Thank you for the suggestion. To better isolate the architectural effect and control for parameter count, we conducted a follow-up experiment using a **Multi-Head Attention (MHA)** variant of the MoE model (changing from q=32, kv=4 to qkv=16). We then compared **G1 (SDPA Elementwise)** and **G2 (v Elementwise)** with *equal parameter budgets* (100M additional parameters each). The results are shown below:\n    \n    \n    | Method | Added Param (M) | Avg PPL | Hellaswag | MMLU | GSM8k | C-eval |\n    | --- | --- | --- | --- | --- | --- | --- |\n    | Baseline | 0 | 6.085 | 72.84 | 57.68 | 52.01 | 60.41 |\n    | SDPA Elementwise G1 | 100 | **5.842** | **74.01** | **59.54** | **54.43** | 61.34 |\n    | v Elementwise G2 | 100 | 5.900 | 73.32 | 58.52 | 53.84 | **61.51** |\n    \n    As observed, when the parameter counts are balanced, the performance gap between G1 and G2 narrows, but **G1 still outperforms G2**. This further supports our hypothesis that the **query-independent sparsity** and **position after SDPA** in G1 contribute meaningfully to its effectiveness, beyond just parameter count.\n    \n\n---\n\nWe sincerely appreciate the reviewer’s constructive feedback, which has helped strengthen the paper. We will incorporate these experiments and discussions into the final version."}, {"review_id": "Ex14YhawjT", "reviewer": "Reviewer_pXLA", "summary": "This work aims to propose an improvement on the self-attention module widely used in LLMs, introducing gating mechanisms to the typical self-attention layer. The work systematically investigates the  impact of gating mechanisms from a range of perspectives, including placing gates across various positions, different granularity, headspecific or shared, etc., which provides a comprehensive comparison between gating mechanisms with these nuances.\nExperimental results  across popular benchmarks indicate that this simple modification can improve model performance and training stability. Particularly, SDPA Output gating can reduce massive activation and attention-Sink, creating more balanced roles for weights and attention scores. Additionally, this gating helps improve the performance on tasks involving context length extension.", "questions": "Can you provide more details in experimental setups? Which architectures are used for the target LLMs (15A2B MoE and 1.7B dense models)?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. This work introduces gating mechanisms within the self-attention layer, accounting for a range of nuances such as positions and granularity.\n2. This work conducts comprehensive empirical comparison on both MoE and dense LLMs under various gating mechanisms, investigating which factor may be more impactful in improving the performance of target LLMs. The findings offer insights for readers to understand the significance of different gating settings in architectural design.\n3. The gating mechanisms can be used to provide the explanation of gains via non-linearity and input-dependent sparsity.\n\nWeaknesses:\n1. The architectures of the target LLMs are limited. It is a challenge to claim whether these findings can be generalized to other architectures such as Llama.\n2. This work emphasizes empirical result analysis from a benchmarking perspective, while offering limited investigation into the underlying causes of performance differences across gating configurations. For instance, SDPA output gatings (e.g., elementwise or headwise) appear more effective in mitigating the attention-sink phenomenon. However, it remains unclear why gating applied after SDPA outperforms other variants such as G2, G3, G4, and G5.\n3. Table 1, Table 2, and Table 3 are quite confusing in terms of the methods. It seems different gating mechanisms are not compared on the same settings. Which methods (positions like G1, G2, G3, etc are not stated) are compared in Table 2? Is G1 the default setting?", "was_revised": false, "final_justification": "", "rebuttal": "> The architectures of the target LLMs are limited. It is a challenge to claim whether these findings can be generalized to other architectures such as Llama.\n> \n- Thank you for the comment. We will include a more detailed description in the revised version. Specifically, our dense model follows the architecture of **Qwen3 1.5B**, and our MoE model adopts the same structure as **Qwen3-30B-A3B**, except that we reduce the number of layers from 48 to 24. More architectural details can be found in the response to your question below.\n- We believe that the Qwen architecture differs only slightly from Llama3, primarily in the design of the FFN and attention components. Since the paper *Massive Activation* has shown that the attention sink phenomenon is widespread across various architectures, and given that our analysis on non-linearity should be applicable to all Transformer-based models, we expect our findings to generalize well.\n- Furthermore, beyond the 28-layer 1.7B model, we have also conducted experiments on a 48-layer 1.7B model with a different width-to-depth ratio, as shown in Table 2.\n\n---\n\n> This work emphasizes empirical result analysis from a benchmarking perspective, while offering limited investigation into the underlying causes of performance differences across gating configurations. For instance, SDPA output gatings (e.g., elementwise or headwise) appear more effective in mitigating the attention-sink phenomenon. However, it remains unclear why gating applied after SDPA outperforms other variants such as G2, G3, G4, and G5.\n> \n- Thank you for the feedback. Indeed, our current analysis primarily focuses on comparing **G1 (SDPA Elementwise)** and **G2 (v Elementwise)**, explaining why both are effective but G1 performs better: while both enhance non-linearity, the sparsity induced by G2 is limited, and its query-independent gating scores are less effective at mitigating the attention sink. Therefore, G2 underperforms G1 overall.\n- As for the other variants—G3 (k Elementwise), G4 (q Elementwise), and G5 (Dense Output)—they do not effectively introduce non-linearity between the $W_v$ and $W_o$ projections. Our ablation studies further show that gating at these positions induces significantly less sparsity, and the resulting models exhibit similar levels of massive activation and attention sink as the baseline without gating.\n- We will include a detailed analysis of these positions and a discussion on their limited effectiveness in the appendix of the revised manuscript.\n\n---\n\n> Table 1, Table 2, and Table 3 are quite confusing in terms of the methods. It seems different gating mechanisms are not compared on the same settings. Which methods (positions like G1, G2, G3, etc are not stated) are compared in Table 2? Is G1 the default setting?\n> \n- Thank you for raising this important point. The design rationale behind Tables 1, 2, and 3 is as follows: **Table 1** aims to identify the most effective gating position (which we found to be G1, i.e., after SDPA). Building on this finding, **Table 2** and **Table 3** further evaluate and analyze the impact of applying gating at G1 under different training scales and configurations. Therefore, all gating results in Tables 2 and 3 are based on the **G1 (SDPA output)** position.\n- To avoid confusion, we will clarify this design choice more explicitly in the next version of the paper.\n- Additionally, as mentioned in the response to Reviewer bMKL, we are including the full ablation across gating positions on the 28-layer 1.7B dense model below, which will be added to the appendix:\n    \n    \n    | Method | Avg PPL | MMLU | GSM8k | Hellaswag | C-eval | CMMLU |\n    | --- | --- | --- | --- | --- | --- | --- |\n    | Baseline | 7.499 | 50.21 | 27.82 | 64.94 | 49.15 | 49.52 |\n    | SDPA Elementwise G1 | 7.404 | 51.15 | 28.28 | 65.48 | 50.72 | 50.72 |\n    | v Elementwise G2 | 7.453 | 50.53 | 27.97 | 65.12 | 49.82 | 49.91 |\n    | k Elementwise G3 | 7.485 | 50.25 | 28.01 | 64.99 | 49.03 | 49.17 |\n    | q Elementwise G4 | 7.466 | 50.33 | 27.89 | 65.04 | 49.15 | 49.33 |\n    | Dense Output G5 | 7.489 | 50.34 | 28.08 | 65.01 | 49.38 | 49.25 |\n\n---\n\n> Can you provide more details in experimental setups? Which architectures are used for the target LLMs (15A2B MoE and 1.7B dense models)?\n> \n- The detailed architectures for the **1.7B-28 layer**, **1.7B-48 layer**, and **15A2B MoE** models are summarized in the table below:\n    \n    \n    | Model | 1.7B-28 layer | 1.7B-48 layer | 15A2B MoE |\n    | --- | --- | --- | --- |\n    | Layers | 28 | 48 | 24 |\n    | Query Heads | 16 | 16 | 32 |\n    | Key / Value Heads | 8 | 8 | 4 |\n    | Head Dimension | 128 | 128 | 128 |\n    | Tie Embedding | Yes | Yes | No |\n    | QK Normalization | Yes | Yes | Yes |\n    | Hidden Size | 2048 | 1536 | 2048 |\n    | FFN Size | 6144 | 4608 | 768 |\n    | Number of Experts | – | – | 128 |\n    | Top-k | – | – | 8 |\n- Additionally, we plan to **open-source the 1.7B models** to facilitate further research and reproducibility within the community."}]

Author rebuttal:

Meta-review: In this work, the authors do a study of gating mechanisms which tend to be almost omnipresent nowadays, especially in the context of LLMs. To do so they introduce a gating in self-attention layers. However, they go way beyond that and study a different range of gating mechanisms (over thirty of them) on where to put the gating, the different granularities of it providing a very complete and comprehensive study. Their change in the mechanism of softmax not only improve the results in different benchmarks, but also show a reduction in both activations and softmax weight in particular tokens (studied before in different terms as 'Attention-Sink' or 'Registers'), which can have other practical advantages such as simplifying the quantization.

The paper initially had some diverging reviews, with 2 strong accepts, 1 accept and 1 borderline reject. Reviewer cPo5 who is an authority on the field found no weaknesses at all with the paper, having only a couple of suggestions which the authors addressed. Similarity, reviewer Wjax finds no weaknesses at all, and like Reviewer cPo5 scores the paper with the highest score. Reviewer pXLA asks for experiments in modern LLMs, and for the detailed hyperparameters. The authors response that they will add experiments in QWen architecture, give the hyperparameters, and promise to open-source one of the models. 

Reviewer bMKL is the only reviewer to initially score the paper in the negative region (Borderline reject). They have some doubts on the experimental section. In particular, they are not convinced on the results because of the varying batch size, to which the authors respond in showing experiments with an increased batch size when increasing the number of training tokens, similar to what was the reviewer's suggestion. The reviewer also is not convinced that a 0.2 improvement in PPL is not significant. The authors and the AC disagree on this point with the reviewer, while the other reviewers found the improvement significant. Finally, the reviewer had some smaller points which were properly addressed by the authors. During the authors-reviewer phase, the reviewer initially engaged with the authors, and after the stage, they increased the score to 4 (Borderline Accept) claiming that the authors 'added additional experiments and discussion during the rebuttal to address my primary concerns'. At this stage, there is a high consensus between the reviewers for the paper to be accepted with scores 6, 6, 5, 4.

Having read the paper and the reviews, I tend to agree with the reviewer's comments. This is a high-quality paper that deserves to be accepted to NeurIPS. In particular:

- This is a very-well executed study with an impressive number of experiments, and as Reviewer Wjax it is indeed at a production level.

- The paper is well-written, enjoyable to read with arguments being supported by tons of evidence.

- Strong empirical results in a different range of architectures, models and token size.

- Important findings.

I think this paper will be received well by the community, and will improve the knowledge on the gating mechanism in LLMs, if not improving the quantization schemes. I would go as far as claiming that this has been one of the top 5 papers I have read this year. The authors are urged to add the parts they mentioned in the rebuttal for the camera-ready version of the paper. Congratulations to the authors for a perfectly executed study!
Decision: Accept (oral)