Compression Beyond the Uncompressed: A Two-Stage Training Recipe for Soft Context Compression in RAG

📄 arXiv: 2609.05152v1 📥 PDF

作者: Shuyu Guo, Shuo Zhang, Zhaochun Ren

分类: cs.CL

发布日期: 2026-09-04

💡 一句话要点

提出DEX-Comp以解决RAG模型上下文压缩效率问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 检索增强生成 上下文压缩 蒸馏训练 强化学习 问答系统

📋 核心要点

  1. 现有的软上下文压缩方法在训练时依赖于未压缩RAG系统的输出，限制了其性能提升的潜力。
  2. 本文提出DEX-Comp，通过两阶段训练方案，首先进行纯蒸馏，然后在未压缩RAG失败的查询上进行强化学习，优化压缩模型。
  3. 在五个开放域问答基准测试中，DEX-Comp实现了16倍的上下文压缩和4到24倍的推理加速，性能超过未压缩RAG基线。

📝 摘要（中文）

检索增强生成（RAG）通过外部知识增强语言模型，但冗长的检索上下文会增加输入长度并降低推理效率。软上下文压缩将每个文档编码为显著更短的嵌入序列。然而，现有方法通常通过蒸馏未压缩RAG系统的输出进行训练，限制了性能。为了解决这一问题，本文提出了DEX-Comp，一个两阶段的训练方案：纯蒸馏阶段仅在未压缩RAG的正确响应上预热压缩模型，随后通过强化学习在未压缩RAG失败的查询上进行硬探索，迫使模型探索更适合压缩表示的计算模式。在五个开放域问答基准测试中，DEX-Comp将检索上下文压缩了16倍，并加速推理4到24倍，同时在各个检索深度上实现了与未压缩RAG基线相当或更优的性能。

🔬 方法详解

问题定义：本文旨在解决现有软上下文压缩方法在训练时依赖未压缩RAG系统输出的问题，这限制了模型的性能提升。

核心思路：DEX-Comp的核心思路是通过两阶段训练方案，首先在未压缩RAG的正确响应上进行预热蒸馏，然后在未压缩RAG失败的查询上进行强化学习，以优化压缩模型的表现。

技术框架：DEX-Comp的整体架构分为两个主要阶段：第一阶段是纯蒸馏，第二阶段是硬探索。纯蒸馏阶段通过未压缩RAG的正确响应初始化压缩模型，而硬探索阶段则专注于未压缩RAG未能正确回答的查询。

关键创新：DEX-Comp的创新在于其两阶段训练策略，尤其是强化学习的应用，使得模型能够在特定失败查询上进行优化，从而更好地适应压缩表示。与现有方法相比，这种设计显著提升了模型的压缩能力和推理效率。

关键设计：在训练过程中，采用了特定的损失函数来平衡蒸馏和强化学习的目标，同时在网络结构上进行了调整，以支持更高效的上下文压缩。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DEX-Comp在五个开放域问答基准测试中表现出色，实现了上下文压缩16倍，推理加速4到24倍，并在各个检索深度上达到了与未压缩RAG基线相当或更优的性能，验证了其有效性。

🎯 应用场景

该研究的潜在应用领域包括智能问答系统、对话生成和信息检索等。通过提高上下文压缩效率和推理速度，DEX-Comp能够在资源受限的环境中实现更高效的知识检索与生成，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

Retrieval-Augmented Generation (RAG) enhances language models with external knowledge, but the lengthy retrieved context inflates the input and degrades inference efficiency. Soft context compression encodes each document into a substantially shorter embedding sequence. However, most existing approaches are trained by distilling outputs from uncompressed RAG systems, inherently limiting their performance relative to the original model. To address this limitation, we propose DEX-Comp, a two-stage training recipe: Pure Distillation warm-starts the compression model on the uncompressed RAG's correct responses only, and Hard Exploration then runs reinforcement learning solely on queries the uncompressed RAG fails, forcing the model to explore computation patterns better suited to compressed representations. On five open-domain QA benchmarks at retrieval depths from top-5 to top-30, DEX-Comp compresses retrieved contexts by $16\times$ and accelerates inference by $4\times$--$24\times$, while achieving performance comparable to or exceeding the uncompressed RAG baseline across retrieval depths. Ablations and evaluations across diverse datasets and backbones further confirm the contribution of each stage and the generalization of our approach.