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

2026-07-15

AI Summary

Existing context compression methods struggle to dynamically retain information useful for subsequent generation in long-context reasoning, often lacking adaptability to inference needs due to premature decisions or reliance on external compressors. This work proposes PReM, a framework that maintains long contexts as layer-wise key-value (KV) memory within the Transformer architecture and enables adaptive context compression and refreshing during generation through a learnable memory selection mechanism, dedicated memory layers, and a memory-trigger token \<m\>. Coupled with a stage-separated training strategy, PReM significantly outperforms strong baselines at both 16× and 32× compression ratios under a 32K-token context length, achieving an effective balance between answer quality and inference efficiency.

Abstract

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.

Problem

• context compression
• long-context inference
• KV cache
• memory management
• evidence preservation

Innovation

• context compression
• KV cache
• memory refresh
• long-context inference
• adaptive memory