Neural Garbage Collection: Learning to Forget while Learning to Reason¶

会议: NeurIPS 2026
arXiv: 2604.18002
OpenReview: https://openreview.net/forum?id=WibVOPxZMf
论文: https://arxiv.org/abs/2604.18002
领域: LLM效率
关键词: KV Cache 驱逐, 强化学习, 思维链推理, 神经垃圾回收, 资源理性

一句话总结¶

提出神经垃圾回收（NGC），将 KV 缓存驱逐建模为语言模型采样的离散动作，直接使用基于任务正确性的结果奖励通过策略梯度端到端联合优化推理生成与记忆遗忘，在 2-3 倍峰值缓存压缩下大幅超越启发式驱逐基线。

研究背景与动机¶

思维链（Chain-of-Thought, CoT）推理使大语言模型在复杂数学与组合问题上取得了显著突破，但随着思考步数的增加，自回归解码过程中累积的 Key-Value (KV) 缓存线性膨胀，导致显存占用与注意力计算开销急剧飙升，成为制约长程推理扩展的核心瓶颈。现有缓解手段大多依赖手工设计的硬编码启发式规则（如基于注意力权重、局部滑动窗口或键向量多样性进行贪心裁剪），或者引入基于特征重构、教师模型蒸馏等辅助代理目标进行离线压缩与提示词总结。这些方法将模型推理与显存管理割裂开来，不仅需要繁琐的人工超参数调优，更忽略了一个根本事实：推理轨迹中大量草稿、子问题求解等临时上下文在得出阶段结论后已失去长程效用。

这种割裂带来的核心矛盾在于：现有的显存裁剪决策是由系统外部根据先验代理目标强加给模型的，而非由模型根据自身后续推理的真实语义需求自适应决定的。代理损失对任务正确性往往欠缺保真度，极易造成推理关键线索的灾难性误删，导致下游任务精度崩塌；而通过端到端强化学习，模型本身恰恰最清楚哪些记忆应当保留。

本文的核心角度是：既然语言模型能够在强化学习的奖励驱动下学会如何思考，它同样有能力自主学会该遗忘什么。核心 idea：将 KV 缓存的淘汰剔除显式建模为语言模型的离散行动，与推理 Token 生成统一在单一的任务正确性结果奖励下，通过无辅助损失的端到端强化学习让模型在学习推理的同时学会垃圾回收。

方法详解¶

整体框架¶

神经垃圾回收（Neural Garbage Collection, NGC）打破了传统语言模型上下文单调递增的环境假设，引入了“生成-驱逐（grow-then-evict）”的循环决策动力学机制。模型在生成长思维链的过程中每隔固定的步长暂停生成，直接复用 Transformer 自身末尾查询对历史前缀键的注意力权重构建块级重要性分布；随后通过带有 Gumbel 噪声的无放回 Top-k 采样选出待保留的缓存块，永久剔除其余条目；在后续解码阶段，模型在经过剪枝的紧凑上下文状态上继续生成。整个决策链路无需任何辅助重构损失或监督微调数据，完全依靠下游任务输出的二值正确性奖励计算组归一化优势，并借助回放注意力掩码（Replay Attention Masks）精确对齐策略梯度更新。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入题目提示词与资源预算"] --> B["周期性自回归推理与动态累积<br/>每生成 delta 个 Token 触发驱逐轮次"]
    B --> C["注意力复用与块级评分<br/>末尾 w 个查询计算 prefix 注意力并聚合为块得分"]
    C --> D["Gumbel-top-k 离散动作采样<br/>计算保持块集合及其精确对数概率"]
    D --> E["回放注意力掩码与端到端联合策略梯度<br/>Token 策略梯度与驱逐策略梯度共享结果奖励"]
    E --> F["输出紧凑缓存与最终推理答案"]

关键设计¶

1. 周期性注意力复用与块级评分：零新增参数构建离散驱逐分布

针对传统显存管理需要额外训练独立检索器或引入重构网络的问题，NGC 直接复用 Transformer 模型已有的预训练注意力权重机制。在每个驱逐周期，模型无需额外参数即可自然衡量前缀键的上下文重要性。具体而言，在第 $\ell$ 层，模型提取生成序列末尾滑动窗口内 $w$ 个查询向量 $\mathbf{Q}^{(\ell)}$，计算其与所有有效前缀键 $\mathbf{K}_{\mathrm{prefix}}^{(\ell)}$ 的缩放点积注意力，并在所有注意力头与滑动窗口内平均得到每个前缀键的位置得分： $$\psi_t = \frac{1}{H \cdot w} \sum_{h=1}^H \sum_{q=1}^w A_{q,t}^{(\ell,h)}$$ 为规避单键（token-level）离散动作空间过大带来的高方差和信用分配（credit assignment）难题，同时兼顾底层 GPU 显存的分页管理效率，NGC 将前缀序列划分为大小为 $b$ 的连续块（block），通过块内有效键得分的平均值构造块级得分 $s_j = \frac{\sum_{t \in \mathcal{B}_j} m_t \psi_t}{\sum_{t \in \mathcal{B}_j} m_t}$。该设计完全避免了引入任何新的模型权重或辅助预测头，使推理参数与驱逐参数天然共享。

2. Gumbel-top-k 离散动作采样：闭式计算子集对数概率

为了在强化学习框架下无偏优化驱逐策略，缓存淘汰不能采用传统的确定性贪心 Top-k，而必须作为从参数化策略中采样的随机动作。直接对组合子集求解分母配分函数在计算上是不可行的。NGC 将保留 $K$ 个块的过程形式化为基于 Gumbel-top-k 的无放回序列采样：向所有候选块的对数几率 $s_j$ 中注入独立同分布的 Gumbel 噪声，并选取扰动值最大的 $K$ 个块保留。该被选块序列 $\bm{\sigma} = (\sigma_1, \dots, \sigma_K)$ 具备严格的闭式对数概率表达： $$\log p(\bm{\sigma} \mid \mathbf{s}) = \sum_{j=1}^K \left[ s_{\sigma_j} - \log \sum_{t \notin \{\sigma_1, \dots, \sigma_{j-1}\}} \exp(s_t) \right]$$ 在工程实现上，模型利用前缀和技巧追踪移除的概率质量，避免了多次重复计算 LogSumExp 的计算冗余，使离散淘汰决策能够无缝嵌入标准策略梯度推导之中。

3. 回放注意力掩码与统一结果奖励驱动：消除脱策偏差并实现端到端梯度回传

在自回归生成中，早期被丢弃的 KV 条目无法被后续 Token 观察到。如果在策略梯度反向传播时盲目采用标准的上三角因果注意力掩码，模型在计算 Token 对数似然时就会访问比生成时更丰富的全局上下文，引发严重的脱策（off-policy）偏差并导致训练迅速崩塌。NGC 提出了每层独立的回放注意力掩码（Replay Attention Masks）机制：在单次并行前向传播中，构建严格复现采样轨迹中 Token 真实可见状态的二维布尔注意力掩码。更重要的是，驱逐决策所依赖的键查询张量在这次回放前向传播中被重新计算，从而使得驱逐动作对数几率直接接入自动微分计算图。训练直接采用 GRPO 架构，Token 级别损失 $\mathcal{L}_{\text{token}}$ 与显存驱逐决策损失 $\mathcal{L}_{\text{mem}}$ 共享完全相同的任务正确性组归一化优势标量 $\hat{A}_i$，没有任何人工构造的 KL 正则或稀疏惩罚项。

4. 渐进式阶梯驱逐课程与感知预算提示：平滑训练动态并赋予资源元认知

在强化学习初期，若立即执行激进的显存驱逐，急剧缩减的上下文会导致模型无法建立基础推理逻辑。NGC 引入了平滑的阶梯式保留率课程调度（Staircase Eviction Curriculum），使每个驱逐周期的淘汰比例在训练过程中分阶段阶梯式提升，并在每个阶段末尾进行线性插值过渡。同时，为了让模型能够提前感知自身面临的硬件资源严苛程度，NGC 提出了“预算感知内部感受（Budget-aware Interoception）”机制：在输入提示词中显式嵌入诸如 <eviction_rate>50%</eviction_rate> 的元数据标签。这赋予了模型在落笔前感知自身显存限制的能力，引导模型在严苛预算下主动压缩推导冗余、在充裕显存下展开细致推演。

损失函数 / 训练策略¶

NGC 采用基于验证奖励的强化学习算法（Dr. GRPO 框架），对一个 Prompt 采样 $G$ 条轨迹。总优化目标由 Token 生成损失与驱逐决策损失直接相加构成： $$\mathcal{L} = \mathcal{L}_{\text{token}} + \mathcal{L}_{\text{mem}}$$ 其中： $$\mathcal{L}_{\text{token}} = -\mathbb{E} \left[ \frac{1}{G} \sum_{g=1}^G \sum_{t=1}^{|o_i|} \log \pi_\theta(o_{i,t} \mid \mathcal{C}_{i,t}) \hat{A}_i \right]$$ $$\mathcal{L}_{\text{mem}} = \sum_{\ell=1}^L \mathcal{L}_{\text{mem}}^\ell, \quad \mathcal{L}_{\text{mem}}^\ell = -\mathbb{E} \left[ \frac{1}{G} \sum_{i=1}^G \frac{1}{|\mathcal{T}_i^\ell|} \sum_{t \in \mathcal{T}_i^\ell} \log \pi_\theta(\bm{\sigma}_{i,t}^\ell \mid \mathcal{H}_{i,t}^\ell) \hat{A}_i \right]$$ 其中优势值 $\hat{A}_i = r_i - \frac{1}{G}\sum_{j=1}^G r_j$ 仅由最终答案正确性给出。对于 DAPO-17k 数学竞赛任务，为防止过短答案在未经历任何驱逐轮次时偶然获得正向奖励而扰乱显存决策梯度，设定了最小序列长度惩罚规则。模型基于 DeepSeek-R1-Distill-Qwen-1.5B 展开训练，优化器采用 AdamW，学习率设为 $5 \times 10^{-6}$，不使用 KL 惩罚，块大小 $b=32$，注意力滑动窗口 $w=5$。

实验关键数据¶

主实验¶

论文在 Countdown 算术组合任务（1024 个留出测试样本）以及竞赛级数学推理数据集 DAPO-17k（在 AMC 2023、AMC 2025、AIME 2025 上测试）上进行了全面评估。下表展示了在每轮固定 50% 驱逐率（对应峰值 KV Cache 压缩 2.4× 至 4.6×）下，NGC 与四种前沿启发式驱逐基线的性能对比：

数据集 / 基准	峰值缓存压缩比	基线最高方法及表现	NGC (本文)	相对基线提升	全量缓存上限 (Full Cache)
Countdown (Pass@1)	2.4×	KeyDiff: 21.2% (SnapKV: 2.1%)	49.6%	+28.4% (翻倍)	54.3%
AMC 2023 (Pass@1)	~3.5×	SnapKV: 22.0%	35.4%	+13.4%	40.0%
AMC 2025 (Pass@1)	~3.8×	SnapKV: 21.3%	32.8%	+11.5%	38.5%
AIME 2025 (Pass@1)	4.6×	KeyDiff: 6.7%	16.7%	+10.0%	20.0%

消融实验¶

在 Countdown 任务上，论文深入拆解了端到端强化学习各个组件的核心贡献。下表对比了针对脱策偏差与驱逐策略梯度的消融表现：

配置 / 变体	准确率 (Pass@1)	训练动态状态	关键机制说明
NGC 完整框架	49.6%	稳定收敛，梯度范数平稳	采用回放掩码消除脱策，联合优化 $\mathcal{L}_{\text{token}} + \mathcal{L}_{\text{mem}}$
仅 Token 对数概率 (Token log-probs only)	35.7%	稳定收敛但上限受限	使用回放掩码纠正脱策，但移除了驱逐决策损失 $\mathcal{L}_{\text{mem}}$
目标 KV Dropout (Targeted KV Dropout)	2.5%	训练在第 100 步左右彻底崩溃	在训练中进行驱逐，但使用普通因果掩码，脱策导致梯度爆炸
全量缓存基线 (No Eviction)	54.3%	稳定收敛（无显存压缩）	不进行任何驱逐，计算与显存开销完全未压缩

关键发现¶

双重策略梯度的不可替代性：移去 $\mathcal{L}_{\text{mem}}$ 导致精度从 49.6% 骤降至 35.7%，表明仅让模型适应残缺上下文是不够的，必须通过结果奖励显式更新驱逐分布，赋予模型主动筛选记忆的能力。
脱策崩塌现象：若不采用回放掩码（Targeted KV Dropout），训练在显存驱逐率提升至约 100 步时出现严重的梯度爆炸与奖励归零，印证了上下文不一致对自回归强化学习训练稳定性的毁灭性打击。
内部感受带来的跨预算泛化：提示词中显式添加 <eviction_rate> 标签后，在超出训练阶段分布的极高压缩比下，模型准确率实现了 8%–13% 的进一步增益。
启发式基线的任务脆性：SnapKV 在 Countdown 上因注意力过度分散近乎失效（2.1%），而在数学题中表现优异；KeyDiff 则在不同任务间排名反复反转，凸显了人工先验无法普遍适配动态多样化推理流的根本缺陷。

亮点与洞察¶

效率作为一种自进化能力：打破了传统将推理能力交由 RL、将系统效率交由手工规则的思维定式，将显存管理动作纳为模型原生输出空间，实现了“越变越聪明的同时越跑越便宜”的闭环。
回放注意力掩码工程设计：利用单次前向回放掩码优雅解决了自回归不可逆剪枝造成的脱策问题，兼具理论严格性与工程轻量性。
元感知元 Token 的自然延展：论文附录进一步展示了在驱逐边界插入强制摘要元 Token $\mathbf{g}$ 的机制，无需额外蒸馏目标即可自发将即将丢弃的历史信息浓缩沉淀至保留节点中。

局限与展望¶

长文本预填充阶段适用性有限：当前主要针对模型自回归推理阶段展开动态剪枝，对于海量外部文档的超长 Prompt Prefill 阶段，尚未提供直接的即时检索或预处理优化。
硬件内核未深度定制：目前在底层仍主要依赖对张量切片的掩码计算，尚未编写深度融合的定长紧凑显存分页 Paged-Kernel，实际吞吐提升受限于现有计算库。
未来方向：探索结合推测解码（Speculative Decoding）与神经垃圾回收的复合加速架构，以及向多模态长视频推理场景中的显存自适应管理延伸。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将 KV 缓存管理完全形式化为与推理生成协同演化的端到端强化学习离散决策。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖组合算术与国家级竞赛级数学，提供了扎实的脱策消融、梯度动态分析及超分布泛化验证。
写作质量: ⭐⭐⭐⭐⭐ 理论框架清晰，将认知科学的“资源理性”与大模型 RLVR 紧密融合，逻辑推导层层递进。
价值: ⭐⭐⭐⭐⭐ 为长思维链大模型突破显存瓶颈开辟了全新的端到端自我演化范式。