Neural Garbage Collection: Learning to Forget while Learning to Reason

TL;DR

Neural Garbage Collection (NGC) trains a reasoning LLM end-to-end, from outcome-based task reward alone, to evict its own KV cache while it reasons. Eviction is treated as a discrete action: at fixed cadence the model pauses, scores cache blocks using its native attention, samples a subset to keep via Gumbel top-k, then continues reasoning on the remaining cache. Token generation and eviction decisions are jointly optimized with policy gradient (Dr. GRPO) using a single binary correctness reward — no SFT, no proxy losses, no learned scoring modules. On Countdown / AMC / AIME, NGC matches the full-cache upper bound at 2–3× peak KV compression and beats heuristic eviction baselines (SnapKV, StreamingLLM); on Countdown it more than doubles the next-best baseline (49.6% vs 21.2%) at 2.4× compression.

Key claims

KV-cache eviction can be optimized end-to-end with the same RL objective that trains reasoning; the model learns to forget purely from outcome-based task reward, without SFT, proxy losses, or a separately trained indexer [Abstract, §1].
Eviction rounds are inserted at a fixed cadence δ during chain-of-thought generation; at each round the model scores cache blocks via its own attention (softmax over a recent window of size w), then samples which blocks to retain via Gumbel top-k [§3, “Method”].
Coarsening the action space to fixed-size blocks (rather than individual tokens) is load-bearing for stable credit assignment; per-layer evictions are chosen independently [§3 “block-level evictions”].
Replay attention masks make policy-gradient log-probability computation exact: naively re-running rollouts under a standard causal mask would be incorrect because the rollout policy actually saw a hole-punched cache; replay masks reproduce visibility exactly and a prefix-sum trick gives numerically stable log-probs for Gumbel top-k [§3, autograd schematic].
“Budget-aware interoception” — prompting the model with its memory budget — generalizes across budgets at inference time and is what lets a single trained model handle multiple compression rates [§3, §5].
An eviction curriculum — start gentle (small eviction fraction ε), end aggressive — is needed for stable RL; without it the policy collapses early [§4 / §5].
Per-task results: on Countdown, NGC reaches 49.6% accuracy at ~2.4× peak-cache compression vs 21.2% for the strongest heuristic baseline at a comparable budget [Results table, §5].
Maintains strong accuracy relative to the full-cache upper bound at 2–3× peak KV compression on Countdown, AMC, and AIME, beating SnapKV / StreamingLLM at matched budgets [Abstract, §5].
Steady-state cache size is bounded analytically: with cadence δ, eviction fraction ε, and L layers, the cache stabilizes at ~L·δ/ε rather than growing unboundedly with reasoning length [§3, “Method”].
The KL term against a reference policy is omitted (following Dr. GRPO / DanceGRPO recipes); the only optimization signal is the verifiable correctness reward [§3].

Method

NGC reframes resource-constrained reasoning as a sequential decision process where both tokens and eviction decisions are discrete actions sampled from the LM’s distribution. The model reasons in chain-of-thought; every δ tokens it enters an eviction round, computes a softmax score over its KV cache using attention over a recent window of width w, groups neighboring keys into blocks, and samples a subset to retain via Gumbel top-k (which yields exact log-probabilities). Reasoning then resumes on the smaller cache.

Training uses Dr. GRPO with group-normalized advantages over multiple rollouts per prompt; the reward is binary task correctness (RLVR). Two engineering pieces make this practical:

Replay attention masks reconstruct, at backprop time, exactly what each rollout token actually attended to (post-eviction), so the policy-gradient log-probs are exact rather than off-policy.
A prefix-sum trick computes the Gumbel-top-k log-probability in O(N) instead of N repeated logsumexp passes.

An eviction curriculum ramps ε from low to target, and budget-aware interoception writes the target budget into the prompt so the same trained model handles many compression rates.

Results

Countdown: NGC 49.6% at ~2.4× peak-cache compression vs 21.2% for the next-best heuristic baseline (SnapKV / StreamingLLM family) at comparable budget — more than 2× the prior best [Results table].
AMC / AIME: maintains strong accuracy “relative to the full-cache upper bound” at 2–3× peak KV compression, outperforming heuristic eviction baselines across the same budget axis [Abstract, §5].
Budget generalization: a single trained model serves multiple compression rates at inference time via budget-aware interoception in the prompt [§5].
Steady-state memory: peak KV is bounded by ~L·δ/ε regardless of chain-of-thought length [§3].

Why it’s interesting

NGC is the cleanest demonstration so far that cache management is a learnable behavior under the same RLVR signal that trains reasoning. Every other entry on LLM Inference Efficiency approaches KV-cache pressure from outside the policy: Inference-Time Hyper-Scaling with KV Cache Compression (DMS) retrofits a per-head learned eviction via distillation against the original LLM as teacher; TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate (and the Revisiting RaBitQ and TurboQuant: A Symmetric Comparison of Methods, Theory, and Experiments rebuttal) compresses values via quantization; 1-bit Bonsai 8B: End-to-End 1-bit Language Models from PrismML compresses weights. NGC trains the policy to want a smaller cache — the eviction decision is sampled from the LM itself, and the only loss is task correctness. This is the same end-to-end-RLVR pattern surveyed in Reasoning RL (the orchestrator-policy version is K2.5 PARL; the diffusion-policy version is HuDA), now applied to the model’s own working memory. It also adds a sixth lever to Inference-Time Scaling: KV-cache eviction policies are themselves learned rather than retrofitted, opening the question of whether DMS-style distillation and NGC-style RL are substitutes or stack.