PReM: Learning What to Preserve and When to Refresh for Context Compression

As long-context inference becomes standard, the memory cost of maintaining large Key-Value (KV) caches grows significantly. Existing compression methods often make static, early decisions about what information to keep, which fails to account for the evolving evidence needs of multi-hop reasoning tasks. This paper asks: can an LLM learn to dynamically manage its own context memory, preserving relevant information and refreshing it as reasoning progresses?

The authors introduce PReM (Preserve and Refresh Memory), a framework that treats context compression as an internal, learned process. Instead of relying on external compressors, PReM maintains the long context as layer-wise KV memory within the model. It uses two primary mechanisms:

1. Memory Selection Layer: A dedicated layer (typically in the middle-to-late stages of the model) that scores context chunks based on their relevance to the current reasoning step.
2. Refresh Token <m>: A special token that the model generates to trigger a memory refresh, allowing it to re-select context chunks whenever the current reasoning step requires different information.

To train this behavior, the authors employ Phase-Separated Refresh Training. This technique splits each generation step into a selection phase (where the model identifies relevant chunks) and a generation phase (where the model produces output conditioned on the updated memory). This ensures that the model learns to align its internal memory state with the specific requirements of the text it is currently generating.

Experiments on 32K-token contexts across various benchmarks (including TriviaQA, SQuAD, and multi-hop datasets like HotpotQA) demonstrate that PReM significantly outperforms static compression baselines. Under 16x and 32x compression ratios, PReM improves average exact match and F1 scores by up to 10.23 and 12.55 points, respectively, compared to the strongest baselines. Notably, a 3B-parameter model using PReM can outperform 7B-parameter models using standard soft-context compression, suggesting that dynamic memory management effectively compensates for smaller model scale.