AdaSplash：自适应稀疏闪光注意力

基于softmax注意力的计算成本限制了transformer在长上下文任务中的应用。其中一种自适应稀疏性方法——$\alpha$-entmax注意力——提供了灵活的数据依赖性替代方案，但现有实现效率低下，未能利用稀疏性获得运行时和内存收益。本文提出AdaSplash，将GPU优化算法的效率与$\alpha$-entmax的稀疏性优势相结合。我们首先引入混合Halley-二分算法，使计算$\alpha$-entmax变换所需的迭代次数降低7倍。随后，我们实现自定义Triton内核，高效处理自适应稀疏性。针对RoBERTa和ModernBERT的文本分类和单向量检索，以及GPT-2的语言建模进行实验，结果表明我们的方法在运行时和内存效率方面显著优于现有$\alpha$-entmax实现。它接近——在某些情况下甚至超过——高度优化的softmax实现如FlashAttention-2，从而实现长上下文训练，同时保持出色的任务性能。