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

Shuyu Guo; Shuo Zhang; Zhaochun Ren

cs.CLSep 4, 2026

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

Authors: Shuyu Guo, Shuo Zhang, Zhaochun Ren

Organizations: Shandong University Qingdao, China · Bloomberg London, United Kingdom · Leiden University Leiden, The Netherlands

Abstract

Retrieval-Augmented Generation (RAG) improves knowledge-intensive generation by conditioning language models on retrieved documents, but processing these documents becomes increasingly expensive as retrieval depth grows. Soft context compression reduces this cost by encoding documents into compact continuous representations that can be precomputed and reused across queries. However, many existing methods train compressed models by distilling from a full-context teacher. When the teacher is wrong, such distillation can reinforce its errors, while teacher imitation provides no direct signal for improving beyond the teacher. We propose DEX-Comp, a two-stage training recipe that separates reliable imitation from targeted exploration. Pure Distillation learns only from teacher-correct questions to mitigate error propagation, while Hard Exploration applies outcome-based reinforcement learning to teacher-failed questions to directly optimize answer correctness. Across five open-domain QA benchmarks and retrieval depths from top- $5$ to top- $30$ , DEX-Comp at $16\times$ compression outperforms all evaluated compression baselines and surpasses the untuned full-context RAG model in average accuracy, while reducing time-to-first-token by $4.4\times$ -- $23.7\times$ . Evaluations across additional datasets and backbones further demonstrate its generalization.

Figures & tables

Figure 1: (a) Representative case. Standard distillation reproduces an incorrect answer from the full-context teacher, whereas Pure Distillation (PD) and PD with Hard Exploration (PD+HE) answer correctly. (b) Outcome distribution. Under greedy decoding on a fixed subset of teacher-failed questions, PD reduces the frequency of reproducing teacher errors, while HE further increases correct answers. (c) Sampling performance. On the same questions, PD improves pass@

1

and pass@

3

over standard distillation, with further gains from HE.

Figure 2: Overview of DEX-Comp. (a) Soft compression. A compressor turns each document into a shorter sequence of soft embeddings by selecting the final-layer hidden states at compression-token positions. The embeddings can be precomputed offline and consumed by the decoder at inference time. (b) Pure Distillation. On teacher-correct queries ( $\mathcal{S}^{+}$ ), KL distillation aligns the compressed model’s response distribution with that of a frozen RAG teacher reading the full retrieved documents. (c) Hard Exploration. Initialized from Stage I, the compressed model trains on teacher-failed queries ( $\mathcal{S}^{-}$ ) using outcome-based GRPO, directly rewarding answer correctness rather than teacher agreement.

Top-5
Method	Comp. Rate	NQ		TriviaQA		HotpotQA		ASQA		PopQA		Avg.
Method	Comp. Rate	CEM	LLM	CEM	LLM	CEM	LLM	CEM	LLM	CEM	LLM	CEM	LLM
Full context	-	58.65	73.67	90.18	89.04	46.45	55.57	68.99	77.00	58.53	57.36	64.56	70.53
LLMLingua-2	$4\times$	49.74	63.59	84.41	83.19	36.43	45.50	59.49	65.30	39.07	39.70	53.83	59.46
xRAG	$128\times$	37.36	51.46	78.06	78.54	27.66	35.27	43.25	50.63	28.15	29.12	42.90	49.00
ICAE	$4\times$	41.59	56.33	77.78	79.27	28.75	40.48	47.68	60.65	38.28	40.24	46.82	55.39

Table 1: Results with Mistral-7B at top-

5/15/30

. Best and second-best scores at each depth are bold and underlined. † : retrained at the corresponding top-

k

because matching checkpoints are unavailable.

Top- $k$	TTFT (ms)			GFLOPs			Memory (MB)
Top- $k$	RAG	DEX-Comp	Speedup	RAG	DEX-Comp	Reduction	RAG	DEX-Comp	Reduction
Top-5	373	84	4.42 $\times$	82,850	12,423	6.67 $\times$	16,100	14,242	1.13 $\times$
Top-15	1,626	131	12.33 $\times$	245,650	22,654	10.84 $\times$	26,887	14,457	1.86 $\times$
Top-30	4,782	201	23.73 $\times$	518,767	38,129	13.61 $\times$	59,328	14,766	4.02 $\times$

Table 2: Inference efficiency: TTFT, computation (GFLOPs), and peak GPU memory (MB).

Initialization	RL subset	NQ	TriviaQA	HotpotQA	ASQA	PopQA	Average
Initialization	RL subset	NQ	TriviaQA	HotpotQA	ASQA	PopQA	Acc	Resilience	Boost
Standard distill.	None	67.57	86.59	50.20	72.68	46.89	64.79	85.25	45.96
PD	None (w/o HE)	70.96	88.74	55.32	74.79	48.95	67.75	87.63	49.56
Standard distill.	Failure (HE)	69.46	86.83	51.85	74.22	49.67	66.41	87.01	48.16
PD	Success	74.41	90.11	57.30	77.43	53.74	70.60	89.58	53.31
PD	Mixed	75.15	90.10	58.13	78.90	55.07	71.47	90.13	54.49

Table 3: Top-

30

ablation. Dataset scores and Acc are LLM-judged accuracy (%). RL subsets contain teacher-correct (Success), teacher-failed (Failure), or both (Mixed) questions.

Configuration	NQ	TriviaQA	HotpotQA	ASQA	PopQA	Avg.	GPU-hours	Max. RL batch
Full context	$73.67$	$89.04$	$55.57$	$77.00$	$57.36$	$70.53$	–	–
Full context + RL	$80.23$	$92.36$	$65.00$	$82.28$	$64.02$	$76.78$	$38$	$16$
PD ( $16\times$ )	$71.04$	$88.25$	$54.72$	$74.18$	$51.41$	$67.92$	$1.5$	–
PD + HE ( $16\times$ )	$76.52$	$89.97$	$57.55$	$79.96$	$58.70$	$72.54$	$13$	$64$

Table 4: Top-

5

LLM-judged accuracy (%) and training cost on eight GPUs. Max. RL batch is the largest per-GPU batch fitting in memory; GPU-hours count total GPU training time.

Figure 3: Resilience Rate (top) and Boost Rate (bottom) across retrieval depths. Subscripts denote the training top- $k$ .

Figure 4: Effects of (a) compression rate and (b) training depth, with the other fixed.

Figure 5: Generalization across (a) backbone families and (b) out-of-distribution datasets.

Figure 6: Compression embeddings: (a) token cosine similarity; (b) top logit-lens tokens.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Assignment
optimizer	AdamW
learning rate	$1\!\times\!10^{-4}$
warmup ratio	$0.05$
weight decay	$0.1$
gradient clipping	$1.0$
epochs	$1$

Appendix

Table 5: Hyperparameters for Pure Distillation.

Hyperparameter	Assignment
optimizer	AdamW
learning rate	$5\!\times\!10^{-6}$
warmup ratio	$0.05$
weight decay	$0.01$
gradient clipping	$1.0$
gradient checkpointing	True

Appendix

Table 6: Hyperparameters for Hard Exploration.

Figure 7: PISCO compression embeddings learned by standard distillation, for the same passage as Figure 6 . (a) Cosine similarity between embeddings and document tokens. (b) Document tokens appearing among the top- $10$ logit-lens vocabulary projections of the memory embeddings. Colors identify embeddings within each model.

Explore similar work

REVA: Reusable Evidence View Aggregation for Context-Efficient RAG Serving
Sep 11, 2026Tuan Nguyen, Qiran Hu, Banruo Liu +3Retrieval-Augmented Generation Pipelines Interaction History
VinUni-Illinois Smart Health Center, VinUniversity, Vietnam · University of Illinois Urbana-Champaign, USA
RAGOCR: Optical Compression of Retrieval-Augmented Text via Visual Representation
Aug 1, 2026Jiayang Yu, Jialun Zhong, Lei ZouKnowledge-Based Visual Question Answering Lingdt-Vl-Ocr
Wangxuan Institute of Computer Technology, Peking University Beijing, China
Fixed RAG Compression Collapses Measured Reader Scaling
Jun 20, 2026Sugam Panthi, Rabab AbdelfattahData Compression Methods Question-Answering Benchmarks
The University of Southern Mississippi

Sep 11, 2026, cs.LGJ/K move · Enter open · S save

REVA: Reusable Evidence View Aggregation for Context-Efficient RAG Serving

Tuan Nguyen, Qiran Hu, Banruo Liu, Khoa D. Doan+2

VinUni-Illinois Smart Health Center, VinUniversity, Vietnam · University of Illinois Urbana-Champaign, USA

Retrieval-Augmented Generation Pipelines Interaction History

Open paper AskPDF arXiv

Retrieval-augmented generation (RAG) improves knowledge-intensive large language model (LLM) applications by conditioning generation on retrieved documents, but longer contexts increase latency, key-value (KV) cache memory, and token cost. Post-retrieval compression can reduce this cost, yet existing compressors often operate independently for each query, rely on auxiliary models or rewriting, and introduce online overhead that can offset the benefit of shorter prompts. We revisit RAG compression from a data-mining perspective by aggregating historical query--document--model interactions into reusable evidence views. We first show that modern compressors have unstable gains over simple truncation and can add substantial inference-time latency. We then propose Reusable Evidence View Aggregation (REVA), a framework that mines the target generator's historical attention traces into a document-keyed, budget-agnostic score store. REVA maps token-level attention to readable word units, aggregates importance across repeated document accesses, and renders budget-specific plain-text views that preserve document order and the standard RAG interface. Across four representative benchmarks and modern LLMs, REVA improves generation quality by 1.0--5.8 points over existing advances, while reducing compression overhead by a factor of 5.3 to 15.6, adding less than 40 ms of latency.

Abstract

Explore similar work

REVA: Reusable Evidence View Aggregation for Context-Efficient RAG Serving

RAGOCR: Optical Compression of Retrieval-Augmented Text via Visual Representation

Fixed RAG Compression Collapses Measured Reader Scaling