DepthWeave-KV：用于长上下文KV缓存压缩的令牌自适应跨层残差分解方法

# DepthWeave-KV：面向长上下文的令牌自适应跨层残差因子化键值缓存压缩
来源：https://arxiv.org/html/2607.06523
Adam Puente Tercero、Nerea Angulo Hijo、Mar Linares Tercero、Julia Barrientos、Ainhoa Miranda、Jesús Olivera

###### 摘要

长上下文语言模型推理日益受到存储键值缓存所需内存带宽和容量的限制，然而现有的压缩方法通常对各个层或令牌施加统一的预算，当词汇线索和语义状态需要不同的保留程度时，会降低检索性能。我们引入了 DepthWeave-KV，一种令牌自适应的缓存压缩方法，它通过共享的低秩通道基将相邻 Transformer 层的键和值状态进行因子化，同时在注意力行为敏感的令牌处保留轻量级的令牌特定残差。DepthWeave-KV 结合了跨深度残差因子化与一个令牌条件深度路由器，该路由器为携带指令和检索关键的令牌分配更高的重建秩，并使用来自注意力输出探针的无校准在线误差跟踪来在生成过程中自适应地调整压缩，无需重新训练基础模型。一个融合的 CUDA 实现联合执行基查找、残差反量化和注意力投影，以减少解码时的内存流量。在 LongBench、Needle-in-a-Haystack、L-Eval 以及长文本问答和摘要基准测试上，DepthWeave-KV 在显著降低内存使用的同时，达到了接近全缓存的任务质量，在先前压缩缓存的基础上提高了平均分数和检索准确率，同时实现了 8.3 倍的 KV 内存缩减和 64K 上下文下每秒 72.8 个令牌的生成速度。

机器学习，ICML

## 1 引言

随着提示从短暂的交互式轮次扩展到文档、对话、代码库和检索增强上下文，长上下文语言模型越来越依赖于高效的推理系统。在此情况下，键值（KV）缓存通常成为解码期间占主导地位的内存驻留状态，限制了批大小、上下文长度和服务吞吐量。先前的工作通过令牌驱逐、令牌合并、量化、低秩注意力和层间缓存共享来减轻这一负担（Li et al., 2024 (https://arxiv.org/html/2607.06523#bib.bib29); Jiang et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib24); Saxena et al., 2024 (https://arxiv.org/html/2607.06523#bib.bib51); Xiang et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib60); Yang et al., 2024a (https://arxiv.org/html/2607.06523#bib.bib66); Yao et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib69)）。然而，长上下文任务暴露了一个持久的矛盾：在局部延续统计下看似不重要的令牌可能在后续检索中变得至关重要，而编码浅层词汇锚点和深层语义抽象的不同 Transformer 层并不需要相同的缓存保真度。因此，对令牌或深度进行统一压缩可能会导致在逐针检索、多跳问答和长文档摘要等任务上出现脆弱的失败。

深度维度的 KV 压缩是一个特别有前景的方向，因为相邻 Transformer 层通常包含相关的缓存结构。MiniCache（Liu et al., 2024a (https://arxiv.org/html/2607.06523#bib.bib31)）表明，利用深度维度的冗余可以在保留大部分模型行为的同时减少 KV 内存。然而，仅靠深度共享对于异构的长上下文工作负载是不够的。与一个指令分隔符、实体提及、引用跨度或承载答案的句子相关联的缓存条目可能需要比一个流畅的延续令牌更忠实的重建；同样，同一个令牌在早期的词汇层和较晚的语义层可能需要不同的处理。最近的研究还强调，激进的 KV 压缩可能会以平均困惑度或短上下文基准测试无法察觉的方式降低检索行为（Chen et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib9); Haverbeck et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib21); Bui et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib7)）。这些观察结果启发了一种同时具备深度感知、令牌自适应并且与注意力行为的在线证据相关联的压缩机制。

我们引入了 DepthWeave-KV，一种用于长上下文 KV 缓存压缩的令牌自适应跨层残差因子化方法。DepthWeave-KV 不是为每一层存储独立的 KV 张量，而是为每个注意力头在相邻 Transformer 层之间编织一组紧凑的共享低秩通道基。层特定的键和值从这些共享基以及稀疏的、令牌特定的残差分量中重建，这些残差分量的秩被动态分配。核心设计原则是，大多数令牌可以共享一个深度局部缓存表示，而检索关键和指令承载的令牌则保留更高保真度的残差，以防止注意力输出发生漂移。

本文做出四项贡献。首先，我们引入了跨深度残差因子化，它为每个注意力头存储共享的基通道，并通过轻量级学习到的残差门来重建层特定的键/值状态。其次，我们添加了一个令牌条件深度路由器，它为显著令牌分配更高的重建秩，同时激进地压缩低显著性延续令牌。第三，我们提出了使用注意力输出探针的无校准在线误差跟踪，使得压缩比能够在生成过程中自适应调整，而无需重新训练基础 LLM。第四，我们提供了一个融合的 CUDA 实现，用于基查找、残差反量化和注意力投影，减少了长上下文工作负载下的解码时间内存流量。

我们在包括 LongBench、Needle-in-a-Haystack、L-Eval、NarrativeQA、Qasper、HotpotQA、MultiFieldQA-en、GovReport、QMSum 和 TriviaQA 在内的已建立的长上下文基准测试上评估了 DepthWeave-KV。在检索、问答和摘要设置中，将 DepthWeave-KV 与完整 KV 缓存及代表性压缩基线（如 StreamingLLM、H2O、SnapKV、PyramidKV、MiniCache、KVSharer、ChunkKV、TailorKV 和 Eigen Attention）进行了比较。结果显示，在内存-质量权衡方面有了显著改善，具有强大的逐针检索行为、降低的注意力输出重建误差，以及长上下文工作负载下改进的解码效率。

参见图注图 1：DepthWeave-KV 概述：不是对所有层和令牌进行统一压缩，该方法在相邻层之间编织共享的深度基，并且仅将残差容量分配给那些注意力行为表明检索或指令重要性的令牌。
## 2 相关工作

KV 缓存压缩已成为长上下文推理中的一个核心系统问题，调查指出缓存驻留和内存带宽是服务大型语言模型的主要瓶颈（Li et al., 2024 (https://arxiv.org/html/2607.06523#bib.bib29); Jiang et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib24); Javidnia et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib23)）。一系列广泛的工作通过根据估计的重要性驱逐或合并令牌来减少缓存大小。StreamingLLM 保留最近的令牌和注意力汇聚令牌，而 H2O 风格的重型击中者策略优先考虑具有较大累积注意力质量的令牌；最近的变体通过语义分块、自适应合并和工作负载感知路由改进了长上下文工作负载的令牌选择（Bui et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib7); Xiang et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib60); Wang et al., 2024 (https://arxiv.org/html/2607.06523#bib.bib55); Akulov et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib4); Yu & Chai, 2025a (https://arxiv.org/html/2607.06523#bib.bib70), b (https://arxiv.org/html/2607.06523#bib.bib71)）。ChunkKV 强调在片段粒度上的语义保留（Xiang et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib60)），TailorKV 结合了多种缓存管理策略以应对异构上下文（Yao et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib69)），而 StructKV 在提示中保留结构锚点（Chen et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib10)）。当令牌显著性可以从注意力统计或提示结构中推断时，这些方法是有效的，但它们通常决定哪些令牌保留在缓存中，而不是每个保留的令牌应该获得多少层特定的重建容量。相反，DepthWeave-KV 为所有令牌保留压缩表示，同时分配令牌条件残差秩，这使得低显著性延续令牌和检索关键锚点可以在同一个压缩框架中共存，而无需硬驱逐边界。

另一类方法压缩缓存的键和值的数值或代数表示。量化和精度分配方法通过低比特格式、令牌特定精度或残差向量量化来减少每个条目的存储（Kumar, 2024 (https://arxiv.org/html/2607.06523#bib.bib27); Zhang et al., 2024 (https://arxiv.org/html/2607.06523#bib.bib74); Slothouber, 2026 (https://arxiv.org/html/2607.06523#bib.bib52); Yang et al., 2024c (https://arxiv.org/html/2607.06523#bib.bib68), b (https://arxiv.org/html/2607.06523#bib.bib67)）。诸如 Eigen Attention 之类的低秩方法将注意力计算投影到紧凑子空间中（Saxena et al., 2024 (https://arxiv.org/html/2607.06523#bib.bib51)），而冗余感知方法利用令牌、头或任务之间的重复结构（Cai et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib8); Corallo et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib11); Patel & Joshi, 2026 (https://arxiv.org/html/2607.06523#bib.bib43); Kriuk & Ng, 2026 (https://arxiv.org/html/2607.06523#bib.bib26)）。系统层面的工作进一步表明，压缩会改变内存访问模式，应与服务基础设施协同设计（Qin et al., 2024 (https://arxiv.org/html/2607.06523#bib.bib46); Xie et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib62); Qiao et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib45)）。DepthWeave-KV 对这些方向是互补的：它使用低秩通道基和量化残差，但关键区别在于因子化是显式地跨相邻 Transformer 层组织的，并由注意力输出误差在线调节，而不是仅由全局压缩目标固定。

深度维度和层共享方法最接近我们的工作。MiniCache 证明了 KV 缓存沿层维度存在大量冗余，并且可以通过跨深度共享缓存状态来进行压缩（Liu et al., 2024a (https://arxiv.org/html/2607.06523#bib.bib31); Haffari et al., 2024 (https://arxiv.org/html/2607.06523#bib.bib20); Liu et al., 2024b (https://arxiv.org/html/2607.06523#bib.bib32)）。KVSharer、层间相似性方法、随机 KV 路由以及深度缓存权衡分析同样表明，并非所有层都需要独立的缓存状态（Yang et al., 2024a (https://arxiv.org/html/2607.06523#bib.bib66); Ma et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib36); Filippova et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib15); Wang, 2026 (https://arxiv.org/html/2607.06523#bib.bib54)）。SpindleKV 和 HCAttention 进一步论证，在激进压缩下，浅层和深层应被区别对待（Tang et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib53); Yang et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib64), 2026 (https://arxiv.org/html/2607.06523#bib.bib65)）。DepthWeave-KV 直接建立在 MiniCache 的深度冗余洞察之上，但用跨深度残差因子化取代了离散的缓存共享：共享基捕获相邻层之间的共同分量，而令牌特定的残差门仅在需要时恢复层局部信息。这种设计避免了将整个令牌或层组强制置于单一保真度级别。

最后，最近的一些研究警告，KV 压缩可以保持平均困惑度，但同时损害检索、推理和指令遵循（Chen et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib9); Haverbeck et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib21); An et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib5)）。ReST-KV 通过层间输出重建和时间平滑来解决这个问题（An et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib5)），而推理和训练导向的工作则研究模型何时自然可压缩或可以变得更缓存高效（Gelberg et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib18); Cai et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib8)）。其他专门的压缩设置包括扩散 LLM 缓存、多模态模型、视觉自回归生成以及通过缓存状态注入知识（Nguyen-Tri et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib41); Zhang et al., 2025 (https://arxiv.org/html/2607.06523#bib.bib75); Qin et al., 2026 (https://arxiv.org/html/2607.06523#bib.bib47); Pustovit, 2026 (https://arxiv.org/html/2607.06523#bib.bib44)）。这些工作促使我们在短上下文语言建模之外进行评估，并支持了我们使用在线注意力输出探针的方法。与依赖大量校准或重新训练的方法不同，DepthWeave-KV 在生成过程中跟踪重建误差，并在不修改基础模型的情况下调整残差秩。

## 3 方法

DepthWeave-KV 用深度局部因子化替换了逐层 KV 存储，该因子化将跨相邻 Transformer 层共享的缓存分量与令牌特定的残差信息分离开来。对于一个具有层 \(\ell \in \{1,\ldots,L\}\)、注意力头 \(h\)、令牌位置 \(t\) 和头维度 \(d\) 的 Transformer，我们将层划分为大小为 \(w\) 的重叠深度窗口 \(\mathcal{W}_m\)。在每个窗口和头内，DepthWeave-KV 维护共享的基通道 \(B^{K}_{m,h}, B^{V}_{m,h} \in \mathbb{R}^{r_b \times d}\)，并通过层特定的混合权重和一个门控残差来重建位于层 \(\ell \in \mathcal{W}_m\) 的令牌 \(t\) 的键/值状态：

\[
\widehat{X}^{Z}_{\ell,h,t} = A^{Z}_{\ell,h,t} B^{Z}_{m,h} + g^{Z}_{\ell,h,t} R^{Z}_{\ell,h,t}, \qquad Z \in \{K,V\}.
\tag{1}
\]

这里 \(A^{Z}_{\ell,h,t} \in \mathbb{R}^{1 \times r_b}\) 是一个紧凑的系数向量，\(R^{Z}_{\ell,h,t} \in \mathbb{R}^{1 \times d}\) 是一个仅以路由秩存储的量化残差，而 \(g^{Z}_{\ell,h,t} \in [0,1]\) 是一个轻量级的残差门。基张量在深度窗口内的所有令牌间共享，而系数和残差保持令牌局部性。这与诸如 MiniCache 和 KVSharer（Liu et al., 2024a (https://arxiv.org/html/2607.06523#bib.bib31); Yang et al., 2024a (https://arxiv.org/html/2607.06523#bib.bib66)）的直接层共享方法不同：相邻层无需重复使用完全相同的缓存状态，而是共享一个低秩通道子空间，其偏差仅在需要时才被恢复。

共享基从第一个压缩前缀块开始在线初始化。对于每个深度窗口，我们在预填充后计算缓存键和值的一个小型随机化低秩草图，然后随着生成的进行使用指数移动平均更新基。更新独立于每个头和序列维度执行。