神经垃圾收集：学会遗忘，同时学会推理

思想链推理推动了语言模型能力的显着进步，但每个推理步骤都会增加 KV 缓存，从而为进一步扩展该范式造成瓶颈。当前的方法使用手工设计的标准代表模型管理这些约束。一种更具可扩展性的方法将使端到端学习完全包含这种设计选择，遵循深度学习中更广泛的模式。毕竟，如果一个模型可以学会推理，为什么它不能学会忘记呢？我们引入了神经垃圾收集（NGC），其中语言模型在学习推理的同时学习遗忘，仅根据基于结果的任务奖励进行端到端训练。当模型推理时，它会定期暂停，决定要逐出哪些 KV 缓存条目，并根据剩余缓存继续推理。通过将思想链中的标记和缓存驱逐决策视为从语言模型中采样的离散动作，我们可以使用强化学习来共同优化模型的推理方式以及它管理自己的记忆的方式：模型驱逐的内容决定了它记住的内容，它记住的内容决定了它的推理，而推理的正确性决定了它的奖励。至关重要的是，该模型完全从单个学习信号（基于结果的任务奖励）中学习这种行为，而无需监督 微调 或代理目标。在 Countdown、AMC 和 AIME 任务中，NGC 在 2-3 倍峰值 KV 缓存大小压缩下保持相对于全缓存上限的较高准确性，并且大大优于逐出基线。我们的结果是迈向更广阔愿景的第一步，其中端到端优化可提高语言模型的能力和效率。