Source-linked AI summary
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 Huang, Suozhi Huang, Dayiheng Liu, Jingren Zhou, Junyang Lin
TL;DR
Prior work widely uses gating, but its specific effects in softmax attention remain insufficiently isolated. This paper compares gating variants across large MoE and dense models, finding that a head-specific sigmoid gate after SDPA improves performance, stability, scaling, and attention behavior through non-linearity and query-dependent sparsity.
Problem
Existing literature rarely isolates gating’s specific contribution in softmax attention, especially when gating is confounded with other architectural factors.
Method
The paper systematically compares 30 gating variants across 15B MoE and 1.7B dense models, varying gating positions, granularity, and computational form.
Results
A head-specific sigmoid gate after SDPA consistently improves performance, training stability, learning-rate tolerance, scaling, and attention behavior.
Takeaways & Limitations
Gated softmax attention can add non-linearity and query-dependent sparsity while mitigating attention sinks and supporting longer-context generalization.
Takeaways & Limitations
The paper does not provide a rigorous theoretical explanation for how attention sinks affect longer-context generalization.
Abstract
from arXiv · showhide
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 'attention sink' and enhances long-context extrapolation performance, and we also release related $\href{https://github.com/qiuzh20/gated_attention}{codes}$ and $\href{https://huggingface.co/QwQZh/gated_attention}{models}$ to facilitate future research.
1 Introduction
The paper systematically studies gating in standard softmax attention, finding that head-specific gating after SDPA most reliably improves performance and training behavior. The authors attribute these gains to added non-linearity and input-dependent sparsity, which also mitigates attention sinks.
- Motivation: Existing work uses gating broadly, but its specific contribution is difficult to assess when combined with other architectural factors.The paper highlights confounding in systems such as Switch Heads and Native Sparse Attention.
- Training behavior: Gating improves training stability, reduces loss spikes, supports larger learning rates, and enhances scaling.These effects are reported for SDPA-output gating under the paper’s experimental settings.
- Mechanisms: SDPA-output gating introduces non-linearity into the low-rank transformation formed by consecutive value and output projections.This increases the expressiveness of the attention mapping.
- Mechanisms: Sparse gating scores add input-dependent sparsity to SDPA outputs and eliminate the attention-sink phenomenon.The mechanism can suppress disproportionately dominant initial-token attention.
- Conclusion: The authors summarize gated attention as improving model performance and behavior while releasing attention-sink-free models for future research.The release is presented as a way to support further investigation.
2 Gated-Attention Layer
The section introduces standard multi-head softmax attention and formalizes gating as a dynamic modulation of intermediate outputs. It then defines the investigated gating positions, granularities, and default sigmoid-based configuration.
- Preliminary: Multi-Head Softmax Attention: Transformer attention begins with an input X that is projected into queries, keys, and values, followed by attention, head concatenation, and a final output layer.The input has sequence length n and model dimension dmodel.
- Preliminary: Multi-Head Softmax Attention: SDPA computes query-key attention scores, normalizes them with softmax, and returns a weighted sum of values.The resulting attention weights are non-negative and sum to 1 across each row.
- Augmenting Attention Layer with Gating Mechanisms: Gating modulates an input Y using scores computed from another input X, learnable parameters Wθ, and an activation function such as sigmoid.The gated output selectively preserves or erases features of Y.
- Augmenting Attention Layer with Gating Mechanisms: The paper examines gating after Q, K, or V projections, after SDPA outputs, and after concatenated multi-head outputs.These positions are labeled G2, G3, G4, G1, and G5, respectively.
- Augmenting Attention Layer with Gating Mechanisms: The study compares headwise and other gating granularities along with elementwise or headwise, head-specific or shared, and additive or multiplicative variants.The exploration is organized around gating position, granularity, and computational form.
- Augmenting Attention Layer with Gating Mechanisms: Unless otherwise specified, experiments use head-specific multiplicative gating with sigmoid activation.The sigmoid is σ(x) = 1/(1+e^-x).
3 Experiments
Experiments compare gated-attention variants in MoE and dense models across benchmarks, perplexity, learning rates, batch sizes, and model configurations. SDPA-output gating consistently performs best or improves performance while supporting more stable training.
- 3 Experiments: Head-specific gating adds fewer than 2M parameters in the MoE-15A2B model while delivering substantial improvements.Sharing gating scores across heads produces smaller benchmark gains than headwise gating.
- 3 Experiments: Multiplicative SDPA-output gating outperforms additive gating, although additive gating still improves over baselines.
- 3 Experiments: Sigmoid activation performs better than SiLU in the most effective gating configuration.Replacing sigmoid with SiLU leads to less improvement.
- 3 Experiments: SDPA-output (G1) and value-layer (G2) gating reduce PPL by more than 0.2 and outperform parameter-expanding baselines.G1 achieves better PPL and benchmark results than G2.
- 3 Experiments: SDPA-output gating consistently benefits dense models across model configurations, training-data settings, and hyperparameters.
- 3 Experiments: Gating improves training stability and enables noticeable gains when increasing the maximum learning rate, whereas baselines encounter convergence issues.Under the 3.5T-token setting, gating largely reduces the loss spike.
4 Analysis: Non-Linearity, Sparsity, and Attention-Sink-Free
The analysis attributes gated attention’s effectiveness to increased non-linearity in attention’s low-rank mapping and sparse, query-dependent modulation of SDPA outputs. These properties are associated with improved performance, reduced massive activations, and mitigation of attention sinks.
- Non-Linearity: RMSNorm applied independently to each attention head before concatenation significantly reduces PPL with almost no additional parameters.This corresponds to the G1-side modification that introduces non-linearity between the value and dense projections.
- Non-Linearity: Gating between the value and dense projections increases the expressiveness of attention’s low-rank linear mapping.The value and output projections can be merged into one low-rank mapping, while gating at G1 or G2 introduces non-linearity between them.
- Sparsity: SDPA output gating has the lowest mean scores and strongest concentration near zero, indicating substantial sparsity.Figure 3 compares SDPA elementwise, value elementwise, and head-shared SDPA gating distributions; SDPA output gating is the sparsest.
- Sparsity: Query-dependent SDPA output gating is more effective than value gating because its scores are computed from the current query rather than past keys and values.Input-independent gating improves over baseline through non-linearity, but its higher gate scores reinforce the importance of input-dependent sparsity.
- Attention-Sink-Free: Head-specific, input-dependent sigmoid gating at SDPA output reduces first-token attention and massive activations, mitigating attention sinks.Shared gating or value-position gating decreases massive activations without reducing first-token attention, while less sparse gating intensifies both effects.
5 Related Works
The related work traces gating from recurrent and feedforward architectures to state-space, linear-attention, and softmax-attention models. It also situates attention sinks alongside redundant-token and massive-activation phenomena.
- Gating Mechanisms: LSTMs, GRUs, and Highway Networks use gates to regulate information flow and support gradient propagation or deep-network training.These early architectures selectively retain or discard information across time steps, layers, or feedforward transformations.
- Gating Mechanisms: State-space models and linear-attention systems use gating modules to control token-mixer information.Examples include FLASH, RetNet, Lightning Attention, and Gated Delta Networks.
- Gated Attention: Prior softmax-attention work applied gating to attention outputs and reported significant performance improvements.The paper positions its systematic comparison within this emerging line of gated softmax-attention research.
- Attention Dynamics: Attention sinks are related to redundant register tokens and excessive attention assigned to tokens with massive activation values.The paper distinguishes these attention-dynamics observations from its finding that output gating can reduce the phenomenon.
6 Conclusion
The paper systematically evaluates gating in standard softmax attention across diverse model scales and variants. It finds that sigmoid gating after SDPA delivers broad improvements while adding non-linearity, input-dependent sparsity, and attention-sink reduction.
- Conclusion: Experiments across 30 variants of 15B MoE and 1.7B dense models trained on up to 3.5T tokens identify SDPA-output sigmoid gating as most effective.The conclusion reports effects on performance, training stability, and attention dynamics.
- Conclusion: SDPA-output gating introduces non-linearity and input-dependent sparsity while eliminating the attention-sink phenomenon.The conclusion connects these mechanisms with the reported improvements in attention behavior.
- Conclusion: Gating facilitates context-length extension, allowing models to generalize to longer sequences without retraining.The paper also releases attention-sink-free models and code for future research.
Limitations
The study analyzes gating through ablation studies, but leaves the broader effects of attention non-linearity and the role of attention sinks in long-context generalization unresolved.
- The broader implications of non-linearity for attention dynamics and overall training remain under-explored.
- The study observes improved long-context extension after eliminating attention sinks but lacks a rigorous theoretical explanation for this generalization effect.
A.1 Switch Head Baselines
The Switch Head experiments separate the value of gating from expert routing and compare configurations with different activated expert counts. Single-expert sigmoid gating remains especially competitive.
- Switch Head experiments test sparse activation by selecting top-k key/value/output experts through learnable sigmoid routing.
- ‘Switch v, 1top1’ is equivalent to v Headwise Gate in Table 1, providing a direct correspondence between the methods.
- The best benchmark scores and PPL were achieved by ‘Switch v 1top1’, analogous to applying sigmoid gating directly to the value-layer output.
- Increasing activated kv experts improved PPL modestly, while gains in overall benchmark performance were less pronounced.
A.2 More Discussion on Sparse Gating Score
Sparse gating substantially reduces SDPA hidden-state magnitudes and increases threshold-based sparsity, while post-gating states remain similar to the ungated baseline.
- 0.71 to 0.05: gating reduced the mean absolute hidden-state value while keeping gated states close to the baseline.The passage attributes the reduction to generally small gating scores and suggests filtering of irrelevant information.
- Gating significantly increased hidden-state sparsity across multiple thresholds.Small mean gating scores naturally push some hidden-state values below the selected thresholds.
- Figure 4 compares mean absolute SDPA output values before and after gating against the baseline.
- Figure 5 reports the proportion of SDPA outputs below 1e-2 and 1e-3 after gating, alongside sparsity estimated from average gating scores.
A.3 Layerwise Massive Activations and Attention Sinks
Layerwise analyses connect gating sparsity with massive activations and attention sinks: SDPA output gating suppresses both, whereas reduced or shared sparsity restores baseline-like behavior.
- Analysis setup: The analysis compares attention sinks, defined here as attention scores assigned to the first token, with massive activations across configurations.
- Layerwise comparisons: Value-layer gating produces activations similar to SDPA gating but leaves residual attention sinks, showing massive activations are not necessary for sinks.
- Sparsity constraints: Shared head gating and sparsity-suppressing activation functions reduce gating sparsity, making massive activations and attention sinks comparable to baseline.
- Scope: The authors suggest sufficient attention sparsity may mitigate massive activations but call for further study in deeper and larger models.
- Sparsity constraints: Different heads require different sparsity, because sharing gating scores increases scores across most layers.
- Alternative intervention: Clipping attention and FFN outputs before residual addition was explored as an alternative constraint on large activations.
- Layerwise comparisons: Baseline massive activations and attention sinks emerge after the 6th layer, while SDPA gating reduces activations and removes observed sinks.
- Sparsity constraints: Figure 7 shows gating-score distributions under different constraints for SDPA output-gating variants.