ReST-KV：基于逐层输出重构与时空平滑的鲁棒 KV Cache 驱逐方法

# ReST-KV：基于层级输出重构与时空平滑的鲁棒 KV Cache 淘汰

来源：https://arxiv.org/html/2605.08840

Yongqi An1,2, Chang Lu3, Kuan Zhu1,2, Tao Yu1,2, Chaoyang Zhao1,2, Hong Wu3, Ming Tang1,2, Jinqiao Wang1,2,4,5

1. 中国科学院自动化研究所基础模型研究中心，中国北京
2. 中国科学院大学人工智能学院，中国北京
3. 电子科技大学，中国成都
4. 武汉人工智能研究院，中国武汉
5. Objecteye Inc.，中国北京

[email protected], [email protected]

###### 摘要

大型语言模型（LLMs）由于键值（KV）缓存的内存需求日益增加，特别是在处理长序列时，其在高效生成推理方面面临越来越大的挑战。现有的淘汰方法通常保留具有高注意力权重的 KV 对，但忽视了令牌移除引起的注意力重分配影响，以及 KV 选择中的时空动态特性。在本文中，我们提出了 ReST-KV，一种鲁棒的 KV 淘汰方法，它结合了层级输出重构和时空平滑，为 KV 缓存淘汰任务提供了更全面的视角。具体而言，ReST-KV 将 KV 缓存淘汰公式化为一个通过高效层级重构来最小化输出差异的优化问题。通过直接建模每个令牌的移除如何影响模型输出，我们的方法自然地捕捉了注意力重分配效应，超越了单纯依赖原始注意力权重的局限性。为了进一步增强鲁棒性，我们设计了指数移动平均平滑来处理时间变化，并采用基于自适应窗口的机制来捕捉空间模式。我们的方法 ReST-KV 在长上下文基准测试上显著提升了性能。它在 LongBench 上比最先进（SOTA）的基线高出 2.58%，在 RULER 上高出 15.2%。此外，ReST-KV 在 Needle-in-a-Haystack 和 InfiniteBench 上也始终优于现有方法，同时在 128k 上下文长度下实现了惊人的 10.61 倍解码延迟降低。代码公开在 https://github.com/an-yongqi/rest-kv，以促进可复现性和进一步研究。

## 1 引言

大型语言模型（LLMs）（Achiam et al., 2023; Anthropic, 2023; Dubey et al., 2024; MistralAI, 2023）极大地推动了自然语言处理（NLP）的发展。这些模型在文档摘要（Zhang et al., 2024a）、多轮对话（Du et al., 2021）、检索增强（Yao et al., 2022）和代码生成（Roziere et al., 2023）等各种任务中实现了突破。最近的模型如 GPT-4（Achiam et al., 2023）、Claude 3.5（Anthropic, 2023）和 Llama-3.1（Dubey et al., 2024）将其上下文长度扩展到了超过 128K 个令牌，从而支持长上下文应用。然而，随着上下文长度的增加，存储 KV 缓存所需的内存迅速增长，在处理较长序列时可能达到数百 GB。因此，在不重新训练的情况下优化推理过程中的 KV 缓存，对于提高效率可扩展性至关重要。

> 图 1：ReST-KV 与现有方法的比较。与忽视注意力重分配的先前方法不同，ReST-KV 考虑其影响以改善 KV 保留。

KV 缓存淘汰通过识别并移除不太重要的 KV 对，是减少内存消耗和提高计算效率的一种有前途的方法（Li et al., 2024a）。当前方法通常依赖于固定的注意力模式（Han et al., 2024; Ge et al., 2023）或使用注意力权重的统计信息（Zhang et al., 2023; Li et al., 2024b; Cai et al., 2024）来估计 KV 对的重要性。然而，如图 1 所示，这些方法仅专注于保留具有高相似性得分的查询-键对，而忽视了移除某些对所引起的注意力重分配效应。这种重分配可能会改变整体注意力景观，导致次优的保留决策和性能下降，尤其是在缓存受限的情况下。

在本文中，我们提出了 ReST-KV，一种鲁棒的 KV 缓存淘汰方法，它考虑了注意力重分配的影响以及 KV 选择中的时空动态。我们重新审视了 KV 缓存淘汰问题，并将其重新公式化为在固定内存约束下保留每层的注意力输出。具体而言，我们衡量移除每个单独 KV 对引起的重构损失，并将其作为淘汰指标：损失越大，KV 对越重要。这种损失隐式地捕捉了移除引起的注意力重分配的影响。此外，我们的经验观察显示，KV 重要性在时间和空间上都存在显著差异。为了进一步提高鲁棒性，我们引入了两种平滑机制：（1）指数移动平均，通过强调更近的 KV 对来建模时间动态；（2）基于自适应窗口的空间平滑方法，通过估计空间动态来调整不同的窗口大小和偏移。

通过在包括 LongBench、RULER、Needle-in-a-Haystack 和 InfiniteBench 在内的广泛下游任务上进行评估，我们证明 ReST-KV 始终优于最先进的基线，特别是在低缓存预算下，并在多轮对话场景中表现出更强的鲁棒性。我们在 LongBench、RULER、Needle-in-a-Haystack 和 InfiniteBench 等具有挑战性的长上下文基准上对 ReST-KV 进行了广泛评估。结果表明，它始终超越最先进的基线，在 LongBench 上特别获得了 2.58% 的提升，在 RULER 上获得了 15.2% 的提升。ReST-KV 在多轮对话中也表现出更大的鲁棒性，在受限缓存预算下效率更高。在解码方面，当与 FlashAttention-2 集成时，它在 128k 上下文长度下实现了 10.61 倍的延迟降低。重要的是，ReST-KV 完全兼容现有的预填充稀疏注意力方法，实现了 2.37 倍的 TTFT（首 token 时间）加速。

总之，我们做出以下贡献：

-   一种新颖的 KV 淘汰公式，将其视为层级输出重构，从而实现一种能够捕捉注意力重分配效应的新重要性指标。
-   一种结合指数移动平均和自适应窗口的时空平滑机制，显著增强了 KV 选择的鲁棒性。
-   广泛的实验表明，ReST-KV 在低缓存预算下优于最先进的基线，并在 128k 上下文长度下将解码延迟降低多达 10 倍。

## 2 相关工作

### 2.1 KV 缓存淘汰

KV 缓存淘汰是一种在不重新训练的情况下优化推理过程中 KV 缓存的突出方法，缓解了长上下文 LLMs 中的内存和延迟问题（Li et al., 2024a）。早期的淘汰方法专注于特定的注意力模式，例如 StreamingLLM（Xiao et al., 2023）和 LM-Infinite（Han et al., 2024），仅保留初始和局部令牌。虽然开发了像 FastGen（Ge et al., 2023）和 RazorAttention（Tang et al., 2024）这样更灵活的方法，但它们仍然依赖于预定义的模式，并存在忽视重要令牌的风险。后续研究引入了淘汰指标来评估 KV 缓存条目的重要性，通常使用注意力权重。例如，H2O（Zhang et al., 2023）使用累积注意力权重，而 SnapKV（Li et al., 2024b）池化最后一个窗口上的平均注意力权重。除了指标改进外，一些研究还探索了非均匀层级和头级预算分配策略。PyramidKV（Cai et al., 2024）和 PyramidInfer（Yang et al., 2024）以金字塔方式分配预算，而 DynamicKV（Zhou et al., 2024）、D2O（Wan et al., 2024）和 CAKE（Qin et al., 2025）根据特定层的信息自适应地分配预算。AdaKV（Feng et al., 2024）根据输出 $\ell_1$ 损失边界调整每个头的预算。

我们的工作侧重于现有淘汰指标的局限性，这些指标主要依赖于源自查询-键交互的注意力权重，并忽视了值向量及时空动态的综合影响。此外，我们的方法完全兼容现有的层级和头级预算分配策略。

> 图 2：ReST-KV 概览。（a）层级输出重构量化每个 KV 对输出误差的影响作为其淘汰指标。（b）两种平滑机制增强鲁棒性：用于时间平滑的指数移动平均和用于空间平滑的基于自适应窗口的方法。

### 2.2 注意力动态

虽然注意力是 Transformer 成功的关键，但由于其二次复杂度，它在长上下文设置中也带来了可扩展性挑战。因此，最近的工作调查了注意力动态——特别是注意力的时空模式和权重重分配——作为一种实现更高效推理的手段。几项研究揭示了结构化的注意力行为。MInference（Jiang et al., 2024）发现了一种“垂直斜杠”模式，其中注意力随着时间的推移在令牌之间逐渐转移，表明令牌重要性的演变。FlexPrefill（Lai et al., 2025）同样识别出预填充期间一致的注意力轨迹。Keyformer（Adnan et al., 2024）检查了 KV 淘汰如何扭曲注意力分布，并提出归一化以减轻这种转变。与上述方法不同，我们通过显式建模注意力重分配和时空动态来重新公式化 KV 缓存淘汰。我们的方法不仅依赖于静态的注意力权重，还捕捉注意力的时间演化和层级偏移，从而实现更鲁棒的重要性估计，并在内存约束下显著改善性能。

## 3 方法论

### 3.1 预备知识

LLMs 通常以自回归方式解码文本，这允许它们生成高质量、上下文连贯的文本。然而，这种解码过程计算成本高昂，因为它涉及高度重复的计算，使其难以应用于实时或大规模场景。KV 缓存是一种广泛认可的技术，通过存储先前计算的键和值来减少冗余计算。在本节中，我们描述了 KV 缓存框架下的注意力计算，为讨论 KV 缓存淘汰奠定基础。为清晰起见，我们关注单个注意力头和层，省略脚注。

在每个解码步骤 $t$，KV 缓存存储先前计算的键和值 $\langle \mathbf{K}_{1:t-1}, \mathbf{V}_{1:t-1} \rangle$ 用于 $X[1:t-1]$，以便在未来步骤中复用。为了方便，我们将 $\mathbf{K}_{1:t-1}$ 记为 $\mathbf{K}_{T-1}$，将 $\mathbf{V}_{1:t-1}$ 记为 $\mathbf{V}_{T-1}$。因此，模型只需要当前令牌 $\mathbf{x}_t$ 来生成 $\mathbf{x}_{t+1}$，而不是完整序列 $X=[\mathbf{x}_1, \dots, \mathbf{x}_t]$。

形式上，在步骤 $t$，查询 $\mathbf{q}_t$、键 $\mathbf{k}_t$ 和值 $\mathbf{v}_t$ 计算如下：

$$ \mathbf{q}_t = \mathbf{x}_t \mathbf{W}_Q, \ \mathbf{k}_t = \mathbf{x}_t \mathbf{W}_K, \ \mathbf{v}_t = \mathbf{x}_t \mathbf{W}_V, \quad (1) $$

其中 $\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V$ 是对应于单个注意力头的 $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ 权重矩阵的分量。当前计算的 $\mathbf{k}_t$ 和 $\mathbf{v}_t$ 将与先前缓存的键和值连接，并用于解码步骤 $t$ 的注意力计算：

$$ \mathbf{K}_T = \operatorname{Concat}\left(\mathbf{K}_{T-1}, \mathbf{k}_t\right), \quad \mathbf{V}_T = \operatorname{Concat}\left(\mathbf{V}_{T-1}, \mathbf{v}_t\right), \quad (2) $$

其中 $\mathbf{K}_T$ 和 $\mathbf{V}_T$ 是解码步骤 $t$ 处的整个键和值序列。令牌 $\mathbf{x}_t$ 在步骤 $t$ 的注意力输出 $\mathbf{z}_t$ 计算为：

$$ \mathbf{z}_t = \operatorname{softmax}\left(\frac{\mathbf{q}_t \mathbf{K}_T^\top}{\sqrt{d_k}}\right) \mathbf{V}_T = \mathbf{A}_t \mathbf{V}_T, \quad (3) $$

其中 $\mathbf{A}_t$ 代表令牌 $\mathbf{x}_t$ 的注意力权重，并被现有方法用于计算淘汰指标。$d_k$ 代表注意力机制中键向量的维度。最后，多头注意力中单个头的输出可以表示为：

$$ \operatorname{MHA}\left(\mathbf{x}_t, \langle \mathbf{K}_T, \mathbf{V}_T \rangle\right) = \mathbf{z}_t \mathbf{W}_O, \quad (4) $$

其中 $\mathbf{W}_O$ 是对应于单个注意力头的输出投影权重矩阵。

### 3.2 层级重构指标

我们将 KV 缓存淘汰重新公式化为在固定内存约束下保留每层的注意力输出分布，自然地捕捉注意力重分配的影响。我们将这一范式正式定义为*层级重构*，这是一个与 Transformer 固有的层级计算流对齐的框架。具体而言，对于单层，子问题表示为：

###### 定义 3.1。给定单层的缓存预算 $B$，任务是从步骤 $t$ 的总缓存条目 $\langle \mathbf{K}_T, \mathbf{V}_T \rangle$ 中选择一系列最多包含 $B$ 个元素的重要 KV 缓存条目 $\langle \hat{\mathbf{K}}_T, \hat{\mathbf{V}}_T \rangle$，目标是最大化原始 MHA 输出的保留。我们使用 $\ell_2$ 距离来计算重构