arXiv:2604.18002cs.LG2026-04被引 1

让大模型自己学会在推理时删无用记忆,提升效率

Neural Garbage Collection: Learning to Forget while Learning to Reason

  • 用强化学习让模型自主决定何时丢弃推理缓存
  • 在2-3倍缓存压缩下仍保持接近全缓存的准确率
  • 仅靠任务奖励信号训练,无需人工设计规则

链式思维推理推动了语言模型能力的显著进步,但每一步推理都会增长键值缓存(KV cache),成为进一步扩展的瓶颈。当前方法依赖人工设计的淘汰策略来管理缓存,而更可扩展的方案应让端到端学习完全替代这一设计选择。既然模型能学会推理,为何不能学会遗忘?我们提出神经垃圾回收(Neural Garbage Collection, NGC),让语言模型在推理过程中自主学习遗忘行为,仅通过任务结果奖励信号进行端到端训练。模型在推理时周期性暂停,判断哪些缓存条目应被清除,再基于剩余缓存继续推理。将思维链中的每个标记和缓存清除决策视为从语言模型中采样的离散动作,使用强化学习联合优化推理与内存管理:删除什么影响记忆,记忆影响推理,推理正确性决定奖励。关键在于,模型完全由单一学习信号——任务结果奖励——驱动,无需监督微调或代理目标。在Countdown、AMC和AIME任务上,NGC在2-3倍峰值缓存大小压缩下保持强准确率,显著优于传统淘汰基线。这为端到端优化同时提升模型能力和效率迈出了第一步。

原文摘要 · Abstract (English)

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.

大模型推理内存管理强化学习链式思维

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。