NestedKV: 嵌套内存路由用于长上下文KV缓存压缩

arXiv cs.CL 论文

摘要

NestedKV是一种无需训练的KV缓存压缩方法,它采用嵌套内存路由和多时间尺度异常评分,提升长上下文语言模型的效率,在RULER和LongBench等基准测试上取得了显著效果。

arXiv:2605.26678v1 公告类型: 新 摘要: 长上下文语言模型受限于键值(KV)缓存的内存占用。现有的免训练KV压缩方法通常依据单一重要性信号(注意力、新近性、逐层分配或键独特性)对令牌进行排序,但当有用上下文具有全局独特性、局部偶发性或即时相关性时,这种方法会变得脆弱。我们提出NestedKV,一种仅基于键的KV缓存压缩方法,受嵌套学习中的连续内存系统启发。NestedKV维护全局、块级和滑动窗口键锚点,通过多时间尺度余弦异常评分对令牌进行评分,并将所得排名与一个免训练的外部学习器结合,该学习器采用头自适应混合和惊喜门控令牌路由。该评分与自适应每头预算配对,无需训练或修改LLM。在Qwen3和Llama-3.2模型上,针对RULER(4k-32k)、LooGLE、LongBench、LongBench-E、InfiniteBench和MMLU-Pro,NestedKV在保留较小时表现最佳。在Qwen3-4B上,当$r=0.75$时,相比KeyDiff,它在RULER上提升高达19.10个百分点,在LongBench上提升19.29个百分点;当$r=0.95$时,它在LongBench上保留37.32,而KeyDiff仅为17.55。
查看原文
查看缓存全文

缓存时间: 2026/05/27 09:08

# NestedKV:面向长上下文KV缓存压缩的嵌套内存路由
来源:https://arxiv.org/html/2605.26678
Hong Chen1, Xiang Liu1, Yubo Gao1, Yuxuan Fan1, Bo Wang1, Yuanlin Chu1, Yuanguo Lin2, Xuming Hu1
1香港科技大学(广州)
2集美大学
\{hchen763,xliu886,ygao704,yfan546,bwang423,ychu763\}@connect\.hkust\-gz\.edu\.cn
xdlyg@jmu\.edu\.cn
xuminghu@hkust\-gz\.edu\.cn

###### 摘要

长上下文语言模型的性能受限于键值(KV)缓存的内存占用。现有的无需训练的KV压缩方法通常根据单一重要性信号——注意力、近因性、逐层分配或键独特度——对标记进行排序,当有用上下文在全局上独特、在局部上偶发或与当前直接相关时,这些方法变得脆弱。我们提出NestedKV,一种受嵌套学习中连续记忆系统启发的仅键KV缓存压缩方法。NestedKV维护全局、分块级和滑动窗口键锚点,通过多时间尺度余弦异常对标记进行评分,并将得到的排序与一个无需训练的外部学习器结合,该学习器使用头自适应混合和基于惊喜的标记路由。该评分配以自适应每头预算,无需训练或修改LLM。在Qwen3和Llama-3.2模型上的RULER(4k–32k)、LooGLE、LongBench、LongBench-E、InfiniteBench和MMLU-Pro基准测试中,NestedKV在保留缓存较小时表现最强。在Qwen3-4B上,当r=0.75时,它在RULER上比KeyDiff提升最多19.10分,在LongBench上提升19.29分;当r=0.95时,它在LongBench上保持37.32分,而KeyDiff仅为17.55分。

NestedKV: Nested Memory Routing for Long\-Context KV Cache Compression

Hong Chen1, Xiang Liu1, Yubo Gao1, Yuxuan Fan1, Bo Wang1, Yuanlin Chu1, Yuanguo Lin2, Xuming Hu1
1香港科技大学(广州)
2集美大学
\{hchen763,xliu886,ygao704,yfan546,bwang423,ychu763\}@connect\.hkust\-gz\.edu\.cn
xdlyg@jmu\.edu\.cn
xuminghu@hkust\-gz\.edu\.cn

## 1引言

长上下文语言模型已成为文档理解、检索增强生成、编码和多轮交互的标准接口。然而,它们的实际部署受到一个日益简单的瓶颈的限制:键值(KV)缓存随上下文长度和批量大小线性增长。对于长提示和高吞吐量服务,即使模型权重固定,这种瞬态内存也可能主导推理成本。因此,越来越多的研究关注无需训练的KV缓存压缩,旨在不微调模型或改变注意力实现的情况下减少缓存内存(Liu等人,2023 (https://arxiv.org/html/2605.26678#bib.bib1);Zhang等人,2023 (https://arxiv.org/html/2605.26678#bib.bib2);Xiao等人,2024 (https://arxiv.org/html/2605.26678#bib.bib3);Li等人,2024b (https://arxiv.org/html/2605.26678#bib.bib4);Cai等人,2024 (https://arxiv.org/html/2605.26678#bib.bib5);Feng等人,2026 (https://arxiv.org/html/2605.26678#bib.bib6);Park等人,2026 (https://arxiv.org/html/2605.26678#bib.bib7))。

大多数现有方法可以理解为为标记重要性选择一个锚点:基于持久重要性假设的过去注意力质量(Liu等人,2023 (https://arxiv.org/html/2605.26678#bib.bib1);Zhang等人,2023 (https://arxiv.org/html/2605.26678#bib.bib2))、近因性和注意力汇(Xiao等人,2024 (https://arxiv.org/html/2605.26678#bib.bib3))、提示末尾附近的观察窗口(Li等人,2024b (https://arxiv.org/html/2605.26678#bib.bib4))、逐层缓存预算(Cai等人,2024 (https://arxiv.org/html/2605.26678#bib.bib5))或键与均值方向的几何差异(Park等人,2026 (https://arxiv.org/html/2605.26678#bib.bib7))。一个补充方向是跨头自适应分配缓存预算而不是均匀分配(Feng等人,2026 (https://arxiv.org/html/2605.26678#bib.bib6))。

参见图注图1:长上下文检索提示上最后64个查询的注意力(Qwen3-4B, RULER niah\_multivalue, N=3,800, 4个针 ⋆1–⋆4)。顶部:注意力质量(对数尺度)。底部:注意力排序压缩器在r=0.50和r=0.85时保留的标记;存活的针为绿色,被驱逐的为红色。图1 (https://arxiv.org/html/2605.26678#S1.F1) 说明了为什么在激进压缩下,基于注意力的视图在结构上不足。甚至在考虑任何具体设计选择之前,注意力信号本身就集中在提示尾部和注意力汇前缀上,而携带答案的标记位于低注意力区域,只接收到极小部分质量。因此,任何根据过去注意力对标记评分的压缩器都会继承这种几何错位,并首先驱逐错误的标记——这个问题在预算缩小时会变得更糟,而不是更好。这促使将压缩评分完全移出注意力空间,进入键流。

这些方法虽然有效,但也暴露了一个共同局限:每种方法通过单一记忆视图压缩缓存。一个标记可能因为在文档中全局异常、因为在某个段落内标志着局部主题转换、或者因为属于将塑造即时生成的近期流而重要。在轻度压缩下,单一统计量可能足够。但在激进压缩或更长上下文下,这些概念会分叉。全局均值可能错过局部事件;局部规则可能过拟合重复块;近期窗口规则可能丢弃早期证据,而这些证据是检索或多步推理所需的。

我们提出NestedKV,一种基于连续记忆视角的标记重要性无需训练的KV缓存压缩方法。遵循嵌套学习的观点,即模型通过具有自我修改更新规则的嵌套记忆系统维护压缩上下文流(Behrouz等人,2026a (https://arxiv.org/html/2605.26678#bib.bib14)),NestedKV直接在缓存键流上维护一个三时间尺度记忆——一个稳定锚点、一个剧情锚点和一个当前锚点——并针对每个尺度通过余弦异常对每个标记进行评分。这三个尺度充当内部学习器:一个标记获得三个排序而不是一个,如果它对任何尺度异常,则被保留。

一个无需训练的外部学习器随后将这两个轴上的内部排序合并。在每个注意力头下,最具区分力的尺度相对于固定先验被加权,使头能专门化不同的时间角色。在每个标记下,三个排序之间的跨尺度不一致被解读为压缩引发的惊喜信号,高惊喜将评分从混合视角平滑引导至最强的个体记忆。这两轴一起在测试时实例化嵌套学习的自我修改压缩器模式,没有可训练参数,也不修改底层LLM。

这个评分仅基于键,并且与优化的注意力内核保持兼容。完整策略与自适应每头内存分配相结合,分离了两个经常纠缠的问题:头内哪些标记具有信息性,以及每个头应获得多少内存。

我们在长上下文基准测试套件(RULER (Hsieh等人,2024 (https://arxiv.org/html/2605.26678#bib.bib18))、LongBench (Bai等人,2024 (https://arxiv.org/html/2605.26678#bib.bib19))、LooGLE (Li等人,2024a (https://arxiv.org/html/2605.26678#bib.bib21))、LongBench-E 和 InfiniteBench (Zhang等人,2024 (https://arxiv.org/html/2605.26678#bib.bib22))以及一个短上下文知识基准测试(MMLU-Pro (Wang等人,2024 (https://arxiv.org/html/2605.26678#bib.bib23))上评估NestedKV,使用Qwen3-4B作为主要冻结模型。主要经验模式是,连续记忆评分在最需要它的地方最为有用:在更高压缩比和更长上下文下,单一锚点最弱。NestedKV在大多数RULER上下文-比例单元格中表现最佳或接近最佳,在激进压缩和更长上下文下增益最明显。在r=0.75时,它还将LongBench平均值从30.77提升到50.06,并且在r=0.25时,在MMLU-Pro上保持在Full KV上限的0.2分以内,而大多数基线下降。

我们的贡献是:

- • 我们将无需训练的KV缓存压缩重新定义为键流上的连续记忆异常检测,赋予标记驱逐一个嵌套学习视角下的有界测试时内存维护解释。
- • 我们提出NestedKV,它使用三个时间尺度的键统计量——稳定、剧情和当前——作为内部学习器,并通过一个无需训练的外部学习器合并它们的每个标记异常排序,该学习器适应每个头和每个标记,后者由压缩引发的惊喜驱动。评分与自适应每头内存分配相结合,无需训练或修改LLM。
- • 我们提供六个基准测试——RULER、LongBench、LooGLE、LongBench-E、InfiniteBench 和 MMLU-Pro——的经验证据,表明多时间尺度评分在激进压缩和长上下文下特别有价值,同时不损害短提示能力。

## 2方法

NestedKV在预填充之后、自回归解码之前压缩KV缓存。它在每个Transformer层独立应用,同时协调该层KV头之间的内存分配。模型参数、注意力函数和保留的值向量不变;该方法仅决定哪些缓存位置保留在有限的测试时内存中。图2 (https://arxiv.org/html/2605.26678#S2.F2)总结了三个组件:缓存键上的嵌套连续记忆状态、每个尺度的异常评分混合成主连续读数,以及一个基于惊喜的引导路由规则,该规则在每个标记的混合读数和最强个体记忆之间进行选择。

参见图注图2:NestedKV概述。左侧(第2.2节 (https://arxiv.org/html/2605.26678#S2.SS2))。缓存键流的三个时间尺度摘要:稳定均值 μ\_s、剧情块均值 μ\_e(i) 和当前滑动窗口均值 μ\_c(i)。中间(第2.3节 (https://arxiv.org/html/2605.26678#S2.SS3)–2.4节 (https://arxiv.org/html/2605.26678#S2.SS4))。每个键产生每个尺度的余弦异常 s\_s(i)、s\_e(i)、s\_c(i),按头归一化并通过头自适应softmax混合为混合评分 s\_b(i)。右侧(第2.4节 (https://arxiv.org/html/2605.26678#S2.SS4))。基于惊喜的引导测量跨尺度不一致:一致标记保留 s\_b(i),不一致标记路由到最大锚点读数,使任何单一异常标志足以保留。最终评分驱动保留/驱逐决策(底行)。

### 2.1 KV压缩作为嵌套内存维护

对于冻结的LLM,预填充的KV缓存是内部记忆状态,模型通过它将上下文流携带到未来的解码步骤。因此,KV压缩要求一个有界的内存策略,而不是一个独立的标记删除规则。对于层 ℓ 和KV头 h,令

M_{ℓ,h} = (K_{ℓ,h}, V_{ℓ,h}) (1)

为完整的预填充内存。NestedKV构造压缩内存

M_{ℓ,h}^{B_h} = C_φ(K_{ℓ,h}, V_{ℓ,h}; B_h), (2)

其中 B_h 是头特定的内存预算,φ 表示固定的NestedKV内存策略。φ 中没有参数被学习;策略由连续记忆状态和下面的逐头分配规则定义。

为简化符号,我们描述一个层和一个KV头,并省略 ℓ, h。令

K = [k_1, ..., k_N] ∈ ℝ^{N×d}, (3)
V = [v_1, ..., v_N] ∈ ℝ^{N×d_v}。

给定预算 B,压缩器返回索引集 S,其中 |S| = B,以及内存 M^B = (K_S, V_S)。所有评分通过归一化键 \hat{k}_i = k_i / ∥k_i∥_2 计算,因此内存策略关注键空间中的方向结构。

参见图注图3:LongBench-Qasper注意力系列探测(Dasigi等人,2021 (https://arxiv.org/html/2605.26678#bib.bib20);Bai等人,2024 (https://arxiv.org/html/2605.26678#bib.bib19))。Q1–Q3关注不同的答案区域(垂直线),而NestedKV在这些分散的位置分配显著度。图3 (https://arxiv.org/html/2605.26678#S2.F3) 显示了使用记忆状态而非单一标记重要性视图的动机。即使在同一个Qasper文档中,不同的下游问题会在不同层激活不同的答案区域。一个承诺于单一时间锚点的压缩器有风险只保留其中一种模式。因此,NestedKV在键流上构造一个连续记忆,这样当一个标记在全局、在其局部剧集内或相对于当前流与众不同时,它就可以被保留。

### 2.2 连续记忆状态

连续记忆系统的观点表明,记忆不应坍缩到单一时间尺度。一个标记可能与整个文档冗余、在其局部剧集内冗余、或与近期流冗余。NestedKV将这三个概念表示为一个连续记忆状态

M(i) = {μ_s, μ_e(i), μ_c(i)} (4)

对于每个缓存的标记 i。

#### 稳定记忆。

稳定组件总结了整个预填充上下文:

μ_s = (1/N) ∑_{j=1}^N \hat{k}_j。 (5)

它捕捉整个上下文中持续的文档级规律性。

#### 剧情记忆。

剧情组件总结了包含标记 i 的局部段落。令 B(i) 为包含 i 的块,块大小为

b = clip(⌊N/32⌋, 128, 256)。 (6)

然后

μ_e(i) = (1/|B(i)|) ∑_{j∈B(i)} \hat{k}_j。 (7)

它捕捉全局总结可能不可见的段落级或轮次级结构。

#### 当前记忆。

当前组件总结了结束于 i 的即时因果流。窗口大小为 W=64,

μ_c(i) = (1/(i - ℓ_i + 1)) ∑_{j=ℓ_i}^i \hat{k}_j, (8)
ℓ_i = max(1, i - W + 1)。

它捕捉短程连续性和即时冗余性。

### 2.3 每个尺度的异常评分

NestedKV在标记的键已经可以从连续记忆状态预测时驱逐它,并在它相对于该状态异常时保留它。因为三个记忆尺度总结不同的时间结构,我们保持它们的异常读数分离,而不是预先坍缩成单一锚点。对于每个缓存的标记 i,每个尺度的异常评分是

a_s(i) = -cos(\hat{k}_i, μ_s), (9)
a_e(i) = -cos(\hat{k}_i, μ_e(i)),
a_c(i) = -cos(\hat{k}_i, μ_c(i))。

较低的 a_k(i) 意味着标记 i 相对于记忆尺度 k 是典型的;较高的 a_k(i) 意味着标记携带相对于该记忆的信息。

相似文章

KV Packet: 免重计算的上下文无关KV缓存用于大语言模型

Hugging Face Daily Papers

KV Packet 提出了一种免重计算的缓存复用框架,用于大语言模型。该框架使用可训练的软标记适配器来弥合上下文不连续性,消除了开销,同时在 Llama-3.1 和 Qwen2.5 上的性能与完全重计算基线相当。