Gated Sparse Attention: Combining Computational Efficiency with Training Stability for Long-Context Language Models
Alfred Shen, Aaron Shen
TL;DR
Gated Sparse Attention (GSA) integrates a gated sparse indexer with adaptive sparsity and dual gating to achieve sub-quadratic long-context attention without sacrificing expressiveness. The architecture delivers a rigorous complexity analysis, expressiveness improvements, and convergence guarantees, while empirical results show substantial throughput gains (12–16x at 128K context) alongside perplexity and downstream task improvements and dramatically reduced attention sinks and training spikes. The combination of a sigmoid-based gated indexer and attention-sparing mechanisms yields strong long-context performance on 1.7B-parameter models trained on 400B tokens, with practical benefits in retrieval and reasoning at very long contexts. Overall, GSA provides a scalable, stable, and effective approach to long-context Transformer modeling, with promising directions for hierarchical indexing and mixture-of-experts integration.
Abstract
The computational burden of attention in long-context language models has motivated two largely independent lines of work: sparse attention mechanisms that reduce complexity by attending to selected tokens, and gated attention variants that improve training sta-bility while mitigating the attention sink phenomenon. We observe that these approaches address complementary weaknesses and propose Gated Sparse Attention (GSA), an architecture that realizes the benefits of both. GSA incorporates a gated lightning indexer with sigmoid activations that produce bounded, interpretable selection scores, an adaptive sparsity controller that modulates the number of attended tokens based on local uncertainty, and dual gating at the value and output stages. We establish theoretical foundations for the approach, including complexity analysis, expressiveness results, and convergence guarantees. In experiments with 1.7B parameter models trained on 400B tokens, GSA matches the efficiency of sparse-only baselines (12-16x speedup at 128K context) while achieving the quality gains associated with gated attention: perplexity improves from 6.03 to 5.70, RULER scores at 128K context nearly double, and attention to the first token, a proxy for attention sinks, drops from 47% to under 4%. Training stability improves markedly, with loss spikes reduced by 98%.