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

2026-09-04

AI Summary

本文针对RAG中因长上下文导致的推理效率降低问题，提出DEX-Comp方法，通过两阶段训练实现软上下文压缩，显著提升推理速度而不损失性能。

Abstract

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.

Problem

• Retrieval-Augmented Generation
• soft context compression
• inference efficiency

Innovation

• Two-Stage Training
• Soft Context Compression
• Pure Distillation
• Hard Exploration
• Retrieval-Augmented Generation