AdaSplash: Adaptive Sparse Flash Attention

要約

トランスのソフトマックスベースの注意の計算コストは​​、長期コンテストタスクへの適用性を制限します。
$ \ alpha $ -Entmaxの注意は例ですが、柔軟なデータ依存の代替品を提供しますが、既存の実装は非効率的であり、スパースを活用してランタイムとメモリのゲインを得ることはありません。
この作業では、GPU最適化されたアルゴリズムの効率と$ \ alpha $ -Entmaxのスパースの利点を組み合わせたAdasplashを提案します。
最初にハイブリッドハレー二等分アルゴリズムを導入し、$ \ alpha $ -Entmax変換を計算するために必要な反復回数が7倍減少しました。
次に、カスタムトリトンカーネルを実装して、適応性のあるスパースを効率的に処理します。
テキスト分類および単一ベクトル検索のためのRobertaおよびModernBertとの実験と、言語モデリングのGPT-2とともに、既存の$ \ Alpha $ -EntMaxの実装と比較して、ランタイムとメモリ効率の大幅な改善が達成されることを示しています。
Flashattention-2などの高度に最適化されたSoftMax実装の効率に近づき、場合によっては上回り、タスクのパフォーマンスを強化しながら長いコンテキストトレーニングを可能にします。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google