Neural Garbage Collection: Learning to Forget while Learning to Reason

Michael Y. Li${}^{\ast }$Michael Y. Li
Stanford University

Jubayer Ibn Hamid
Stanford University

Emily B. Fox
Stanford University

Noah D. Goodman
Stanford University

Abstract

Chain-of-thought reasoning has driven striking advances in language model 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 language model 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 language model, 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 language models.

1 Introduction

A central lesson of deep learning is that desirable properties emerge from end-to-end optimization pressure rather than being built in manually [Sutton, 2019]as argued by Sutton in 2019. Yet, when it comes to making models more efficient, practitioners rely on deep domain knowledge to design architectures [Ainslie et al., 2023, Dao and Gu, 2024], systems-level optimizations [Dao et al., 2022, Kwon et al., 2023], and inference-time heuristics [Li et al., 2024]. We ask whether efficiency itself can be treated as a model capability. Under this perspective, the same end-to-end optimization that drives capability can also drive efficiency, enabling an additional dimension of self-improvement [Silver and Sutton, 2025] in which models become more capable but also more efficient.

A setting where this approach is particularly promising is the chain-of-thought reasoning paradigm [Zelikman et al., 2022]introduced by Zelikman and colleagues that has driven dramatic advances in model capability [Guo et al., 2025, OpenAI, 2024]: by thinking longer before answering, models can solve problems that would otherwise be out of reach. However, these thinking traces rapidly grow the KV cache, creating a bottleneck to further scaling along this dimension. But paying this cost is not fundamental to reasoning itself — after all, humans reason over long horizons despite severely limited working memory: instead, it is an artifact of how we manage the KV cache. Much of the KV cache consists of transient information: intermediate steps that become irrelevant as reasoning progresses or scratch work on sub-problems where only the answer matters downstream. This suggests that

Neural Garbage Collection

Figure 1: Neural Garbage Collection: Learning to Forget while Learning to Reason. (top) As the language model (LM) generates its chain-of-thought, it periodically enters an eviction round: (1) the LM scores KV cache entries via softmax, (2) samples cache evictions using those scores via Gumbel top-$k$k, and (3) samples the next token conditioned only on the pruned cache. Since both tokens and cache evictions are discrete actions sampled from the LM, we can jointly train the LM to reason and manage its memory via end-to-end reinforcement learning from outcome-based task reward alone (e.g., binary correctness) — no proxy objectives, no supervised finetuning. (bottom left) NGC substantially outperforms baseline eviction methods at a 50% eviction rate per round, corresponding to 2.4x reduction in peak KV cache size for Countdown and 4.6x for AIME 2025. (bottom right) To give the LM explicit awareness of its resource constraint before it reasons, we include the eviction rate in the prompt during training — a technique we call budget-aware interoception. This technique improves generalization across cache budgets at test time.

actively managing the KV cache during reasoning need not degrade performance, and may enable models to reason over longer horizons than current resource constraints permit.

Cache management, so far, relies on fixed heuristics or proxy objectives. Heuristic eviction policies remove cache entries based on criteria such as attention weights or recency [Li et al., 2024, Park et al., 2025, Xiao et al., 2024]. Compression methods learn compact representations of the model’s KV cache using hand-tailored training objectives like reconstruction or distillation losses [Monea et al., 2025, Zweiger et al., 2026], and summarization-based approaches use elaborate prompts that specify what the model should preserve [Vajipey et al., 2025]. Relying on fixed rules or proxy objectives for these decisions is anachronistic for reasoning models, whose capabilities emerge from end-to-end optimization against outcome-based task reward. And since the contents of the KV cache arise from the model’s own reasoning process, the model is naturally positioned to manage its own cache. If a model can learn to reason, why can’t it learn to forget?

To that end, we introduce Neural Garbage Collection (NGC), in which a language model learns to forget while learning to reason. As the model generates a chain-of-thought, it periodically pauses, computes a softmax over its KV cache entries, samples which to evict, and continues reasoning conditioned on the pruned cache. Since this eviction decision is discrete, prior work resorts to differentiable relaxations [Łańcucki et al., 2025] or separate supervised training for those decisions using auxiliary objectives [Chen et al., 2026, Liu et al., 2025a]. Our key insight is that, because we already train reasoning models against task reward with reinforcement learning — where tokens are discrete actions sampled from the LM — we can cast cache eviction as another discrete action sampled from the LM and optimize it within the same framework. A single outcome-based task reward jointly trains the model to reason and manage its own memory end-to-end: what the model evicts shapes what it remembers, what it remembers shapes its reasoning, and the correctness of that reasoning trains both the eviction decisions and reasoning tokens that produced it. Our key methodological contribution is to show that both types of decisions can be optimized from a single learning signal — the outcome-based task reward — without auxiliary objectives or separate training stages. This follows the tabula rasa spirit of AlphaZero [Silver et al., 2017]Silver and colleagues: end-to-end optimization pressure alone guides how the model reasons and what it forgets.

We evaluate NGC on Countdown [Gandhi et al., 2024, 2025], and on math tasks AMC, and AIME. On Countdown, NGC more than doubles the accuracy of the next-best baseline (49.6% vs 21.2%) at a 2.4x reduction in peak cache size. We then validate NGC at larger scale by training on DAPO-17k [Yu et al., 2025]. On mathematical reasoning tasks (AMC, AIME), we show that NGC maintains strong performance at 2-3x reductions in peak KV cache size and outperforms all baseline cache management methods.

2 Related work

Language Model Self-Improvement. A growing body of work explores language models that self-improve. STaR [Zelikman et al., 2022] demonstrated that LMs can bootstrap reasoning ability by fine-tuning on self-generated chains-of-thought filtered by correctness. Subsequent work extended this paradigm to more general latent reasoning and iterative self-improvement, including Quiet-STaR [Zelikman et al., 2024], ReST [Gulcehre et al., 2023], and Self-Rewarding Language Models [Yuan et al., 2024]. The reinforcement learning from verifiable rewards (RLVR) paradigm [Guo et al., 2025, OpenAI, 2024] showed that self-improving reasoning scales: models trained on outcome supervision can acquire sophisticated reasoning capabilities. NGC is inspired by this line of work, but targets a different dimension of self-improvement: rather than improving task performance, the model learns to improve its own efficiency. This reframes efficiency as a capability: models typically become more capable at the cost of being slower and more expensive to run, whereas a model that improves its efficiency, in principle, becomes cheaper as it improves. Crucially, many forms of resource management (e.g., eviction, routing) naturally take the form of discrete decisions, making them directly amenable to the same reinforcement learning with verifiable rewards (RLVR) framework used to train modern reasoning models.

Resource Rationality, Bounded Rationality, Metacognition. A long tradition in cognitive science and AI studies intelligent behavior under computational constraints, viewing agents as trading off task performance against the cost of computation [Horvitz, 1989, 1990, Russell and Wefald, 1991, Simon, 1955]. Resource rationality [Lieder and Griffiths, 2020] formalizes this perspective by treating cognition as the optimization of task performance subject to resource costs. This work has primarily used resource rationality as a normative framework to explain behavior, rather than as a training objective to produce it in language models. A related tradition studies metacognition—the ability of an agent to monitor and regulate its own cognitive processes. Recent work explores metacognitive prompting and metacognitive reuse, where models reflect on prior reasoning traces to extract reusable strategies [Didolkar et al., 2025]. NGC teaches a model

to regulate its own current cognitive state by deciding what information to retain under a memory budget. This makes NGC a form of online, prospective metacognition: the model learns not just to reason, but to manage the memory on which its reasoning depends.

Sparse attention, alternative architectures, and conditional computation. Several lines of work improve efficiency by changing how transformers use compute. Sparse attention methods reduce the cost of attending over long contexts by selecting only a subset of keys per query [Chen et al., 2026, Liu et al., 2025a, Yuan et al., 2025]. These methods use attention scores to determine which entries are read at a given step, but they do not reduce the size of the cache itself; memory still grows without bound. NGC uses a related scoring mechanism but with fundamentally different semantics: its decisions determine which entries are removed indefinitely, permanently reshaping the context available to future tokens, and providing a practical mechanism for keeping the KV cache constant size. A related line of work uses conditional computation to improve efficiency by routing tokens to only a subset of model parameters, as in mixture-of-experts (MoE) architectures [Fedus et al., 2022, Liu et al., 2025a, Shazeer et al., 2017]. These models show that routing decisions can themselves be learned, but they target a different resource and training regime: MoEs sparsify parameter compute through architectural routing whereas NGC sparsifies memory during inference by deciding what information to retain under an explicit budget. State-space models [Dao and Gu, 2024, Gu and Dao, 2023] achieve efficiency through different computational primitives, but require pre-training from scratch and cannot fully leverage the thriving hardware and software ecosystem specific to transformers. By contrast, NGC requires no modification to the base transformer architecture, repurposing the existing parameters to perform eviction, and can be applied directly during post-training.

Most crucially, what distinguishes NGC from prior approaches is how its discrete decisions are trained. In methods such as DeepSeek Sparse Attention [DeepSeek-AI, 2025], the indexer, which selects which tokens the model attends to, and the main model are not trained under a unified objective. Instead, the indexer is first warm-started to imitate dense attention via a KL-divergence loss against the attention scores, and its input is then explicitly detached from the computational graph throughout sparse training: the indexer continues to receive signal only from the KL-divergence term, while the main model is optimized only by the language modeling loss [DeepSeek-AI, 2025]. The indexer is a separate module trained by the designer to recognize importance according to a proxy criterion, then installed into the system while the main model optimizes around it. NGC instead trains eviction under a unified objective: the same reward signal that shapes reasoning also shapes how it remembers. Moreover, NGC does not require a separate warm-up stage. Likewise, conditional-computation methods such as mixture-of-experts learn routing decisions, but their routers are typically trained with standard backpropagation-based objectives and load-balancing schemes rather than from downstream task reward [Fedus et al., 2022]. In contrast, NGC treats discrete choices that modulate memory use as a first-class action and optimizes them with reinforcement learning.

KV cache compression. KV cache compression methods reduce memory by evicting or compressing cache entries during generation. Static methods apply fixed eviction rules: retaining only the most recent tokens [Xiao et al., 2024] or tokens with high attention scores [Li et al., 2024]. More recent methods use richer proxies for importance such as diversity [Park et al., 2025] or optimization objectives to compress the cache [Eyuboglu et al., 2026, Mu et al., 2024, Zweiger et al., 2026]. Breadcrumbs [Monea et al., 2025] and Memento [Kontonis et al., 2025] compress the KV cache during reasoning, but via a fundamentally different mechanism. Breadcrumbs trains a compressed student policy to match an uncompressed teacher via token-level KL divergence, distilling from the teacher’s rollouts. Memento teaches models to produce compressed summaries via SFT on reasoning traces interleaved with summaries generated by a strong teacher prompted with hand-designed rubrics. Both methods inject inductive bias about what good compression looks like—Breadcrumbs via the frozen teacher’s reasoning style, Memento via the prompt author’s rubrics—rather than letting task reward determine what to remember. NGC takes a conceptually different approach and

jointly learns how to manage memory and reason, end-to-end from outcome-based task reward alone. This requires no annotation pipeline, no teacher model, no supervised data, and no hand-designed objectives.

3 Problem setting: resource use as a learned capability

Today, practitioners manage resource constraints through design choices that span the entire modeling stack: architecture, systems-level optimizations, and even scaling laws that account for inference time cost [Hoffmann et al., 2022, Sardana et al., 2025]as discussed in prior work. We ask whether efficiency can instead be learned—treating it as a capability the model acquires through end-to-end training.

We model resource-constrained generation as a sequential decision process of horizon $T$T, in which the model makes two kinds of decisions at each step: a task action $a_t$a sub t (e.g. a generated token) and a resource-management action $u_t$u sub t (e.g., which KV cache entries to evict). Let $z_t\in \mathcal{Z}$z sub t in Z denote the model’s internal computational state at step $t$t, and let

the joint actions a t and u t are sampled from the policy pi parameterized by theta and phi given state z t

where $\theta$theta governs task behavior and $\varphi$phi governs resource allocation; these parameters can be shared or tied but this is a design choice independent of the general formulation.

A trajectory is

tau is the sequence of states and actions from time one to T

We optimize expected task reward subject to a resource budget $B$B:

we maximize the expected task reward R of tau, subject to the expected resource consumption C of the state sequence being less than or equal to B

where $R$ is a task reward (e.g., answer correctness) and $C:{\mathcal{Z}}^T\to {\mathbb{R}}_{\ge 0}$C measures resource consumption (e.g., the mean fraction of KV cache entries retained over the trajectory). The resource-management decision $u_t$u sub t is typically discrete, which makes it naturally amenable to the reinforcement-learning paradigm already used to train reasoning models. We instantiate this framework for KV-cache management in the next section.

4 Neural Garbage Collection

In this section, we describe NGC, our method for training a language model to jointly reason and manage its KV cache via end-to-end reinforcement learning using purely an outcome-based task reward. The main deviation from standard RLVR settings is that our action space consists of both eviction decisions and tokens. Our key methodological contribution is to show that both types of decisions can be optimized from a single learning signal — the outcome-based task reward — without auxiliary objectives or separate training stages.

4.1 State space and environment dynamics: grow-then-evict cycle

In this section, we define the state space and environment dynamics of our reinforcement learning formulation. In language modeling, the state corresponds to the model’s context window—represented by its KV cache—while the environment dynamics dictate how this context evolves over time. In standard RLVR, environment transitions are strictly monotonic: the LM samples the next token conditioned on all previously generated tokens, steadily growing the KV cache. In NGC, the environment dynamics fundamentally differ because the model actively modifies its own state space by managing its own KV cache, in addition to generating tokens.

Specifically, we introduce a “grow-then-evict” dynamic. Periodically, at a set of discrete eviction steps, the LM selects cache entries to keep. Every $\delta$delta tokens, an eviction round fires: the model produces scores for all cache entries and samples a $(1-ϵ)$one minus epsilon fraction of them to keep, where $ϵ\in (0,1]$epsilon between zero and one is the eviction rate, and

permanently evicts the rest; we will refer to $\delta$delta as the eviction cadence. Both prefill and generation tokens are eligible for eviction and we evict the same fraction (though not the same entries) in each of the $L$L layers of the transformer. The first eviction round fires when the total number of tokens in the cache including prefill tokens reaches $\delta$delta. Between eviction rounds, the cache grows as usual as the model generates additional tokens conditioned on the surviving entries. Under these dynamics, the maximum cache size before an eviction round converges to the steady state value $L\frac{\delta }{\epsilon }$L times delta over epsilon; this is independent of the number of reasoning tokens generated. We formalize this in Section A.2 of the Appendix.

4.2 Extending the action space: block-level evictions via existing attention mechanism

Having formalized the state space, we now describe how we extend the LM’s action space to include cache evictions. At each eviction round, the action is a subset of cache entries to keep. Language models do not natively support this type of action since their action space consists only of tokens. In the rest of this section, we describe how we use the transformer’s existing attention mechanism to parameterize cache evictions, enabling the model to act on its own KV cache — even though it was never trained, through pre-training or supervised fine-tuning (SFT), to do so.

Parameterizing eviction actions To perform KV cache evictions, the LM must be able to evaluate the utility of its cache entries—a task it was never explicitly trained to perform. Fortunately, the transformer’s native attention mechanism is already well-suited for this role. First, after standard pre-training, attention weights naturally encode a prior over which past entries are relevant for ongoing computation. This provides an effective initialization, which has been shown to be essential for stable reinforcement learning [Gandhi et al., 2025]Gandhi and colleagues. Second, re-purposing this existing machinery introduces no additional architectural overhead. We therefore parameterize eviction decisions directly with the model’s attention scores, setting $\varphi =\theta$phi equals theta and introducing no new parameters.

Concretely, at each eviction step, we score the remaining (prefix) keys according to how much they are attended to by the $w$w most recent queries. For layer $\ell$ell, we use the queries ${\mathbf{Q}}^{(\ell )}$Q of ell associated with the $w$w most recent tokens, compute attention weights from these queries to the prefix keys,

The attention matrix A is the softmax of the product of Q and the transpose of K prefix, divided by the square root of d sub h

and average across heads and recent queries to obtain a scalar importance score for each prefix key:

The importance score psi sub t is the average of the attention weights across heads and recent queries

In our experiments, we use $w=5$w equals five, so scoring incurs only a small, constant memory overhead. This construction is inspired by sparse attention methods [Liu et al., 2025a]Liu and colleagues, which also use attention scores to decide which keys to read when decoding a token. NGC uses the same primitive but with fundamentally different semantics: rather than deciding which entries are read at the current step, these scores decide which entries are permanently removed from memory. Moreover, instead of requiring a separate quantized KV cache, used to decide which actual KV cache entries get used for a decode step, and additional scoring parameters, we re-purpose the attention mechanism into a parameter-free eviction policy.

Coarsening the action space via block-level evictions Evicting at the individual key level creates a credit assignment problem: the marginal contribution of any single key is difficult to isolate. Moreover, the importance of adjacent keys is highly correlated: semantically coherent units such as sub-computations typically occupy contiguous spans of tokens in the chain-of-thought. Consider a set of surviving keys

of length $T$T. We coarsen the action space by working at the block level and partitioning the $T$T keys into $N=\lceil T/b\rceil$N equals the ceiling of T over b blocks of contiguous tokens of size $b$b, where the final block may be smaller if $b\nmid T$b does not divide T ($b$b does not divide $T$T). Let $m_t$m sub t be a mask indicating whether a key $t$t is valid (e.g., non-padding); we omit the layer index for simplicity. We aggregate the per-key scores $\psi \in {\mathbb{R}}^T$psi in R T to block-level scores by averaging over valid keys in a block:

yielding $\mathbf{s}\in {\mathbb{R}}^N$s in R N, one score per block. This dramatically reduces the size of the action space.

While the focus of our work is on the algorithmic properties of NGC rather than hardware optimizations, blockwise selection broadly aligns with hardware principles: GPUs have higher throughput for contiguous memory accesses, and block-level [Yuan et al., 2025]Yuan and colleagues eviction maps naturally onto paged KV cache systems such as vLLM [Kwon et al., 2023]Kwon and colleagues, where the cache is partitioned into fixed-size pages of contiguous keys.

4.3 Efficiently sampling eviction actions with Gumbel-top-k

We now define a procedure for tractably sampling eviction actions for training. KV cache eviction is conventionally performed using deterministic top-k rules [Li et al., 2024, Xiao et al., 2024]. We instead formulate cache eviction as a stochastic action sampled from the LM. This formulation enables us to leverage policy gradient methods for unbiased gradient estimates: stochastic sampling with exact log-probabilities is precisely the structure these methods require.

Concretely, we need to be able to sample a size-$K$K subset of kept cache entries in a way that exposes a well-defined log-probability of the selected set. Initially, this seems challenging since subset selection is combinatorial and the partition function is intractable. We formulate this task as sequential sampling without replacement using the Gumbel-top-k trick [Kool et al., 2019, Vieira, 2014]. Conveniently, sampling can be done in parallel: we sample a size-$K$K subset by adding i.i.d. Gumbel noise to the block logits and keep the $K$K largest perturbed values. The log-probability of the selected block sequence $\sigma =({\sigma }_1,\dots ,{\sigma }_K)$sigma equals sigma one through sigma K admits a closed form

where indices range over blocks.

Each term naively requires recomputing the partition function over the remaining blocks, i.e., computing a logsumexp $K$K times. We avoid this via a prefix sum trick that tracks the cumulative fraction of removed probability mass.

4.4 Policy optimization: token and eviction policy gradients

With the states, actions, and sampling procedure established, we now formalize the policy optimization objective. In standard RLVR, the policy gradient applies solely to the token space. Here, we define a unified objective over a joint action space consisting of both tokens and discrete cache evictions. Concretely, we detail how the single, outcome-based task reward in RLVR can be used to update the model to jointly optimize both its chain-of-thought reasoning and how it manages it own KV cache.

Given a prompt $x$x, we sample $G$G trajectories $\{{\tau }_i{\}}_{i=1}^G$tau sub i, from i equals one to G from ${\pi }_{\theta }$pi theta, where each trajectory ${\tau }_i=\{(o_{i,t},{\mathbf{\sigma }}_{i,t}){\}}_{t=1}^{|{\tau }_i|}$tau i, defined as pairs of token generation o and eviction decision sigma at each timestep t interleaves token generation $o_{i,t}$o sub i t with eviction decisions ${\mathbf{\sigma }}_{i,t}$sigma sub i t, where ${\mathbf{\sigma }}_{i,t}=\varnothing$sigma is empty at timesteps $t$t that do not correspond to an eviction round. Each trajectory is scored with a binary task reward $r_i\in \{0,1\}$r sub i in zero, one computed from the generated tokens $\{o_{i,t}\}$the set of o i t.

Figure 2: Replay masks enable efficient end-to-end RL training. During reasoning, the language model samples eviction decisions that dynamically change its KV cache. We can perform the policy gradient update efficiently and correctly using replay attention masks. The eviction decisions can be captured by attention masks that reproduce the visibility patterns over previous tokens induced by the eviction decisions. We then use these masks to compute log-probabilities in parallel using a single forward pass over all the tokens. (top) The model evicts cache entries at fixed intervals. (<textcolor{red}{round 1}: evicts $t_1,t_3$; <textcolor{orange}{round 2}: evicts $t_2,t_5$). For illustrative purposes, we show the schematic for a single layer, but in practice we evict separately across layers. (bottom) The resulting replay mask: row $t$ marks which keys were visible when token $t$ was generated. A single forward pass with this mask reproduces the next-token distributions from generation in parallel. Without the replay mask, token log-probabilities are computed under a richer context than the model actually saw during generation, introducing a systematic off-policyness that causes training collapse.

Following GRPO [Guo et al., 2025]G R P O, we compute group-normalized advantages:

The advantage A sub i hat is the reward r sub i minus the average reward of the group G

The only training signal is task accuracy.

We optimize ${\pi }_{\theta }$pi theta using Dr. GRPO [Liu et al., 2025c]D R G R P O. For the token-level actions, we have a token-level policy gradient:

The token loss is the negative expectation of the average over the group and the sum over tokens of the log probability of token o sub i t given context C sub i t, multiplied by the advantage A sub i hat

Here, ${\mathcal{C}}_{i,t}$C sub i t denotes the context available at step $t$t of trajectory $i$i (i.e., the surviving KV cache entries).

For the eviction decisions, we have an eviction decision level policy gradient

L mem is the sum over layers of L mem ell, where L mem ell is the negative expectation of the average over rollouts and eviction rounds of the log probability of the retained subset sigma, given context H, multiplied by the advantage estimate A hat.

where $L$L is the number of transformer layers, ${\mathcal{T}}_i^{\ell }$T sub i ell is the set of eviction rounds at layer $\ell$ell for rollout $i$i that fire before sequence termination, ${\sigma }_{i,t}^{\ell }$sigma i t ell is the subset of blocks retained at round $t$t, and ${\mathcal{H}}_{i,t}^{\ell }$H i t ell is the context used to make eviction decisions at the $\ell$ell-th layer (layer-$\ell$ell queries and alive keys at round $t$t); with slight abuse of notation, we write ${\pi }_{\theta }({\sigma }_{i,t}^{\ell }\mid {\mathcal{H}}_{i,t}^{\ell })$pi theta of sigma given H to denote the probability of the retained subset given by Gumbel top-k in Equation 2. The inner mean averages over eviction rounds within an individual rollout. We use the layer index to emphasize that we make separate eviction decisions for each KV entry in each layer rather than making the same decision for all layers for a given token.

Crucially ${\hat{A}}_i$A hat i is the same for both ${\mathcal{L}}_{\text{token}}$L token and ${\mathcal{L}}_{\text{mem}}$L mem: both learning to generate reasoning tokens and learning to evict come from the same learning signal. The total objective is:

The total loss L equals L token plus L mem

We again emphasize that although we call $\mathcal{L}$L a loss, it is not an auxiliary loss (e.g., an ${\ell }_1$L one sparsity penalty or a KL divergence between teacher and student attention weights [DeepSeek-AI, 2025]as seen in work by DeepSeek AI): it introduces no independent training signal beyond the task reward already present in ${\hat{A}}_i$A hat i. It is a scalar whose gradient is the correct policy gradient estimate for both token and eviction decisions. The resource constraint in our setting enters by construction (evicting a fixed amount of cache each time) rather than through auxiliary losses.

Because eviction decisions are scored by the model’s own queries and keys ($\varphi =\theta$phi equals theta), the policy gradient from ${\mathcal{L}}_{\text{mem}}$L mem flows into every parameter that shapes $\mathbf{Q}$Q and $\mathbf{K}$K—which is every parameter the LM uses to reason. Token generation and cache management are therefore treated as a single policy optimized under a single outcome reward.

4.5 Computing rollout log-probabilities efficiently with replay masks

To compute the policy gradients defined above, we need to evaluate the log-probabilities of the exact trajectories sampled during rollouts. This section introduces replay attention masks, a mechanism to efficiently compute these probabilities.

The key challenge is that because NGC dynamically modifies its own context through cache eviction, it breaks an assumption in standard RLVR: that the model attends to all previous tokens when generating a new token. Concretely, when the model generates a new token, it does not attend to cache entries it has already evicted. Therefore, naively computing log-probabilities under a standard causal attention mask would be incorrect and create a form of off-policyness [Liu et al., 2025b, Zheng et al., 2025]as discussed in recent literature. We use replay attention masks to correctly and efficiently compute Equation 3 (i.e., the log-probability of the generated tokens) during training. Specifically, we “replay” the rollout log-probabilities under a per-layer attention mask that replicates the exact visibility pattern over past tokens that was induced by the eviction decisions (See Figure 2). A single forward pass over those tokens with these replay masks gives us the correct log-probabilities.

Another subtle technical challenge is ensuring that gradients flow from ${\mathcal{L}}_{\text{mem}}$L mem; see Figure 9 for a schematic of the desired autograd graph. Eviction decisions at a given round are scored using the queries and keys at that eviction round. If those tensors were simply cached during rollouts and passed directly into Equation 4, they would be numerically correct but detached from the computation graph, and no gradient from ${\mathcal{L}}_{\text{mem}}$L mem would reach $\theta$theta. We recompute these inputs in the replay forward pass, so that the values are on the autograd graph. Using these inputs, we then re-compute the eviction log-probabilities via Equation 4, analogously to how we recompute per-token log-probabilities for ${\mathcal{L}}_{\text{token}}$L token. The resulting eviction log-probabilities are therefore live on the autograd graph, and gradients flow from ${\mathcal{L}}_{\text{mem}}$L mem through the eviction decisions into $\theta$theta.

Crucially, this replay step introduces no additional overhead beyond standard RL training, which already requires a forward pass over sampled trajectories. We only have to store a negligible amount of additional state (a binary mask per layer and per eviction round) to construct the replay masks.

4.6 Eviction Rate Curriculum

Finally, we introduce a curriculum to ensure that our model can be trained stably to both optimize its own reasoning and cache management. Intuitively, evicting large amounts of the KV cache from the start of training can be destabilizing to the model. During training, we smooth the transition to more aggressive eviction rates by slowly increasing the eviction rate using a staircase curriculum over the retention rate $p_0$p zero (i.e., percentage of kept entries). More formally, let $\{{\rho }_0,{\rho }_1,\dots ,{\rho }_K\}$the set of rho values from zero to K be a sequence of retention rates with ${\rho }_0\ge {\rho }_1\ge \cdots \ge {\rho }_K$rho zero greater than or equal to rho one, and so on, and let $\Delta$delta denote the number of training steps per stage. At training step $t$t, the current stage index is $\ell =\min (\lfloor t/\Delta \rfloor ,K)$ell, defined as the minimum of the floor of t over delta, and K. Within stage $\ell <K$ell less than K, we apply a linear blend toward the next level during the final fraction $\alpha \in (0,1)$alpha between zero and one of the stage. We set $\alpha =0.6$alpha equals zero point six in our experiments. Let $s=(t\mathrm{mod}\Delta )/\Delta$s denote the fractional progress within the stage. Then:

At the final stage $\ell =K$ell equals K, we set $p_0(t)={\rho }_K$p zero of t equals rho K. See the rightmost panel of Figure 5 for an illustration of this schedule.

5 Experiments

We begin with a controlled study on Countdown, validating and ablating core design choices, and then evaluate on standard math competition benchmarks.

5.1 Shared experimental setup

Model and training recipe. We train DeepSeek-R1-Distill-Qwen-1.5B [Guo et al., 2025]Guo and colleagues using on-policy Dr. GRPO with a staircase eviction curriculum. We use block size $b=32$ and window size $w=5$. Training hyperparameters are standard; see Section A.1 in the Appendix for details.

Baselines. We compare NGC against four inference-time eviction heuristics applied to a model trained with the full KV cache: SnapKV [Li et al., 2024]Snap K V (attention weights), KeyDiff [Park et al., 2025]Key Diff (key diversity), KNorm [Devoto et al., 2024]K Norm (key norm statistics), and StreamingLLM [Xiao et al., 2024]Streaming L L M (sliding window with attention sinks).

Metrics. We report accuracy on held-out test problems. At evaluation time, we replace Gumbel-top-$k$Gumbel top k sampling with greedy top-$k$top k selection, retaining the $K$K highest-scored blocks deterministically.

To measure memory savings, we report the average peak KV cache reduction: the factor by which a method reduces the peak KV cache size relative to no eviction. Concretely, let $p_i$p i denote the prompt length, $c_i^{\text{base}}$c base i the completion length under the no-eviction baseline, and $c_i^{\text{method}}$c method i the completion length under the method being evaluated. Let $\text{peak}(p,c,ϵ,\delta )$peak of p, c, epsilon, delta denote the maximum number of KV entries held in memory at any point during generation for a sequence with prompt length $p$, completion length $c$, eviction rate $ϵ$, and cadence $\delta$. The average peak KV cache reduction is then

where the expectation is over prompts. A value of $2x$two x means the method uses half the peak memory of the no-eviction baseline.

5.2 Controlled analysis on Countdown

Setup. We train DeepSeek-R1-Distill-Qwen-1.5B [Guo et al., 2025]Guo and colleagues for 250 steps on Countdown [Gandhi et al., 2024]Gandhi and colleagues, a combinatorial arithmetic task requiring a model to combine randomly drawn numbers via basic operations to reach a target number. We use a maximum completion length of 1024 tokens, eviction cadence 256, 32 unique prompts per update step, and 16 rollouts per prompt. Our training set contains 327,680 problems (3 and 4 input numbers); our test set contains 1024 held-out problems. Countdown is a demanding testbed since DeepSeek-R1-Distill-Qwen-1.5B achieves near-zero accuracy before training; NGC must jointly learn to reason about this task and manage memory from scratch. It has also become a standard setting for evaluating LM reasoning at scales accessible to academic labs [Gandhi et al., 2024, 2025, Pan et al., 2025]as established in prior work.

5.2.1 Main results

Figure 3: Comparing NGC against KV cache eviction baselines. (left) NGC significantly outperforms baseline cache eviction methods. Error bars correspond to 1 SE and we report pass@1 accuracies. All methods evict 50% of entries at each eviction round, corresponding to 2.4x reduction in peak KV cache size. The dashed horizontal line corresponds to a model trained and evaluated with full cache. (right) We also test NGC’s generalization to other eviction rates (smaller cache sizes) that it was not explicitly trained for; the vertical dashed line indicates the minimum cache size during training. NGC pareto-dominates all other methods even at higher compression rates.

NGC outperforms all baselines on test-set. We report accuracy on 1024 held-out test problems. In Figure 3, we compare NGC against several baseline methods at a 50% eviction rate per eviction round (corresponding to $2.4x$two point four x reduction in peak cache size). NGC achieves 49.6% accuracy, substantially outperforming all heuristic methods. This shows that NGC learns to reason and manage memory jointly, a more challenging task than learning to reason alone.

NGC generalizes to unseen eviction rates. In Figure 3 (right), we evaluate NGC across a range of eviction rates more aggressive than those seen during training. NGC pareto-dominates all baselines, and its performance degrades gracefully beyond the training distribution until a very low cache size.

Figure 4: Ablating NGC design choices. We use two ablations to characterize the properties of end-to-end training. Targeted KV dropout evicts cache entries during RL but ignores the off-policyness introduced by evictions. Token log-probs only corrects for off-policyness via replay masks but drops the eviction decision policy gradient term. Both significantly under-perform NGC.

5.2.2 Ablations

To understand the benefits of end-to-end training, we consider two ablations.

End-to-end training is essential. Our first ablation evaluates whether exposing the model to cache eviction during training enables the model to handle evictions at test time. Concretely, Targeted KV dropout evicts cache entries during training using NGC’s scoring mechanism and applies the same curriculum, but optimizes only Equation 3equation three with a standard causal mask—ignoring the off-policyness that eviction introduces; see Section 4.5. Token log-probs only corrects for this off-policyness using replay masks but drops the eviction policy gradient term ${\mathcal{L}}_{\text{mem}}$L mem. Both significantly under-perform NGC (2.5% and 35.7% vs. 49.6%), confirming that correct off-policy handling and end-to-end training of both reasoning and eviction are essential for performance.

Training dynamics. Figure 5 shows reward and gradient norm over training. NGC improves steadily but more slowly than training without eviction (no evict). Targeted KV dropout initially improves at low eviction rates but collapses during training around step 100, coinciding with a spike in gradient norm. This illustrates the training instability created by off-policyness.

5.3 Budget-aware Interoception

If resource use is part of the task, then intuitively the model could benefit from “sensing” its budget constraints before reasoning. We ask whether explicitly conditioning the model on the eviction rate in the prompt improves its accuracy and generalization to various compression rates. We call this budget-aware interoception: like organisms sensing internal state (such as hunger) to maintain homeostasis, the model perceives its memory budget before it reasons.

Setup. To implement this, we simply modify the prompt during training and testing: we sample a single eviction rate $\rho$rho for a group of $G$G rollouts, as before, and simply append the same structured tag to every prompt in that group:

<eviction_rate>ρ%</eviction_rate>

Figure 5: NGC training dynamics on Countdown. (left) Reward over training: NGC improves steadily but more slowly than training without any eviction (no evict) but eventually reaches a similar level of accuracy. Targeted KV Dropout increases initially at a low eviction rate but collapses after around step 100 due to off-policyness. (center) Targeted Dropout exhibits exploding gradient norms coinciding with its reward collapse, whereas NGC and the no-eviction baseline remain stable throughout training. (right) The staircase eviction rate curriculum, shared across both NGC and targeted KV dropout, which gradually increases the percentage of KV cache evicted to 50% over $\sim$100 steps. Vertical dashed teal lines in the left and center panels mark the steps where the eviction rate begins to increase.

This approach takes inspiration from the “difficulty-conditioned” conjecturer from Poesia et al. [2024]Poesia and colleagues, 2024. At inference time, we vary the eviction rate and compare against a standard NGC model trained under identical conditions but without the eviction-rate tag in the prompt.

Figure 6: Budget-aware interoception. By training and evaluating the model with the eviction rate $\rho$rho in the prompt <eviction_rate>$\rho$rho%</eviction_rate>, we see softened performance degradation as peak KV cache size decreases and generalization to stricter budgets. At aggressive budgets, this gives an 8-13% boost in performance. The dashed vertical line represents the minimum cache size used during training.

Results. Figure 6 reports accuracy as a function of the peak cache size. The interoceptive model matches or exceeds standard NGC across the full training distribution, and the gap widens at more aggressive eviction rates which require generalization beyond the training distribution. The result also illustrates a broader

Figure 7: Pass@k across math reasoning benchmarks. NGC consistently outperforms all inference-time eviction baselines across AMC 2023, AMC 2025, and AIME 2025; all methods evict 50% at each eviction round, corresponding to a 3-5x peak cache size reduction.

Figure 8: NGC pass@32 across varying cache size reductions. For moderate reductions in max cache size 2-3x, NGC can still maintain relatively strong accuracy relative to an upper bound that is trained with no eviction and uses no eviction during inference.

principle: conditioning on resource constraints in the prompt is a simple way to enable a model to have a meta-awareness of its own computational constraints.

5.4 Mathematical reasoning experiments

To test whether the same approach generalizes to broader reasoning tasks, we now evaluate NGC on standard math competition benchmarks.

Setup. We train DeepSeek-R1-Distill-Qwen-1.5B on DAPO-17k [Yu et al., 2025]Yu and colleagues, 2025, a large-scale mathematical reasoning dataset spanning a broad range of problem types and difficulties. We evaluate on AMC 2023, AMC 2025, and AIME 2025 to assess whether NGC generalizes beyond its training distribution. We use the same baseline eviction approaches as before on top of a model trained without eviction. We train with an eviction cadence of 350 tokens and a maximum response length of 1050, using 256 prompts per update step and 8 rollouts per prompt for a total of 469 steps. At test time, we use top_p = 0.95top p equals 0.95 and temperature 0.6 and a max completion length of 3850 tokens.

Results. Figure 7 reports pass@k [Chen et al., 2021]pass at k at a 3-5x cache reduction size (corresponding to 50% eviction rate at every eviction round) across AMC 2023, AMC 2025, and AIME 2025. NGC consistently outperforms baselines across all three benchmarks; heuristic baselines can degrade to near-zero accuracy while NGC maintains a reasonable performance. Moreover, heuristic baselines are inconsistent across datasets. The relative ranking of KeyDiff and SnapKV shifts across benchmarks: SnapKV fails catastrophically on Countdown but performs reasonably well on math reasoning tasks. This inconsistency illustrates why a fixed proxy for importance is fragile — it is task-dependent in ways that are difficult for a practitioner to anticipate. NGC’s consistently strong performance across tasks shows that end-to-end training allows the model to adapt to each task and discover what information to keep in a task-dependent manner. In Figure 8, we compare the pass@32 performance across different cache reduction sizes. We find that at more moderate (2-3x) cache reduction sizes, NGC can come close to matching the upper bound (no-eviction ceiling).

6 Conclusion

We introduced Neural Garbage Collection, a framework in which a language model jointly learns to reason and manage its own KV cache by optimizing outcome-based task reward alone. Experiments on Countdown and DAPO-17k show that NGC outperforms heuristic eviction baselines. NGC is a first step towards a broader vision: as models become more capable, they can also learn to use their own resources more efficiently. Under this view, efficiency is a behavior the model can acquire through the same end-to-end optimization pressure that drives capability. This could unlock the possibility of inverting the usual tradeoff between capability and cost.

7 Acknowledgements

We especially thank Kanishk Gandhi and Aditya Cowsik for detailed comments. We also thank Hengyuan Hu, Neil Band, Omar Shaikh, Thomas Chen, and Cocolab members for helpful discussions as always. This work was supported in part by ONR Grant N00014-22-1-2110, NSF Grant 2205084, and the Stanford Institute for Human-Centered Artificial Intelligence (HAI). EBF is a Biohub, San Francisco, Investigator.

References