DoSA、近似なしでスパース注意のTop-K選択を高速化

概要

スパース注意の狙いは、長いコンテキストから関連性の高い少数のキーだけを参照し、推論コストを抑えることです。しかし実際のシステムでは、必要なキーを見つけるためのIndexerが別のボトルネックになる可能性があります。候補となるすべてのキーについて、複数のIndexerヘッドのスコアを計算してからTop-Kを選ぶ構成では、最終的に少数しか使わない場合でも探索段階の負荷が残ります。

DoSA（Dots Sparse Attention）は、この探索を条件付きにするロスレスな方法です。候補をすべて完全評価するのではなく、どれほど良いスコアを仮定しても現在のTop-K境界を越えられない候補を、完全計算の前に除外します。

中核となる仕組み

• 正確な値ではなく、棄却に十分な証拠を使う：候補の取り得る最大スコアが、現在の選択閾値を下回ると証明できれば、正確なスコアを求める必要はありません。
• 種集合から閾値の下界を得る：まず種集合を使って、現在の勝者が与える閾値の下界を作ります。この下界がスクリーニングの基準になります。
• 正負の寄与を別々に評価する：非負のヘッド投影と残差の共同上界によって正の寄与を抑え、負のヘッドについては活性値に基づく集約上界で減算量の大きさを制限します。これらを組み合わせ、候補がまだ選択対象になり得るかを判定します。
• 完全計算後の処理も削減する：残った候補を正確にスコアリングした後も、厳密なスコアによるフィルタリングで書き戻しと最終選択の負荷を減らします。
• Top-Kの意味を変更しない：論文はこの証明手続きと、元の計算と同じTop-Kを得る条件を定式化しています。近似検索による速度向上ではありません。

結果と意義

GPU実装は英語、中国語、コードのワークロードで評価されています。報告された結果では、完全なIndexerがSGLangに対して最大2.74倍、LiteTopKに対して最大1.69倍高速化しました。エンドツーエンドでは、GLM-5.2で時間到達初回トークンが最大1.94倍、DeepSeek-V3.2で最大1.50倍改善しています。削減対象はスコア計算だけでなく、候補の書き戻しやTop-K選択の処理にも及びます。

DoSAの重要な点は、スパース性を最終的な読み出しだけでなく、候補発見の段階にも適用したことです。最終的に少数のトークンしか参照しないなら、その少数を探す工程でも、結果に影響しない計算を避けられる可能性があります。

一方、効果は常に一定とは限りません。論文は、表現の幾何学的な関係と選択候補間のマージンが、証明の有効性を左右すると分析しています。上界が十分に鋭く、選択された候補と棄却候補の差が大きいほど、早期除外を適用しやすくなります。報告された速度向上は、評価された実装、モデル、ワークロードにおける最大値です。それでも、Top-Kの正確さを維持しながら「不要であることを証明してから計算する」という設計は、長文脈推論の実装最適化に有力な方向を示しています。

出典：Hugging Face Daily Papers