Source-linked AI summary
TRACER: Per-Tool Context Retention for LLM Agents via Consequence-Attributed Reinforcement Learning
Ziqi Lin, Ye Wu, Mengying Yang, Xu Liu, Yizhou Liu, Qiang Ke, Qin Guo
TL;DR
Long-horizon data agents accumulate large contexts, while existing compression strategies often ignore the downstream consequences of removing individual tool outputs. TRACER treats retention as a sequential per-tool reinforcement-learning problem with consequence-aware credit assignment, reducing token consumption while maintaining task success across compressors and transfer settings.
Problem
Existing compression strategies allocate retention without accounting for downstream tool re-invocations, creating a compression–consequence gap for long-horizon data agents.
Method
TRACER uses a query-conditioned REINFORCE policy for per-tool retention ratios and a learned outcome model for single-tool counterfactual credit assignment.
Results
TRACER preserves task success while reducing token consumption across compressor backends, with 18–25% savings on held-out LOCA-bench environments.
Takeaways & Limitations
Consequence-aware per-tool context retention transfers across agent backbones and compressor architectures and improves efficiency on held-out environments.
Takeaways & Limitations
The outcome model is not assumed to provide a formal causal guarantee; its usefulness is evaluated empirically.
Abstract
from arXiv · showhide
Enterprise data agents answer business queries by chaining many tool calls over multiple reasoning steps, routinely accumulating hundreds of thousands of context tokens per session. Existing compression strategies typically allocate retention budgets without accounting for the downstream consequences of removing individual tool outputs. Aggressive compression may therefore trigger costly tool re-invocations that offset the initial savings. We call this the compression--consequence gap. To close it, we propose TRACER, which formulates compression as a sequential per-tool decision problem. A lightweight REINFORCE policy assigns query-conditioned retention ratios using only information available at each compression event. Its consequence-aware objective jointly accounts for task success, total token consumption, and post-compression tool re-invocations. To improve credit assignment, TRACER uses a learned outcome model to compare the predicted consequences of the selected retention ratio with those of fully retaining each tool output. On held-out production queries across three compressor backends, TRACER reduces total token consumption by 29--46% relative to keeping all context while maintaining comparable or higher task success. Compared with a tool-type-conditional static policy, TRACER provides an additional 15--18% of token savings. Interventional rollouts show that the learned per-tool credit scores correlate with measured single-tool consequences. The learned policy also yields positive savings when transferred across agent backbones and compressor architectures, and reduces token consumption by 18--25% on five held-out LOCA-bench environments. These results demonstrate the value of consequence-aware, per-tool context retention for improving the efficiency of long-horizon language agents.
1 Introduction
Enterprise data agents accumulate massive contexts because tool outputs persist across long tool-call chains. TRACER addresses the resulting compression–consequence gap with query-conditioned, per-tool retention and consequence-aware reinforcement learning.
- Motivation: Hundreds of thousands of context tokens can accumulate because agents retain outputs across five to ten tool calls, exceeding windows and budgets.A single SQL result can exceed 10,000 tokens, and diluted attention can degrade reasoning.
- Motivation: Existing compressors apply uniform strategies despite heterogeneous downstream importance across tool outputs.Schema lookups may be repeatedly referenced, whereas permission checks may never be revisited.
- Compression–consequence gap: The compression–consequence gap arises when removing information saves tokens initially but triggers costly downstream tool re-invocations.These recovery costs can compound as compression recurs across a session.
- TRACER: TRACER dynamically assigns a retention ratio to each tool output through a reinforcement-learning framework.Its reward combines task success, token efficiency, and post-compression tool re-invocations.
- TRACER: A learned outcome model supplies denser per-tool credit by predicting consequences through single-tool counterfactuals.This addresses the sparse credit-assignment problem of one episode-level reward across many tool decisions.
- Contributions: TRACER formalizes compression as a sequential per-tool decision problem whose reward explicitly prices downstream re-invocations.The framework also uses dynamically determined per-tool retention ratios and a learned attributor.
2 Related Work
Prior context-compression methods operate at granularities from tokens to interaction steps, while prior agentic RL methods optimize action policies or penalize redundant calls. TRACER instead trains a context policy and assigns credit at the individual tool-output level.
- Context compression: Context compression spans token pruning, interaction-step dropping, tiered storage, and learned decisions about when or how to compress.Examples include perplexity-based pruning, counterfactual observation–action retention, and MemGPT’s tiered storage.
- RL for language agents and credit assignment: Prior language-agent RL methods assign credit at step or turn level, exploit episode structure, or penalize redundant tool calls.These approaches optimize the agent’s action policy.
- RL for language agents and credit assignment: TRACER freezes the action policy and trains only the context policy, using excess tool calls as a sparse credit signal.Its counterfactual attribution isolates individual tool-output decisions.
3 Method
TRACER embeds a per-tool retention policy inside the agent loop and optimizes compression using task success, token usage, and re-invocation consequences. Counterfactual outcome modeling provides directional credit for individual tool decisions, while measured recall penalties supply a coarser signal.
- Policy and execution: At each compression event, TRACER reads per-tool features, emits heterogeneous retention ratios, applies a fixed compressor, and continues agent reasoning.The trajectory records success, token usage, and post-compression tool re-invocations for online REINFORCE updates.
- 3.1 Problem Formulation: Each session is an episode, and each compression event is a decision point whose action assigns retention ratios to current tool outputs.A discarded output can force a future re-invocation, coupling decisions across events.
- 3.2 Consequence-Aware Objective: The objective combines session-level task success and token efficiency with event-local penalties for post-compression re-invocations.Reference call counts are measured from matched keep-all timelines rather than assumed.
- 3.2 Consequence-Aware Objective: The measured recall signal contains noise from natural call-count variation between rollouts.The paper describes a later provenance upgrade intended to remove this noise through causal matching.
- Policy parameterization: Retention ratios use a Beta distribution whose bounded support matches the action interval, and evaluation uses its deterministic mean.The policy is a two-layer MLP producing per-tool means for output-instance slots.
- Credit assignment: REINFORCE combines a uniform penalty floor with learned per-tool directionality from an outcome model.The floor remains a coarse penalty signal even if attribution degrades.
- Counterfactual Credit Assignment: The outcome model predicts per-tool re-invocations, token totals, and success probability while modeling cross-tool coupling.Single-tool counterfactuals compare the chosen ratio with full retention; savings exceeding downstream cost encourage stronger compression.
4 Experiments
TRACER is evaluated on a production data-analysis agent across multiple compressor backends, held-out queries, heuristic baselines, attribution ablations, and transfer settings. It reduces token use while preserving success and limiting tool re-invocations, with learned retention patterns reflecting tool regenerability and transferring across backends and backbones.
- Experimental setup: The evaluation uses a production data-analysis agent with three primary compressor backends, a held-out AGORA-style backend, 120 stratified queries, and five seeds.The agent uses Claude Sonnet 4.6 with a 200K-token context window; evaluation separates 80 training queries from 40 held-out queries.
- Main comparison: Fixed baselines increase total consumption by 3–24% under truncation and Self-Info, with precall re-invocation rates of at least 0.07.These results show that uniform compression can trigger recovery behavior that offsets or exceeds initial savings.
- Main comparison: TRACER saves 29–46% of tokens relative to keep-all while maintaining task success across all compressors.Summarization saves 46%, Self-Info 29%, truncation 31%, and the AGORA-style backend 34%; re-invocation rates remain at or below 0.025.
- Main comparison: TRACER saves an additional 15–18% beyond the tool-type-conditional baseline, whose savings range from 11–28%.The results attribute the residual improvement to query-conditioned adaptation because the same tool type can require different retention across queries.
- Attribution ablation: Attribution improves task success by 4.8, 3.5, and 2.2 percentage points under summarization, Self-Info, and truncation, respectively.It also reduces iso-success token ratios by 3.4%, 4.1%, and 4.7%, while aggregate re-invocation probability decreases by 13%.
- Validation: The learned attributor’s predicted consequences correlate with interventions at Spearman ρ = 0.66, recover 82% of the true top-three tools, and match 84% of credit directions.These measurements validate the usefulness of the per-tool consequence signal without establishing a formal causal guarantee.
- Behavioral analysis: Under summarization, execute_sql receives the highest converged retention ratio at 0.65 and resolve_date_range the lowest at 0.48, matching their regenerability costs.Under truncation, the ratio spread narrows to 0.52–0.60 because positional compression cannot exploit fine-grained ratio differences.
- Transfer: Zero-shot transfer retains 18–40% savings across Qwen-3.7-Max and GPT-5.5, while target-specific AGORA training reaches 34% savings at matched success 0.65.GPT-5.5 preserves 78–87% of training-backbone gains and Qwen-3.7-Max preserves 62–76%; the learned importance structure transfers across backends without fine-tuning.
5 Conclusion
TRACER treats context compression as a sequential per-tool decision problem and uses counterfactual credit to attribute downstream consequences. The framework saves tokens across compressor backends while preserving task success and transfers across backbones and architectures.
- TRACER treats context compression as a sequential per-tool decision problem.
- TRACER saves tokens across four compressor backends while preserving task success.
- A two-channel credit scheme attributes consequences to individual tools through single-tool counterfactuals, lifting success on semantic compressors.
- TRACER generalizes to held-out queries, transfers zero-shot across agent backbones and compressor architectures, and achieves 18–25% savings on LOCA-bench.
B Full Training Algorithm
The training algorithm combines session-level rewards with event-local and model-based attribution. Its credit favors compression for low-consequence outputs while protecting outputs whose removal is predicted to cause downstream cost.
- Training schedule: Phase 0 keep-all rollouts establish Tbase, while later eligible policy events train the outcome model after a 20-episode warmup.
- Credit design: The immediate-saving term rewards lower requested retention, while the counterfactual term subtracts only predicted harm relative to fully retaining a tool.
- Credit design: The net credit favors compression for low-consequence outputs and protects outputs with predicted downstream cost.
- Evaluation: The model is evaluated empirically through attribution ablations and single-tool interventional rollouts rather than assumed to provide a formal causal guarantee.
- Reward structure: The shaped surrogate uses an event-local recall signal, while the terminal base reward is defined only at the episode’s final event.
E Provenance-Based Attribution
The implementation describes per-tool event recording, bounded retention actions, outcome-model attribution, and several reference baselines. Provenance matching is presented as an upgrade path, while the reported experiments retain the measured window-based signal.
- Provenance attribution: Provenance attribution tags removed spans with event, tool, and content-signature information for later matching to tool calls.
- Provenance attribution: The provenance-based penalty and targeted floor use matched costs, removing the time-window mixing identified in the main paper.
- Scope: The reported experiments retain the measured window-based signal; provenance matching is an upgrade path rather than a replacement.
- Implementation: TRACER records ordered tool-output instances, applies bounded retention ratios, and updates policy and outcome models from event-level records.
- Outcome model: The outcome model predicts re-invocation counts, log token cost, and auxiliary success probability, with only the first two entering predicted cost.
- Baselines: The reference baselines include keep-all, uniform-0.5, recency, token-proportional, and tool-type-conditional policies.
- Documentation: The appendix restates held-out metrics and implementation defaults rather than introducing new experimental claims.
- Runtime: The implementation uses a CPU-based policy and outcome model with approximately 102K trainable parameters, while external agent and compressor calls dominate wall-clock cost.
H.5 LOCA-bench Training Configuration
LOCA-bench evaluation trains TRACER on 10 of 15 environments and evaluates it on the other five across three context-window settings and five seeds.
- TRACER is trained for five iterative-RL rounds on 10 of LOCA-bench’s 15 environments and evaluated on the remaining 5.
- Evaluation aggregates the 96K, 128K, and 256K context-window settings over 5 seeds.
I Difficulty Tiers and Evaluation Boundary
The evaluation uses stratified production-query splits and distinguishes training penalties, event-level recall statistics, and token-ratio definitions. It also specifies separate uncertainty constructions and paired comparisons for different analyses.
- Difficulty tiers: 120 production queries span three difficulty tiers, with 80 for training and 40 held out for evaluation.The held-out composition is stratified by difficulty, and held-out policies are deterministic.
- Evaluation metrics: The reported recall statistic is an event-local re-invocation rate averaged within records and then across valid query–seed records.It excludes the λrecall training factor and is not summed across events.
- Uncertainty and tests: Table 3 reports native point estimates with state-clustered percentile-bootstrap intervals that preserve paired interventions within each held-out state.These intervals are distinct from the IQM intervals used for method-comparison token summaries.
- Uncertainty and tests: Method comparisons use two-sided paired Wilcoxon signed-rank tests after explicit observation-key alignment, with the pooled attribution-ablation result reported as p = 0.008.Unmatched records are not positionally paired, and unreported per-baseline p-values are not assigned to Table 1.
- Evaluation metrics: Attribution token savings use the iso-success token ratio, while the general token ratio remains T/Tkeep-all.The success conditioning applies to the attribution comparison rather than every token ratio in the paper.
J.4 Reproducibility of Randomness
Randomness controls make the released training choices reproducible across production-query RL runs, while the related-work distinction places TRACER among methods that freeze the agent policy and optimize context retention.
- Reproducibility of randomness: Production-query RL runs use seeds 0–4 and seed Python, NumPy, PyTorch, query sampling, policy initialization, exploration, and replay-buffer sampling.Frozen evaluation uses the deterministic policy mean, while LOCA additionally propagates the seed through PYTHONHASHSEED.
- Policy distinction: TRACER differs from cited credit-assignment and tool-call approaches by freezing the agent policy and training a context policy.Its counterfactual credit compares the selected retention ratio with a single-tool fully retained alternative.
L Cost-Benefit Analysis
The cost-benefit analysis compares recurring token savings with upfront training cost using a simple break-even calculation. Its scenario range is price-sensitive and excludes several deployment costs and risks.
- Break-even model: Break-even time equals training cost divided by daily token-cost savings: dbreak-even = Ctrain / (V T0(1 − r)p).The calculation uses training cost, keep-all tokens, observed token ratio, token price, and daily volume.
- Scenario calculation: Approximately 24 days, or 3.4 weeks, results from Ctrain = $2,800, T0 = 85,000 tokens, r = 0.541, p = $3/10^6 tokens, and V = 1,000 sessions/day.An effective price near $1.7 per million gives a 42-day break-even, showing that the stated 3–6 week range reflects pricing sensitivity.
- Scope and limitations: The calculation credits only input-token reduction and excludes summarizer costs, scorer compute, output pricing, infrastructure, retries, latency, maintenance, distribution shift, and task-success value.It is therefore not a net-profit, ROI, permanence, or non-degradation guarantee, and deployments must recompute it from billed costs.
M Released Data, Code, and Reproduction Boundary
The released supplement supports inspection, testing, and reruns against compatible external backends, but proprietary production artifacts and raw evaluation inputs prevent exact end-to-end reproduction.
- Released code: The code supplement includes the reference policy, feature construction, outcome model, training loop, compressor interfaces, reward functions, statistical utilities, and evaluation protocols.It also includes offline unit tests and scripts for held-out, intervention, transfer, AGORA-style, and LOCA-bench evaluations.
- Released data: The released query file documents the nominal 120-query composition, difficulty labels, 80/40 split, and 7/19/14 held-out stratification.Redacted or empty targets are marked unscorable, so the file is not a drop-in replacement for the private evaluation set.
- Reproduction boundary: The supplement omits the proprietary production agent, data contents, raw prompts and tool outputs, trajectories, checkpoints, matched table inputs, LOCA-bench, and commercial endpoints.The API client specifies an integration contract rather than implementing the production backend.
- Reproduction boundary: The artifact supports source inspection, offline semantic tests, and reruns against a separately configured compatible backend, but not exact end-to-end replay of production experiments.Required external backends, checkpoints, trajectories, or benchmark artifacts are absent.