Neural Garbage Collection: Learning to Forget while Learning to Reason
作者: Michael Y. Li, Jubayer Ibn Hamid, Emily B. Fox, Noah D. Goodman
分类: cs.LG
发布日期: 2026-04-20
💡 一句话要点
提出神经垃圾回收(NGC)，通过端到端学习实现语言模型在推理过程中自主遗忘，提升效率。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 神经垃圾回收 语言模型 链式思考 强化学习 内存管理 端到端学习 KV缓存 推理效率
📋 核心要点
- 链式思考推理的语言模型面临KV缓存增长瓶颈，限制了模型扩展。
- 提出神经垃圾回收(NGC)，使模型通过强化学习自主决定遗忘哪些信息，优化内存管理。
- 实验表明，NGC在保持准确率的同时，显著压缩了KV缓存大小，优于传统驱逐策略。
📝 摘要（中文）
链式思考推理显著提升了语言模型的能力，但每个推理步骤都会增加KV缓存，成为扩展这种模式的瓶颈。现有方法使用手工设计的标准来管理这些约束。一种更具扩展性的方法是让端到端学习完全包含这种设计选择。毕竟，如果模型可以学会推理，为什么不能学会忘记呢？我们引入神经垃圾回收（NGC），其中语言模型在学习推理的同时学习忘记，仅从基于结果的任务奖励进行端到端训练。当模型推理时，它会定期暂停，决定要驱逐哪些KV缓存条目，并继续以剩余缓存为条件进行推理。通过将链式思考中的token和缓存驱逐决策视为从语言模型中抽样的离散动作，我们可以使用强化学习来联合优化模型如何推理以及如何管理自己的内存：模型驱逐的内容会影响它记住的内容，它记住的内容会影响它的推理，而推理的正确性决定了它的奖励。至关重要的是，模型完全从单一的学习信号（基于结果的任务奖励）中学习这种行为，而无需监督微调或代理目标。在Countdown、AMC和AIME任务中，NGC相对于完整缓存的上限保持了强大的准确性，同时实现了2-3倍的峰值KV缓存大小压缩，并且大大优于驱逐基线。我们的结果是朝着更广泛的愿景迈出的第一步，在该愿景中，端到端优化驱动了语言模型的能力和效率。
🔬 方法详解
问题定义：现有链式思考推理的语言模型在推理过程中，每个步骤都会增加KV缓存，导致内存占用迅速增长，成为扩展此类模型的瓶颈。传统方法依赖于人工设计的规则来管理缓存，缺乏灵活性和适应性，难以在不同任务和场景中取得最佳效果。
核心思路：论文的核心思路是让语言模型自主学习如何管理自己的内存，即学习“忘记”哪些信息。通过将缓存驱逐决策视为模型行动的一部分，并使用强化学习进行端到端训练，模型可以根据任务奖励来优化推理和内存管理策略。这种方法避免了人工设计的规则，使模型能够更好地适应不同的任务需求。
技术框架：NGC框架包含一个语言模型，该模型在推理过程中会定期暂停，并决定要驱逐哪些KV缓存条目。整个过程可以分为以下几个阶段：1) 模型进行推理，生成token序列；2) 模型暂停推理，根据当前状态选择要驱逐的KV缓存条目；3) 模型根据剩余缓存继续推理，直到任务完成；4) 根据任务结果计算奖励，并使用强化学习算法更新模型参数。模型将token生成和缓存驱逐都视为离散动作，并使用策略梯度方法进行训练。
关键创新：NGC的关键创新在于将内存管理问题转化为一个可学习的决策过程，并使用强化学习进行端到端优化。与传统方法相比，NGC无需人工设计规则，而是让模型自主学习如何平衡推理能力和内存效率。这种方法具有更强的适应性和扩展性，可以应用于不同的任务和模型架构。
关键设计：NGC的关键设计包括：1) 使用语言模型作为策略网络，输出token和缓存驱逐决策；2) 使用基于结果的任务奖励作为强化学习的信号，引导模型学习有效的内存管理策略；3) 将缓存驱逐视为离散动作，可以使用策略梯度方法进行训练；4) 通过调整奖励函数，可以控制模型的内存使用量和推理准确率。
🖼️ 关键图片
📊 实验亮点
实验结果表明，NGC在Countdown、AMC和AIME等任务上，能够在保持较高准确率的同时，实现2-3倍的KV缓存压缩。与基于启发式规则的驱逐基线相比，NGC在性能上取得了显著提升，证明了端到端学习在内存管理方面的有效性。例如，在特定任务上，NGC在缓存压缩2倍的情况下，准确率仅下降了不到1%，远优于其他驱逐策略。
🎯 应用场景
NGC技术可应用于各种需要长程推理和记忆的语言模型应用，例如对话系统、问答系统、代码生成等。通过降低内存占用，NGC可以使这些模型在资源受限的设备上运行，并提高大规模部署的可行性。此外，NGC还可以促进对语言模型内部记忆机制的理解，为未来的模型设计提供指导。
📄 摘要（原文）
Chain-of-thought reasoning has driven striking advances in language model capability, yet every reasoning step grows the KV cache, creating a bottleneck to scaling this paradigm further. Current approaches manage these constraints on the model's behalf using hand-designed criteria. A more scalable approach would let end-to-end learning subsume this design choice entirely, following a broader pattern in deep learning. After all, if a model can learn to reason, why can't it learn to forget? We introduce Neural Garbage Collection (NGC), in which a language model learns to forget while learning to reason, trained end-to-end from outcome-based task reward alone. As the model reasons, it periodically pauses, decides which KV cache entries to evict, and continues to reason conditioned on the remaining cache. By treating tokens in a chain-of-thought and cache-eviction decisions as discrete actions sampled from the language model, we can use reinforcement learning to jointly optimize how the model reasons and how it manages its own memory: what the model evicts shapes what it remembers, what it remembers shapes its reasoning, and the correctness of that reasoning determines its reward. Crucially, the model learns this behavior entirely from a single learning signal - the outcome-based task reward - without supervised fine-tuning or proxy objectives. On Countdown, AMC, and AIME tasks, NGC maintains strong accuracy relative to the full-cache upper bound at 2-3x peak KV cache size compression and substantially outperforms eviction baselines. Our results are a first step towards a broader vision where end-to-end optimization drives both capability and efficiency in language models.