Neural Garbage Collection: Learning to Forget while Learning to Reason

Abstract: Chain-of-thought reasoning has driven striking advances in LLM 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 LLM 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 LLM, 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 LLMs.

• The paper demonstrates an end-to-end joint optimization where LLMs learn both reasoning and adaptive forgetting by scoring and evicting KV cache entries using policy gradient reinforcement learning.
• The methodology employs block-level evictions via the Gumbel-top-k trick and replay attention masks to achieve hardware-efficient memory management and reliable gradient updates.
• Empirical results show that Neural Garbage Collection outperforms baseline eviction methods, achieving 49.6% accuracy on arithmetic reasoning and robust performance under aggressive cache reductions.

Neural Garbage Collection: Learning to Forget while Learning to Reason

Motivation and Problem Formulation

Chain-of-thought (CoT) reasoning in LLMs fundamentally expands their problem-solving horizon, but every additional reasoning step exacerbates memory bottlenecks via unbounded growth of the KV cache. Existing approaches address this by imposing hand-crafted eviction heuristics or proxy objectives, but these methods lack the adaptivity and scalability observed in end-to-end optimization. The paper "Neural Garbage Collection: Learning to Forget while Learning to Reason" (2604.18002) proposes reframing efficiency as a learnable capability, extending self-improvement beyond reasoning to efficient memory management.

Methodology: End-to-End Joint Optimization of Reasoning and Forgetting

NGC operationalizes resource-constrained LM generation as a sequential decision process. At designated eviction rounds, the model uses its internal attention mechanism to score KV cache entries, samples a subset to retain via Gumbel-top-$k$, prunes the cache, and then continues generation conditioned only on the surviving cache. Both tokens and eviction decisions are treated as discrete sampled actions, enabling their joint end-to-end optimization using policy gradient RL from outcome-based reward alone.

Figure 1: NGC operationalizes dynamic memory management via periodic cache eviction rounds, scored and sampled by the LM itself, enabling simultaneous learning-to-reason and learning-to-forget.

The method capitalizes on block-level evictions, aligning decision granularity with contiguous token spans, which improves credit assignment and is hardware-efficient. The Gumbel-top-$k$ trick is employed to realize stochastic subset selection with tractable log-probabilities, critical for unbiased RL updates. Replay attention masks are introduced to reconstruct the per-layer visibility induced by prior evictions, guaranteeing accurate log-prob computation and correct gradient routing in the autograd graph.

The optimization objective comprises two principal terms: token generation and memory management, both with group-normalized advantage, yielding a unified policy gradient. Notably, both are updated solely via the task reward, with no auxiliary losses or supervision.

Experimental Validation and Analysis

Experiments span arithmetic reasoning (Countdown) and mathematical competitions (AMC, AIME), employing DeepSeek-R1-Distill-Qwen-1.5B as the backbone. Baseline comparisons include SnapKV (attention-weighted), KeyDiff (diversity-weighted), KNorm (norm-based), and StreamingLLM (sliding window) eviction heuristics, each applied inference-time to models trained without eviction.

Key empirical results:

• On Countdown, NGC achieves 49.6% accuracy at 2.4x peak cache reduction, exceeding the next-best baseline (21.2%) by more than double.
• On AMC/AIME, NGC maintains robust pass@$k$ at 3-5x cache reduction, consistently outperforming all baselines.

  Figure 2: NGC decisively outperforms baseline cache eviction methods at fixed eviction rates, and generalizes gracefully to stricter budgets unseen during training.

Ablation studies indicate that both end-to-end training and correct replay mask handling are essential for stability and performance. Applying targeted KV dropout during RL but neglecting off-policy effects causes gradient explosion and reward collapse.

Figure 3: Ablation demonstrates severe accuracy degradation when end-to-end optimization or correct masking are removed, confirming the necessity of both contributions.

NGC’s learning dynamics are distinguishable: it improves more slowly than models trained without eviction but ultimately achieves comparable accuracy, whereas targeted dropout collapses with increasing eviction rates due to systematic off-policyness.

Figure 4: Training dynamics reveal NGC's robustness and stability versus targeted KV dropout, which collapses with increasing eviction rates.

Budget-Aware Interoception

To further enhance adaptivity, the paper introduces budget-aware interoception, conditioning the model explicitly on the eviction rate via prompt tag insertion. This enables improved generalization to unseen cache constraints and softens performance degradation at aggressive compression rates.

Figure 5: Conditioning on eviction rate in the prompt enhances generalization, yielding higher accuracy as the memory budget becomes more stringent.

Scaling to Mathematical Reasoning

NGC applied to large-scale mathematical datasets (DAPO-17k) demonstrates generalization across AMC and AIME benchmarks. Across varied tasks and cache reduction sizes, NGC shows universal superiority to all baseline eviction strategies; heuristics are task-dependent and fragile, whereas NGC’s strategy adapts dynamically to each domain.

Figure 6: NGC shows consistent superiority across math reasoning benchmarks at substantial cache reduction ratios.

For moderate cache reductions (2-3x), NGC’s accuracy remains close to that of the upper bound model with no cache limit, highlighting its compression efficacy.

Figure 7: NGC maintains strong pass@32 even as peak cache size reduces, approaching no-eviction model performance at moderate compression.

Implications and Theoretical Perspectives

NGC advances the paradigm of self-improving LMs, demonstrating that efficiency can be learned as a capability co-optimized with reasoning. By making resource allocation a native action, models are incentivized to align memory management with task-relevant retention. The approach unifies cache management and reasoning under a single RL reward, contrasting prior methods relying on proxy objectives, distillation, or rigid heuristics.

Practically, NGC unlocks trajectories where models can reason over longer horizons without memory constraints, directly benefiting high-throughput inference and enabling deployment in resource-constrained environments. Theoretically, it opens avenues for meta-cognitive modeling—agents regulating their own cognitive state—within deep learning frameworks. The system can be extended to other resource allocation regimes, e.g., compute budgeting and conditional parameter routing.

Speculatively, integrating NGC-like mechanisms in future transformer or state-space models could drive more general adaptive computation, enabling not just longer context but principled trade-offs between capability and cost. This approach may catalyze research into RL-based conditional computation, meta-awareness, and algorithmic resource rationality within LLMs, synchronizing with trends in continual reasoning and self-improving architectures.

Conclusion

Neural Garbage Collection demonstrates that LMs can learn efficient memory management jointly with reasoning, achieved solely through outcome-based reward optimization. The mechanism outperforms all baseline cache eviction strategies, proving its adaptivity and robustness across arithmetic and mathematical reasoning. This work substantiates the vision that efficiency is a behavior which can be acquired the same way as capability and lays groundwork for future models in which resource use is governed by learned policies, not prescribed heuristics.