OjaKV: 上下文感知的在线低秩KV缓存压缩

# OjaKV：上下文感知的在线低秩KV缓存压缩
来源：https://arxiv.org/html/2509.21623
Yuxuan Zhu¹, David H. Yang¹\*, Mohammad Mohammadi Amiri¹, Keerthiram Murugesan², Tejaswini Pedapati², Pin-Yu Chen²  
¹伦斯勒理工学院 ²IBM研究院  
\{yangd13, zhuy27, mamiri\}@rpi.edu [email protected] \{keerthiram.murugesan, pin-yu.chen\}@ibm.com  

###### 摘要
大型语言模型不断扩展的长上下文能力受到一个显著的内存瓶颈的制约：自回归生成所需的键值（KV）缓存。这一瓶颈十分显著；例如，处理一个32K token的提示词、批大小为4的Llama-3.1-8B模型，其KV缓存需要约16 GB，这一大小超过了模型本身的权重。虽然通过低秩投影进行KV缓存压缩是一个有前景的方向，但现有方法依赖静态的、离线学习的子空间，在数据分布变化时表现不佳。为克服这些限制，我们提出OjaKV，一个融合了策略性混合存储策略与在线子空间自适应的新型框架。首先，OjaKV认识到并非所有token对压缩同等重要；它保留关键token的全秩形式，为注意力机制提供高保真锚点。其次，对于大多数中间token，它采用低秩压缩，通过Oja算法在线主成分分析逐步调整投影基。该过程在预填充阶段进行一次全面更新，并在解码阶段进行轻量级周期性更新，确保子空间始终与不断演变的上下文对齐。至关重要的是，我们的框架完全兼容现代注意力模块如FlashAttention。实验表明，在高压缩比下，OjaKV能保持甚至提升准确率。特别地，OjaKV在需要复杂推理的极长上下文基准测试中取得了最显著的增益，凸显了在线子空间自适应在动态追踪上下文变化中的重要性。此外，我们的方法与token选择方法兼容，可实现复合内存节省。这些结果确立了我们的混合框架作为一种实用、即插即用的解决方案，用于实现无需模型微调的内存高效长上下文推理。代码已开源：https://github.com/zzbright1998/OjaKV

OjaKV：上下文感知的在线低秩KV缓存压缩
Yuxuan Zhu¹†††*表示同等贡献。*, David H. Yang¹\*, Mohammad Mohammadi Amiri¹, Keerthiram Murugesan², Tejaswini Pedapati², Pin-Yu Chen²  
¹伦斯勒理工学院 ²IBM研究院  
\{yangd13, zhuy27, mamiri\}@rpi.edu [email protected] \{keerthiram.murugesan, pin-yu.chen\}@ibm.com  

## 1 引言
大型语言模型（LLM），如GPT-4o（OpenAI et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib6)）和Deepseek-R1（DeepSeek-AI et al., 2025 (https://arxiv.org/html/2509.21623#bib.bib7)），在包括编程（Nam et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib8)）、数学（Setlur et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib9)）和开放式文本生成（Kumichev et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib10)）等多个领域展现出卓越性能。然而，随着模型能力和上下文长度的扩展，GPU内存成为推理的关键瓶颈。内存占用主要源于两个方面：(i) 模型权重，像Llama-3.1-8B这样的模型本身就需要16 GB；(ii) 提示词预填充和自回归解码过程中使用的键值（KV）缓存。例如，以float16精度处理一个32K token的提示词、批大小为4的Llama-3.1-8B，其KV缓存额外消耗16 GB，与模型权重本身的大小相当。这种巨大的内存消耗使得长上下文推理在除高端硬件外的所有设备上都难以实现。为缓解这一挑战，已有多种方法被提出以优化KV缓存内存使用（Shi et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib11)）。这些方法可分为四类：(1) **量化**，以较低精度（如8位）存储键和值（Hooper et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib12); Liu et al., 2024c (https://arxiv.org/html/2509.21623#bib.bib13)）；(2) **Token选择**，基于注意力分数或启发式显著性度量修剪或合并被认为不重要的token（Xiao et al., 2023 (https://arxiv.org/html/2509.21623#bib.bib14); Li et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib15); Zhang et al., 2023 (https://arxiv.org/html/2509.21623#bib.bib16)）；(3) **卸载**，将KV缓存传输到CPU内存，在解码过程中选择性回传（Tang et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib18); Sun et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib19); Zhu et al., 2025 (https://arxiv.org/html/2509.21623#bib.bib17)）；以及(4) **低秩近似**，将键和值投影到更低维的子空间（Saxena et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib20); Lin et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib21)）。我们的工作专注于第四个方向。通过将每个键和值向量从维度\(d\)（例如\(d=128\)）压缩到\(r\)（例如\(r=96\)），低秩方法可以减少\(((1-r/d) \times 100)\%\)的缓存内存，同时保持模型准确率，为长上下文推理带来显著的节省。虽然token选择已成为广泛采用的策略，但我们提供了理论支持，表明低秩投影和token驱逐是兼容的，它们结合使用时能产生复合收益，实现更大的内存减少。

现有的低秩方法主要分为两类。(1) **权重分解**技术直接将查询、键和值的线性投影权重（\( \boldsymbol{W}_q, \boldsymbol{W}_k, \boldsymbol{W}_v \)）分解为低秩矩阵，从而缓存已压缩的中间状态（Chang et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib22)）。然而，这种方法往往会导致明显的准确率下降。(2) **基于投影**的技术从校准数据集中学习固定的正交投影基（\( \boldsymbol{U}_q, \boldsymbol{U}_k, \boldsymbol{U}_v \)）。这些基随后用于压缩KV缓存，在注意力计算时进行重建（Saxena et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib20); Lin et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib21)）。虽然有效，但这些静态基隐含地假设推理提示词遵循与校准数据相同的分布。在实践中，分布转移（例如从对话到代码生成）会导致近似效果恶化，损害生成质量。

为解决这些限制，我们提出OjaKV，一个基于两个核心原则运行的KV缓存压缩新框架。我们的第一个关键见解是，对所有token进行均匀压缩并非最优。受注意力汇聚（attention sinks）（Xiao et al., 2023 (https://arxiv.org/html/2509.21623#bib.bib14)）和SnapKV（Li et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib15)）发现的启发，OjaKV采用一种混合存储策略，策略性地将那些具有高重建误差的关键token从低秩投影中排除。这保留了它们的全秩保真度，为注意力机制创建稳定的锚点，并形成一个显著更强的性能基线。其次，对于剩余的中间token，我们采用Oja的增量主成分分析（PCA）（Oja, 1997 (https://arxiv.org/html/2509.21623#bib.bib23)）进行在线子空间自适应。该机制在预填充阶段对选择出的显著token进行一次全面更新，随后在解码阶段执行轻量级周期性更新。这确保了低秩基能够以极小的开销持续适应不断演变的上下文。我们的框架完全兼容现代注意力模块，如FlashAttention（Dao et al., 2022 (https://arxiv.org/html/2509.21623#bib.bib24)），确保了在实际长上下文推理中的实用性。

我们在lm-eval-harness（Biderman et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib25)）中的多项选择基准上评估OjaKV，与先前研究保持一致（Saxena et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib20); Lin et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib21)）。此外，据我们所知，我们首次在面向生成的長上下文任务上评估基于投影的低秩KV缓存压缩方法的性能，使用了LongBench（Bai et al., 2023 (https://arxiv.org/html/2509.21623#bib.bib26)）和RULER（Hsieh et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib27)）。我们还在长推理和生成任务AIME（Balunović et al., 2025 (https://arxiv.org/html/2509.21623#bib.bib38)）上评估了我们的方法。在所有设置下，OjaKV在等效压缩比下均表现出优于静态低秩基线的性能。总之，我们的贡献有三点。第一，我们引入了一个混合低秩KV缓存压缩框架OjaKV，它结合了选择性全秩存储策略与上下文感知的在线子空间自适应。第二，我们的设计与FlashAttention兼容，确保了对现代长上下文推理管道的实用性。第三，我们首次在具有挑战性的面向生成的基准上对低秩KV压缩进行了全面评估，超越了简单的多项选择任务。

## 2 相关工作
近期关于减少LLM推理内存占用的工作催生了各种KV缓存压缩策略，包括量化、token修剪、卸载和子空间压缩。我们的方法建立在低秩近似的基础上，并扩展为在线、上下文自适应的形式。下面我们回顾相关方法。

### 2.1 KV缓存压缩
KV缓存的内存开销推动了四种主要压缩策略（Shi et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib11)）。(1) **量化**：这些方法降低存储精度（例如降至4位或2位）。如KVQuant（Hooper et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib12)）和KIVI（Liu et al., 2024c (https://arxiv.org/html/2509.21623#bib.bib13)）在最小质量损失下实现了显著压缩。(2) **Token选择**：该类方法仅保留必要token。StreamingLLM（Xiao et al., 2023 (https://arxiv.org/html/2509.21623#bib.bib14)）保留“注意力汇聚”，而SnapKV（Li et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib15)）基于重要性分数选择token。(3) **卸载**：诸如Quest（Tang et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib18)）和ShadowKV（Sun et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib19)）等方法将缓存存储在CPU内存中，在计算时选择性重载相关部分。(4) **低秩近似**：这些方法将键和值投影到更低维的子空间（Saxena et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib20); Lin et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib21)）。我们的方法属于最后一类，并且与上述其他方法**正交**，可实现叠加的内存节省。附录A.9 (https://arxiv.org/html/2509.21623#A1.SS9)提供了分析及实验，表明我们的方法可以与基于token驱逐的方法结合使用。

### 2.2 注意力机制的低秩近似
低秩结构被广泛应用于压缩。Palu（Chang et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib22)）和ReCalKV（Yan et al., 2025 (https://arxiv.org/html/2509.21623#bib.bib28)）通过分解权重来缓存压缩状态，但通常伴随准确率下降。EigenAttention（Saxena et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib20)）和MatryoshkaKV（Lin et al., 2024 (https://arxiv.org/html/2509.21623#bib.bib21)）使用从校准数据导出的基（\(\boldsymbol{U_k}, \boldsymbol{U_v}\)）来投影激活。然而，这些方法依赖**静态**基，在分布转移下性能会下降。虽然像多头部潜在注意力（Multi-Head Latent Attention）（Liu et al., 2024a (https://arxiv.org/html/2509.21623#bib.bib36),b (https://arxiv.org/html/2509.21623#bib.bib37)）这样的架构原生集成了低秩结构，但需要训练新模型。我们的方法通过**在线子空间自适应**解决了这些限制，在推理过程中持续更新投影矩阵，提供即插即用的、上下文感知的近似。

### 2.3 在线主成分分析
在自回归解码中，对完整历史重新运行SVD在计算上是不可行的。在线PCA算法通过增量更新子空间来解决这个问题。扰动技术（Gu and Eisenstat, 1994 (https://arxiv.org/html/2509.21623#bib.bib29)）提供了准确性，但对于高维数据流来说资源消耗过大。增量SVD（Brand, 2002 (https://arxiv.org/html/2509.21623#bib.bib30)）提供了一个折中方案，但对于实时LLM推理而言仍然成本高昂。相比之下，像Oja规则（Oja, 1997 (https://arxiv.org/html/2509.21623#bib.bib23)）这样的随机方法通过梯度更新高效地优化方差。我们提出的方法OjaKV建立在此基础之上。通过应用Oja规则动态更新键和值的子空间，它实现了快速、自适应且无需校准的近似。

## 3 预备知识：低秩注意力
我们首先描述标准注意力机制以及它如何适用于低秩KV缓存压缩。为简单起见，我们关注单个注意力头。设头部维度为\(d_h\)，输入序列为\(\boldsymbol{X} \in \mathbb{R}^{n \times d}\)。标准注意力首先使用投影矩阵\(\boldsymbol{W}_q, \boldsymbol{W}_k, \boldsymbol{W}_v \in \mathbb{R}^{d \times d_h}\)将输入投影为查询、键和值表示：
\[
\boldsymbol{Q} = \boldsymbol{X}\boldsymbol{W}_q, \quad \boldsymbol{K} = \boldsymbol{X}\boldsymbol{W}_k, \quad \boldsymbol{V} = \boldsymbol{X}\boldsymbol{W}_v,
\]
其中\(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V} \in \mathbb{R}^{n \times d_h}\)。KV缓存存储\(\boldsymbol{K}\)和\(\boldsymbol{V}\)供解码时使用。注意力分数\(\boldsymbol{A}\)随后计算为查询与键之间的缩放点积，接着进行softmax操作。注意力头的输出\(\boldsymbol{O}\)通过将这些注意力分数应用于值矩阵\(\boldsymbol{V}\)来计算：
\[
\boldsymbol{A} = \text{softmax}\left( \boldsymbol{Q}\boldsymbol{K}^T / \sqrt{d_h} \right) \in \mathbb{R}^{n \times n}, \quad \boldsymbol{O} = \boldsymbol{A}\boldsymbol{V} \in \mathbb{R}^{n \times d_h}.
\]

### 3.1 低秩KV缓存近似
低秩近似的核心思想是将\(\boldsymbol{K}\)和\(\boldsymbol{V}\)投影到更低维的子空间。我们定义两个正交基：用于键和查询的\(\boldsymbol{U}_k \in \mathbb{R}^{d_h \times r_k}\)，以及用于值的\(\boldsymbol{U}_v \in \mathbb{R}^{d_h \times r_v}\)，其中\(r_k, r_v \ll d_h\)是压缩所需的目标秩。这些基满足\(\boldsymbol{U}_k^{\mathsf{T}}\boldsymbol{U}_k = \boldsymbol{I}_{r_k}\)和\(\boldsymbol{U}_v^{\mathsf{T}}\boldsymbol{U}_v = \boldsymbol{I}_{r_v}\)。低秩基\(\boldsymbol{U}_k\)和\(\boldsymbol{U}_v\)从小型校准数据集初始化（详见附录A.3 (https://arxiv.org/html/2509.21623#A1.SS3)）。