上下文蒸馏作为潜在记忆管理

# 上下文蒸馏作为潜在记忆管理
来源：https://arxiv.org/html/2605.28889
Ziyang Zheng1,2, Zeju Li1, Xiangyu Wen1, Jianyuan Zhong1, Junhua Huang2, Lei Chen2, Mingxuan Yuan2, Qiang Xu1

1香港中文大学，2华为诺亚方舟实验室 通讯作者：qxu@cse\.cuhk\.edu\.hk (https://arxiv.org/html/2605.28889v1/mailto:email@domain)

###### 摘要

上下文蒸馏将上下文信息压缩到模型参数中，然而现有方法通常忽略了多个蒸馏后的潜在记忆应如何存储、检索以及在没有oracle设定的情况下安全激活。我们将上下文蒸馏建模为一个潜在记忆管理问题。我们将每个上下文蒸馏为一个独立的LoRA适配器，形成一个模块化的记忆库，从而实现显式的记忆选择。给定一个查询，我们的框架检索候选记忆，将查询路由到最合适的适配器，并使用Self-Gating机制决定是否激活潜在记忆。为了提高效率，我们进一步引入缓存共享以减少推理过程中的管理开销。实验表明，我们的方法在有检索的情况下显著优于基线方法，而Self-Gating通过停用不必要的潜在记忆提高了鲁棒性。

上下文蒸馏作为潜在记忆管理

Ziyang Zheng1,2, Zeju Li1, Xiangyu Wen1, Jianyuan Zhong1, Junhua Huang2, Lei Chen2, Mingxuan Yuan2, Qiang Xu1
1香港中文大学，2华为诺亚方舟实验室
通讯作者：qxu@cse\.cuhk\.edu\.hk (https://arxiv.org/html/2605.28889v1/mailto:email@domain)

## 1 引言

LLM主要通过将相关信息放入上下文窗口来适应文档、任务和用户 (Brown et al., 2020 (https://arxiv.org/html/2605.28889#bib.bib23); Lewis et al., 2020 (https://arxiv.org/html/2605.28889#bib.bib24))。这种上下文内适应是灵活的，但本质上是暂时的：相同的信息必须为每个查询重新读取，并且长提示会增加延迟 (Tay et al., 2022 (https://arxiv.org/html/2605.28889#bib.bib25))、内存使用以及生成不稳定性 (Zhao et al., 2021 (https://arxiv.org/html/2605.28889#bib.bib26); Liu et al., 2024 (https://arxiv.org/html/2605.28889#bib.bib27))。这些局限性引出了一个核心问题：上下文信息能否从临时文本转换为持久的模型记忆？如果可以，那么如何存储、检索并安全地激活许多这样的记忆？

上下文蒸馏 (CD) (Cao et al., 2025 (https://arxiv.org/html/2605.28889#bib.bib1); Wang et al., 2024 (https://arxiv.org/html/2605.28889#bib.bib2); Ye et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib9); Zhang et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib16); Charakorn et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib14), 2025 (https://arxiv.org/html/2605.28889#bib.bib15)) 为问题的第一部分提供了自然的答案。CD并非在每次推理时都依赖长上下文进行条件生成，而是训练模型在未显式看到上下文的情况下重现其基于上下文的行为，从而将文本上下文压缩为存储在模型参数中的潜在记忆，供后续使用。

参考图题Figure 1: 累积上下文蒸馏与我们的方法对比。累积式蒸馏将所有上下文压缩成一个适配器，而我们的框架将上下文存储为模块化记忆：为每个查询检索、路由并门控潜在记忆。然而，问题的第二部分在很大程度上仍未解决。实际部署不仅需要将上下文写入参数，还需要管理一组参数化的记忆。现有的CD方法通常通过假设能oracle式地访问相关的蒸馏记忆 (Ye et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib9); Zhang et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib16); Snell et al., 2022 (https://arxiv.org/html/2605.28889#bib.bib20); Caccia et al., 2025 (https://arxiv.org/html/2605.28889#bib.bib30))，或通过将新上下文持续写入单个适配器的累积范式 (Wang et al., 2024 (https://arxiv.org/html/2605.28889#bib.bib2); Cao et al., 2025 (https://arxiv.org/html/2605.28889#bib.bib1)) 来绕过这个问题。

累积范式如图1 (https://arxiv.org/html/2605.28889#S1.F1) 所示，概念简单：每个新上下文被吸收到当前参数状态中，因此所得模型应隐式包含所有先前观察到的上下文。其结果是，新蒸馏的上下文可能会覆盖较早的上下文 (Shenfeld et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib10); Li and Hoiem, 2017 (https://arxiv.org/html/2605.28889#bib.bib29); Kirkpatrick et al., 2017 (https://arxiv.org/html/2605.28889#bib.bib28))，导致严重的遗忘和先前潜在记忆的不可靠激活。此外，由于所有上下文都被压缩成一个单一状态，模型缺乏一种显式机制来在记忆之间进行选择，或者当查询不需要上下文知识时将其停用。因此，在现实的非oracle设定中，累积式蒸馏是脆弱的，因为系统必须同时确定使用哪个潜在记忆以及是否应激活任何记忆。

在本文中，我们将上下文蒸馏建模为一个*潜在记忆管理*问题，其中蒸馏后的上下文被维护为持久的模块化记忆。具体来说，我们将每个上下文蒸馏到一个专用的LoRA适配器中，形成一个记忆库。如图1 (https://arxiv.org/html/2605.28889#S1.F1) 所示，我们的框架针对给定查询检索候选记忆，将其路由到最合适的适配器，并采用Self-Gating机制决定是否激活潜在记忆，从而避免在上下文无关的查询上产生有害使用。为了提高效率，我们进一步引入缓存共享以在蒸馏期间分摊上下文编码，并减少推理时的检索和门控开销。在现实设定下的实验表明，我们的方法大幅优于累积式记忆基线方法，而Self-Gating和缓存共享则提高了鲁棒性和效率。

我们的贡献有三点：

- • 我们将上下文蒸馏形式化为潜在记忆管理，将关注点从单上下文的内部化转移到潜在记忆的存储、检索和激活。
- • 我们提出了一个模块化适配器记忆库，结合两阶段潜在检索和Self-Gating。
- • 我们引入了缓存共享蒸馏，在保持下游性能的同时实现高效的适配器切换。

## 2 背景

### 2.1 KV缓存压缩

KV缓存压缩旨在将长上下文压缩成更短的潜在令牌，通常达到15×15\\times\-30×30\\times的压缩比。先前的工作 (Eyuboglu et al., 2025 (https://arxiv.org/html/2605.28889#bib.bib13); Liu and Qiu, 2025 (https://arxiv.org/html/2605.28889#bib.bib18); Li et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib19)) 首先使用压缩器φ\\phi压缩上下文cc，并应用上下文蒸馏目标：

DKL\(π\(⋅∣q,c\)∥π\(⋅∣q,φ\(c\)\)\),D\_\{\\mathrm\{KL\}\}\\\!\\left\(\\pi\(\\cdot\\mid q,c\)\\,\\middle\\\|\\,\\pi\(\\cdot\\mid q,\\phi\(c\)\)\\right\),\(1\)其中qq表示查询。不同的工作对φ\\phi的选择各不相同。Cartridges (Eyuboglu et al., 2025 (https://arxiv.org/html/2605.28889#bib.bib13)) 和 Latent Context Compilation (Li et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib19)) 使用带有LoRA的基础模型压缩上下文，而C3 (Liu and Qiu, 2025 (https://arxiv.org/html/2605.28889#bib.bib18)) 使用一个小型基础模型作为上下文压缩编码器。

### 2.2 上下文蒸馏

上下文蒸馏旨在通过最小化以下损失将上下文内化到模型参数中：

DKL\(πθ\(⋅∣q,c\)∥πθ\+Δθ\(⋅∣q\)\)。D\_\{\\mathrm\{KL\}\}\\\!\\left\(\\pi\_\{\\theta\}\(\\cdot\\mid q,c\)\\,\\middle\\\|\\,\\pi\_\{\\theta\+\\Delta\\theta\}\(\\cdot\\mid q\)\\right\)。\(2\)Snell et al. (2022 (https://arxiv.org/html/2605.28889#bib.bib20)) 首次提出通过蒸馏将上下文信息内化到模型参数中。类似地，Prompt baking (Bhargava et al., 2024 (https://arxiv.org/html/2605.28889#bib.bib17)) 提出将提示信息烘焙到模型参数中。Charakorn et al. (2025 (https://arxiv.org/html/2605.28889#bib.bib15), 2026 (https://arxiv.org/html/2605.28889#bib.bib14)) 采用元学习思想，使用超网络促进蒸馏。最近，OPCD (Ye et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib9)) 和 Opsdl (Zhang et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib16)) 采用在策略蒸馏进一步增强蒸馏能力。Caccia et al. (2025 (https://arxiv.org/html/2605.28889#bib.bib30)) 首次讨论了CD与RAG之间的协同作用。然而，他们仅用检索到的段落增强oracle潜在记忆，而不是管理潜在记忆本身。除了内化单个上下文外，如何管理潜在记忆在很大程度上仍未解决。

参考图题Figure 2: 我们训练阶段和推理阶段的概览。训练：教师模型预计算目标文档的KV缓存以供训练期间重用，而学生使用基础模型计算前缀KV缓存，以实现管理所需的Hot-swap属性。推理：给定一个查询，我们首先检索最相关的潜在记忆Δθi\\Delta\\theta\_\{i\}，然后使用Self-Gating决定是继续使用θ\+Δθi\\theta\+\\Delta\\theta\_\{i\}还是回退到θ\\theta。由于训练时引入的缓存共享蒸馏，记忆管理仍然是高效的。

### 2.3 累积式蒸馏

累积式上下文蒸馏 (Wang et al., 2024 (https://arxiv.org/html/2605.28889#bib.bib2); Cao et al., 2025 (https://arxiv.org/html/2605.28889#bib.bib1)) 通过顺序地将每个新上下文蒸馏到其参数中来更新模型：

θi=argminθDKL\(πθi−1\(⋅∣q,ci\)∥πθ\(⋅∣q\)\)，\\theta\_\{i\}=\\arg\\min\_\{\\theta\}D\_\{\\mathrm\{KL\}\}\\big\(\\pi\_\{\\theta\_\{i\-1\}\}\(\\cdot\\mid q,c\_\{i\}\)\\parallel\\pi\_\{\\theta\}\(\\cdot\\mid q\)\\big\),\(3\)目标是逐渐将所有之前的上下文内化到最新的模型θn\\theta\_\{n\}中，从而允许使用最新的适配器进行推理，而不是依赖oracle式的选择来选取合适的适配器。然而，这种顺序写入可能导致严重的灾难性遗忘 (Shenfeld et al., 2026 (https://arxiv.org/html/2605.28889#bib.bib10); Li and Hoiem, 2017 (https://arxiv.org/html/2605.28889#bib.bib29); Kirkpatrick et al., 2017 (https://arxiv.org/html/2605.28889#bib.bib28))。除了遗忘问题，我们进一步论证，在理想化的递归蒸馏假设下，累积式蒸馏的行为类似于将基础模型在累积上下文\{c1,...,ci\}\\\{c\_\{1\},\\ldots,c\_\{i\}\\}条件下的蒸馏结果：

πθi∗\(y∣q\)=πθ\(y∣q,\[c1,...,ci\]\)。\\pi\_\{\\theta\_\{i\}^\{\*\}\}\(y\\mid q\)=\\pi\_\{\\theta\}\(y\\mid q,\[c\_\{1\},\\ldots,c\_\{i\}\]\)。\(4\)我们在附录A (https://arxiv.org/html/2605.28889#A1) 中提供了更多细节。这一视角表明，累积式CD继承了教师模型对不相关累积上下文的敏感性，以及在长上下文输入下的性能退化。因此，即使在oracle测试设定下，性能也可能恶化。

## 3 方法论

### 3.1 缓存共享上下文蒸馏

给定一个上下文流C=\{c1,c2,...,ct\}\\mathcal\{C\}=\\\{c\_\{1\},c\_\{2\},\\dots,c\_\{t\}\\\}，我们冻结基础模型fθf\_\{\\theta\}，并为每个上下文cic\_\{i\}训练一个独立的PEFT模块，例如LoRA，Δθi\\Delta\\theta\_\{i\}。学生模型θ\+Δθi\\theta\+\\Delta\\theta\_\{i\}被训练来模仿教师模型，即显式以cic\_\{i\}为条件的基础模型。对于一组查询\{q1,...,qn\}\\\{q\_\{1\},\\dots,q\_\{n\}\\\}，我们最小化KL散度：

Δθi∗=arg⁡minΔθi1n∑j=1nDKL\(πθ\(yj\|qj,ci\)∥πθ\+Δθi\(yj\|KVqj\)\)，\\begin\{split\}\\Delta\\theta\_\{i\}^\{\*\}=\\mathop\{\\arg\\min\}\_\{\\Delta\\theta\_\{i\}\}\\frac\{1\}\{n\}\\sum\_\{j=1\}^\{n\}D\_\{\\text\{KL\}\}\\big\(&\\pi\_\{\\theta\}\(y\_\{j\}\|q\_\{j\},c\_\{i\}\)\\\\ &\\parallel\\pi\_\{\\theta\+\\Delta\\theta\_\{i\}\}\(y\_\{j\}\|\\text\{KV\}\_\{q\_\{j\}\}\)\\big\)，\\end\{split\}