# MemoryWalker: Stop Training Agents on Contexts They Never Saw

arXiv: https://arxiv.org/abs/2609.00865 · HTML: https://mlsanity.com/paper/2609.00865/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 36 of its 49 references that are in the map.
- First posted on arXiv 1 Sep 2026.

## Builds on (most cited first, 30 of 36)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning](https://mlsanity.com/paper/2503.09516/) (2025)
- [MemGPT: Towards LLMs as Operating Systems](https://mlsanity.com/paper/2310.08560/) (2023)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Improving language models by retrieving from trillions of tokens](https://mlsanity.com/paper/2112.04426/) (2021)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [Measuring and Narrowing the Compositionality Gap in Language Models](https://mlsanity.com/paper/2210.03350/) (2022)
- [MuSiQue: Multihop Questions via Single-hop Question Composition](https://mlsanity.com/paper/2108.00573/) (2021)
- [SnapKV: LLM Knows What You are Looking for Before Generation](https://mlsanity.com/paper/2404.14469/) (2024)
- [Group Sequence Policy Optimization](https://mlsanity.com/paper/2507.18071/) (2025)
- [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [Understanding the planning of LLM agents: A survey](https://mlsanity.com/paper/2402.02716/) (2024)
- [LLMLingua: Compressing Prompts for Accelerated Inference of Large Language Models](https://mlsanity.com/paper/2310.05736/) (2023)
- [MEM1: Learning to Synergize Memory and Reasoning for Efficient Long-Horizon Agents](https://mlsanity.com/paper/2506.15841/) (2025)
- [Memory-R1: Enhancing Large Language Model Agents to Manage and Utilize Memories via Reinforcement Learning](https://mlsanity.com/paper/2508.19828/) (2025)
- [Fact, Fetch, and Reason: A Unified Evaluation of Retrieval-Augmented Generation](https://mlsanity.com/paper/2409.12941/) (2024)
- [ReSearch: Learning to Reason with Search for LLMs via Reinforcement Learning](https://mlsanity.com/paper/2503.19470/) (2025)
- [AgentFold: Long-Horizon Web Agents with Proactive Context Management](https://mlsanity.com/paper/2510.24699/) (2025)
- [WideSearch: Benchmarking Agentic Broad Info-Seeking](https://mlsanity.com/paper/2508.07999/) (2025)
- [Reweighted Wake-Sleep](https://mlsanity.com/paper/1406.2751/) (2014)
- [A Survey of LLM-based Deep Search Agents: Paradigm, Optimization, Evaluation, and Challenges](https://mlsanity.com/paper/2508.05668/) (2025)
- [REDSearcher: A Scalable and Cost-Efficient Framework for Long-Horizon Search Agents](https://mlsanity.com/paper/2602.14234/) (2026)

## Abstract (arXiv)

Production agent harnesses such as Claude Code and Qwen-Agent compress context during rollout, but training under compression creates a conditioning problem: every eviction branches the effective history, so the learning object is a tree rather than a sequence. Existing linearizations either retain the rightmost path, causing time-travel leakage, or replay a depth-first traversal, causing train-inference mismatch. We introduce two exact, gradient-equivalent corrections: LogitTree, a segmented K-forward traversal, and a packed 4D attention mask. LogitTree requires K+1 backward passes; the 4D mask requires a custom kernel and white-box eviction records. We also propose SDCC (Self-Distillation for Conditioning Consistency), a single-backward-pass variational relaxation. At each eviction, it minimizes forward KL between the compressed student and a stop-gradient teacher on the reconstructed pre-eviction prefix. A residual per-junction KL of epsilon_KL gives an O(sqrt(epsilon_KL)) bound on the train-deployment total-variation gap. SDCC also applies to black-box harnesses. On seven web-search benchmarks with TC-RAG, AgentFold, MemexRL, Claude Code, and OpenCode, naive training inflates the train-rollout log-probability gap, especially on eviction-heavy batches. The exact methods stay at the no-compression floor, and SDCC substantially closes the gap, with lower logit drift and higher rollout rewards.
