LKV：通过端到端学习多头预算与 Token 选择优化大模型 KV 缓存淘汰机制

# LKV：面向大模型 KV Cache 驱逐的头级预算与 Token 选择的端到端学习

来源: https://arxiv.org/html/2605.06676 Yifan HaoChao WangRui ZhangDi HuangJiaming GuoXing HuZidong DuQi GuoYunji Chen

###### 摘要

大型语言模型（LLMs）中的长上下文推理受限于 Key-Value (KV) 缓存内存的线性增长。现有的 KV 缓存压缩范式根本性地受限于启发式方法：启发式预算分配依赖统计先验而非任务目标，导致资源分配不当；而启发式选择则依赖耦合的查询-键交互或静态归纳偏置（如注意力汇聚点）。为解决这一局限性，我们引入了 LKV（Learned KV Eviction，学习型 KV 驱逐），将 KV 压缩表述为一个端到端的可微优化问题。LKV 集成了 LKV-H 以学习任务优化的全局预算，并集成 LKV-T 以在无需实例化注意力矩阵的情况下推导固有的 KV 重要性。这种设计绕过了启发式代理，使压缩严格对齐任务目标。大量评估表明，LKV 在 LongBench 和 RULER 基准测试中，在高压缩率下均取得了最先进的性能。特别是在 LongBench 上，LKV 仅保留 15% 的 KV 缓存即可实现近乎无损的性能。关键在于，我们的分析确定了学习型预算是保真度的主要驱动因素，证明了数据驱动的分配对于克服手工启发式方法的局限性至关重要。

机器学习

## 1 引言

处理长上下文对于大型语言模型（LLMs）的应用至关重要，如仓库级代码分析（Zhang et al., 2023a (https://arxiv.org/html/2605.06676#bib.bib40)）和多文档问答（Liu et al., 2024a (https://arxiv.org/html/2605.06676#bib.bib21)）。然而，处理这些大量输入带来了 prohibitive 的内存墙（Gholami et al., 2024 (https://arxiv.org/html/2605.06676#bib.bib13); Dao et al., 2022 (https://arxiv.org/html/2605.06676#bib.bib6)）。随着模型编码长序列，Key-Value (KV) 状态的累积引发了即时瓶颈（Pope et al., 2023 (https://arxiv.org/html/2605.06676#bib.bib25); Kwon et al., 2023 (https://arxiv.org/html/2605.06676#bib.bib18)），这严重限制了最大可处理的上下文长度，并大幅降低了内存受限硬件上的推理吞吐量（Sheng et al., 2023 (https://arxiv.org/html/2605.06676#bib.bib27)）。

为了克服这一障碍，KV Cache Eviction（KV 缓存驱逐）已成为一种关键技术（Liu et al., 2023 (https://arxiv.org/html/2605.06676#bib.bib22); Zhang et al., 2023b (https://arxiv.org/html/2605.06676#bib.bib42)）。其核心原理涉及动态识别并丢弃非必要的 token，同时保留关键信息，从而减轻序列长度增加带来的内存增长。通过激进地压缩 KV 缓存，驱逐技术显著减少了高吞吐量服务和设备端部署等应用的内存占用，使模型能够在有限的内存预算内，在扩展的上下文中维持高性能生成（Ge et al., 2023 (https://arxiv.org/html/2605.06676#bib.bib12); Li et al., 2024a (https://arxiv.org/html/2605.06676#bib.bib19)）。

然而，当前的方法仍受限于预算分配和 token 选择方面相互关联的局限性。首先，早期方法如 H2O（Zhang et al., 2023b (https://arxiv.org/html/2605.06676#bib.bib42)）和 SnapKV（Li et al., 2024a (https://arxiv.org/html/2605.06676#bib.bib19)）主要采用均匀预算，忽略了注意力头的显著异质性，其中关键信息是稀疏且非均匀分布的。其次，虽然先进策略试图解决这一问题，但它们被困于僵化的先验或昂贵的代理指标中。PyramidKV（Cai et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib5)）强制执行静态的层间衰减，无法适应中间检索头。同时，像 Ada-KV（Feng et al., 2025b (https://arxiv.org/html/2605.06676#bib.bib9)）和 D2O（Wan et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib33)）这样的自适应方法依赖计算注意力分数来分配预算或选择 token。这造成了推理时的循环依赖：识别重要组件需要执行我们旨在避免的昂贵查询-键交互（$O(t^2)$ 复杂度），从而阻碍了真正的与查询无关的加速（Adnan et al., 2024 (https://arxiv.org/html/2605.06676#bib.bib1)）。第三，也是最关键的，优化这些代理指标（注意力权重）而非任务目标，往往需要手动保障措施以确保稳定性。对于预算分配，方法通常强制实施通用约束（例如，最小预算阈值）（Feng et al., 2025b (https://arxiv.org/html/2605.06676#bib.bib9); Fu et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib11)）以防止过度修剪特定头。对于 token 选择，它们依赖固定的“最近窗口”或“注意力汇聚点”（Xiao et al., 2023 (https://arxiv.org/html/2605.06676#bib.bib36); Feng et al., 2025a (https://arxiv.org/html/2605.06676#bib.bib8)）来保持流畅性。虽然这些通用先验起到了有效的稳定器作用，但它们施加了规定性结构，限制了模型自主发现偏离手工规则的复杂稀疏模式的能力。

我们认为，最佳的 KV 压缩应是一种端到端学习的能力，而非一系列启发式规则的集合。任务目标本身应直接指导宏观层面预算分配（例如，区分关键检索头与流式头）和微观层面 token 效用（例如，平衡远距离锚点与最近局部上下文）的优化，而不是依赖手动规则。

我们介绍了 **LKV**（Learned KV Eviction），据我们所知，这是第一个将 KV 缓存压缩重构为统一的、端到端可微优化问题的框架。

> **图 1：** KV 预算分配策略（15% 保留率）。(a) SnapKV: 均匀。(b) PyramidKV: 层间衰减。(c) DuoAttention: 刚性二分分类（检索 vs. 流式）。(d) Ada-SnapKV: 层内自适应但具有均匀层先验。(e) LKV（本文）: 学习的全局细粒度策略，优化任务目标且无刚性先验。

如图 1 (https://arxiv.org/html/2605.06676#S1.F1) 所示，LKV 重塑了所有 KV 头的资源分配。不同于均匀策略（图 1a）、强制层间衰减或严格二分法的刚性结构先验（图 1b-c），或依赖嘈杂代理指标的启发式适应（图 1d），LKV 学习了一种全局优化、无约束的策略，纯粹根据需求将预算导向最关键的头部。

**LKV-H**（全局学习型预算）通过展平头结构实现全局竞争，允许每个头通过可学习的嵌入竞争总预算。这打破了“金字塔”陷阱：如图 1d 所示，如果任务需要，LKV 自由地向更深层分配高预算，而 Ada-KV 仍受限于静态衰减。该分配完全从数据中学习，不依赖启发式或保障超参数。

同时，**LKV-T**（无矩阵选择）通过轻量级网络直接从 KV 状态预测 token 效用。这种设计实现了带有可忽略线性 $O(t)$ 开销的无矩阵选择，避免了昂贵的 $O(t^2)$ 注意力计算。值得注意的是，LKV 在预填充阶段逐层操作，有效地驱逐 token 以降低峰值内存消耗。

我们通过自蒸馏（Zhang et al., 2019 (https://arxiv.org/html/2605.06676#bib.bib41)）端到端地训练 LKV。在 LongBench（Bai et al., 2024 (https://arxiv.org/html/2605.06676#bib.bib3)）和 RULER（Hsieh et al., 2024 (https://arxiv.org/html/2605.06676#bib.bib15)）上的大量实验表明，LKV 显著优于最先进的基线。我们的分析进一步揭示，**学习型全局预算**对性能提升的贡献大于选择策略本身，证明了精确的头级资源分配的重要性。

我们的贡献总结如下：

- 我们提出了 LKV，这是第一个将 KV 驱逐转化为端到端可微学习问题的框架，消除对手工启发式、最小预算或强制最近窗口的需求。
- 我们引入了 LKV-H，一种揭示固有头重要性的全局预算机制，以及 LKV-T，一种用于高效推理的无矩阵选择器。
- 我们在长上下文基准测试中取得了 SOTA 性能，并提供了新的见解：学习型预算是压缩质量的主导因素。

## 2 相关工作

### 2.1 基于驱逐的 KV 缓存压缩

基于驱逐的 KV 缓存压缩主要关注两个方面：预算分配策略和 token 选择机制。

#### 预算分配策略

KV 缓存预算的分布显著影响模型保真度。早期方法通常在所有层和头之间采用均匀分配，例如 H2O（Zhang et al., 2023b (https://arxiv.org/html/2605.06676#bib.bib42)）和 SnapKV（Li et al., 2024a (https://arxiv.org/html/2605.06676#bib.bib19)）。认识到层在信息处理中的贡献不均（Wang & Tu, 2020 (https://arxiv.org/html/2605.06676#bib.bib35); Skean et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib29)），最近的工作引入了非均匀策略。PyramidKV（Cai et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib5)）和 PyramidInfer（Yang et al., 2024 (https://arxiv.org/html/2605.06676#bib.bib39)）采用金字塔策略，优先为较浅的层分配更高的缓存预算。类似地，D2O（Wan et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib33)）根据层间注意力密度动态调整预算。过渡到更细粒度，Ada-KV（Feng et al., 2025b (https://arxiv.org/html/2605.06676#bib.bib9)）调整每层内不同头之间的预算分布。进一步扩展这种灵活性，CAKE（Qin et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib26)）基于时空注意力模式将缓存分配视为级联的切蛋糕任务。虽然这些方法依赖统计或基于规则的启发式，但我们的 LKV-H 模块将宏观预算视为端到端可微优化问题，实现绕过刚性层间先验的资源全局竞争。

#### Token 选择机制

一旦预算确定，模型必须决定保留哪些 token。像 StreamingLLM（Xiao et al., 2023 (https://arxiv.org/html/2605.06676#bib.bib36)）这样的静态策略优先保留初始“注意力汇聚点”和最近 token（Wang et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib34); Huang et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib16)）以稳定解码。启发式动态方法通常通过耦合的查询-键交互来评估 token 重要性；例如，H2O（Zhang et al., 2023b (https://arxiv.org/html/2605.06676#bib.bib42)）跟踪累积注意力分数，而 SnapKV（Li et al., 2024a (https://arxiv.org/html/2605.06676#bib.bib19)）利用局部观察窗口来识别关键簇。最近的进展引入了感知价值的指标，如 CriticalKV（Feng et al., 2025c (https://arxiv.org/html/2605.06676#bib.bib10)），考虑投影值状态的范数以限制输出扰动。此外，基于学习的方法如 DuoAttention（Xiao et al., 2024 (https://arxiv.org/html/2605.06676#bib.bib37)）和 PruLong（Bhaskar et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib4)）试图通过蒸馏或优化来识别检索头。然而，大多数现有机制需要实例化注意力矩阵以评估重要性。相比之下，LKV-T 采用无矩阵选择过程，从固有潜在特征中预测 token 效用，消除了对昂贵依赖查询的注意力计算的需求。

### 2.2 其他 LLM 效率范式

我们的工作与其他范式正交，包括稀疏注意力（例如，Quest（Tang et al., 2024 (https://arxiv.org/html/2605.06676#bib.bib31)），MInference（Jiang et al., 2024 (https://arxiv.org/html/2605.06676#bib.bib17)））、提示压缩（例如，LLMLingua-2（Pan et al., 2024 (https://arxiv.org/html/2605.06676#bib.bib24)），500xCompressor（Li et al., 2024b (https://arxiv.org/html/2605.06676#bib.bib20)））和 KV 量化（例如，KIVI（Liu et al., 2024b (https://arxiv.org/html/2605.06676#bib.bib23)），AQUA-KV（Shutova et al., 2025 (https://arxiv.org/html/2605.06676#bib.bib28)））。与这些专注于选择性加载、输入缩小或精度降低的方法不同，LKV 减少基数，并可与它们集成以获得互补增益。

## 3 方法

> **图 2：** LKV 概述。左侧（LKV-H）：从头嵌入学习全局预算比例 $\mathbf{r}$。中间（LKV-T）：通过 Soft-TopK 执行可微的、与查询无关的 token 选择。右侧：通过对冻结教师模型进行自蒸馏实现端到端优化。

### 3.1 预备知识与问题表述

考虑一个具有 $L$ 层的仅解码器大型语言模型（LLM）。我们将表述重点放在 **Key-Value (KV)** 头上，因为它们构成了内存瓶颈。设 $H$ 为每层的 KV 头数量。这一通用设置自然地适用于多头注意力（MHA）和分组查询注意力（GQA），其中多个查询头可能共享单个 KV 头（Vaswani et al., 2017 (https://arxiv.org/html/2605.06676#bib.bib32); Ainslie et al., 2023 (https://arxiv.org/html/2605.06676#bib.bib2)）。

给定输入 token 序列 $x_1, \dots, x_t$，模型以自回归方式生成下一个 token $x_{t+1}$。

#### KV 缓存瓶颈

在步骤 $t$ 的自回归生成期间，模型维护一个 KV 缓存 $\mathcal{C}_t = \{ (\mathbf{K}^{(l,h)}_{\leq t}, \mathbf{V}^{(l,h)}_{\leq t}) \}_{l=1,h=1}^{L,H}$ 以避免冗余计算。这里，$\mathbf{K}, \mathbf{V} \in \mathbb{R}^{t \times d}$ 代表第 $l$ 层第 $h$ 个 KV 头的堆叠键和值状态。随着序列长度 $t$ 的增加，$\mathcal{C}_t$ 的内存足迹线性增长，对系统吞吐量施加了显著的**内存墙**。针对此缓存的查询 $\mathbf{q}_t$ 的标准注意力输出计算如下：

$$
\mathbf{o}_t = \text{Softmax}\left(\frac{\mathbf{q}_t (\mathbf{K}_{\leq t})^\top}{\sqrt{d}}\right) \mathbf{V}_{\leq t}. \quad (1)
$$

#### 问题表述

我们的目标是通过为所有 KV 缓存项学习二元保留掩码 $\mathbf{M} \in \{0,1\}^{L \times H \times t}$ 来压缩 $\mathcal{C}_t$。关键在于，此掩码在 KV 头级别操作，确保任何驱逐决策一致地应用于共享该 KV 缓存的所有查询头。我们定义一个目标全局保留比率 $R \in (0,1)$，它确定了步骤 $t$ 的总 KV 预算为 $B_{\text{total}} = \lfloor R \cdot L \cdot H \cdot t \rfloor$。我们将 LKV 表述为一个端到端优化问题，以