Compression Beyond the Uncompressed: A Two-Stage Training Recipe for Soft Context Compression in RAG
Shuyu Guo, Shuo Zhang, Zhaochun Ren
cs.CL
2026-09-04
DEX-Comp distills only correct uncompressed RAG answers, then runs GRPO on failures. At 16× compression it matches or beats the teacher on five QA sets, with up to 24× faster TTFT.
RAG prepends retrieved documents to the query, so latency grows with context and the window can overflow. Soft compression encodes each document into a short sequence of continuous embeddings, which carry more per slot than extractive or abstractive hard compression. Most prior recipes distill the uncompressed RAG's outputs, so the student inherits computation patterns shaped by reading discrete tokens. Theory says the compressed slots can hold more; practice has stayed below the uncompressed teacher.
Prior evaluations also stop at top-5. Production retrievers often return top-15 or top-30, and shallow-k numbers do not transfer.
DEX-Comp is a two-stage recipe. Compressor and decoder share a Mistral-7B backbone with separate LoRA adapters. Compression runs offline: each document is concatenated with mi=⌊Li/τ⌋ compression tokens, and the last-layer hidden states at those positions become the cached embedding. At inference the decoder reads only those vectors plus the query.
Stage I is Pure Distillation. The training set is split by whether the teacher's answer matches the gold label, S⁺ versus S⁻. Only S⁺ is used. The student minimizes KL between the teacher's response distribution and its own, with gradients flowing through both decoder and compressor. Failed teacher answers are held out so their mistakes are not cloned into the initialization.
Stage II is Hard Exploration. GRPO runs only on S⁻. Each query draws a group of rollouts, scored 0/1 against the gold answer, with group-normalized advantages and a KL tether to the PD checkpoint. On queries the teacher already misses, imitation has no correct target, so the student has to find computation that fits compressed inputs.
Evaluation covers NQ, TriviaQA, HotpotQA, ASQA, and PopQA at retrieval depths top-5 through top-30 on 128-token Wikipedia-KILT chunks. Splade-v3 retrieves, DeBERTa-v3 reranks. Metrics are Containment Exact Match and Gemini 3 Flash as judge. Baselines include uncompressed Mistral-7B, LLMLingua-2, xRAG, ICAE, COCOM, and PISCO, with PISCO retrained at top-30 for a fair deep-k comparison.
With matched train/test depth and 16× compression, DEX-Comp beats uncompressed RAG on every dataset and depth (paired permutation test, p<0.05). Average CEM/LLM: 68.38/72.54 vs 64.56/70.53 at top-5, 67.83/72.88 vs 63.94/70.54 at top-15, 66.03/72.02 vs 62.91/69.49 at top-30. PISCO at the same 16× and top-30 training reaches only 57.90 CEM. A single top-30-trained model still beats the uncompressed average across depths, with PopQA at top-5/15 the exception.
Efficiency is measured at batch 8 on one RTX PRO 6000. At top-30, TTFT falls from 4782ms to 201ms (23.73×), GFLOPs drop 13.61×, and peak memory from 59328MB to 14766MB (4.02×). Top-5 already has 4.42× TTFT. Compression rate 16 is best; 32 already matches the uncompressed baseline.
Both stages matter. SFT on all teacher outputs lands at 64.79 average LLM; PD alone at 67.75; PD plus RL on correct-only or mixed slices at 70.60 and 71.47, still below the full 72.02. Resilience Rate gains are larger, so compression also acts as a denoiser. The pattern holds across backbone families and on BioASQ, CovidQA, and FEVER.
This is the first query-independent soft compressor shown to beat uncompressed RAG across several realistic retrieval depths. For systems that can cache the corpus once, 16× shorter context and up to 24× faster first token come with a two-to-four point CEM lift, not a quality trade. The recipe is not tied to one compressor architecture.
The uncompressed RAG was never given task-specific RL; the authors say long-context training was too expensive. The student received GRPO on short compressed inputs, the teacher did not, so some of the "win" may be "short inputs are easier to RL" rather than "compressed embeddings are better." They leave a cleaner isolation to future work.
Offline compression and storage of the whole corpus will hurt if the collection churns. Evaluation is open-domain QA correctness only: faithfulness, citations, adversarial inputs, and long-form generation are untested. Soft vectors are opaque, so a wrong answer is hard to trace to a source span.