Agent-Omit: Adaptive Context Omission for Efficient LLM Agents

AI summary

Overview

Research area: Efficient large language model (LLM) agents and agentic reinforcement learning (RL), specifically context management during multi-turn agent–environment interaction (arXiv:2602.04284v2, cs.AI).

Technical level: Intermediate. The paper combines a quantitative empirical study, a two-stage training pipeline (supervised fine-tuning plus reinforcement learning), and a short theoretical analysis, so readers will get the most from it with some familiarity with LLM agents, chain-of-thought reasoning, and RL fine-tuning.

Scope: This paper argues that the usefulness of an agent's thoughts and observations varies by interaction turn, and it introduces a framework called Agent-Omit that trains agents to selectively omit redundant thoughts and observations.

What This Paper Is About

LLM agents solving multi-turn tasks accumulate large contexts because they generate reasoning "thoughts" at every step and stack every environment "observation" from previous turns. Existing efficiency methods compress or prune the whole interaction trajectory uniformly. The authors show that not all turns contribute equally — some thoughts and observations are redundant, while others (particularly early planning and late-turn observations) are critical — and they build a training framework that teaches an agent to decide, turn by turn, what to omit.

Key Contributions

1. A unified thought-and-observation analysis framework. The authors quantitatively characterize how thoughts and observations affect both agent effectiveness (task accuracy) and efficiency (token cost) across turns, showing that selective omission can improve efficiency without sacrificing effectiveness.

2. The Agent-Omit framework. A two-stage training approach: (a) Agent Omission Behavior Synthesis, which builds single-turn and multi-turn omission cold-start data for fine-tuning, and (b) Omit-Aware Agentic Reinforcement Learning, which adds a dual sampling mechanism (full and partial trajectories) and a tailored omission reward.

3. Theoretical analysis. The authors prove that the deviation of the learned omission policy from the optimal omission policy is upper-bounded by the KL-divergence between them (Theorem 5.2, "Bounded Omission Error"), building on a Semantic Lipschitz Continuity assumption (Assumption 5.1).

4. Extensive evaluation. Experiments on five agent benchmarks show Agent-Omit-8B achieves performance comparable to seven frontier LLM agents, and that it delivers the best effectiveness–efficiency trade-off among seven efficient LLM agent construction methods.

Main Findings

• Token cost is dominated by thought and observation, not actions. On WebShop using Qwen3-8B, thoughts account for 45.1% of token cost and observations 52.2%, while actions account for only 2.7%.

• Thoughts and observations have different cost profiles over time. Thoughts are "front-loaded," with heavy token consumption in early turns (e.g., Turns 1–2) for high-level planning, whereas observations grow linearly due to stacking, creating a heavy context burden in later turns.

• Accuracy contributions diminish over turns. Using Monte Carlo rollouts, the authors find early thoughts and observations are critical for high Pass@8 accuracy, but the accuracy gain diminishes rapidly as interaction progresses, with later turns often falling below the Pass@1 baseline.

• Omission helps only when applied selectively. Omitting intermediate reasoning steps improved accuracy and reduced token usage, but omitting thoughts in the initial or final turns was detrimental. Observations were most omissible in intermediate turns (Turns 2, 4, 6); in later turns omitting observations caused a significant accuracy drop and the agent generated excessive reasoning tokens trying to bridge the information gap.

• Agent-Omit-8B-RL is the strongest variant. It scored 26.56 Pass@1 / 4,356 avg tokens on DeepSearch, 23.57 / 8,764 on WebShop, 87.00 / 7,328 on TextCraft, 84.36 / 6,643 on BabyAI, and 18.45 / 9,643 on SciWorld. By comparison, Agent-Omit-8B-SFT scored 22.25 / 6,153, 14.43 / 11,376, 72.00 / 11,125, 72.25 / 12,265, and 8.54 / 12,256, respectively. The paper reports Agent-Omit-8B-RL as achieving the highest Pass@1 on WebShop, TextCraft, BabyAI, and SciWorld, while remaining competitive on DeepSearch.

• RL both improves accuracy and efficiency. The authors state that the RL stage not only further enhanced model accuracy over the initial SFT but also actively improved agent efficiency; token statistics for closed-source baselines (OpenAI o3/o4-mini) were unavailable due to API restrictions.

• Beats efficient-agent baselines. On Qwen3-8B, Agent-Omit-8B-RL is reported as achieving the highest Pass@1 and lowest average token cost across all five benchmarks compared with thought management (Thinking-Retention, DEPO, Tool-Light), observation management (Observation-Mask, DeepMiner), and thought-and-observation management (MEM-Agent, ReSum) methods.

• Ablation results. Both SFT and RL stages contribute to gains; in SFT, single-turn omission data is the dominant factor; in RL, partial trajectory sampling matters more than full trajectory sampling, and the omission reward is the primary factor for token reduction. Changing the reward reweighting factor μ from 0.2 (to 0.1 or 0.3) yields a worse effectiveness–efficiency trade-off.

• Omission behavior at inference. The trained agent adaptively omits an average of 3 to 4 turns of redundant thoughts or observations per trajectory, with omission frequency peaking during intermediate turns (turns 3–10).

Methodology in Plain English

The authors first run controlled experiments on WebShop with Qwen3-8B, deliberately deleting one thought or one observation at a specific turn and letting the agent finish the trajectory, then comparing accuracy and token cost against a base run (averaged over 200 randomly sampled test cases, for turns 1 through 12). The turns where token cost dropped without hurting accuracy were treated as "omittable."

To teach the behavior, they synthesize cold-start data in two levels. Single-turn omission samples teach the format — the agent emits an empty thought inside <think></think> tags, or outputs a strategy command like <omit_tool_response_N_...> to trigger removal of a historical observation set. Multi-turn samples replace original thoughts and observations with these omission behaviors so the agent learns to keep reasoning coherent under a shortened history. Full-parameter fine-tuning is then done on these datasets with standard language modeling loss and a loss mask applied to all environment observations.

Because supervised fine-tuning depends on synthetic data, the second stage uses reinforcement learning. A key issue the authors identify is that once an observation is omitted, the agent can no longer see the context that justified the omission, so standard post-hoc trajectory training makes the omission policy unlearnable. Their fix is dual sampling: a full trajectory y (the complete multi-turn interaction, used for efficiency and final success) and partial trajectories y' (each turn where an omission was triggered, using the pre-omission context). Each full trajectory yields a variable number p(y) of partial trajectories. Rewards combine a task reward based on answer correctness with an omission reward equal to the ratio of saved thought tokens plus saved observation tokens to total tokens — set to zero if the task reward is zero, to prevent reward hacking. Training uses Group Relative Policy Optimization (GRPO) with a KL penalty, μ = 0.2, and follows Verl for the loss mask. The theoretical section formalizes the resulting error bound.

Why This Matters

The work reframes agent context management as a turn-level, learned decision rather than a uniform compression heuristic. If thoughts and observations can be dropped exactly where they are redundant, agents can produce the same answers at lower token cost, which directly affects inference expense and energy use. The paper notes that reducing token consumption and computational cost may help lower the energy footprint of deploying LLM agents.

Real-world applications (drawn from the domains the paper evaluates and cites):

• Deep search / information retrieval: agents that call search engines for knowledge-intensive queries (DeepSearch).
• Web navigation and e-commerce: agents that browse sites to extract product attributes and make purchases (WebShop).
• Digital games and long-horizon planning: Minecraft-inspired recipe crafting and planning tasks (TextCraft).
• Embodied control and robotics: instruction following and navigation in partially observable grid-worlds (BabyAI).
• Scientific discovery: hypothesis testing and experiment design in physical simulations (SciWorld).

Industry relevance: The models used are small (Qwen3-4B and Qwen3-8B backbones), and Agent-Omit-8B-RL is reported as competitive with frontier models on several benchmarks at far lower token cost. Training the Qwen3-8B model required 8 NVIDIA A100 GPUs and the 4B model 4 GPUs, which is within reach of many labs and companies. Lower token cost per agent trajectory translates into cheaper deployment of long-running agentic applications, and the general domain of search assistants and embodied decision-making is commercially active (one author group is affiliated with Didichuxing Co. Ltd).

Future Directions

• Scale the omission data synthesis pipeline to large-scale pre-training, rather than relying on a cold-start fine-tuning stage.
• Extend the training paradigm to larger-size LLMs — the current work is validated only on Qwen3-4B and Qwen3-8B.
• Reduce the additional training samples required during RL, which the authors identify as a limitation on scalability.
• Test adaptivity outside the five evaluated environments, which span DeepSearch, WebShop, TextCraft, BabyAI, and SciWorld; whether the learned omission policy transfers to other interaction structures is not established in the reported results.

Target Audience

Researchers and engineers working on LLM agents, agentic reinforcement learning, and inference-cost reduction will benefit most. It is also relevant to practitioners deploying long-horizon tool-using agents where token cost and latency matter, and to readers interested in how context-compression behavior can be learned rather than hand-specified. Readers without background in RL fine-tuning may find the dual sampling and reward design sections harder going, but the empirical analysis in Section 3 is accessible on its own.

Authors’ abstract

Managing agent context (e.g., thought and observation) during multi-turn agent-environment interactions is an emerging strategy to improve agent efficiency. However, existing studies treat the entire interaction trajectories equally, overlooking the thought necessity and observation utility varies across turns. To this end, we first conduct quantitative investigations into how thought and observation affect agent effectiveness and efficiency. Based on our findings, we propose Agent-Omit, a unified training framework that empowers LLM agents to adaptively omit redundant thoughts and observations. Specifically, we first synthesize a small amount of cold-start data, including both single-turn and multi-turn omission scenarios, to fine-tune the agent for omission behaviors. Furthermore, we introduce an omit-aware agentic reinforcement learning approach, incorporating a dual sampling mechanism and a tailored omission reward to incentivize the agent's adaptive omission capability. Theoretically, we prove that the deviation of our omission policy is upper-bounded by KL-divergence. Experimental results on five agent benchmarks show that our constructed Agent-Omit-8B could obtain performance comparable to seven frontier LLM agent, and achieve the best effectiveness-efficiency trade-off than seven efficient LLM agents methods. Our code and data are available at https://github.com/usail-hkust/Agent-Omit.