让检索增强生成更快更准,压缩上下文16倍仍保持高精度。
Compression Beyond the Uncompressed: A Two-Stage Training Recipe for Soft Context Compression in RAG

- 分两阶段训练:先蒸馏正确答案,再强化学习纠错场景。
- 压缩16倍,推理速度提升4到24倍,效果不降反升。
- 适合需要高效推理的开放域问答系统部署。
检索增强生成(RAG)通过外部知识提升语言模型性能,但长篇检索上下文导致输入过长,降低推理效率。软上下文压缩将每篇文档编码为更短的嵌入序列。然而,现有方法多通过蒸馏未压缩RAG系统的输出进行训练,性能受限于原始模型。为此,本文提出DEX-Comp,一种两阶段训练方案:纯蒸馏阶段仅使用未压缩RAG的正确响应初始化压缩模型;硬探索阶段则仅对未压缩RAG失败的查询进行强化学习,迫使模型探索更适合压缩表示的计算模式。在五项开放域问答基准上,检索深度从前5到前30,DEX-Comp实现16倍上下文压缩,推理加速4至24倍,且在所有检索深度下表现与或优于未压缩基线。消融实验及跨数据集、骨干网络的评估进一步验证了各阶段贡献和方法泛化能力。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。