OjaKV：文脈対応型オンライン低ランクKVキャッシュ圧縮

Abstract

大規模言語モデルの拡張長コンテキスト能力は、自己回帰生成に必要なキー・バリュー（KV）キャッシュという重要なメモリのボトルネックによって制約されています。このボトルネックは非常に大きく、例えば、Llama-3.1-8Bモデルがバッチサイズ4で32Kトークンのプロンプトを処理する場合、そのKVキャッシュには約16GBが必要であり、これはモデルの重みのサイズを上回ります。低ランク射影によるKVキャッシュ圧縮は有望な方向性ですが、既存手法は静的でオフライン学習された部分空間に依存しており、データ分布の変化に対して性能が低下します。これらの制約を克服するために、我々は戦略的なハイブリッド記憶ポリシーとオンライン部分空間適応を統合する新しい枠組みであるOjaKVを提案します。まず、OjaKVは、圧縮において重要なトークンはすべて同等ではないことを認識し、重要な最初および直近のトークンをフルランクのまま保持することで、注意のための高忠実度なアンカーを維持します。次に、圧倒的多数の中間トークンに対しては、Ojaのアルゴリズムを用いたオンライン主成分分析により射影基底を逐次的に適応させることで、低ランク圧縮を適用します。この適応は、プロンプトのプレフィリング中に包括的な更新を行い、復号中には軽量な周期的更新を行うことで構成され、部分空間が変化し続けるコンテキストに整合した状態を保つことを保証します。重要な点として、我々の枠組みはFlashAttentionのような最新の注意モジュールと完全に互換です。実験の結果、OjaKVは高い圧縮率においてゼロショット精度を維持、あるいは改善することが示されています。特にOjaKVは、複雑な推論を必要とする非常に長いコンテキストのベンチマークで最大の改善を達成しており、コンテキストの変化を動的に追跡する上でのオンライン部分空間適応の重要性を強調しています。これらの結果は、モデルの微調整を不要とし、メモリ効率の高い長コンテキスト推論のための実用的なプラグアンドプレイ解として、我々のハイブリッド枠組みを確立するものです。