Abstract
スパース・アテンションは、長い文脈（ロングコンテキスト）の学習における中核的なボトルネックである、トランスフォーマの二次コストを軽減する手段として提案されてきました。有望な研究の流れとして、
alpha-entmax attention（
alpha-エントマックス・アテンション）があります。これは、softmax の微分可能なスパース代替であり、入力に依存したスパース性を可能にする一方で、正規化項 au を計算するために必要な計算オーバーヘッドのため、softmax に後れを取ってきました。本論文では、この制約を解決する AdaSplash-2 を導入します。新規のヒストグラムベースの初期化によって、 au を計算するのに必要な反復回数を、典型的には 1〜2 回にまで削減します。要点は、注意スコアの粗いヒストグラムをその場で計算し、オンチップ SRAM に保存することで、より正確な初期化を実現し、迅速な順伝播および逆伝播の計算を可能にする点です。さらに、ゼロブロックを低オーバーヘッドでスキップする、スパース性を考慮した GPU 実装と組み合わせることで、AdaSplash-2 は、ブロックスパース性が中程度〜高い場合（例：>60%）に、FlashAttention-2 と比べてステップあたりの学習時間を一致または改善します。このような状況は、長い文脈長でしばしば発生します。下流タスクにおいては、効率的な
alpha-entmax attention で学習したモデルは、短い文脈長では softmax のベースラインと同等の性能を達成し、長い文脈設定では大幅な向上を実現します。
