MemoryWalker: Stop Training Agents on Contexts They Never Saw
Zinco J, Xunjie Zhu, Shen Huang, Zhenyi Wang, Pengjun Xie, Jieping Ye
Abstract
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 epsilonKL gives an O(sqrt(epsilonKL)) 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.
Create a lesson
Related papers
A Common Measure of Communication for Speech Brain-Computer Interfaces
Dulhan Jayalath, Benjamin Ballyk, Oiwi Parker Jones
Graph Machine: Towards Better Pretraining via Edges
Lintai Hou
The Implications of Linguistic Illegibility for LLM Security
James Mickens
Post-Training Language Models for Gold-Medal Performance in Coding Competitions
Aleksander Ficek, Sean Narenthiran, Mehrzad Samadi et al.
UE5M3 FP4 Block Scaling for Stable Language Model Pretraining
Robert Hu, Carlo Luschi, Paul Balanca
Cliff: Learning Process Rewards from the First Mistake
Peixuan Han, Runhui Wang, Ketan Ramaneti et al.