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.CL, cs.IR, cs.LG
2026-08-12
AdaMem allocates RAG memory tokens by passage relevance, beating uniform-allocation baseline OSCAR by 14.6% at 64x compression with 4x lower latency.
Retrieval-augmented generation systems compress retrieved passages into continuous memory embeddings to cut inference cost. Existing soft-compression methods like OSCAR use a relevance score only to rank or drop passages, then give every surviving passage the same number of memory tokens regardless of how useful it actually is for the question. When retrieval pools are large and the memory budget is tight, this flat allocation wastes capacity that should go to the passages that actually matter.
AdaMem uses one query-conditioned compressor to produce, in a single forward pass, both continuous memory embeddings for each passage and a relevance score, read off the final hidden state at a dedicated <RERANK> token appended to every passage. Scores are standardized within the retrieval pool, then converted into a fractional memory-budget share for each passage via a softmax with temperature τ, and rounded to integers with largest-remainder rounding; passages scoring low enough get zero tokens and are effectively dropped. τ controls how sharply the budget concentrates: as τ approaches infinity the rule collapses to OSCAR's uniform split, as τ approaches zero it approaches rerank-then-truncate. The authors prove this softmax rule is the exact solution to a logarithmic-utility budget allocation problem under Kelly's proportional-fairness framework, and a first-order approximation of the reverse-water-filling optimum from classical bit-allocation theory; the theory predicts optimal τ should shrink as compression gets more aggressive, which the experiments confirm (optimal τ drops from 1.0 at 16x to 0.5 at 64x). Training runs in three stages: query-independent autoencoding pretraining, query-dependent joint training with the relevance head initialized by cross-encoder distillation, and a final fine-tuning stage on a new dataset, KILT-SCR, with a combined generation and ranking loss.
At the standard 16x compression rate with 25 retrieved passages, AdaMem beats OSCAR, the closest uniform-allocation baseline sharing the same architecture, across all six QA benchmarks:
| Compression | Metric | AdaMem | OSCAR | Full context |
| 16x | TriviaQA Match | 84.0 | 82.4 | 84.4 |
| 16x | PopQA Match | 61.6 | 58.4 | 60.2 |
| 64x | TriviaQA Match | 83.3 | 77.1 | 84.4 |
| 64x | PopQA Match | 59.6 | 49.8 | 60.2 |
The average relative gain is 3.4% at 16x and grows to 14.6% at 64x, with a 9.8-point (19.7%) jump on PopQA. The bigger story is stability. Between 16x and 128x compression, AdaMem's quality drops by only about 2%, versus more than 7% for OSCAR over the same range. AdaMem at 32x already beats OSCAR at 16x, meaning it reaches better quality with roughly half the memory tokens. Against the strongest hard-compression baseline, LLMLingua-2, AdaMem wins by 9.7% on average while using an 8x smaller context. On latency, AdaMem at 32x matches full-context quality at 4x lower latency and 3.6x less compute.
For RAG systems that retrieve large pools, 25 or more documents, under a tight memory budget, AdaMem is a near-free upgrade. It reuses the relevance score that existing pipelines already compute for ranking and spends it on allocation instead of discarding it after selection. Because the authors train PISCO, OSCAR, and AdaMem on identical data and hyperparameters, the only variable is the allocation rule, which makes the reported gains attributable to the method rather than to confounds. Teams already running a query-conditioned soft compressor could plausibly bolt this on without redesigning the pipeline.
The optimality proof rests on an assumed logarithmic-utility model for how passage value scales with token count; whether that assumption actually holds is not independently verified. Experiments cover only English open-domain QA with one fixed compressor-decoder pair, a Llama-3.2-1B compressor and a Mistral-7B-Instruct decoder, so it's unclear whether the gains hold at different model scales. BioASQ is the one benchmark where a hard-compression method, LLMLingua-2, beats AdaMem outright, and the paper doesn't explain why the advantage disappears for biomedical QA. AdaMem's peak memory is also somewhat higher than OSCAR's, since it must encode the full candidate bank before allocating, an overhead that matters more at lower compression rates.