PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shuyu Guo, Shuo Zhang, Zhaochun Ren

arXiv:2609.05152v1cs.CL

TL;DR

RAG’s retrieved contexts can be costly, while soft compression methods remain limited by distillation from uncompressed systems. DEX-Comp combines correct-response distillation with failure-focused reinforcement learning, achieving efficient performance comparable to or exceeding uncompressed RAG across retrieval depths.

  • Problem

    Retrieved RAG contexts increase inference cost, while existing soft compression methods are typically bounded by the behavior of uncompressed RAG through response distillation.

  • Method

    DEX-Comp uses Pure Distillation on correctly answered teacher examples, followed by Hard Exploration through reinforcement learning on queries the uncompressed model fails.

  • Results

    Across five open-domain QA benchmarks and top-5 to top-30 retrieval depths, DEX-Comp achieves performance comparable to or exceeding uncompressed RAG while compressing contexts by 16× and accelerating inference by 4×–24×.

  • Takeaways & Limitations

    DEX-Comp extends query-independent soft compression beyond the uncompressed RAG baseline across multiple retrieval settings and generalizes across backbone models and data distributions.

  • Takeaways & Limitations

    The evaluation focuses on QA benchmarks and answer-correctness metrics, leaving faithfulness, grounding, and robustness under adversarial or long-form generation for future work.

Abstract

from arXiv · show

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.

1 Introduction

RAG improves knowledge-intensive tasks by retrieving external documents, but long contexts increase inference cost and can exceed the model’s context window. DEX-Comp addresses soft compression’s training limitation with targeted distillation and exploration, achieving efficient performance comparable to or exceeding uncompressed RAG.

  • Retrieved contexts increase RAG inference latency and can exceed the underlying LLM’s context window.
  • Soft compression encodes documents into shorter continuous embedding sequences, but existing methods are typically bounded by uncompressed RAG behavior through response distillation.
  • DEX-Comp combines Pure Distillation on correct uncompressed-RAG responses with Hard Exploration on queries the uncompressed model fails.The two stages provide a competent initialization and encourage behaviors suited to compressed representations.
  • 16× compression and 4×–24× inference acceleration accompany performance comparable to or exceeding uncompressed RAG across retrieval depths from top-5 to top-30.The evaluation covers five open-domain QA benchmarks and reports stronger results than prior soft compression methods.

2 Related Work

Prior work reduces RAG computation through inference-time KV compression or pre-inference hard and soft context compression. Within soft compression, existing approaches explore varied compression mechanisms and training objectives, while DEX-Comp targets the underexplored combination of reliable imitation and targeted exploration.

  • Context compression reduces RAG’s effective decoder computation through inference-time KV compression or pre-inference context shortening.
  • Hard compression extracts salient tokens or summarizes documents, whereas soft compression encodes context into continuous representations.
  • Soft compression research spans auto-encoding, cross-modal projection, attention-based aggregation, reconstruction or continuation pretraining, and answer-oriented supervision.
  • Prior methods also study joint end-to-end optimization, adaptive compression-ratio selection, and combinations of hard and soft compression.
  • Prior work largely uses imitation-based or globally optimized objectives without explicitly combining imitation with exploration for soft compression.DEX-Comp addresses this limitation by separating reliable imitation from targeted exploration on challenging cases.

3 Method

DEX-Comp combines an offline document compressor with a compression-aware decoder and trains them in two stages. Pure Distillation initializes the model from teacher-correct queries, while Hard Exploration uses teacher-failed queries for reinforcement-learning-based exploration.

  • 3.1 DEX-Comp Architecture: DEX-Comp replaces each retrieved document with cached soft embeddings consumed by an autoregressive decoder.The compressor encodes documents offline, and the decoder generates responses from the cached embeddings and query.
  • 3 Method: The two-stage design separates competent initialization from adaptation to computation patterns suited to compressed representations.The coupled compressor and decoder lack known oracle representations or parameters, motivating staged optimization.
  • 3.2 Stage I: Pure Distillation: Pure Distillation trains only on queries answered correctly by the uncompressed teacher, aligning compressed and uncompressed response distributions.The teacher reads original tokens while the student reads compressed embeddings, with training based on Kullback–Leibler divergence.
  • 3.3 Stage II: Hard Exploration: Hard Exploration applies Group Relative Policy Optimization exclusively to teacher-failed queries, sampling response groups and scoring them against gold answers.The resulting clipped policy-gradient update is regularized toward the Pure Distillation checkpoint.
  • 3.3 Stage II: Hard Exploration: Focused exploration on teacher-failed queries enables DEX-Comp to match or exceed the uncompressed teacher across standard RAG benchmarks.The paper motivates this restriction as a way to discover behaviors not captured by teacher predictions.

4 Experimental Setup

The experiments evaluate DEX-Comp on five open-domain QA datasets across retrieval depths from top-5 to top-30 against uncompressed and compression baselines. Evaluation uses both containment exact match and LLM-as-judge, with a Mistral-7B backbone and offline query-independent compression.

  • 4 Experimental Setup: Five datasets cover factoid, multi-hop, ambiguous long-form, and long-tail open-domain question answering.The benchmarks are Natural Questions, TriviaQA, HotpotQA, ASQA, and PopQA.
  • 4.2 Baseline Methods: The study compares DEX-Comp with uncompressed RAG and representative hard- and soft-compression baselines.The baselines include LLMLingua-2, xRAG, ICAE, COCOM, and PISCO.
  • 4 Experimental Setup: Evaluation spans retrieval depths from top-5 to top-30, including a retrained PISCO baseline at top-30 for fair comparison.The wider range is intended to address prior work’s emphasis on shallow retrieval settings.
  • 4.3 Evaluation: CEM checks whether the ground-truth answer appears in the prediction, while LLM-as-judge assesses correctness against the ground truth.Gemini 3 Flash2 is used as the judge model, and CEM is retained as a reproducible metric.
  • 4 Experimental Setup: All methods use Mistral-7B inference, while DEX-Comp uses offline query-independent document encoding and separate LoRA adapters for its compressor and decoder.The implementation uses deterministic decoding, with preprocessing and retrieval handled through BERGEN, Splade-v3, and DeBERTa-v3.

5 Experimental Results

DEX-Comp surpasses the uncompressed RAG baseline at 16× compression across datasets and retrieval depths while substantially reducing inference cost. It also generalizes across retrieval depths when trained only at top-30, with limited exceptions.

  • 5.1 Main Results: At 16× compression, DEX-Comp surpasses the uncompressed RAG baseline on every dataset and retrieval depth with statistically significant gains.Average CEM / LLM improvements are +3.82 / +2.01, +3.89 / +2.34, and +3.12 / +2.53 at k=5, 15, and 30.
  • 5.1 Main Results: DEX-Comp is reported as the first query-independent soft compression method to surpass uncompressed RAG across multiple realistic retrieval depths.The paper characterizes this as a new state of the art among comparable compression methods.
  • 5.1 Main Results: A single model trained at top-30 outperforms uncompressed RAG at k=5, 15, and 30 on average, except on PopQA at top-5 and top-15.Average score improvements are +0.34 / +0.29, +1.73 / +1.27, and +3.12 / +2.53 at k=5, 15, and 30.
  • 5.2 Computational Efficiency: At top-30, DEX-Comp reduces peak GPU memory by 4.02×, computational cost by 13.61×, and TTFT by 23.73×.Efficiency was measured with batch size 8 on a single RTX PRO 6000 Blackwell GPU.
  • 5.2 Computational Efficiency: DEX-Comp requires a one-time offline corpus-compression pass, with compressed embeddings fixed until documents are updated.For dynamic corpora, incremental recompression can be applied only to updated documents.

6 Analysis

The analysis attributes DEX-Comp’s gains to stronger robustness and information use, complementary training stages, parameter choices, and generalization across models and domains. Embedding analyses indicate specialization between local and global document information.

  • 6.2 What makes DEX-Comp effective?: DEX-Comp surpasses uncompressed RAG on Resilience Rate and Boost Rate when training and inference depths match.The larger Resilience Rate gains suggest compression acts as an implicit denoising mechanism.
  • 6.2 What makes DEX-Comp effective?: Removing either Pure Distillation or Hard Exploration causes consistent performance drops, while teacher-correct distillation and teacher-failed-query reinforcement learning perform better.The ablation identifies the two stages as complementary and jointly responsible for the improvements.
  • 6.3 Training Parameters: Accuracy improves as compression rate decreases; CR = 16 performs best, while CR = 32 matches uncompressed RAG.Training and inference depths perform best when matched, but larger training k generalizes better across depths.
  • 6.4 Generalization: DEX-Comp surpasses corresponding uncompressed RAG baselines across backbone families and matches or exceeds them on BioASQ, CovidQA, and FEVER.These results indicate robustness across domains and tasks.
  • 6.5 How do compression embeddings work?: Learned embeddings show spatial specialization, with some embeddings focusing on document regions and others capturing global information.This indicates that DEX-Comp allocates limited embedding capacity between local and global representations.

7 Limitations

DEX-Comp has two stated limitations: offline, query-independent compression requires corpus-wide precomputation and storage, and the evaluation does not isolate compression from task-specific reinforcement learning or test broader generation qualities.

  • Offline compression requires precomputing and storing compressed embeddings for the corpus.The cost is amortized across queries but may become significant for frequently updated or very large-scale corpora.
  • The current setup does not strictly isolate gains from compression versus task-specific reinforcement learning.The authors leave a more controlled comparison for future work.
  • Evaluation focuses on QA benchmarks and answer-correctness metrics rather than faithfulness, grounding, or robustness in adversarial and long-form generation.The authors identify these broader settings as directions for future investigation.

8 Conclusion

The paper concludes that DEX-Comp combines Pure Distillation with Hard Exploration and narrows the gap between compressed and uncompressed RAG. Across five benchmarks and retrieval settings, it preserves or exceeds uncompressed-RAG performance while substantially reducing context and inference costs.

  • DEX-Comp combines Pure Distillation with Hard Exploration in a two-stage soft context compression recipe.
  • 16× context compression and over 20× faster inference accompany performance comparable to or exceeding uncompressed RAG across five benchmarks at top-30 retrieval.
  • Further analyses confirm contributions from both stages and generalization across backbone models and data distributions.
  • The authors position DEX-Comp as a step toward understanding soft context compression and narrowing the compressed–uncompressed RAG gap.

A LLM-as-Judge Prompt

The appendix prompt defines an LLM-as-judge accuracy evaluation: given a question, accepted golden answers, and a generated answer, the judge returns a binary correctness decision.

  • The prompt is used to query a judge model for the LLM-as-judge accuracy metric.
  • The judge must respond only with Yes or No.
  • An answer is correct when it matches any one of the provided golden answers.

B Pure Distillation Hyperparameters

Pure Distillation trains on the teacher-correct subset S+, whose membership is determined by whether the ground truth appears in the model prediction; its hyperparameters are listed in Table 4.

  • Pure Distillation is performed on the teacher-correct subset S+.
  • Subset correctness is determined by whether the ground truth is contained within the model prediction.
  • Table 4 lists the core hyperparameters for Pure Distillation.

C Hard Exploration Hyperparameters

Hard Exploration initializes from the Pure Distillation checkpoint and applies GRPO to the teacher-failed subset. Its core hyperparameters are provided in Table 5.

  • Hard Exploration initializes from the Pure Distillation checkpoint.
  • GRPO is run on the teacher-failed subset S−.
  • Table 5 lists the core hyperparameters for Hard Exploration.
Loading 2609.05152v1…