OjaKV: コンテキストに意識したオンラインで低ランクな KV Cache 圧縮

arXiv:2509.21623v2 Announce Type: replace-cross Abstract: 大規模言語モデルの長文コンテキスト対応は、自己回帰生成に必要不可欠なキー値 (KV) Cache による重要なメモリボトルネックに制約されています。このボトルネックは重大です；例えば、Llama-3.1-8B モデルが 32K トークンのプロンプトをバッチサイズ 4 で処理する場合、その KV Cache には約 16GB が必要となり、これはモデル自体の重みを超えます。低ランク投影による KV Cache 圧縮は有望な方向ですが、既存の方法は静的かつオフラインで学習された基底空間に依存しており、データ分布の変化に対してパフォーマンスが低下します。これらの限界を超えるため、戦略的混合保存政策とオンライン基底空間適応を組み合わせた新たなフレームワークである OjaKV を提案します。まず、OjaKV はすべてのトークスが等しく重要ではないことを認識しており、重要な最初のトークンと最新のトークンをフルランクで保持することで、注意機構において高い忠実度のアンカーを維持します。次に、ほとんどの中間トークンについては、Oja のアルゴリズムを用いたオンライン主成分分析によって投影基底を漸動的に適応させることで、低ランク圧縮を適用します。この適応はプロンプトのプリフイルリング中に包括的な更新を行い、デコーディング中は軽量な定期的な更新を行い、基底空間が変化し続けるコンテキストに常に同期することを保証します。至关重要的是、私たちのフレームワークは FlashAttention など現代的な注意モジュールと完全に互換性があります。実験结果表明、OjaKV は高圧縮率においてゼロショットの精度を維持し、あるいは改善しています。特に、OjaKV は複雑な推論を必要とする非常に長いコンテキストベンチマークにおいて最も顕著な性能向上を達成しており、オンライン基底空間適応が動的にコンテキストの変化を追跡する重要性を強調しています。これらの結果は、モデルのファインチューニングを伴わないまま、メモリ効率の高い長文コンテキスト推論のための実用的なプラグアンドプレイソリューションである我々のハイブリッドフレームワーク確立しました。

arXiv:2509.21623v2 Announce Type: replace-cross Abstract: The expanding long-context capabilities of large language models are constrained by a significant memory bottleneck: the key-value (KV) cache required for autoregressive generation. This bottleneck is substantial; for instance, a Llama-3.1-8B model processing a 32K-token prompt at a batch size of 4 requires approximately 16GB for its KV cache, a size exceeding the model's weights. While KV-cache compression via low-rank projection is a promising direction, existing methods rely on a static, offline-learned subspace that performs poorly under data distribution shifts. To overcome these limitations, we introduce OjaKV, a novel framework that integrates a strategic hybrid storage policy with online subspace adaptation. First, OjaKV recognizes that not all tokens are equally important for compression; it preserves the crucial first and most recent tokens in full-rank, maintaining high-fidelity anchors for attention. Second, for the vast majority of intermediate tokens, it applies low-rank compression by incrementally adapting the projection basis using Oja's algorithm for online principal component analysis. This adaptation involves a comprehensive update during prompt prefilling and lightweight periodic updates during decoding, ensuring the subspace remains aligned with the evolving context. Crucially, our framework is fully compatible with modern attention modules like FlashAttention. Experiments demonstrate that OjaKV maintains or even improves zero-shot accuracy at high compression ratios. In particular, OjaKV achieves its strongest gains on very long-context benchmarks that require complex reasoning, highlighting the importance of online subspace adaptation in dynamically tracking context shifts. These results establish our hybrid framework as a practical, plug-and-play solution for memory-efficient long-context inference without requiring model fine-tuning.