🤖 AI Summary
To address inaccurate pattern localization and degraded generalization in long sequences, this work identifies that conventional Softmax attention—enforcing dense, normalized attention distributions—is susceptible to interference from irrelevant tokens, leading to attention dispersion and representation collapse. We propose ASEntmax, a sparse attention mechanism built upon α-entmax: it introduces a learnable, adaptive temperature parameter enabling dynamic interpolation between sparsity and density; additionally, we design a novel positional encoding scheme to enhance long-range modeling. Seamlessly integrated into standard Transformer architectures, ASEntmax achieves significant improvements over Softmax, Sparsemax, and fixed-temperature α-entmax baselines across multiple long-context generalization tasks. Notably, it is the first method to enable end-to-end adaptive optimization of sparse attention parameters, combining theoretical rigor with strong empirical performance.
📝 Abstract
Transformer-based architectures traditionally employ softmax to compute attention weights, which produces dense distributions over all tokens in a sequence. While effective in many settings, this density has been shown to be detrimental for tasks that demand precise focus on fixed-size patterns: as sequence length increases, non-informative tokens accumulate attention probability mass, leading to dispersion and representational collapse. We show in this paper that sparse attention mechanisms using $alpha$-entmax can avoid these issues, due to their ability to assign exact zeros to irrelevant tokens. Furthermore, we introduce Adaptive-Scalable Entmax (ASEntmax), which endows $alpha$-entmax with a learnable temperature parameter, allowing the attention distribution to interpolate between sparse (pattern-focused) and dense (softmax-like) regimes. Finally, we show that the ability to locate and generalize fixed-size patterns can be further improved through a careful design of position encodings, which impacts both dense and sparse attention methods. By integrating ASEntmax into standard transformer layers alongside proper positional encodings, we show that our models greatly outperform softmax, scalable softmax, and fixed-temperature $alpha$-entmax baselines on long-context generalization.