トークン・スパース・アテンション：インタリーブされたトークン選択による効率的な長文脈推論

要旨：注意の二次計算量は、大規模言語モデルにおける長文コンテキスト推論の中核的なボトルネックであり続けている。先行する加速手法は、構造化されたパターンにより注意マップを疎にするか、特定の層でトークンを恒久的に追い出すかのいずれかである。しかし、前者は無関係なトークンを保持し得るし、後者はトークン重要度の層／ヘッドごとのダイナミクスにもかかわらず、取り返しのつかない初期の判断に依存してしまう可能性がある。本論文では、Token Sparse Attention（トークン疎注意）を提案する。これは、注意中に各ヘッドの Q, K, V を縮小したトークン集合に圧縮し、その後出力を元のシーケンスへと解凍する、軽量かつ動的なトークンレベル疎化メカニズムである。これにより、次の層でトークン情報を再検討できるようになる。さらに、Token Sparse Attention は、トークン選択と疎注意の交点に新しい設計上の論点をもたらす。提案手法は Flash Attention を含む密な注意の実装と完全に互換であり、既存の疎注意カーネルとシームレスに組み合わせられる。実験結果は、Token Sparse Attention が一貫して精度とレイテンシのトレードオフを改善し、128K コンテキストにおいて精度低下 1% 未満で最大 imes3.23 の注意速度向上を達成することを示している。これらの結果は、動的でインターリーブ（交互挿入）されたトークンレベル疎化が、スケーラブルな長文コンテキスト推論のための補完的かつ有効な戦略であることを示している。