OjaKV动态压缩大模型注意力缓存,提升长文本推理效率。
OjaKV: Context-Aware Online Low-Rank KV Cache Compression
- 分层存储:保留关键首尾token全精度,中间token用在线更新的低秩投影压缩。
- 在32K上下文下压缩比达16倍,零样本准确率不降反升。
- 无需微调,可直接接入FlashAttention,适合长文本生成场景。
大型语言模型的长上下文能力受限于自回归生成所需的键值(KV)缓存带来的显著内存瓶颈。例如,处理32K token提示、批量大小为4时,Llama-3.1-8B模型的KV缓存需约16GB内存,超过模型权重本身。尽管低秩投影压缩是可行方向,但现有方法依赖静态离线学习的子空间,在数据分布变化时表现不佳。为此,我们提出OjaKV,结合策略性混合存储与在线子空间自适应。首先,识别出并非所有token同等重要,保留关键首尾token的完整秩,作为注意力的高保真锚点。其次,对大多数中间token,采用增量式更新投影基底的低秩压缩,利用Oja算法进行在线主成分分析。该更新包括预填充阶段的全面更新和解码阶段的轻量周期更新,确保子空间始终与上下文演变同步。关键的是,框架与FlashAttention等现代注意力模块完全兼容。实验表明,OjaKV在高压缩比下维持甚至提升零样本准确率,尤其在需要复杂推理的超长上下文基准上表现最优,凸显在线子空间自适应的重要性。结果证明,该混合框架是一种无需模型微调的实用、即插即用的长上下文推理内存优化方案。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。