PReM: Learning What to Preserve and When to Refresh for Context Compression
For LLMs handling long contexts, PReM addresses the limitation of static compression by dynamically adapting memory to reasoning needs, improving efficiency without sacrificing quality.
PReM introduces a context-compression framework that maintains long context as internal layer-wise KV memory, learning what to preserve and when to refresh it. It outperforms strong baselines under 16x and 32x compression with 32K-token contexts, balancing answer quality and inference efficiency.
Efficient long-context inference is not only about reducing memory cost, but also about keeping useful contextual evidence accessible as generation proceeds. However, existing compression-oriented approaches, such as key-value (KV) cache compression and context compression, often either make an early decision about which contextual information to keep or rely on an external compressor. Such designs make it difficult to adapt the compressed context to the evidence needed by later reasoning steps. This paper introduces PReM (Preserve and Refresh Memory), a context-compression framework that maintains the long context as the model's internal layer-wise KV memory and learns what to preserve and when to refresh it. Specifically, PReM uses a dedicated memory layer to make memory-selection decisions, and a special memory token <m> to trigger refreshes during generation. To train this behavior, PReM introduces Phase-Separated Refresh Training, aligning memory selection with memory-conditioned generation while preserving continuity across refreshes. Experiments with 32K-token contexts show that PReM outperforms strong baselines under both 16x and 32x compression, while maintaining a favorable balance between answer quality and inference efficiency.