PaperlayerView exact revision on arXiv

Source-linked AI summary

MemoryWalker: Stop Training Agents on Contexts They Never Saw

Zinco J, Xunjie Zhu, Shen Huang, Zhenyi Wang, Pengjun Xie, Jieping Ye

arXiv:2609.00865v1cs.LGcs.CL

TL;DR

Context compression turns agent rollouts into conditioning trees, so sequential replay can score tokens under histories different from those used during generation. The paper introduces exact tree-consistent replay and SDCC, a one-backward-pass relaxation; exact methods recover the no-compression drift floor, while SDCC substantially narrows the gap across evaluated harnesses and tasks.

  • Problem

    Evictions branch rollout histories, but existing sequential replay schemes condition training on paths different from those used during generation.

  • Method

    The paper proposes exact LogitTree and packed 4D attention traversals, plus SDCC, which distills original pre-eviction behavior into compressed contexts.

  • Results

    Across seven web-search benchmarks and multiple white-box and black-box harnesses, exact replay restores the no-compression drift floor, while SDCC substantially narrows the gap and improves rollout rewards.

  • Takeaways & Limitations

    Conditioning-tree replay provides a consistent training target for edited-context agents, while SDCC offers a lower-cost option that also works with black-box harnesses.

  • Takeaways & Limitations

    The packed 4D formulation requires white-box eviction records, custom attention support, and substantial memory and computational resources.

Abstract

from arXiv · show

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.

1 Introduction

Context compression makes post-training inconsistent because evictions turn each rollout into a conditioning tree, while common replay schemes score tokens under histories different from those used during generation. The paper formalizes this problem and proposes exact tree-consistent methods plus SDCC, a one-backward-pass relaxation.

  • Motivation: Long-horizon agent trajectories make retaining full interaction histories computationally expensive and can introduce noise, invalidated branches, and diluted goals.These pressures motivate bounded-context memory editors during rollout.
  • Problem: Memory editors rewrite interaction history before decoding, so training may lack the context under which rollout tokens were originally generated.The mismatch affects post-training methods including SFT and RL.
  • Problem: Each eviction branches the conditioning history, making the learning object a tree rather than a sequence.Subsequent tokens use the compressed replacement, while earlier tokens retain the original prefix.
  • Pitfalls: Naive-Compressed causes time-travel leakage, whereas Naive-Full causes train–inference conditioning mismatch by using histories that are respectively too short or too long.Both flattenings assign rollout tokens contexts different from those used during generation.
  • Solutions: LogitTree and packed 4D attention provide exact gradient-equivalent tree traversals, while SDCC offers a training-efficient relaxation.The exact methods restore conditioning; SDCC distills behavior under the original context into the compressed context.
  • Evaluation: Across seven web-search benchmarks and multiple memory editors, the methods reduce train–inference inconsistency and logit drift while improving rollout rewards over naive compressed-stream training.The evaluation covers three white-box and two black-box memory editors.

2 Related Work

Prior agent-RL pipelines commonly treat edited contexts as ordinary sequential data, while memory-management methods optimize editing itself. This paper addresses the distinct unresolved question of which context the training gradient should use after rollout history has been rewritten.

  • Harness context compression: Production harnesses use bounded windows, summaries, paging, or retrieval to rewrite physical interaction history during rollout.Examples include Claude Code, Qwen-Agent, MemexRL, MEMGPT, TC-RAG, and StackPlanner.
  • Agent RL over edited contexts: Recent agent-RL methods fine-tune language models on rollout trajectories, typically treating edited context as ordinary sequential data.The related work includes policy-gradient RL and agentic extensions.
  • Research gap: Learnable memory-management methods optimize offloading and retrieval, but do not determine which context should condition training gradients after rewriting.This distinguishes memory optimization from conditioning consistency.

3 Pitfalls: Two Default Ways to Train on an Edited Rollout

Edited rollouts form conditioning trees because each eviction preserves the pre-edit generation path while continuing from a compressed spine. The two natural sequence flattenings fail in opposite directions, producing measurable train–inference logit drift.

  • 3.1 Train–inference logits drift: The train–inference logits drift averages the difference between rollout-time and training-time token log probabilities over loss-carrying response tokens.With no compression, the contexts coincide and the statistic has a 0.014 numerical floor.
  • 3.2 Setup: the live view and the trajectory tree: Each eviction forks the history into a generation-time leg retaining the evicted span and a compressed spine from which rollout generation continues.With K evictions, every token has a root-to-leaf generation path and a corresponding prefix in the final compressed walk.
  • 3.3 The two pitfalls and the conditioning invariant: The final compressed walk and full physical trace are natural representations, but neither preserves every token’s generation-time conditioning path.The compressed walk applies edits too early; the physical trace retains edited-away content too long.
  • 3.3 The two pitfalls and the conditioning invariant: Figure 1 contrasts the staircase of live contexts with two flattenings: one omits previously visible tokens, while the other revives evicted tokens.Their logit-difference distributions have opposite signs, with µcomp ∈[−4.0, −1.6] and µfull ∈[+0.2, +0.7].
  • 3.3 The two pitfalls and the conditioning invariant: Naive-Compressed scores tokens on the final compressed path, allowing future edits to propagate backward as time-travel leakage.In the example, y6 and y7 are scored as if y5 had never existed; one leaf contributes ∆comp = −22.8 nats.
  • 3.3 The two pitfalls and the conditioning invariant: Naive-Full scores post-eviction tokens on the physical prefix, exposing them to content unavailable during rollout and deployment.The matched example gives ∆full = +18.5 nats, opposite in sign and comparable in magnitude to Pitfall A.
  • 3.3 The two pitfalls and the conditioning invariant: Conditioning consistency requires training to use only the edits fired through each token’s decoding step, not future edits.Both naive paths violate this invariant whenever an edit changes a target’s evaluation prefix.

(d) SDCC

SDCC keeps the compressed rollout as a student and aligns it with teacher behavior reconstructed under pre-eviction contexts. It approximates exact conditioning consistency with one backward pass while avoiding explicit replay of every tree branch.

  • SDCC: SDCC pulls the single compressed walk toward teacher legs representing the original pre-eviction conditioning paths.The student remains the compressed forward, while teacher contexts are reconstructed for distillation.
  • SDCC: The exact materializations traverse each branch separately or pack the same branches into one masked forward pass, while SDCC retains the compressed forward.This positions SDCC as the training-efficient option among the three consistent materializations.
  • SDCC: SDCC is a soft alternative to exact tree replay: the exact methods have zero conditioning bias, whereas SDCC provides an O(√ε_KL) relaxation.The relaxation avoids explicitly materializing every branch.

4 Exact Solutions: Two Ways to Walk the Context Tree

The paper restores rollout-time conditioning by traversing the conditioning tree exactly, either through separate branch sequences or a packed 4D attention mask. These methods are gradient-equivalent, but trade computational cost against systems requirements.

  • Core idea: The exact solutions score each target on the unique branch whose prefix matches its rollout-time live view, eliminating conditioning bias.Both approaches enforce conditioning-invariance across the K+1 branches of the tree.
  • LogitTree: LogitTree materializes the K+1 root-to-leaf branches as separate sequences and masks each target onto exactly one matching branch.Standard causal attention then restricts each target to predecessors on its own leg.
  • Packed 4D attention mask: The packed 4D mask keeps branches in one sequence while preventing cross-branch attention and assigning gap-free logical positions.Each query consequently attends only to the root-to-leaf path under which it was generated.
  • Equivalence: The packed 4D mask and LogitTree compute identical training gradients under dense softmax attention.The theorem identifies both constructions as equivalent implementations of the same conditioning-tree traversal.
  • Cost and deployment: LogitTree requires K+1 backward passes and incurs a 5–20× wallclock overhead, while the 4D formulation requires one pass but white-box harness and model access.The 4D method also faces high memory and computational complexity for very long contexts.

5 SDCC: Self-Distillation for Conditioning Consistency

SDCC approximates exact conditioning-tree walks by distilling the live-context policy into the compressed-context policy at diverging leaves. It retains one gradient-carrying backward pass while providing a bound on the resulting behavioral gap.

  • Motivation and role: SDCC aligns the compressed replay policy with the original live-context policy only at diverging leaves, rather than materializing the full conditioning tree.The live-context evaluation is a stopped-gradient teacher, while the compressed-context evaluation is the gradient-carrying student.
  • Local alignment pairs: A diverging leaf is a response token whose original live prefix differs from its prefix in the final compressed walk.For example, an evicted token can remain in the live prefix used to decode a later response but be absent from the compressed prefix.
  • Procedure: SDCC runs one ordinary compressed-walk student pass, evaluates reconstructed live prefixes without gradients, and adds a KL penalty at diverging leaves.The penalty is forward KL, so the stopped-gradient teacher distribution acts as the fixed soft target.
  • Objective: The SDCC loss combines the ordinary compressed-stream task loss with a leaf-gated forward-KL penalty weighted by λ.With no differing contexts, the leaf set is empty and the loss equals the usual task loss; λ = 0 recovers Naive-Compressed.
  • Guarantee: For each diverging leaf, the residual KL εp bounds the teacher–student next-action distribution gap through a behavioral Pinsker bound.When εp = 0, SDCC matches the exact walk’s next-token distribution at that leaf; otherwise it remains an explicitly bounded approximation.

6 Experiments

Experiments across models, harnesses, recomputation schemes, and web-search settings show that naive replay increases conditioning drift, while exact methods recover the no-compression floor and SDCC narrows the gap.

  • Experimental setup: Experiments span Qwen3-4B and Qwen3.7-Air, three white-box editors, two black-box agents, and five logit-recomputation schemes.The evaluation uses live web tools and seven web-search benchmarks, with pooled training data from REDSEARCHER and ASEARCHER.
  • Grand result matrix: The no-compression control has logdiff = 0.014, while Naive-Compressed rises 30.5× to 0.366 and Naive-Full rises 9.7× to 0.203 as eviction increases.Exact methods remain near the floor: 4D moves 3.7× and ends at 0.022, SDCC moves 6.5× to 0.071, and LogitTree is flat at 0.012–0.013.
  • Grand result matrix: Exact walks recover the drift floor, with LogitTree at 0.012–0.013 and 4D at 0.006–0.022 versus the no-compression anchor at 0.014.Naive-Compressed reaches 0.366 on AgentFold, approximately 26× the floor; independently trained rows need not match exactly despite gradient equivalence.
  • Structural diagnostics: On untrained Qwen3-4B, Naive-Compressed consistently under-places tokens and Naive-Full-leak consistently over-places them across all three white-box harnesses.The predicted signs hold for a majority of tokens in every cell, indicating structural failure modes rather than effects that reinforcement learning can simply absorb.
  • SDCC dynamics: On the AgentFold cell, SDCC’s mean logdiff falls 55% over the logged window and crosses below the same-window Naive-Compressed trace.This is a single-editor, single-scale comparison without a confidence interval; learning impact is assessed by EM rather than logdiff alone.
  • Black-box transfer: For black-box harnesses, SDCC reaches 37.5 average EM on Claude Code and 36.9 on OpenCode, while observed logdiff remains 0.015 and 0.013, respectively.Claude Code and OpenCode expose neither eviction spans nor compression depth, so evaluation uses downstream EM and observed maximum context length.
  • Large-scale evaluation: On WIDESEARCH, LogitTree raises Pass@1 from 4.5% to 5.0% and Pass@4 from 7.0% to 10.0% under Claude Code.The corresponding relative improvements are 11.1% and 42.9%; row- and item-level F1 improve by 16.1% and 7.0%.

7 Conclusion and Future Work

The paper frames edited rollouts as conditioning trees, presents exact replay through LogitTree and a 4D mask, and introduces SDCC as a single-backward approximation. Future work targets broader compression policies, more harnesses, and sparse or linear-attention architectures.

  • Conclusion: Conditioning trees expose a train–inference mismatch in editor-based agent reinforcement learning, where naive replay produces measurable logit drift and degraded rollout reward.The paper identifies Naive-Compressed and Naive-Full as opposite-direction failures and evaluates the issue across multiple settings.
  • Conclusion: LogitTree and the 4D attention mask provide exact conditioning-consistent replay, while SDCC substantially narrows the gap with a single-backward objective.The exact methods restore the no-compression drift floor; SDCC is the training-efficient approximation.
  • Future work: Future work includes broader compression policies, a wider range of harnesses, and exact tree-consistent replay for sparse and linear-attention architectures.The stated examples include dropping tool observations and using latent summaries.

Part I Position in the agentic-RL literature

The paper positions itself against prior harness-RL work that trains on physical trajectories and claims to formalize the conditioning invariant, characterize its violations, and propose a soft remedy with a provable bias bound.

  • Prior work: Published QA-oriented harness-RL baselines such as Search-R1 and ReSearch train on the physical trajectory H_t.The paper also situates itself alongside emerging stack-based and learned-memory RL approaches for long-horizon agents.
  • Positioning: The authors claim their work is the first to formalize the conditioning invariant, characterize failure modes from violating it, and propose a soft remedy with a provable bias bound.This positioning is explicitly qualified as a claim made to the authors’ knowledge.

Part II Method supplements: pitfalls, structure, proofs, variational derivation, convergence

Context compression turns rollout training into a conditioning-tree problem: naive serializations either apply edits too early or retain evicted context. The paper formalizes the tree and compares exact traversal with SDCC’s cheaper approximation.

  • Pitfalls: Naive-Compressed scores targets on the final compressed walk, so earlier targets can be conditioned on summaries created only after their generation.This is labeled time-travel leakage.
  • Pitfalls: A worked eviction rollout shows two opposite failures: Naive-Compressed loses context needed by earlier targets, while Naive-Full retains context already discarded.The example is illustrative rather than measured; its measured counterparts are reported separately.
  • Tree structure: The trajectory tree contains one root, K junctions, R leaves, and K+1 branches; K governs exact traversal cost, while D identifies SDCC’s corrected leaves.Each leaf also carries an eviction mass measuring tokens present at generation but absent from the final-walk prefix.
  • Pitfalls: Naive-Full scores targets on the physical depth-first tape, retaining evicted spans and creating stale-context leakage.The training prefix can strictly exceed the live view available during rollout.
  • Methods: The 4D mask and LogitTree restore live conditioning exactly, whereas SDCC matches compressed students to stop-gradient teachers only at diverging leaves.The exact methods require K+1 backward passes or specialized masking; SDCC uses one backward pass but gives up exactness, with an O(√ε_KL) bias bound.

B.4 Per-leaf divergence and the two pitfalls

Per-leaf divergence arises when the live-view prefix differs from the final compressed-walk prefix. The two naive methods fail in opposite directions, while branch-aware gradients restore the conditioning invariant for SFT and RL.

  • B.4 Per-leaf divergence and the two pitfalls: Future edits applied retroactively make the final compressed prefix shorter than the live prefix that generated an earlier target, producing Pitfall A.This mismatch affects both supervised scoring and RL importance ratios.
  • B.4 Per-leaf divergence and the two pitfalls: Physical-prefix replay preserves already-evicted content, so Naive-Full trains on information unavailable to the deployed policy, producing Pitfall B.In RL, the resulting advantage baseline is evaluated on the wrong observation and the bias accumulates with eviction events.
  • B.4 Per-leaf divergence and the two pitfalls: 0.014 is the no-compression logdiff floor, while Naive-Compressed reaches 0.024–0.059 on average and 0.23–0.29 on eviction-heavy batches.The latter peak is reported as approximately 20× the baseline and grows with eviction density.
  • B.4 Per-leaf divergence and the two pitfalls: LogitTree and the 4D mask place gradients on each target’s live branch, making rollout and training conditioning identical for SFT and RL.The 4D construction additionally requires exact logical visibility, reassigned positions, decoupled masks, boundary-safe normalization, and a shared mask.

C.3 Materialization equivalence (theorem 1)

Theorem 1 establishes that packed 4D masking and segmented LogitTree are two materializations of the same conditioning-tree computation. They therefore produce identical logits and gradients, subject to the mask construction assumptions.

  • C.3 Materialization equivalence: The 4D-masked union forward and K-segment LogitTree produce the same logits at every non-evicted position.The equivalence assumes the logical mask, position reassignment, mask separation, boundary-safe normalization, and shared rollout/training mask specified for the 4D method.
  • C.3 Materialization equivalence: Both methods compute the same leaf-summed loss, so their gradients coincide up to floating-point association order.Their difference is computational organization: one structured attention call versus K+1 separate forwards.
  • C.3 Materialization equivalence: SDCC is derived as a forward-KL conditioning-consistency regularizer whose correctness depends on the KL residual’s zero set and behavioral control.The derivation does not require a tight ELBO.
  • C.3 Materialization equivalence: The forward KL is chosen because reverse-KL optimization would require sampling continuations from the compressed context and can have unbounded score-function variance on the surplus side.The deficit side corresponds to the time-travel failure of Naive-Compressed.

D.1 Probabilistic model and the conditioning invariant

The paper models compressed rollouts as conditioning trees and defines consistency as matching each token’s training context to the context delivered at deployment. SDCC uses a forward-KL residual to relax exact consistency while providing convergence and deployment-gap guarantees under stated assumptions.

  • Probabilistic model: At each diverging leaf, the clean context zp is the pre-eviction prefix, while compressed context xp removes spans active at that junction.The action y is the next generated token, and y⋆ is an oracle outcome whose likelihood can encode reward or supervised targets.
  • Conditioning invariant: The conditioning invariant requires training to use the same edit history that was available when each rollout token was generated.The invariant is a target achieved through training rather than a modeling assumption.
  • SDCC relaxation: SDCC minimizes forward KL from a stop-gradient teacher conditioned on zp to a student conditioned on xp, using shared parameters and diverging leaves.The teacher is the deployment-target distribution, while the student is the amortized policy evaluated on compressed context.
  • SDCC relaxation: Zero SDCC residual is equivalent to satisfying the conditioning invariant at every diverging leaf.The residual is the per-leaf KL between the teacher and student action distributions.
  • Guarantees: O(√ε_KL) bounds the deployment total-variation gap when ε_KL is the maximum per-leaf SDCC residual.The bound is unconditional because the harness redelivers the clean pre-eviction context at inference.
  • Guarantees: SDCC’s single-backward objective has λ-tunable slack, coincides with the exact objective on the zero-invariant set, and can collapse at excessively large β.The convergence guarantees additionally require assumptions including a restrictive local PL condition; without it, convergence to a critical point remains but the linear rate is lost.

E.2 Harnesses and methods

The experimental harnesses rewrite context through white-box or black-box editors, while the methods differ in how they reconstruct or approximate the conditioning tree. Exact methods trade compute or infrastructure for fidelity, whereas SDCC preserves single-backward training and supports black-box harnesses.

  • Harnesses: White-box editors expose compression as trainable tool calls, with TC-RAG popping stack envelopes and other editors using different triggers and tree shapes.TC-RAG removes the oldest envelope and has few heavy-span junctions, while model-triggered editors can leave the tree unbranched when untrained.
  • Harnesses: Claude Code and OpenCode are black-box editors whose internal sliding-window and auto-compact operations expose rendered transcripts but not eviction spans or summary provenance.A separate platform pipeline detects prefix divergences and uses the previous payload as the teacher prefix.
  • Methods: Naive-Compressed replays the rendered compressed walk, Naive-Full reinjects evicted spans, and LogitTree splits each rollout into K+1 live-view branches.The exact branch split preserves the context active during each layer’s generated tokens.
  • Methods: The packed 4D mask places all branches in one sequence and admits only keys live when each token was decoded, with branch-specific position ids.Its gradient is identical to the LogitTree construction.
  • Methods: SDCC keeps the compressed student forward, reconstructs stop-gradient teachers by reinserting evicted spans, and applies the KL only at diverging leaves.The teacher and student share weights, and the direction is fixed from compressed student to pre-eviction teacher.
  • Cost and applicability: LogitTree costs K additional backward passes and 5–20× per-step overhead, whereas 4D uses one backward pass but requires arbitrary masks, row-wise position ids, and serving-time equivalence checks.SDCC retains a single backward pass, needs only detectable eviction events, and supports black-box harnesses at an O(√ε_KL) residual.

F.5 Training curves

The training curves diagnose optimization behavior and conditioning drift rather than serving as standalone reward comparisons. Exact methods remain at the no-compression drift floor, while SDCC reduces drift over time and naive compression does not.

  • Scope: Training curves are descriptive diagnostics, while Table 1 remains the source of final numerical results.White-box curves cover 4B RL runs; black-box reward traces cover Claude Code and OpenCode over their first 40 optimizer steps.
  • Reward: Reward rises by 0.07–0.39 from first drawn point to peak across all fifteen white-box cells, but noisy curves cannot separate methods by eye.The normalized-progress axis aligns endpoints despite unequal training lengths, so endpoint contrasts would partly reflect differing windows.
  • Conditioning drift: Ten of fifteen cells sit at 0.93–1.04× the control in the fixed comparison window, while eviction-heavy harnesses separate Naive-Compressed from the floor.The fixed window ends at iteration 58, with control median 0.0138 over n=448.
  • SDCC dynamics: The SDCC correction activates selectively: KL-positive steps have median logdiff 0.081, versus 0.015 when KLSDCC is zero.The KL term is nonzero on batches containing live folds and zero elsewhere, consistent with per-leaf gating.
  • SDCC dynamics: SDCC’s KL magnitude scales with eviction density, from ≈10^-5 on MemexRL to 2.9 × 10^-3 on AgentFold.The intermediate TC-RAG value is 1.1 × 10^-4; the reported session and fold rates accompany these editor-specific magnitudes.
  • Conditioning drift: SDCC’s mean logdiff falls 55% over the logged AgentFold window and ends below Naive-Compressed, approaching the no-compression floor of 0.0135.Naive-Compressed shows no comparable directional trend; this comparison uses one editor at one scale.

G Implementation note: branch-replicated packing vs. the physical-union view

The implementation executes the 4D construction through branch replication, while the paper presents a compact physical-union mask view. Under Proposition 4’s conditions, both views produce identical target logits and gradients.

  • Views: The main text defines the 4D mask as row-wise visibility over the physical union, while implementation replicates shared-trunk tokens per branch.Each replicated token receives a hidden state from its branch-local causal prefix before sequences are packed.
  • Equivalence: Under dense softmax attention and Proposition 4’s five conditions, physical-union and branch-replicated constructions yield identical per-target logits and gradients.The physical-union formulation supports the compact mask definition and proof, whereas branch replication is used by the training loop.
  • Packing cost: The physical union length L is a lower bound because trunk replication increases the actual packed length to Npacked.Table 1’s “max tok.” column reports the true packed input size rather than only L.
Loading 2609.00865v1…