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

2025-05-10

AI Summary

Existing studies lack a systematic analysis of how gating mechanisms enhance soft (softmax) attention. This work presents the first systematic investigation into the impact of gating position on attention performance. We propose head-specific sigmoid gating applied at the softmax attention output, enabling query-dependent nonlinear sparse modulation. This design effectively mitigates the “attention sink” problem. Evaluated on both a 15B MoE and a 1.7B dense architecture trained on 3.5T tokens, our method significantly improves multi-task benchmark performance, training stability, learning-rate robustness, and long-context extrapolation capability. Ablation studies demonstrate that output-level gating consistently outperforms input- or intermediate-layer gating variants. The proposed approach is lightweight, architecture-agnostic, and requires no additional parameters beyond the gating weights. All code and pre-trained models are publicly released on Hugging Face and GitHub.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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 '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.

Problem

• Investigates gating effects on softmax attention variants
• Improves performance via head-specific sigmoid gating
• Mitigates attention sink and enhances long-context extrapolation

Innovation

• Head-specific sigmoid gate after SDPA
• Non-linearity in softmax attention
• Query-dependent sparse gating scores