Neural Garbage Collection: Teaching LLMs to Forget to Reason Better

TL;DR

Neural Garbage Collection (NGC) is a paradigm shift that reframes KV cache management—traditionally a systems or heuristic challenge—as a learned capability. By using Reinforcement Learning from Verifiable Rewards (RLVR), NGC trains models to decide which parts of their "memory" are no longer needed for reasoning. It achieves massive context compression (up to 5x) while outperforming traditional human-designed pruning rules.

The "Losing Memory" Problem in Reasoning

As models "think longer" via Chain-of-Thought (CoT), their KV cache swells. This creates a hard ceiling on reasoning depth due to GPU memory limits. Traditionally, we solve this with "garbage collection" rules: Keep the most recent tokens or keep tokens with the highest attention scores.

However, the authors argue these heuristics are anachronistic. If a model is smart enough to solve a complex math problem, why shouldn't it be smart enough to manage its own memory? Heuristics are rigid; they don't know that a specific intermediate calculation is now useless, but the final sum is vital for the next five steps.

Methodology: The Grow-then-Evict Cycle

NGC integrates memory management directly into the model's policy. The process follows a "Grow-then-Evict" cycle:

1. Reasoning: The model generates $\delta$ tokens.
2. Scoring: The model uses its internal attention mechanism to score all current KV blocks.
3. Action: It samples a subset of blocks to keep and permanently deletes the rest.
4. Repeat: Reasoning continues, conditioned only on the remaining "pruned" cache.

The Replay Mask Breakthrough

A major technical hurdle in training this is off-policyness. If you train a model on its final trace, it "sees" tokens that were actually evicted during the generation. To fix this, the authors developed Replay Attention Masks. These masks ensure that during the training's forward pass, the model's gradients correctly reflect exactly what it saw (and didn't see) at every step.

Figure 1: The model samples eviction decisions just like it samples tokens, using a single reward signal to optimize both.

Experiments: Superior to Heuristics

The researchers tested NGC on Countdown (arithmetic reasoning) and high-level math (AMC/AIME).

• Vs. Heuristics: Methods like SnapKV (attention-based) and StreamingLLM (recency-based) often failed on Countdown because they deleted critical early numerical values. NGC learned to keep exactly what was necessary.
• Scaling and Generalization: Even when pushed to more aggressive compression rates than seen during training, NGC's performance degraded gracefully, whereas heuristic baselines often hit a "performance cliff."

Figure 2: NGC (Blue) maintains significantly higher accuracy than all heuristic benchmarks at a 50% eviction rate.

Budget-aware Interoception

One of the most fascinating findings was Budget-aware Interoception. By including the target eviction rate in the prompt (e.g., <eviction_rate>50%</eviction_rate>), the model "senses" its memory constraints. This allowed it to adjust its reasoning strategy—perhaps becoming more concise—to fit the available memory budget, leading to an 8-13% boost in accuracy at high compression levels.

Critical Analysis & Future Outlook

NGC proves that Efficiency is a Behavior. The key takeaways are:

• Unified Reward: You don't need a separate loss for "smart pruning." The standard "is the answer right?" signal is sufficient to teach the model to prune its cache effectively.
• Parameter-Free: NGC re-uses existing attention weights, adding no new parameters to the model.

Limitations: Currently, NGC requires periodic pauses, which might slightly increase latency despite the memory savings. Future work could move toward "asynchronous" garbage collection where the model prunes while it predicts.

In the era of "Infinite Context," NGC suggests that the path forward isn't just bigger GPU memory, but smarter, more "self-aware" models that know what to forget.