[2604.18002] Neural Garbage Collection: Learning to Forget while Learning to Reason

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.

# link_28mojulp8ma.pdf ## Metadata - PDFFormatVersion=1.7 - IsLinearized=false - IsAcroFormPresent=false - IsXFAPresent=false - IsCollectionPresent=false - IsSignaturesPresent=false - Author=Michael Y. Li; Jubayer Ibn Hamid; Emily B. Fox; Noah D. Goodman - Creator=arXiv GenPDF (tex2pdf:a6404ea) - Custom.DOI=https://doi.org/10.48550/arXiv.2604.18002 - Custom.License=http://arxiv.org/licenses/nonexclusive-distrib/1.0/ - Custom.PTEX.Fullbanner=This is pdfTeX, Version 3.141592653-2.6-1.40.28 (TeX Live 2025) kpathsea version 6.4.1 - Custom.arXivID=https://arxiv.org/abs/2604.18002v1 - Producer=pikep

saved by

• Devin Plumb

related reading

• [2005.11401] Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasksarxiv.org
• [2305.18290] Direct Preference Optimization: Your Language Model is Secretly a Reward Modelarxiv.org
• [2412.06769] Training Large Language Models to Reason in a Continuous Latent Spacearxiv.org
• [2201.11903] Chain of Thought Prompting Elicits Reasoning in Large Language Modelsarxiv.org
• [2510.13786] The Art of Scaling Reinforcement Learning Compute for LLMsarxiv.org
• [2506.06266] Cartridges: Lightweight and general-purpose long context representations via self-studyarxiv.org
• Training Large Language Models to Reason in a Continuous Latent Spacearxiv.org
• [2606.10346] Reasoning or Memorization? Direction-Aware Diversity Exploration in LLM Reinforcement Learningarxiv.org
• [2603.07267] How to Steal Reasoning Without Reasoning Tracesarxiv.org
• [2602.19362] LLMs Can Learn to Reason Via Off-Policy RLarxiv.org
• [2501.00663] Titans: Learning to Memorize at Test Timearxiv.org
• Unlocking the Working Memory of Large Language Models for Latent Reasoningarxiv.org