LISA: 効率的な長文脈推論のための線形インデックス付きスパース注意機構
長い思考連鎖推論モデルにおける自己注意の二次複雑性問題に対処するため、本論文ではLISA（線形インデックス付きスパース注意）を提案する。これはプラグアンドプレイで動作し、線形注意とLightning Indexerを並列に統合し、ゲート機構で融合することで推論複雑性をO(n²)からO(nM)に削減する。DeepSeek蒸留Qwenモデルでの実験では、16Kトークンコンテキストで50%の推論高速化、AIMEやMATH-500などのベンチマークで平均5.6%の性能向上を達成した。
近年、DeepSeek-R1に代表される長い思考連鎖推論モデルは、テスト時スケーリングのパラダイムの下で推論コンテキスト長をますます長くしており、モデルはより複雑な推論タスクを処理できるようになっている。しかし、標準の自己注意機構はO(n²)の計算複雑性を持つため、長いシーケンスでは推論コストが急激に増大し、長文脈CoT推論の実運用が制限されている。この問題に対処するため、本研究ではLISA（Linear-Indexed Sparse Attention）を提案する。これは、ゼロからの事前学習を必要としないプラグアンドプレイの注意置換モジュールである。LISAの設計は巧妙で効率的である。元のモデル内に2つの軽量コンポーネントを並列に統合する。（1）線形注意モジュールはO(n)の時間複雑性で長距離記憶を提供し、コンテキスト内のグローバルな依存関係を捕捉する。（2）Lightning Indexerは全コンテキストから重要度上位M個のトークンを選択してスパース自己注意に入力する。これら2つのブランチはゲート機構によって融合され、nトークンを生成する推論複雑性をO(n²)からO(nM)（M << n）に削減する。訓練は2段階のパイプラインで行われる。第1段階では、線形注意を統合して長距離依存関係を捕捉し、スライディングウィンドウ注意機構と組み合わせて知識蒸留により最適化し、凍結された教師モデルの完全な自己注意分布に近似させる。第2段階では、Indexerを導入して静的なスライディングウィンドウ機構を置き換え、より広いコンテキストからの動的なトークン選択を可能にする。Indexerは新しいヘッド別KLダイバージェンス損失を用いて訓練され、その選択動作を教師モデルの注意パターンに合わせる。DeepSeek蒸留Qwenモデルを用いた実験では、LISAは16Kトークンコンテキストで50%の推論高速化を達成し、AIMEやMATH-500などの推論ベンチマークで平均5.6%の性能向上を示した。この研究は、長文脈推論の効率的な展開に新たな解決策を提供し、長い思考連鎖推論モデルの実世界応用を促進することが期待される。