RestoreKV：在激进查询无关KV缓存驱逐下恢复完整缓存行为 RestoreKV: Recovering Full-Cache Behavior Under Aggressive Query-Agnostic KV Cache Eviction
用8个可学习恢复token生成上下文条件化紧凑缓存，弥补KV驱逐损失。
前置知识
KV Cache（键值缓存）
Transformer自回归生成时，为避免重复计算历史token的Key/Value，把每层每个KV头的K、V张量缓存下来，长度随序列线性增长，是长上下文推理显存与带宽的主要瓶颈。
本文所有压缩、驱逐、恢复的对象都是KV cache，理解其结构（层×头×序列维度）才能看懂预算公式与restore cache的生成。
Query-agnostic KV cache eviction（查询无关KV驱逐）
在context prefill阶段、尚未看到未来查询时，就根据重要性评分一次性压缩缓存并复用于任意后续请求。与之相对的是query-dependent方法，必须等查询已知才能压缩。
本文的核心设定就是query-agnostic：缓存构造不能看查询，restore pass也必须在查询前完成，这决定了方法的可行性与效率优势。
LoRA（低秩适配）
冻结原模型权重W，只训练注入的低秩增量$\Delta W=(\alpha/r_{\text{LoRA}})B_\phi A_\phi$，用极小可训练参数（本文0.4%）适配下游任务，推理时可合并或开关。
RestoreKV的restore pass正是靠LoRA注入额外变换生成恢复缓存，且仅在restore阶段激活、之后关闭，理解LoRA才能理解为何query-time零额外成本。
Self-distillation（自蒸馏）
用同一个（或同结构）模型作为teacher蒸馏student：teacher用完整缓存生成答案token分布，student用恢复缓存逼近该分布，损失常为KL散度，teacher梯度截断不更新。
RestoreKV不依赖人工标注，仅用full-cache teacher的答案分布作为离线蒸馏目标训练restore机制，这是其通用性与免task-specific调优的关键。
RoPE（旋转位置编码）
通过对Query/Key施加按位置旋转的编码来注入相对位置信息，注意力分数取决于query与key之间的相对旋转角度，因此token的位置偏移会改变注意力分布。
restore token插入在context之后会改变后续query的相对RoPE偏移，作者专门用Table B的偏移对照实验排除了位置偏移本身带来增益的混淆。
研究动机
现有query-agnostic KV缓存驱逐方法在大语言模型长上下文推理中被广泛采用：在context prefill阶段一次性压缩缓存并复用于任意后续查询。但其本质都是决定保留哪些原始KV对，性能在预算紧张时急剧崩溃。以Qwen3-4B上的KVzip为例，RULER-4K准确率从预算比r=0.1的80.1暴跌到r=0.05的38.2；最终查询token对非sink上下文KV的平均注意力质量从完整缓存的2.73%骤降到0.56%，36层中7层低于0.1%，模型几乎看不见上下文。仅靠选出更好子集无法在激进压缩下避免这种退化。
本文的目标是本文目标是在不增加query-time的KV内存与解码开销前提下，缩小激进压缩与完整缓存之间的性能差距。具体地，作者希望保持基座重要性评分器和驱逐规则不变，仅用约0.4%可训练参数（restore-token嵌入与LoRA适配器）学习一个跨上下文共享的恢复机制，在驱逐前为每个新上下文一次性生成紧凑的、上下文条件化的恢复缓存$C_{res}$，与保留的原始上下文KV对组合后仍满足相同总预算$B=\lfloor rTLH \rfloor$。目标是在Qwen3-4B 5%预算下把KVzip在RULER-4K的38.2显著提升，并保持每上下文仅0.03-0.04秒的一次性开销。
与已有工作不同的是，本文切入点是互补恢复视角，区别于已有selection-based方法只决定哪些原始KV对存活。与KV-Distill（重写选中token表征）或Cartridges（用自学习替换整个前缀缓存）不同，RestoreKV保留绝大部分被选中的原始KV对不加修改，只在总预算内预留$nLH$个槽位给生成的恢复缓存。关键前提洞察是：被驱逐丢失的信息是上下文特定的，但生成其紧凑补集的机制可跨上下文共享。LoRA仅在单次restore pass激活，之后查询处理与解码完全使用冻结的原始backbone——这是与每上下文重新拟合的AM-fast等方法在效率上的本质区别。
核心方法
整体思路是：在context prefill完成后、eviction发生之前，插入一个轻量的恢复阶段。作者引入$n=8$个可学习的restore-token嵌入$E=[e_1,\dots,e_n]\in\mathbb{R}^{n\times d}$，让它们以因果方式attend到完整KV缓存$C$，通过单次LoRA-adapted前向传播生成上下文条件化的恢复缓存$C_{res}=\text{Restore}_{\theta,\phi}(E|C)$。恢复缓存占据固定预算$B=\lfloor rTLH \rfloor$中的一小部分$nLH$个槽位，剩余$B-nLH$个槽位由基座evictor用保留的原始上下文KV填充，两者Concat拼接成预算匹配的最终缓存$\tilde C$，满足$|\tilde C|=B$。训练阶段通过冻结的full-cache teacher做自蒸馏，只更新$E$和LoRA参数$\phi$（约0.4%参数），优化token级对称KL散度。推理时对每个新上下文生成一次$C_{res}$即关闭LoRA。
核心创新点是互补恢复而非更好选择：虽然被驱逐丢失的信息因上下文而异，但生成其紧凑补集的机制可跨上下文共享学习。这与所有selection-based方法的本质区别在于——它们学习的是哪些原始KV对该存活这一离散选择问题，结果缓存始终是原始缓存的子集；而RestoreKV学习的是一个生成式变换：用极少量（8个）恢复token attend全缓存后产出的$nLH$个KV对，是被驱逐前能感知完整上下文的信息浓缩。消融证明收益主要来自attention侧适配（fixed embedding+LoRA达71.9，而仅学习embedding只到42.1），且这些恢复状态以极低的一次性开销（0.03-0.04s）换取了对query-time预算的完整保留。
方法步骤详情
步骤一·恢复缓存生成：prefill产出含$TLH$个KV对的完整缓存$C$后，将$n=8$个可学习嵌入$E$置于位置$T+1,\dots,T+n$，以因果注意力访问$C$，单次LoRA前向得$C_{res}$，每投影应用$\Delta W=(\alpha/r_{\text{LoRA}})B_\phi A_\phi$（$r_{\text{LoRA}}=8,\alpha=16$）。步骤二·预算匹配：从总预算$B$预留$nLH$给$C_{res}$，剩余由基座evictor填充，$\tilde C=\text{Concat}(\text{Evict}(C,s,B-nLH),C_{res})$，保留KV保持原RoPE相位。步骤三·自蒸馏训练：teacher用$C$对查询$q$生成答案$y$，student用$\tilde C$评估同一答案，最小化token平均对称KL $L=\frac{1}{2M}\sum_i[\text{KL}(p_{full}^i\|p_{res}^i)+\text{KL}(p_{res}^i\|p_{full}^i)]$，仅更新$E$与$\phi$，$r\sim U(0.025,0.25)$。步骤四·推理：每上下文一次性生成$C_{res}$并关闭LoRA。
技术新颖性
技术新颖性体现在三方面。其一，设计点是互补而非替代：保持基座重要性评分器与驱逐规则完全不变，保留大部分被选原始KV对，仅预留小预算给生成式补集，与重写整个缓存的方法形成对比。其二，restore pass的工程极简性：只在prefill后额外做8个位置的单次前向，LoRA权重仅84MB常驻、与上下文长度无关，后续查询与解码完全复用冻结backbone，因此在query-time零额外KV内存与解码成本。其三，训练范式：通过full-cache自蒸馏而非task-specific微调，restore机制跨多个预算比统一学习（$r\sim U(0.025,0.25)$），一个checkpoint即可支持多种压缩率，且teacher答案仅用于定义离线蒸馏目标、不破坏query-agnostic的缓存构造。
实验结果
主实验Fig.3跨4个backbone（Qwen3-0.6B/4B/8B与Llama-3.1-8B）和4基准，预算越紧增益越大。Qwen3-4B r=0.05：KVzip 38.2→73.2（+35.0），KVzip+ 51.6→70.7（+19.1）。Table 1跨5种驱逐方法（KVzip/KVzip+/ContrastKV/SnapKV/H2O），60组配对改进59组，证明与选择策略无关的通用插件性。Fig.4对比AM-fast：RestoreKV+在RULER-4K与最紧LongHealth预算更优，4K缓存构造0.74s vs 9.68s约13×加速。消融Table 2：仅学embedding到42.1，fixed embedding+LoRA到71.9（恢复96%增益），限定q/k/v达72.4（仅4.0M参数），证明attention侧适配是主因。Table 3中全缓存生成（73.2）优于已驱逐缓存生成（64.4）；Table 4中$n=8$最优、$n=1$已达65.3。KVPress Qwen3-8B 16×压缩达86.4 RULER准确率。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| RULER-4K（Qwen3-4B, r=0.05） | 字符串匹配准确率 | RestoreKV 73.2 | KVzip 38.2 | +35.0 |
| RULER-4K（Qwen3-4B, r=0.05） | 字符串匹配准确率 | RestoreKV+ 70.7 | KVzip+ 51.6 | +19.1 |
| KVPress RULER（Qwen3-8B, 16×压缩） | RULER准确率 | RestoreKV+ 86.4 | — | 在16×压缩下接近full-cache |
| LongBench 16任务均值（Llama-3.1-8B, r=0.0625） | 任务指标均值 | RestoreKV 37.7 | KVzip 33.5 | +4.2 |
| SCBench 9任务均值（Qwen3-4B, r=0.05, 约104K上下文） | 任务指标均值 | RestoreKV 25.6 | KVzip 24.1 | +1.5 |
| 4K上下文缓存构造时间（Qwen3-4B） | 秒 | RestoreKV+ 0.74s | AM-fast 9.68s | 约13×加速 |
局限与改进
作者承认的局限：跨evictor迁移有限——Table C中SnapKV训练的checkpoint配KVzip推理时r=0.05仅42.0，远低于KVzip训练的73.2，说明restore机制适应训练时的驱逐行为；轻度预算（r=0.2）增益微小，方法在温和压缩下收益有限；训练范围需覆盖紧预算，Table A显示$r\sim U(0.05,0.5)$在r=0.05仅51.5。我观察到的局限：最大评测上下文仅32K，而训练上限15K，超长上下文（100K+）行为未充分验证（SCBench 104K平均增益仅约1-2点，相对温和）；仅测试到8B参数规模，70B级模型是否同构未知；restore token数固定为8而非自适应，可能对不同长度/任务非最优；仅query-agnostic设置，未触及query-aware场景；训练需per-model、per-evictor各约2小时，部署成本随方法数线性增长。
独立分析的弱点
弱点一：增益高度依赖紧预算。在r=0.2温和预算下多数设置增益≤2点（如KVzip r=0.2从91.4到93.5），此时restore cache占用$nLH$预算反而挤占保留槽，建议引入预算自适应开关或动态$n$。弱点二：跨evictor泛化弱。Table C表明训练与推理evictor需匹配，改进方向是设计evictor-agnostic的restore目标或元学习跨方法共享表征。弱点三：超长上下文与超大规模未验证。最长仅32K、最大仅8B，restore机制在100K+/70B+的扩展性存疑，建议补充长上下文与MoE/大模型实验。弱点四：restore token数固定为8，缺乏自适应机制——不同上下文复杂度、不同预算下的最优$n$可能不同，可改为预算/困惑度自适应。弱点五：仅限query-agnostic，未与query-aware方法（如LookaheadKV）结合，存在拓展空间。
未来方向
作者隐含方向：扩展到学习型评分器（已在Fig.A验证Fast KVzip从46.7升到79.3，证明可组合），以及推广到更长上下文与更多backbone。基于成果可延伸：其一，跨evictor统一restore机制——当前需per-evictor训练，可探索元学习或evictor-条件化适配器，让一个checkpoint适配多种驱逐规则；其二，自适应restore token数——按上下文长度/预算/困惑度动态选择$n$，平衡开销与精度；其三，与query-aware方法结合，在restore pass中融合轻量查询前瞻信号；其四，扩展到更大模型（70B+/MoE）和超长上下文（128K+）验证可扩展性；其五，将restore思路迁移到其他压缩后复用场景，如prompt cache、检索增强的缓存复用、多模态KV缓存。
复现评估
复现评估良好。项目页公开（paper.pnu-cvsp.com/RestoreKV/）。训练细节充分：6,248个context-query对（LongAlpaca 2,488、PG-19 2,260、Tulu-3 FLAN 1,500），平均4,236 tokens；$n=8$ restore token，LoRA rank 8、$\alpha=16$、dropout 0，AdamW（$\beta=(0.9,0.999)$, wd 0.01），学习率$2\times10^{-4}$，cosine+50 warmup，5,000步，单卡RTX PRO 6000约2小时。预算采样$r\sim U(0.025,0.25)$。评测用标准基准并遵循KVPress协议。种子敏感性Table G三seed标准差<0.57，而r=0.05平均增益34.6点，变异约60×更小。算力门槛低（单卡2小时）。主要不确定性是代码与checkpoint是否完全开源，以及基线依赖的特定实现版本。
论文图表
Qwen3-8B在KVPress基准上的曲线：基座驱逐方法在温和压缩时接近full-cache，但压缩比增大时急剧偏离，RestoreKV大幅缩小差距。
直观说明为什么需要恢复而不仅是更好选择，是motivation的核心动机图。
附RestoreKV到学习型评分器Fast KVzip，16×压缩下从46.7升到79.3，超过普通KVzip（59.0）。
证明与学习型评分器兼容，保留其单遍评分效率优势。
训练采样范围$U(0.025,0.25)$在紧预算下最优（73.2），扩大到$U(0.05,0.5)$在r=0.05仅51.5，说明需暴露于紧预算。
揭示训练范围对紧预算性能的敏感性。
仅施加+8位置偏移而无恢复状态的KVzip在r=0.05为36.9，接近标准KVzip 38.2，远低于RestoreKV 73.2。
排除RoPE位置偏移本身导致增益的混淆因素。
SnapKV训练的checkpoint配KVzip推理r=0.05仅42.0，KVzip训练达73.2，跨evictor迁移有限但略优于无恢复基线（38.2）。
揭示restore机制对训练时驱逐行为的适应性，是重要局限。
Llama-3.1-8B与Qwen3-8B在16任务上，r=0.0625时平均分分别33.5→37.7、28.6→32.1，紧预算增益更明显。
将主结论扩展到更广任务套件（3,750例，5K-15K上下文）。
平均上下文约104K的9任务，RestoreKV跨4个预算比均改进KVzip均值，如r=0.05从24.1到25.6，证明泛化到训练长度约10倍的长上下文。
验证在超长上下文（训练仅15K）下的泛化能力。