Hybrid Gated Attention

The paper proposes a Hybrid Gated Attention (HyGA) framework to extend the effectiveness-efficiency Pareto frontier of gated attention. HyGA contains three types of gating strategies that leverage diverse information from multiple stages of attention, and collaboratively build element-wise/head-wise gating from multiple perspectives, capturing intra-head and cross-head information interactions. The framework provides multi-source modulation signals, enabling more comprehensive control over information flow and improving the representational capacity of attention. The authors also introduce low-rank matrix decomposition and learnable attention sink to further enhance training efficiency and stability. Experiments on widely-used benchmarks with different backbones show that HyGA comprehensively improves both training loss and various downstream performances compared with Gated attention, achieving the best performance at different computation costs.

The paper addresses limitations of existing gated attention (Qiu et al. 2026b), which applies element-wise gating to Scaled Dot-Product Attention (SDPA) output. The authors identify three key limitations of the original gated attention:

(a) the original gated attention relies on the raw input X for gating, ignoring other information that is a good supplement to the current gating strategy.

(b) Element-wise gating introduces considerable additional parameters and computational overhead, leaving substantial room for improving gating efficiency.

(c) existing gated attention still exhibits a non-negligible BOS-token sink ratio, motivating us to further mitigate the attention-sink phenomenon for stable training.

HyGA extends the original Gated attention with three gating mechanisms:

• X-gate: the original gate that takes the raw input X before the attention layer as input

• H-gate: the proposed gate that takes the output H after SDPA as input

• C-gate: the proposed cross-head gate, which captures the inter-head connections and provides head-level reweighting

The framework adopts a gate fusion strategy to get our final hybrid gated attention output, avoiding undesired dominating gates to ensure smooth training. It also implements HyGA with learnable attention sink to provide a double guarantee for training stability.

The H-gate is formulated as: Hi′ = Hi ⊙ σ(XWi + SiLU(Hi Wid)Wiu), where Wid ∈ Rd×dint and Wiu ∈ Rdint ×d are the down-projection and up-projection gating matrices for the i-th head. The H-gate uses a 2-layer MLP form with SiLU activation to enhance nonlinear modulation and to enable low-rank parameter compression by adjusting the intermediate dimension.

The authors explain the functional complementarity: The former mainly captures the intrinsic features of input tokens before attention, while the latter mainly captures the features after contextual interaction through attention.

They choose fused rather than separate gates because "the multiplicative coupling form... multiplies two gates together. Since the sigmoid activation function has a value range of (0, 1), each gate generally acts as a suppressive modulator of the information flow. Multiplying too many gates may therefore lead to overly strong suppression, which may weaken gradient propagation and limit effective representation learning. The fused form first adds the pre-activation gate logits and then applies the activation function, allowing different factors to jointly control the information flow."

For efficiency, they apply low-rank factorization to both X-gate and H-gate: Hi′ = Hi ⊙ σ(SiLU(X W̄id)W̄iu + SiLU(Hi Wid)Wiu), where the maximum rank of each gating matrix can be controlled by adjusting the intermediate dimension.

The C-gate is formulated as: Hi′ = Hi ⊙ σ(SiLU(X W̄id)W̄iu + SiLU(Hi Wid)Wiu) ⊙ σ (Broadcast((HWc)i)), where Wc ∈ Rhd×h denotes the transformation matrix for computing the cross-head gate for h heads based on all heads' elements.

The C-gate provides one gate score for all elements of one head because: "a) the C-gate is supplementary to the above element-wise X-gate and H-gate and thus should not bring in much additional computation, and b) the cross-head interactions are supposed to provide coarse-grained inter-head reweighting."

The mechanism may suppress heads with lower contribution under the current input, emphasize more important heads, and dynamically allocate information flow across heads according to the current context.

The authors note that when using gated attention alone, its mitigation of the sink ratio is still not perfect, and massive activations occasional occur in production-scale training. They implement HyGA with learnable attention sinks inspired by GPT-OSS (Agarwal et al. 2025), finding that adding learnable attention sinks on top of HyGA can further reduce the sink ratio and effectively alleviate massive activations, even with slight loss advantages.

MoE-5B (MLA backbone, 500B tokens): HyGA achieves an average score of 42.23 across 14 benchmarks compared to 40.70 for Gated attention. Notable improvements include ARC (53.34→57.53), MBPP+ (35.19→40.21), MATH (14.95→17.20), and GSM8K (28.51→29.87). The training loss advantage is approximately 0.012 at 60k step.

Qwen3-0.6B (GQA backbone, 200B tokens): HyGA achieves an average of 30.56 across 6 benchmarks compared to 29.44 for Gated attention and 27.53 for original GQA. HyGA attains a training loss approximately 0.008 lower than that of Gated GQA.

• H-gate: introducing H-gate yields clear improvements in most benchmarks covering different capabilities, achieving significantly better average performance.

• Learnable attention sink: the cooperation with learnable attention sink could also bring in a slight improvement on the average score.

• C-gate with gate fusion: adding C-gate based on X+H gates further reduces the final training loss by 0.004. The sequential additions provide average performance gains: +0.24%→+1.21%→+1.53%.

With low-rank compression (dint = 32), HyGA utilizes only approximately 26% of the gating parameters required by the original Gated attention baseline, while achieving slightly better performance on training loss and better overall performance on downstream tasks. The paper demonstrates that HyGA does extend the effectiveness-efficiency Pareto frontier of the original gated attention.

HyGA with learnable sink substantially reduces the BOS-token's attention score (especially in the last few layers) and shows a marked reduction in massive activations. For Qwen3, both the baseline and Gated attention exhibit obvious loss spikes in training, whereas HyGA does not under the same recommended learning rate, suggesting HyGA is more robust to larger learning rates and provides more stable optimization.

The paper summarizes three main contributions:

1. "We propose HyGA, which jointly adopts three hybrid gating strategies to provide both head-wise and element-wise gate scoring calculated from different factors. Equipped with learnable sink, HyGA achieves more stable training."

2. We explore different low-rank matrix factorization settings in our hybrid gates to extend the effectiveness-efficiency Pareto frontier of Gated attention methods.

3. Overall, HyGA achieves significant improvement compared to baselines with different backbones and model settings, shedding light on more effective, efficient, and stable gated attention modules in practice.

The authors plan to scale HyGA to larger models and evaluate its effectiveness at greater scales and investigate its generalizability across different attention backbones, such as linear/sparse attention architectures.

Based on the paper, here are the specific improvements I can implement in an AI system, and what the improved system can do:

Improvements:

1. Add a post-attention gating branch (H-gate) that takes the attention output (after contextual mixing) as input, in addition to the pre-attention input (X-gate). This provides complementary modulation signals—one capturing intrinsic token features, the other capturing contextually integrated features—leading to more precise control over information flow.

2. Add a cross-head gating mechanism (C-gate) that computes a single gate score per head based on all heads' outputs. This dynamically reweights heads according to the current input, suppressing low-contribution heads and emphasizing important ones, enabling adaptive inter-head resource allocation.

3. Fuse the three gates additively before applying the sigmoid activation rather than multiplying them. This avoids over-suppression from multiplicative gating (since sigmoid outputs are <1), preserving gradient flow and representation learning capacity.

4. Apply low-rank factorization (down-projection to intermediate dimension, then up-projection) to both X-gate and H-gate matrices, with an intermediate dimension of 32. This reduces gating parameters to 26% of the original gated attention while maintaining or improving performance.

5. Add a learnable attention sink token to the attention mechanism, complementing the gating. This further reduces BOS-token attention sink ratio and eliminates massive activations, providing double training stability.

What the improved AI system can do:

• Achieve better downstream performance across reasoning (ARC: +4.2 points), coding (MBPP+: +5.0 points), math (MATH: +2.3 points), and general benchmarks (average +1.5 points over gated attention, +3.0 points over standard attention) at the same compute budget.

• Train faster and more stably—reduces training loss by 0.012 at 60k steps, eliminates loss spikes even at higher learning rates, and avoids massive activation outliers that can destabilize large-scale training.

• Use fewer gating parameters (74% reduction) while performing better, enabling deployment on memory-constrained hardware or scaling to larger models with the same memory footprint.

• Adaptively allocate attention resources—the system can dynamically emphasize or suppress specific heads based on the current input context, improving representational capacity for diverse tasks (e.g., switching between code reasoning and mathematical computation modes).

• Scale more reliably—the combination of learnable sinks and hybrid gating makes the system robust to production-scale training (500B+ tokens), where vanilla gated attention still exhibits occasional massive activations.

• Generalize across architectures—works with both MLA (MoE-5B) and GQA (Qwen3-0.6B) backbones, suggesting applicability to other attention variants (linear, sparse) with minimal adaptation.

Gated attention is an effective approach to mitigate attention sinks and enhance the representational capacity of attention. To further extend its effectiveness-efficiency Pareto frontier, we propose a Hybrid Gated Attention (HyGA) framework that contains three types of gating strategies. Specifically, these gates leverage diverse information from multiple stages of attention, and collaboratively build element-wise/head-wise gating from multiple perspectives, capturing intra-head and cross-head information interactions. Through our hybrid gating components, HyGA could provide multi-source modulation signals, enabling more comprehensive control over information flow and improving the representational capacity of attention. We also introduce low-rank matrix decomposition and learnable attention sink to further enhance training efficiency and stability. In experiments, we evaluate HyGA on widely-used benchmarks based on different backbones. The experimental results show that our HyGA comprehensively improves both training loss and various downstream performances compared with Gated attention. HyGA has also been verified to achieve the best performance at different computation costs, with comprehensive model analyses for better understanding. The proposed HyGA sheds light on a more effective, efficient, and stable attention mechanism.