Abstract
The computational cost of softmax-based attention in transformers limits their applicability to long-context tasks. Adaptive sparsity, of which $\alpha$-entmax attention is an example, offers a flexible data-dependent alternative, but existing implementations are inefficient and do not leverage the sparsity to obtain runtime and memory gains. In this work, we propose AdaSplash, which combines the efficiency of GPU-optimized algorithms with the sparsity benefits of $\alpha$-entmax. We first introduce a hybrid Halley-bisection algorithm, resulting in a 7-fold reduction in the number of iterations needed to compute the $\alpha$-entmax transformation. Then, we implement custom Triton kernels to efficiently handle adaptive sparsity. Experiments with RoBERTa and ModernBERT for text classification and single-vector retrieval, along with GPT-2 for language modeling, show that our method achieves substantial improvements in runtime and memory efficiency compared to existing $\alpha$-entmax implementations. It approaches -- and in some cases surpasses -- the efficiency of highly optimized softmax implementations like FlashAttention-2, enabling long-context training while maintaining strong task performance.
Abstract (translated)
基于softmax的注意力机制在变压器中的计算成本限制了其在长上下文任务中的适用性。自适应稀疏性，例如$\alpha$-entmax注意就是此类方法的一个例子，提供了一种灵活的数据依赖性替代方案，但现有的实现效率低下，并没有利用稀疏性来获得运行时间和内存的节省。在这项工作中，我们提出了AdaSplash，它结合了GPU优化算法的效率与$\alpha$-entmax的稀疏性优势。首先，我们引入了一个混合Halley-二分法算法，使得计算$\alpha$-entmax变换所需的迭代次数减少了7倍。然后，我们通过自定义Triton内核有效地处理适应性稀疏性问题。使用RoBERTa和ModernBERT进行文本分类和单向量检索的实验，以及GPT-2用于语言建模，证明了我们的方法在运行时间和内存效率方面相比现有的$\alpha$-entmax实现取得了显著改进。我们的方法接近于——并在某些情况下超过了——高度优化的softmax实现（如FlashAttention-2）的效率，这使得能够进行长上下文训练的同时保持强大的任务性能。
URL
https://arxiv.org/abs/2502.12082
