💬 NLP

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

PReM is a novel context-compression framework that dynamically learns to preserve and refresh internal layer-wise KV memory via a dedicated memory layer and special tokens, enabling efficient long-context inference without relying on external compressors or making premature retention decisions.

Original authors: Bohan Yu, Lei Shen, Chenxi Zhou, Chen Han, Junlin Liu, Wenbo Su, Yu Cheng, Bo Zheng

Published 2026-07-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Bohan Yu, Lei Shen, Chenxi Zhou, Chen Han, Junlin Liu, Wenbo Su, Yu Cheng, Bo Zheng

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive, complex mystery, but instead of a single clue, you are handed a library containing 32,000 pages of documents. Your brain (or in this case, a computer program called a Large Language Model) has to read through all of it to find the specific facts needed to answer a question. The problem is that holding that entire library in your "working memory" at once is incredibly heavy and slow. It's like trying to carry a library in a backpack while running a race; eventually, you get tired, and you might drop the very book you need.

To fix this, scientists have tried two main tricks so far. The first is "compression," where they try to shrink the library down to a single, tiny summary before you even start reading. The second is "caching," where they keep the whole library but try to only look at the most popular pages. But both of these have a flaw: they make a decision about what to keep before you start solving the puzzle. If the mystery changes halfway through and you suddenly need a page from the back of the book that was thrown away, you're stuck. You can't go back and grab it because the decision was made too early. This paper introduces a new way to handle this, suggesting that instead of a static summary, we need a memory system that can update itself as the story unfolds.

Enter PReM (Preserve and Refresh Memory), a clever new framework that treats the model's memory like a dynamic, living notebook rather than a static filing cabinet. The core idea is simple but powerful: don't just decide what to keep once; decide what to keep right now, and be ready to swap things out the moment the story demands it.

Here is how PReM works, using a fun analogy. Imagine you are a detective solving a case with a team of assistants (the different layers of the AI). In the old methods, the team would read the whole case file, pick out their favorite 10 pages, and lock the rest in a vault. If you later needed a page from the vault, too bad.

PReM changes the rules. It keeps the whole case file, but it organizes it into small "chunks" (like chapters). As the detective starts writing the solution, a special "Memory Token" (let's call it a Magic Bookmark, written as ⟨m⟩) acts as a signal. When the detective writes this bookmark, it's like saying, "Wait! I need to rethink what's on my desk."

At that exact moment, a dedicated "Memory Manager" (a special layer in the AI) wakes up. It looks at the current question and the Magic Bookmark, then quickly scans all the chapters. It decides: "Okay, for this next step of the story, we definitely need Chapter 3 and Chapter 12. We can forget the rest for now." It then swaps the old chapters on the desk with the new ones, keeping the important details sharp and clear while compressing the unneeded ones into a tiny summary. This happens during the writing process, not before.

The paper shows that this "refresh" mechanism is a game-changer. In tests using 32,000-token contexts (that's a lot of text!), PReM managed to compress the memory by 16 times and even 32 times. Despite having so much less information "on the desk" at any one time, it actually answered questions better than models that tried to keep the whole library open. For example, on a 7-billion-parameter model, PReM improved the accuracy of answers by up to 12.55 points compared to the best existing methods.

The researchers also discovered something interesting about where this memory manager should live. They tested placing the "Memory Manager" in the early, middle, or late parts of the AI's brain. They found that the middle and late layers were much better at deciding what to keep, while the early layers were too slow to learn this skill. It's like having a senior detective (the middle layers) make the call on what evidence is relevant, rather than a rookie (the early layers).

One of the most surprising findings is that PReM is so good at this "learned" memory management that a smaller model (3 billion parameters) using PReM could outperform larger models (7 billion parameters) that used other compression tricks. This suggests that knowing when to refresh your memory is just as important as how big your brain is.

The paper also rules out a few ideas. It shows that simply using the AI's natural attention (looking at what it's currently thinking about) isn't enough to decide what to keep; you need a specific, trained system to do the selection. It also shows that you can't just compress the text once at the beginning and hope for the best; the "static" approach of keeping the same compressed summary throughout the whole answer leads to worse results, especially for complex, multi-step questions.

In short, PReM suggests that the future of efficient AI isn't just about shrinking data, but about building a memory that breathes. It learns to preserve the evidence it needs for the current step of the conversation and knows exactly when to hit the "refresh" button to swap in new evidence for the next step. By doing this, it keeps the AI fast and efficient without making it forget the clues it needs to solve the mystery.

Technical Summary: PReM (Preserve and Refresh Memory)

1. Problem Statement

Efficient long-context inference in Large Language Models (LLMs) faces a fundamental tension between reducing computational costs and maintaining access to necessary contextual evidence. As context length increases, the Key-Value (KV) cache grows, significantly increasing decoding latency and memory usage.

Existing approaches to mitigate this, such as KV-cache compression and context compression, suffer from two primary limitations:

Static Evidence Retention: Most methods make a single, early decision on which context to retain. This static compression fails to adapt to the evolving needs of multi-step reasoning, where different reasoning steps may require different pieces of evidence.
Decoupled Evidence Selection: Methods relying on external compressors separate the compression process from generation. This introduces additional computation and prevents the model from tailoring evidence selection dynamically to specific reasoning steps.

The paper argues that effective compression must go beyond simple size reduction; it must dynamically preserve useful evidence and refresh it when later reasoning steps require different information.

2. Methodology: PReM

PReM (Preserve and Refresh Memory) is a context-compression framework that maintains long contexts as internal, layer-wise KV memory within the LLM. Instead of relying on external compressors, PReM learns to decide what to preserve and when to refresh based on the model's internal state.

Core Mechanisms

Layer-Wise KV Memory Representation: The long context is encoded as layer-wise KV states ( $M_\ell$ ) stored within the model. These states are partitioned into contiguous chunks.
Dedicated Memory Layer: A specific layer ( $\ell_m$ $ℓ_{m}$ ) is designated to make memory-selection decisions. This layer uses a dedicated selection query projection matrix ( $W^c_\ell$ $W_{ℓ}^{c}$ ) to score context chunks.
- Chunk Selection: The model scores chunks based on the current query (from the question or a refresh token).
- Preserve-and-Pool: Selected chunks are preserved at token-level granularity. Unselected chunks are compressed into a single averaged KV representative.
The Memory Token $\langle m \rangle$ : A special token serves as an explicit internal signal for refreshing memory. When the model generates $\langle m \rangle$ , the memory layer re-evaluates the context chunks, selects new relevant evidence, and updates the compressed KV memory for subsequent layers. This allows the model to control when to refresh based on the current reasoning step.

Training: Phase-Separated Refresh Training

To train the model to align memory selection with generation, PReM introduces a Phase-Separated Refresh Training strategy. Each reasoning step is split into two distinct forward passes:

Memory-Selection Phase: The model processes the input (question or $\langle m \rangle$ ) to compute chunk scores and select the top- $k$ evidence chunks. This phase is supervised with a Memory Selection Loss ( $L_{ms}$ ) that encourages the model to rank evidence chunks higher than non-evidence chunks.
Memory-Conditioned Generation Phase: The model generates the answer segment ( $y_s$ ) conditioned on the newly refreshed memory state resulting from the selection phase.

To ensure consistency between training and inference, the training objective includes boundary cross-entropy terms ( $L_{bd}$ ). These terms connect the selection phase to the first answer token and the end of an answer segment to the next $\langle m \rangle$ token, ensuring the model learns the autoregressive links between memory updates and generation.

3. Key Contributions

Formulation: The paper reframes context compression as a learned memory preservation and refresh problem, eliminating the need for external compressors.
Architecture: PReM introduces a dedicated memory layer and a special $\langle m \rangle$ token to enable dynamic, step-specific memory updates within the LLM's internal KV states.
Training Strategy: The proposal of Phase-Separated Refresh Training, which decouples memory selection from generation during training to ensure the model learns to condition generation on the specific memory state required for that step.
Empirical Validation: Comprehensive experiments demonstrating that learned, dynamic memory updates outperform static compression and full-context prompting.

4. Experimental Results

Experiments were conducted on 3B and 7B variants of Qwen2.5-Instruct using 32K-token contexts across six benchmarks (TriviaQA, SQuAD, NaturalQuestions, 2WikiMQA, HotpotQA, MuSiQue).

Performance vs. Baselines: Under both 16× and 32× compression ratios, PReM outperformed eight strong baselines, including KV-cache methods (SnapKV, StreamingLLM, CAKE) and context compression methods (LongLLMLingua, EXIT, ICAE).
- On the 3B backbone at 16× compression, PReM improved average Exact Match (EM) and F1 by 5.30 and 5.15 points respectively over the strongest baseline.
- On the 7B backbone at 32× compression, improvements reached 8.20 (EM) and 9.89 (F1) points.
Comparison to Full Context: PReM outperformed direct 32K-token prompting (without compression) on both 3B and 7B models, suggesting that learned memory selection is more effective than exposing the model to the full unfiltered context.
Small Model Superiority: The 3B PReM model outperformed released 7B soft context-compression baselines (e.g., Activation Beacon), indicating that learned preservation and refresh can compensate for smaller backbone scales.
Efficiency: PReM reduced peak GPU memory significantly (e.g., ~19.57GB vs. 30.0GB for full memory) while maintaining competitive inference latency. The use of $\langle m \rangle$ tokens allows for a controllable trade-off between decoding computation (refresh frequency) and memory savings.

5. Significance and Claims

The paper claims that PReM demonstrates that effective long-context inference requires generation-coupled memory updates rather than static context shortening. By maintaining the context as internal KV memory and learning to refresh it dynamically, the model can adaptively access the specific evidence needed for each reasoning step.

The authors emphasize that their approach addresses the limitations of static retention and decoupled selection, showing that a model can learn to "know what to preserve and when to refresh," leading to superior performance in multi-hop reasoning tasks and better efficiency-quality trade-offs compared to existing compression techniques. The work suggests that internal memory management, trained jointly with generation, is a viable path for scaling LLMs to longer contexts without prohibitive computational costs.