Abstract
標準的な注意機構における二次的な計算複雑性は、長い文脈（ロングコンテキスト）のシナリオにおいて、LLMの厳しいスケーラビリティのボトルネックとなっています。Full Attention（FA）とSparse Attention（SA）を組み合わせたハイブリッド注意機構は潜在的な解決策を提供しますが、既存手法は通常、異なるタスクに応じて変動する検索要求を考慮できない静的な割り当て比率に依存しています。さらに、ヘッド単位の動的スパース性はしばしば、深刻な計算負荷の不均衡や同期の長いテールを引き起こし、自己回帰的デコーディング中のハードウェア加速を妨げます。このギャップを埋めるために、本研究では層レベルで注意計算を動的に最適化する、文脈対応型の枠組みであるFlux Attentionを提案します。凍結済みの事前学習済みLLMに軽量なLayer Routerを統合することで、本手法は入力コンテキストに基づき、各層をFAまたはSAへ適応的にルーティングします。この層ごとのルーティングにより、高い忠実度での情報検索を維持しつつ、連続的なメモリアクセスを確保することで、理論上の計算削減を実際のウォールクロック速度向上へと変換します。パラメータ効率の高いアプローチとして、本枠組みは8\timesA800 GPUでの学習をわずか12時間のみ必要とします。複数の長文脈および数学的推論ベンチマークにまたがる大規模な実験により、Flux Attentionはベースラインモデルと比べて、推論性能と推論速度の優れたトレードオフを達成しており、プリフィル段階とデコード段階の両方で最大2.8\timesおよび2.0\timesの速度向上を示します。
