AdaMem: Adaptive Memory Token Allocation for Soft Compression in Retrieval-Augmented Generation

Artem Sakhno; Grigorii Davydenko; Omar Zoloev; Julia Belikova; Andrey Savchenko; Maksim Makarenko

cs.CLAug 12, 2026

AdaMem: Adaptive Memory Token Allocation for Soft Compression in Retrieval-Augmented Generation

Authors: Artem Sakhno, Grigorii Davydenko, Omar Zoloev, Julia Belikova, Andrey Savchenko, Maksim Makarenko

Abstract

Retrieval-augmented generation (RAG) improves language models with retrieved evidence, but processing many long passages is costly and can introduce distracting information. Soft compression addresses this challenge by encoding passages as compact sequences of continuous memory embeddings before generation. However, existing methods typically assign each retained passage an identical number of memory embeddings, irrespective of its query-specific relevance. To address this, we propose AdaMem, a relevance-guided soft-compression framework that maps learned passage-relevance estimates to a query-dependent allocation of a fixed memory-token budget. A shared query-conditioned compressor produces both continuous passage memories and relevance scores in a single pass; a deterministic allocation rule assigns more memory tokens to higher-scoring passages and can omit low-scoring ones. Across six open-domain QA benchmarks, AdaMem consistently outperforms OSCAR (the closely matched soft-compression baseline that uses uniform allocation) as well as other soft-compression methods at matched memory budgets. Under standard 16 $\times$ compression, AdaMem improves sub-string match by up to 3.2 points (5.5%) over uniform allocation baseline, with an average relative gain of 3.4%; under aggressive 64 $\times$ compression the average relative gain grows to 14.6%, with a maximum of 9.8 points (19.7%) on PopQA. AdaMem matches the answer quality of the uncompressed at up to 4 $\times$ lower inference latency than full context baseline. AdaMem retains an efficiency profile comparable to the uniform-compression baseline, while achieving up to $4\times$ lower inference latency than full-context inference. Thus, relevance-guided memory allocation is particularly effective when retrieval pools are large and the available memory budget is tight.

Figures & tables

Figure 1: TriviaQA match as a function of the compression rate on 25 retrieved documents.

Figure 2: Overview of AdaMem. (a) A shared compressor jointly produces passage memory representations and relevance scores. (b) The scores determine a query-dependent allocation of the fixed budget $B$ via softmax and largest-remainder rounding. (c) The allocated memories are projected, ordered by relevance, and passed to a LoRA-adapted decoder for answer generation.

Method	TriviaQA	HotpotQA	NQ	PopQA	ASQA	BioASQ	Compr. rate
No Context	71.9	32.3	41.7	26.5	48.4	29.4	–
Full Context	84.4	51.5	66.3	60.2	70.9	33.1	–
RECOMP	76.0	37.0	50.0	34.8	55.9	29.0	200
LLMLingua-2	80.7	43.3	60.9	44.3	63.0	31.2	2
LLMLingua	76.2	37.2	55.0	39.0	58.4	27.2	2
xRAG	26.9	13.9	9.3	11.2	11.9	8.8	147

Table 1: Match (%) performance on six QA benchmarks with 25 retrieved documents. Hard- (RECOMP, LLMLingua, LLMLingua-2) and soft-compression (xRAG, PISCO) baselines are shown at their respective rates; OSCAR (uniform) and AdaMem (relevance-aware) share one compressor and are reported at

16

32

64\times

, with

32

64\times

at inference only. Per column, the top-3 compression methods are shaded green (darkest = best). The bootstrap-estimated standard deviations for TriviaQA, HotpotQA, NQ, PopQA, ASQA, and BioASQ are 1.0, 1.1, 1.1, 1.1, 1.5, and 1.2, respectively. Generation temperature:

\tau=1

Figure 3: Match (%) versus retrieval depth. AdaMem benefits from deeper retrieval more consistently than OSCAR, particularly under stronger compression.

Figure 4: Memory-token allocation by document relevance rank at $64\times$ compression. The dashed line denotes OSCAR’s fixed $16\times$ allocation.

Figure 5: Quality-latency tradeoff on TriviaQA. AdaMem shows similar latency as OSCAR and matches Full Context quality at $4.0\times$ lower latency.

Training data	TriviaQA	NQ	HotpotQA
Unfiltered	80.3	62.8	42.6
Filtered	82.8	63.3	44.5
Difference	+2.5	+0.6	+2.0

Table 2: Effect of filtered fine-tuning data on Match (%) at compression rate 16. Both models are trained for one epoch.

Training configuration	TriviaQA	NQ	HotpotQA
w/o Stage 1	71.30	44.45	29.35
w/o Stage 2	81.35	61.65	43.10
w/o Stage 3	81.05	59.30	42.70
Full three-stage training	82.80	63.30	44.50

Table 3: Effect of the three-stage training procedure at

16\times

compression with

k=25

Figure 6: Optimal allocation temperature on TriviaQA across compression rates. The analytically predicted schedule ( $\tau\propto B(\rho)$ , anchored at $16\times$ ) follows the empirical optimum and consistently outperforms a fixed $\tau=1$ .

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Rate	Lat. (ms)	Thr. (q/s)	TFLOPs	Mem. (GB)
Full Context	$1\times$	667.4	1.50	59.9	53.5
OSCAR	$16\times$	156.1	6.41	15.0	39.6
	$32\times$	149.3	6.70	13.4	39.0
	$64\times$	142.9	7.00	12.5	38.7
	$128\times$	136.4	7.33	12.1	38.5
AdaMem	$16\times$	192.4	5.20	21.0	56.4

Appendix

Table 4: Inference efficiency on TriviaQA with

k=25

: per-query latency, throughput, total compute, and peak GPU memory across inference-time compression rates.

Figure 7: LLM-as-a-judge accuracy under increasingly strong inference-time compression. AdaMem consistently retains higher answer quality than OSCAR as the memory budget decreases.

Stage	Examples	Epochs
Stage 1	2,000,000	1
Stage 2	500,000	1
Stage 3 (KILT-SCR)	196,916	3

Appendix

Table 5: Training schedule for the three-stage procedure.

RAG Context Compression

See all 12 papers

Sep 4, 2026

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

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.

Shuyu Guo, Shuo Zhang, Zhaochun Ren

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

RL for RAG LLM Inference Acceleration RAG Context Compression Retrieval-Augmented Generation Open-Domain QA Knowledge Distillation+1

Ask PDF

Sep 10, 2026

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

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.

Tuan Nguyen, Qiran Hu, Banruo Liu +3

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

LLM Inference Efficiency Retrieval-Augmented Generation RAG Context Compression

Ask PDF

Sep 25, 2026

HasMem: Hard-Origin Adaptively Softened Memory for Long-Term LLM Agents

Text-based memory and context compression support reuse of past interactions. Resizing continuous memory changes the input to a frozen LLM, coupling capacity allocation with readout. We propose Hard-Origin Adaptively Softened Memory (HasMem). Frozen hard-prompt embeddings provide a verifiable initial state. A controller adjusts memory widths, a Writer re-encodes resized entries, and Reader and Global provide readout adaptation and cross-turn state. On all

535

questions in a reconstruction probe derived from the Multi-Session Chat (MSC) development split, the main configuration achieves lexical F1 of

95.3

(

+4.4

percentage points) at

93.6\%

of the hard reference's framed memory positions. With approximately matched per-question target body budgets, six configurations at mean per-entry retention around

0.83

0.91

exceed rule-based re-encoding by

8.0

23.6

exact-match (EM) percentage points. With fixed model parameters and rule target width ratio

0.75

, Global's EM gain passes a user-level exact paired test with Bonferroni correction over eight comparisons. On all

500

LongMemEval-S questions, local lexical F1 rises from the hard reference's

3.4

8.9

, and answer negative log-likelihood (NLL) falls from

12.257

5.274

. F1 gains accompany lower EM on both evaluations.

Zihong He, Junxiao Shen, Chen Liang +1

The Hong Kong University of Science and Technology (Guangzhou) · University of Bristol · The Hong Kong University of Science and Technology

Memory Management for LLM Agents Long-Term Conversational Memory Agent Memory Management Memory-Augmented Language Model Agents

Ask PDF

Abstract

RAG Context Compression

Explore similar work

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

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

HasMem: Hard-Origin Adaptively Softened Memory for Long-Term LLM Agents