大规模Cartridge：在海量文档集合上训练模块化KV缓存

# 大规模文档集合上的模块化 KV 缓存训练

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

Momchil Hardalov Gonzalo Iglesias Adrià de Gispert
Amazon AGI \{momchilh, gjii, agispert\}@amazon\.com

###### 摘要

大型语言模型能够对长上下文进行推理，但预填充数百万个 token 的成本极高，而其中大量内容在不同查询之间保持静态不变。Cartridges\(Eyuboglu et al\., 2025 (https://arxiv.org/html/2606.04557#bib.bib8)\) 通过将文档集合提炼为可复用的键值（KV）缓存来解决这一问题，从而在保持准确性的同时消除了预填充开销。该方法的一个关键局限在于：cartridge 是单体式、不可组合的——将整个集合编码到单个 KV 块中无法扩展，而将独立训练的 cartridge 简单混合则会导致性能崩溃至接近随机水平。我们提出了*Cartridges at Scale*（CAS），这是一个可扩展多 cartridge 学习的训练框架，具备动态干扰项混合机制，以及一个在 GPU 与持久存储之间轮换数百个单文档 cartridge 的内存高效预算管理器。我们的方法可扩展至超过百万 token 的集合，在相近 token 预算下比单体式 cartridge 提升 10–31 个百分点。即使在高压缩率下，Oracle cartridge 的准确率与完整上下文学习的差距也仅在 2–6 个百分点以内。配合检索进行 cartridge 选择后，CAS 在消耗 3–4× 更少提示 token 的同时，达到或超过传统 RAG 的准确率。

大规模 Cartridges：在大型文档集合上训练模块化 KV 缓存

Momchil Hardalov Gonzalo Iglesias Adrià de Gispert
Amazon AGI \{momchilh, gjii, agispert\}@amazon\.com

## 1 引言

参见图注图 1：LongHealth 上单文档训练 Cartridges 的性能\(Adams et al\.\(2024 (https://arxiv.org/html/2606.04557#bib.bib14)\)，文档数：20 份患者病历，\∼\\sim232K token\)。*Oracle*——提示中仅加载正确患者的 Cartridge；*Full*——加载全部 cartridge。详见第 §4 (https://arxiv.org/html/2606.04557#S4) 节，*训练机制*。

大型语言模型（LLM）能够对上下文窗口中的信息进行推理\(Snell et al\., 2025 (https://arxiv.org/html/2606.04557#bib.bib24); OpenAI, 2024 (https://arxiv.org/html/2606.04557#bib.bib27), 2025 (https://arxiv.org/html/2606.04557#bib.bib28); DeepSeek\-AI, 2025 (https://arxiv.org/html/2606.04557#bib.bib25); Comanici et al\., 2025 (https://arxiv.org/html/2606.04557#bib.bib29); Anthropic, 2025 (https://arxiv.org/html/2606.04557#bib.bib30); Agarwal et al\., 2025 (https://arxiv.org/html/2606.04557#bib.bib23)\)，然而这一能力的代价随上下文长度线性增长——无论是延迟（预填充时间）还是内存（KV 缓存大小）均如此。对于需要持续访问文档集合的应用场景——如企业知识库、患者病历和法律语料库——为每次查询反复处理相同文档的成本极高。检索增强生成（RAG，Lewis et al\.\(2020 (https://arxiv.org/html/2606.04557#bib.bib31)\); Guu et al\.\(2020 (https://arxiv.org/html/2606.04557#bib.bib70)\)）通过在查询时仅检索最相关的文本片段来缓解这一问题，但每个检索到的片段都是原始文档的*碎片*，迫使模型在不完整信息上进行推理。

此前已有工作提出通过 token 淘汰或合并来缩减 KV 缓存大小 Zhang et al\.\(2023 (https://arxiv.org/html/2606.04557#bib.bib34)\); Kang et al\.\(2024 (https://arxiv.org/html/2606.04557#bib.bib37)\); Liu et al\.\(2026 (https://arxiv.org/html/2606.04557#bib.bib35)\)，或将文本压缩为软 Mu et al\.\(2023 (https://arxiv.org/html/2606.04557#bib.bib64)\); Cheng et al\.\(2024 (https://arxiv.org/html/2606.04557#bib.bib32)\); Louis et al\.\(2025 (https://arxiv.org/html/2606.04557#bib.bib33)\) 或离散 Jiang et al\.\(2023 (https://arxiv.org/html/2606.04557#bib.bib38)\); Pan et al\.\(2024 (https://arxiv.org/html/2606.04557#bib.bib39)\) token 表示。这些方法往往需要代价高昂的运行时预填充、依赖独立的编码器，或在处理长文档时丢失信息 Łajewska et al\.\(2025 (https://arxiv.org/html/2606.04557#bib.bib40)\)。近来，研究人员开始探索将文档编码为紧凑、可复用的表示，在推理时无需额外计算代价即可加载到模型 KV 缓存中 Devoto et al\.\(2025 (https://arxiv.org/html/2606.04557#bib.bib42)\); Kim et al\.\(2026 (https://arxiv.org/html/2606.04557#bib.bib43)\); Zweiger et al\.\(2026 (https://arxiv.org/html/2606.04557#bib.bib44)\)。

*Cartridges* Eyuboglu et al\.\(2025 (https://arxiv.org/html/2606.04557#bib.bib8)\) 提出通过上下文蒸馏为每个文档集合训练少量键值向量，在实现 10–100× 压缩的同时保留模型回答编码内容相关问题的能力。然而，该方法依赖于每个集合的单一整体式 cartridge，无法扩展至现实世界的文档集合：将所有文档加载到一个缓存中会用无关 token 膨胀前缀，而信息密集型内容（如财务报表、临床病历、技术规格说明书）在高压缩比下不可避免地会产生信息损失。此外，单一整体式 cartridge 对未来的变更缺乏灵活性：更新或新增单份文档需要重新编码整个 KV 缓存。

自然而然的替代方案——每份文档一个 cartridge——却暴露出更为根本性的缺陷：如图 1 (https://arxiv.org/html/2606.04557#S1.F1) 所示，*将独立训练的 cartridge 混合会导致模型崩溃*——每个 cartridge 在 oracle 场景下（仅加载一个正确的 cartridge）表现良好，但同时加载*所有 cartridge* 时性能骤降至接近随机水平，原因在于模型从未学会在独立训练的 KV 前缀之间进行选择性注意。

在本工作中，我们提出 CAS，一个用于在包含数百份文档、数百万 token 的集合上训练和部署 cartridge 的框架。CAS 提出了一种基于动态 cartridge **轮换与混合**的训练机制，该机制消除了独立训练的性能差距（图 1 (https://arxiv.org/html/2606.04557#S1.F1)），同时在当前 GPU 内存约束下保持可行性。我们进一步证明，CAS 可与密集检索自然结合用于 cartridge 选择，提供了一条匹配或超越文本 RAG 效率的实用部署路径。

我们的贡献如下：

- •我们证明，独立训练的 cartridge 在推理时会产生灾难性干扰，导致性能崩溃至接近随机水平；而使用干扰项 cartridge 进行联合训练可消除这一退化现象。
- •我们提出 CAS，一个可扩展的训练框架，支持 GPU↔\\leftrightarrow 持久存储换出、优化器卸载和优先级轮换，能够在固定 GPU 内存预算内对数百个 cartridge 进行联合训练。
- •我们提出通过按长度比例采样和批量多问题生成来改进自学习数据生成，在提升事实覆盖率的同时将合成成本降低高达 20× 。
- •我们证明，单文档 cartridge 在五个多样化基准测试中始终以最高 30 个百分点的优势超越单一整体式 cartridge，压缩容忍度因内容类型而异——短技术文档在 100× 压缩下接近无损，而表格密集的财务文件最多可接受 5% 的相对性能下降。
- •我们证明，Cartridge RAG 在使用最多 4× 更少提示 token 的情况下，达到或超越 Text RAG 的准确率，将压缩 KV 缓存的高效性与密集检索的选择性相结合。

参见图注
\(a\) 自学习。对于每份文档 $d_i$，我们使用 $M_Q$ 采样一组问题，随后收集答案以及目标教师模型 $M_A$ 的 logits。

参见图注
\(b\) 训练。每份文档 $d_i$ 被分配一个专属 cartridge。*预算管理器*在 GPU 上仅保留 $B$ 个 cartridge，并每 $R$ 步轮换一次 cartridge 池。

参见图注
\(c\) 推理。选择函数从存储中激活 $k$ 个 cartridge。所选 cartridge 被拼接后前置于查询 token，再进行解码。

图 2：Cartridges at Scale（CAS）端到端流水线。

## 2 Cartridges at Scale（CAS）训练

#### 背景。

我们的工作基于 Eyuboglu et al\.\(2025 (https://arxiv.org/html/2606.04557#bib.bib8)\) 的 Cartridge 框架。Cartridge $Z \in \mathbb{R}^{L \times p \times d \times 2}$ 是一个紧凑的可训练 KV 缓存，将包含 $n_{\mathcal{C}}$ 个 token 的上下文编码为固定长度 $p \ll n_{\mathcal{C}}$ 的前缀，注入模型的 $L$ 个注意力层（维度为 $d$），实现 10–100× 的压缩比，同时可直接集成到现有推理服务中。Cartridge 通过上下文蒸馏目标进行训练，训练数据通过*自学习*生成：将语料库切分为子语料库 $\tilde{\mathbf{c}}$，LLM 使用多样化的种子提示为每个子语料库生成合成问答对\(Eyuboglu et al\., 2025 (https://arxiv.org/html/2606.04557#bib.bib8)\)。损失函数最小化*教师*（上下文中包含子语料库的模型）与*学生*（使用 Cartridge 增强的同一模型）之间的 KL 散度。所有模型参数保持冻结——仅 $Z$ 中的键值向量接收梯度。

### 2\.1 多 Cartridge 联合训练

尽管 Eyuboglu et al\.\(2025 (https://arxiv.org/html/2606.04557#bib.bib8)\) 讨论了通过拼接 KV 缓存来组合独立训练的 Cartridge，但这种组合方式从未经过明确测试。图 1 (https://arxiv.org/html/2606.04557#S1.F1) 表明，混合*孤立*训练的缓存——独立训练且不感知推理时可能共存的其他 Cartridge——会导致模型崩溃。此外，朴素地扩展至非常大的 Cartridge 规模也不可行，因为在大型文档集合上训练的内存需求极高。[^1]

[^1]: 一个 1K token 的 Qwen3\-8B Cartridge 在反向传播时占用约 551 MiB（500K token 约需 274 GiB）。梯度和 Adam 动量会使 Cartridge 的内存占用扩大 4×。

为解决这些局限，我们提出 CAS，一个支持同时使用 $N$ 个 cartridge 的训练和推理框架（图 2 (https://arxiv.org/html/2606.04557#S1.F2)）。

#### 混合可见性训练。

混合可见性训练是学习 KV 表示的关键，它使冻结模型能够在存在干扰项的情况下选择性地关注相关 cartridge，直接针对多 cartridge 推理场景进行优化。为此，我们对 KL 损失进行如下改进：设 $Z^* \in \{Z_1, \ldots, Z_N\}$ 为样本 $\mathbf{x}$ 对应的相关 cartridge，$\mathcal{Z}_\text{pool}$ 为 cartridge 池，$\mathcal{Z}_\text{dist} = \mathcal{Z}_\text{pool} \setminus Z^*$ 为干扰项候选集，$k \sim \mathcal{U}(k_\text{min}, k_\text{max})$，则损失变为：

$$\tilde{Z} = \begin{cases} Z^* & \text{以概率 } P_{\text{iso}} \\ Z^* \cup S,\; S \sim \mathcal{U}(\mathcal{Z}_{\text{dist}}, k) & \text{以概率 } 1 - P_{\text{iso}} \end{cases} \tag{1}$$

$$\mathcal{L}(Z) = \sum_{(\mathbf{x}, \tilde{c}) \in \mathcal{D}} \sum_{i=1}^{|\mathbf{x}|} D_{\text{KL}}\!\left(F(\cdot \mid \tilde{c} \oplus \mathbf{x}_{<i}) \;\Big\|\; \cdots\right)$$

---

在 Cartridge 设置中，文档上下文被编码在训练好的 KV 缓存前缀中，不包含在文本提示里——模型仅接收系统提示和用户问题。在 Oracle（文本）基线中，完整的文档文本被前置于用户消息。表 15 (https://arxiv.org/html/2606.04557#A8.T15) 列出了所用的确切提示。

| 基准测试 | 系统提示 + 用户消息格式 |
|---|---|
| T2\-RAGBench / FinQA | 系统：你是一位通过基于简单语法构建数学公式来回答财务问题的专家。- 任务：根据给定上下文，为问题提供公式答案。指南：1. 答案类型：公式可以是一个数字，或以下之一：add(f1, f2)、subtract(f1, f2)、multiply(f1, f2)、divide(f1, f2)、exp(f1, f2)、greater(f1, f2)。2. 推理：仔细分析上下文，特别注意表格。3. 最终答案：" FORMULA "用户：{context}问题：{question} |
| LongHealth | 系统：请参考患者病历回答用户的问题。选择唯一最佳选项，并严格按选项文字作答。将答案包裹在：此处为正确选项文字 中。用户：{question} A. {option_a} B. {option_b} C. {option_c} D. {option_d} E. {option_e} |
| QASPER | 系统：你是一位回答科学论文相关问题的研究助手。尽量简洁地作答，只给出答案，无需解释。如问题无法从论文中得到回答，请说"Unanswerable"。对于是/否问题，回答"Yes"或"No"。将答案包裹在：… 中。用户：{question} |
| QuALITY | 系统：你是一位仔细阅读并回答长篇文章多项选择题的读者。阅读文章后选择唯一最佳答案选项，将答案字母（A、B、C 或 D）包裹在 answer 标签中。示例：B 用户：{question} A. {option_a} B. {option_b} C. {option_c} D. {option_d} 用答案字母（A、B、C 或 D）作答。 |
| TechQA | 系统：你是一位专注于 IT 基础设施、软件产品和企业系统的专业技术支持助手。尽你所能回答技术问题，简明扼要，基于事实。将答案包裹在：YOUR_ANSWER 中。用户：{question} |

表 15：各基准测试所用评估提示。花括号中的变量在评估时替换为基准测试数据。

## 附录 I 自学习数据合成详情

#### 问题生成模型。

我们使用 GPT\-OSS 120B Agarwal et al\.\(2025 (https://arxiv.org/html/2606.04557#bib.bib23)\) 作为问题生成器（$M_Q$）。该模型在系统提示中接收完整文档上下文，每次调用批量生成 20 个多样化问题。答案生成器（$M_A$）为目标模型本身（Qwen3\-8B），确保蒸馏信号反映学生自身的输出分布。

#### 种子提示类型。

每次合成调用从五类中随机选择一种种子提示类型：*结构化*（请求将信息整理为 JSON、YAML 或其他格式）、*摘要*（请求总结特定章节）、*问答*（事实回忆和推理题）、*用例*（实际下游应用任务）以及*创意*（开放式讨论提示）。这种多样性确保 cartridge 在多种交互模式上训练，而不仅限于事实性问答。

#### 多问题提示格式。

在批量模式下，问题生成器接收如下指令（以*问答*类型为例）：

> 生成 {n} 个多样化问题，测试对上述语料库信息的掌握程度。每个问题应涵盖语料库的不同事实、细节或方面。变换风格：混合事实回忆、比较、推理和细节导向的问题。在每个问题中包含具体细节（编号、名称、标题、日期、数值等），以明确所问内容。仅输出一个 JSON 字符串数组，例如 \["问题 1", "问题 2", ...\]。不要输出其他文本、markdown 围栏或解释。

结构化 JSON 输出格式支持可靠解析；解析失败的响应将被丢弃（通常占调用次数的 $<5\%$）。

#### 采样轮次与温度。

对于每个数据集，我们针对每种种子提示类型（问答、结构化、摘要……）分别运行 4 轮独立采样。