PReM：学习保留什么与何时刷新，实现高效上下文压缩

背景

长上下文推理的高效性不仅关乎降低内存成本，还在于随着生成过程的推进，保持有用的上下文证据可访问。然而，现有的压缩方法，例如键值（KV）缓存压缩和上下文压缩，往往要么过早决定保留哪些上下文信息，要么依赖外部压缩器。这类设计使得压缩后的上下文难以适应后续推理步骤所需的证据，导致模型在长序列生成中容易丢失关键信息或产生冗余。因此，需要一种能够动态调整上下文保留与刷新的机制。

核心内容

PReM（Preserve and Refresh Memory）是一种上下文压缩框架，它将长上下文维护为模型内部的逐层KV记忆，并学习哪些内容应该保留以及何时应该刷新。具体来说，PReM使用一个专门的内存层（memory layer）来做出记忆选择决策，并引入一个特殊的内存标记 <m> 来在生成过程中触发刷新。为了训练这种行为，PReM提出了相位分离刷新训练（Phase-Separated Refresh Training），该方法在保持刷新间连续性的同时，将记忆选择与记忆条件化的生成对齐。在32K token上下文的实验中，PReM在16倍和32倍压缩率下均优于强基线，同时保持了答案质量与推理效率之间的良好平衡。

关键要点

• 核心挑战：现有压缩方法（如KV缓存压缩、上下文压缩）通常静态决定保留内容，无法动态适应推理过程中的证据需求。
• PReM框架：利用模型内部的逐层KV记忆来维护长上下文，通过一个专用内存层学习“保留什么”和“何时刷新”。
• 特殊标记：使用特殊token <m> 在生成过程中触发记忆刷新。
• 训练方法：相位分离刷新训练，将记忆选择与记忆条件生成对齐，同时确保刷新间的连续性。
• 实验效果：在32K token上下文中，16倍和32倍压缩率下均优于强基线，在答案质量和推理效率之间取得良好平衡。

意义与影响

PReM提出了一种新颖的上下文压缩范式，将压缩决策从单向的早期选择转变为动态的、与生成过程耦合的机制。这使得模型能够在长序列推理中自适应地保留关键信息并适时刷新过时内容，从而提升长上下文任务的表现。其内置的内存层和特殊标记设计提供了一种可端到端训练的方法，无需外部压缩器，简化了系统集成。PReM在显著压缩KV缓存的同时保持甚至提升答案质量，对构建低成本、高效率的长上下文推理系统具有重要价值，尤其在需要处理大量上下文（如文档问答、多轮对话、代码生成）的场景中潜力巨大。