AdaSplash让注意力计算更快更省内存,适合长文本任务。
AdaSplash: Adaptive Sparse Flash Attention
- 用混合算法加速α-entmax计算,迭代次数减少7倍
- 自研Triton内核实现动态稀疏性,显著降低内存和耗时
- 在长序列上接近甚至超越FlashAttention-2效率
基于softmax的注意力计算开销限制了Transformer在长序列任务中的应用。自适应稀疏性(如α-entmax)提供数据驱动的灵活替代方案,但现有实现效率低,未充分利用稀疏性带来的性能优势。本文提出AdaSplash,结合GPU优化算法与α-entmax的稀疏特性。首先引入混合Halley-二分法,使α-entmax变换的迭代次数减少7倍;其次开发定制Triton内核以高效处理自适应稀疏性。在RoBERTa、ModernBERT文本分类与单向量检索,以及GPT-2语言建模任务上的实验表明,相比现有α-entmax实现,本方法在运行时间和内存效率上均有显著提升,接近甚至在某些场景下超越高度优化的FlashAttention-2,支持长序列训练并保持优异任务性能。
原文摘要 · Abstract (English)
The computational cost of softmax-based attention in transformers limits their applicability to long-context tasks. Adaptive sparsity, of which $α$-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 $α$-entmax. We first introduce a hybrid Halley-bisection algorithm, resulting in a 7-fold reduction in the number of iterations needed to compute the $α$-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 $α$-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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。