Neural Garbage Collection 论文解读：让 LLM 学会遗忘
今天来聊一篇来自斯坦福大学的最新工作 —— Neural Garbage Collection (NGC)，一个让语言模型边推理边学习遗忘的框架。核心思想极其简洁：如果模型能学会推理，为什么不能学会管理自己的记忆？
一、问题：CoT 推理的内存瓶颈
Chain-of-Thought (CoT) 推理极大提升了 LLM 的能力，但也带来了一个严重的副作用 —— KV Cache 爆炸。推理过程中每一步生成的 token 都会累积到 KV Cache 中，长推理链意味着巨大的内存占用。
现有解决方案的问题：
- SnapKV：基于注意力权重启发式保留重要 token
- KeyDiff：基于 key 的多样性评分
- KNorm：基于 key 范数统计量
- StreamingLLM：滑动窗口 + attention sinks
这些方法的共同缺陷：都是人工设计的启发式规则，不是从任务目标中学到的。它们不知道哪些信息对最终答案真正重要。
二、核心洞察：效率是可学习的能力
NGC 的核心主张 —— 效率应该像推理一样被端到端学习。
"If a model can learn to reason, why can't it learn to forget?"
关键思路：将 KV Cache 的驱逐决策视为离散动作，与 token 生成一样从模型中采样，然后用同一个强化学习信号联合优化。
这个闭环非常优雅：
- 模型驱逐什么 → 决定它记住什么
- 它记住什么 → 塑造它的推理
- 推理的正确性 → 决定奖励信号
- 奖励信号 → 反向优化驱逐策略和推理策略
三、NGC 方法详解
3.1 增长-驱逐循环 (Grow-Then-Evict)
NGC 采用周期性的驱逐策略：
- 每生成 δ 个 token，触发一次驱逐轮次
- 模型对所有当前 KV Cache 条目打分
- 保留 (1−ε) 比例的条目，永久删除其余
- 驱逐在每一层独立进行
数学上，峰值 Cache 大小收敛到一个稳态值：
C* = L · (δ / ε)
这个值与推理总长度无关 —— 无论模型思考多久，Cache 大小始终有界。
3.2 用注意力机制打分（零额外参数）
NGC 的一个精妙设计：不引入任何新参数，直接复用 Transformer 自身的注意力机制来评估 KV 条目的重要性。
具体做法：
- 取最近 w=5 个 query
- 计算它们对所有 prefix key 的 softmax 注意力分数
- 跨 heads 和 queries 平均 → 每个 token 的重要性分数 ψ_t
为了降低信用分配难度并配合硬件友好的内存访问模式，key 被分组为大小 b=32 的连续块，块级分数是块内有效 key 分数的平均值。
3.3 Gumbel-Top-k 随机采样
驱逐不是确定性的 top-k，而是随机动作 —— 这是为了支持无偏的策略梯度估计。
- 对块分数加入 i.i.d. Gumbel 噪声
- 保留噪声后分数最高的 K 个块
- 选中子集的 log-probability 有闭式解，可用前缀和技巧高效计算
3.4 联合优化：一个奖励信号，两个梯度
NGC 使用 Dr. GRPO（Group Relative Policy Optimization）进行训练，核心创新是将 token 生成和驱逐决策统一到一个优化目标下：
L = L_token + L_mem
- L_token：标准 RLVR 的 token 级策略梯度
- L_mem：驱逐决策的策略梯度（每层独立求和）
关键：两者使用同一个 advantage Â_i —— 即最终任务正确性相对于组内平均的偏差。这意味着驱逐决策直接从任务成败中学习，而不是依赖任何代理目标。
由于驱逐分数来自模型自身的 Q 和 K，L_mem 的梯度会流入几乎所有模型权重 —— 模型学会的不只是"何时驱逐"，还包括如何生成更容易区分重要性的 Q/K 表示。
3.5 Replay Attention Masks
标准 RLVR 假设所有先前 token 都可见。NGC 打破了这个假设：被驱逐的 key 不再可见。为了正确计算 log-probability，NGC 使用 Replay Attention Masks 精确复现驱逐决策导致的每层可见性模式。
四、实验结果
4.1 Countdown 任务
| 方法 | pass@1 准确率 | 峰值 Cache 压缩比 |
|---|---|---|
| StreamingLLM | 0.2% | — |
| KNorm | 0.1% | — |
| KeyDiff | 21.2% | — |
| SnapKV | 7.8% | — |
| NGC (ours) | 49.6% | 2.4× |
NGC 在 2.4× Cache 压缩 下，准确率是次优基线（KeyDiff）的 2.3 倍。
4.2 AIME 2025（数学竞赛）
| 方法 | pass@32 准确率 | 峰值 Cache 压缩比 |
|---|---|---|
| StreamingLLM | 0.0% | — |
| KNorm | 0.0% | — |
| KeyDiff | 7.1% | — |
| SnapKV | 10.7% | — |
| NGC (ours) | 21.4% | 4.6× |
在 4.6× 压缩 的极端条件下，NGC 仍然达到 21.4% 的 pass@32，是 SnapKV 的 2 倍。
4.3 AMC + AIME 综合表现
NGC 在 2–3× 峰值 KV Cache 压缩 下，保持了强劲的推理性能，全面超越所有启发式驱逐基线。
4.4 关键发现
- 无监督微调、无代理目标、无预热阶段 —— 纯 outcome-based reward 驱动
- 模型学会了主动遗忘：驱逐决策不是随机的，而是与推理内容高度相关
- 驱逐策略在不同层之间差异化：不同层学会保留不同类型的信息
五、为什么这个方法很重要？
范式转变：从系统优化到能力学习
传统上，KV Cache 管理是系统工程师的工作 —— 设计启发式、调阈值、做工程优化。NGC 把它变成了模型的内在能力。
这与 DeepSeek-R1 的 RLVR 范式一脉相承：
- R1 证明了推理能力可以通过纯 RL 从可验证奖励中学到
- NGC 证明了效率也可以 —— 同一个框架，同一个信号，同时优化能力和资源
对本地部署的启示
对于本地模型用户（LMStudio/Ollama），NGC 的意义：
- 更长的 CoT 推理：同样的 GPU 内存可以支持更长的思考链
- 更大的上下文窗口：2-5× 的 Cache 压缩意味着等效上下文扩展
- 无需硬件升级：通过训练而非工程优化获得效率提升
局限与未来方向
- 当前仅在特定推理任务（Countdown、数学竞赛）上验证，通用任务效果待验证
- 训练需要可验证奖励（RLVR），不适合开放式生成任务
- 驱逐决策增加了推理时的计算开销（虽然内存大幅减少）
- 代码尚未开源（截至 2026-04-28）
六、总结
Neural Garbage Collection 是一篇概念上极其优雅的论文。它没有提出新的架构，也没有设计复杂的系统 —— 它只是问了一个简单的问题：
如果模型能学会推理，为什么不能学会遗忘？
然后它证明了答案是：当然可以。而且遗忘和推理可以一起学，用同一个奖励信号，在同一个优化框架里。
NGC 代表了一种更广泛的范式：将效率视为可学习的能力，而非工程约束。在这个范式下，模型不只是变得更聪明，还变得更"自知" —— 知道什么值得记住，什么可以放手。
论文: arXiv:2604.18002v1
作者: Michael Y. Li, Jubayer Ibn Hamid, Emily B. Fox, Noah D. Goodman (斯坦福大学)
关键词: KV Cache Eviction, Reinforcement Learning, Chain-of-Thought, Efficiency