Neural Garbage Collection: Teaching LLMs that "To Reason is to Forget"

TL;DR

The "Bitter Lesson" of AI teaches us that end-to-end optimization usually beats hand-crafted heuristics. Neural Garbage Collection (NGC) applies this to the KV cache bottleneck. Instead of using fixed rules to decide which past tokens to keep, NGC trains the model to choose what to forget, using the same "pass/fail" reward signal it uses to learn math. The result? A model that reasons better with 1/3 the memory.

Problem: The KV Cache Tax on Thinking

Chain-of-thought (CoT) reasoning is the engine of LLM intelligence, but it comes with a steep "cache tax." Every extra step of reasoning adds tokens to the KV cache, eventually hitting hardware limits or slowing down inference.

Current efficiency methods are essentially "outsourced" logic:

1. Heuristics: "Keep the first and last 20% of tokens." (e.g., StreamingLLM)
2. Proxies: "Compress this cache to look like a summary." (e.g., Gist tokens)

The authors argue that if a model is smart enough to solve AIME math problems, it is smart enough to know which of its own previous "scratchpad" steps are now useless junk.

Methodology: The "Grow-Then-Evict" Cycle

NGC transforms the Transformer into an active memory manager. The process follows a rhythmic cycle:

1. Grow: The model generates $\delta$ tokens (the "cadence").
2. Score: The model uses its own Attention Weights—repurposed as an importance metric—to rank KV blocks.
3. Evict: Using the Gumbel-top-k trick, the model stochastically selects which blocks to kill. This stochasticity is key; it allows the model to explore different "forgetting strategies" during Reinforcement Learning.
4. Condition: Future tokens are generated only with the survivors.

The Secret Sauce: Replay Attention Masks

Standard RL training (like PPO or GRPO) assumes the model sees the entire history. But in NGC, the model generates token $T_{100}$ while only seeing a pruned version of tokens $T_{\mathrm{1..99}}$. To train this correctly without "training-inference mismatch," the authors introduced Replay Attention Masks. These masks ensure that during the training backprop, the gradient "knows" exactly what the model didn't see, preventing the training collapse observed in naive attempts.

Experiments: Breaking the Accuracy-Memory Tradeoff

The researchers tested NGC on Countdown (a logic game) and DAPO-17k (competitive math).

• Countdown Performance: While heuristic methods like SnapKV or StreamingLLM collapsed when memory was cut by 60%, NGC maintained nearly full accuracy. It "learned" to keep the specific intermediate numbers vital for the final calculation.
• Math Benchmarks (AIME/AMC): NGC consistently outperformed all baselines. Interestingly, the baselines (like SnapKV) were inconsistent—performing well on math but failing on Countdown. NGC, by contrast, adapts its forgetting strategy to the task at hand.

Budget-Aware Interoception

One of the most fascinatng findings was "Interoception." By putting a tag like <eviction_rate>50%</eviction_rate> in the prompt, the model learns to "sense" its memory limit. Like a human child told they have only one minute to finish a test, the model adjusts its reasoning style to be more concise when the cache budget is tight.

Critical Insight: Efficiency as a Capability

NGC represents a shift in how we view AI optimization. Efficiency is no longer just a "System Engineering" problem solved with better CUDA kernels; it is a "Machine Learning" problem.

Limitations:

• Training Cost: RL is notoriously compute-heavy compared to simple distillation.
• Architecture Specificity: While it uses standard Transformer blocks, the "Grow-then-Evict" cadence is a new hyperparameter to tune.

Future Outlook: Imagine a model that doesn't just manage KV cache, but manages its own "compute budget"—learning when to use a small "draft" model versus a large "reasoning" model, all trained via a single end-to-end reward. NGC is the first step toward that future of autonomous, resource-rational AI.