Xinghao Wang Pengyu Wang Dong Zhang Chenkun Tan Shaojun Zhou Zhaoxiang Liu Shiguo Lian Fangxu Liu Kai Song Xipeng Qiu

要約

大規模言語モデル（LLM）のコンテキスト長を拡張することは、大きな利点をもたらす一方で、計算コストが非常に高くなる。このコストの主な原因は、シーケンス長に対してO(N²)の計算複雑性を持つ自己注意機構であり、メモリ使用量およびレイテンシの面で大きなボトルネックとなっている。幸いなことに、注意行列は特に長いシーケンスにおいてしばしばスパースであるため、最適化の余地がある。ブロックスパース注意（Block-sparse attention）は、シーケンスをブロックに分割し、一部のブロックに対する計算をスキップすることで、そのような問題に対応する有望な手法として登場している。しかし、この手法の効果は、元となる注意パターンに強く依存しており、ブロックレベルでのスパース性が最適にならない場合がある。例えば、あるブロック内のクエリに対する重要なキートークンが多数の他のブロックに散在している場合、計算の重複が生じる。本研究では、注意機構の置換性（permutation properties）を活用してブロックレベルのスパース性を向上させ、LLMのプレフィル（prefilling）段階における計算効率を高める、即座に統合可能な「置換ブロックスパース注意（Permuted Block-Sparse Attention, PBS-Attn）」を提案する。我々は、実世界の長文コンテキストデータセットを用いた包括的な実験を行い、PBS-Attnが既存のブロックスパース注意手法を常に上回るモデル精度を達成するとともに、完全注意（full attention）ベースラインとほぼ同等の性能を示した。独自に開発した置換型FlashAttentionカーネルを活用することで、PBS-Attnは長文コンテキストのプレフィル処理において、エンドツーエンドで最大2.75倍の高速化を実現し、実用的な有効性を確認した。コードは以下のURLから公開されている：https://github.com/xinghaow99/pbs-attn