Don't Lose the Thread: Empowering Long-Horizon LLM Agents with Cognitive Resource Self-Allocation

AI summary

Overview

Research area: Large language model (LLM) agents — specifically context/working-memory management and agentic reinforcement learning for long-horizon web-research tasks.

Technical level: Advanced. The paper combines a formal multi-episode trajectory formulation, a DAPO-derived RL objective, and a full SFT + RL training pipeline; the empirical sections are readable on their own, but the method sections assume familiarity with ReAct-style agents and policy-gradient methods.

Scope (one sentence): The paper proposes CORAL, a two-tool, agent-callable working-memory system (Memory Management and Context Optimization) plus a Multi-episode ARPO training algorithm, and evaluates them on GAIA, five multi-hop Knowledge QA benchmarks, and BrowseComp.

What This Paper Is About

LLM agents that solve long tasks by repeatedly thinking, acting, and observing accumulate a bloated context full of verbose tool outputs and failed attempts, which dilutes attention and degrades planning — the paper calls this "cognitive overload," drawing an explicit analogy to human working-memory limits and cognitive load theory. Existing fixes prune or compress context on fixed schedules or when the window is about to overflow, triggered by external rules rather than the agent's own judgment. CORAL's goal is to let the agent decide for itself when to save a checkpoint of its progress and verified facts, and when to wipe its context and resume from only that checkpoint.

Key Contributions

1. CORAL framework. A callable toolset added to a ReAct loop: a Memory Management (MM) tool that stores task progress (a plan plus pending/completed sub-goals) and verified facts, and a Context Optimization (CO) tool that clears the conversation history and restarts a new episode from only the system prompt, the user question, and the current memory state. Both tools are invoked autonomously by the agent, not by heuristics or external controllers.

2. Multi-episode trajectory formalism. The paper formalizes a trajectory as a sequence of N episodes, where each episode after the first begins from (s, q, o_{i-1,M_{i-1})} — the memory state returned by the previous CO call — rather than from the full history.

3. Multi-episode Agentic Reinforced Policy Optimization (Multi-episode ARPO). An extension of DAPO to multi-episode trajectories: the reward is computed from the final episode (which contains the answer) and propagated backwards to every preceding episode, all episodes from all sampled trajectories for a question are pooled into one set for group-based advantage normalization, so credit assignment crosses episode boundaries and checkpointing policy becomes learnable.

4. Two-stage training recipe. Behavior cloning via SFT on episode-split trajectories (observations masked out of the loss), followed by Multi-episode ARPO. Ablations isolate the contributions of mathematical reasoning data and of retaining failed trajectories.

Main Findings

• GAIA headline number: CORAL with an 8B model (Qwen3-8B + RL) reaches 42.8% average accuracy on the 103-question GAIA text-only validation set (39 Level 1, 52 Level 2, 12 Level 3), versus 19.4% for the Qwen3-8B ReAct baseline, 20.4% for prompting-only CORAL, and 41.6% for SFT alone.

• Gains concentrate on the hardest levels, not the average. CK-Pro-8B has a higher overall score (43.7% vs. 42.8%), but its advantage is entirely on Level 1 (56.4% vs. CORAL's 46.2%); CORAL leads on Level 2 (45.7% vs. 42.3%) and far exceeds it on Level 3 (19.8% vs. 8.3%). The paper attributes the Level 1 gap to the overhead of maintaining a working memory when clutter never became binding.

• Cheaper, not more expensive. Over all 824 GAIA evaluation trajectories (103 questions × 8 samples) with Qwen3-8B, CORAL RL uses 136K input tokens per trajectory versus 1,078K for ReAct (87.4% fewer), 15.2 LLM calls versus 45.3, and a peak context of 12.7K tokens versus 30.6K. The paper attributes this to two compounding mechanisms: CO resets cap context length (attacking the quadratic cost of re-reading history each turn), and an agent that has written down its plan and facts stops re-searching what it already knows, cutting turns roughly threefold.

• Training teaches better checkpointing, not more resetting. The prompting-only agent calls CO 4.3 times per trajectory; the trained RL agent calls it 0.21 times, while calling MM 4.1 times. The paper reads this as evidence that an agent whose checkpoints capture the right state never accumulates the clutter that would justify a reset.

• RL reduces variance. Multi-episode ARPO cuts run-to-run standard deviation on overall GAIA accuracy from 3.5% to 1.9% and raises the worst of eight runs from 36.9% to 39.8%. The 1.2-point headline gain from SFT (41.6%) to RL (42.8%) decomposes into +2.2 on Level 2 and +4.2 on Level 3, against a slight Level 1 regression. Level 3 variance remains high (roughly 9–10%) for both models.

• Strong on a capable proprietary model without training. Claude Sonnet 4 with prompting-only CORAL averages 65.0 on GAIA, comparable to OpenAI Deep Research at 67.4. On BrowseComp (first 150 questions, Claude Sonnet 4, identical search and browse tools), CORAL reaches 30.7% versus ReAct's 22.7% — an 8.0-point gain with no training at all, the largest gain of any benchmark measured.

• Transfers to unseen multi-hop QA. The same Qwen3-8B RL checkpoint trained only on synthesized web questions plus math, with no exposure to these datasets, scores 71.0 (NQ), 70.4 (HotpotQA), 79.2 (2WikiMultiHopQA), 43.2 (Musique), 79.2 (Bamboogle), above Search-o1, R1-Searcher, DeepResearcher, and SimpleDeepSearcher on Musique in particular.

• Benefit tracks capability rather than parameter size. Across five backbones from 7B to proprietary scale with prompting only, the strongest models gain most (QwQ-32B +5.8, Claude Sonnet 4 +4.8); Qwen3-8B is neutral; Qwen2.5-7B-Instruct loses 3.9 points because it cannot reliably emit JSON for four tools instead of two — an obstacle training removes, with SFT lifting that model to 34.4% against its 17.5% ReAct baseline.

• Failed trajectories are useful training data. At matched volume (2,030 conversations), correct-only data beats a random subsample (39.3% vs. 36.5%), but the full 3,682-conversation pool beats correct-only by a further 2.3 points despite the extra 1,652 conversations being predominantly failures. The explanation: with observation tokens masked, behavior cloning supervises the process, so a trajectory that plans and checkpoints well before a late error still demonstrates the target skills.

• Mathematical reasoning data transfers. Removing the 241 math trajectories costs 1.7 points (41.6% → 39.9%), concentrated on Level 2 (43.5% → 40.4%), even though the math data involves no browsing, tool calls, or memory management.

Methodology in Plain English

The setup is a standard ReAct agent with two environment tools (Web Search, Web Browse). CORAL adds two more tools the agent can call at its own discretion. The first, Memory Management, is a scratchpad the agent fills in and edits: a plan with sub-goals marked pending or done, plus facts it has verified. The second, Context Optimization, is a hard reset — the agent throws away the conversation history and continues in a fresh episode that starts from only the system prompt, the question, and whatever is in the scratchpad.

Training happens in two stages. First, behavior cloning: the authors used Claude Sonnet 4 with the toolset to generate 2,932 web-agent trajectories, of which 1,642 (56%) reached the correct answer and 1,290 (44%) did not; they kept both, added 241 math trajectories from NuminaMath and OpenMathReasoning, for 3,173 raw trajectories that expand into 3,682 training conversations after splitting at each CO reset. Because each episode begins from a different initial context, episodes are trained as separate examples rather than as one flat conversation — but they all come from a single uninterrupted run of one model on one question, never stitched from different runs. Loss is computed only on the agent's own thoughts and tool calls, masking out observation tokens.

Second, reinforcement learning: Multi-episode ARPO takes the SFT checkpoint and samples groups of trajectories per question. Only the final episode contains the answer, so the reward (a format check times an LLM-judge answer check) is computed there and copied back to every earlier episode of that trajectory; all episodes across all sampled trajectories for the question are then pooled and normalized together for the advantage estimate. This means an episode that produced no answer of its own still receives credit or blame for the final outcome — the mechanism that makes "when to checkpoint" learnable rather than a fixed rule.

Evaluation samples 8 trajectories per question at temperature 0.7 and top-p 0.8, reporting Pass@1 (Avg@8); Pass@8, Maj@8, and standard deviations are in the appendix.

Why This Matters

Research impact. The paper reframes context management as something the agent learns to do rather than something a wrapper does to the agent. Its two concrete levers — making the reset decision a supervised/RL-learned action, and propagating trajectory-level reward across episode boundaries so that action can be optimized — are directly reusable by anyone training multi-turn agents. The finding that accuracy and token efficiency improve together (87.4% fewer input tokens alongside more than double the accuracy of the ReAct baseline) undercuts the usual assumption that better attention management is a cost paid for accuracy.

Real-world applications:

• Deep-research and web-investigation assistants that must run dozens of search hops without losing the thread.
• Enterprise document and multi-source retrieval agents where earlier findings must survive long tool-call chains.
• Cost-sensitive agent deployments, where per-trajectory input tokens and LLM calls dominate serving cost (136K vs. 1,078K tokens; 15.2 vs. 45.3 calls).
• Any long-running autonomous workflow where agents must resume from a summarized state — customer-support escalation chains, multi-step data analysis, or coding agents working across many files.

Industry relevance. The paper's most immediately deployable result needs no training at all: bolting the two tools onto a strong proprietary model yielded an 8.0-point gain on BrowseComp and a GAIA average (65.0) near OpenAI Deep Research (67.4). Combined with the 7–8B result at 42.8%, the method offers a route to competitive long-horizon performance at small model sizes and low token budgets.

Future Directions

• Closing the Level 1 gap. The working-memory overhead costs accuracy on short questions where clutter never became binding (CK-Pro-8B's 56.4% vs. 46.2% on Level 1). The paper does not propose a mechanism for making memory maintenance conditional on task difficulty; learning when not to use the tools is an open problem.

• Recovering from resets. The paper notes that a reset also discards the history that lets an agent repair an earlier mistake (Appendix D). Since the trained agent resets only 0.21 times per trajectory, whether a repair mechanism — or a partially reversible reset — would help on tasks with early wrong commitments is left open.

• Answer selection. The gap between Pass@8 and Avg@8 (e.g., 71.8% vs. 41.6% for SFT on GAIA) indicates substantial unused capability; better verification or selection over sampled trajectories is an obvious lever the paper flags but does not pursue.

• Weak base models. Qwen2.5-7B-Instruct loses 3.9 points under prompting-only CORAL because it cannot reliably emit JSON for four tools; SFT raises it to 34.4%, but how far down the capability scale the toolset can be pushed is not established. The paper also points to Appendix E for a discussion of limitations, which is not included in the provided text.

Target Audience

Researchers and engineers working on LLM agents, agentic reinforcement learning, or context engineering — particularly those building long-horizon web-research or tool-use systems and those concerned with inference cost. Readers need comfort with ReAct-style agent loops and policy-gradient terminology to follow Sections 2–3, but the experimental results (Tables 1–4) are legible to a broader applied-ML audience. Practitioners who only want the deployment takeaway can read Section 4.2–4.4 and skip the training details.

Authors’ abstract

Agents powered by large language models (LLMs) have demonstrated remarkable progress in solving complex reasoning tasks. However, LLM agents often falter on long-horizon tasks due to cognitive overload, as their working memory becomes cluttered with expanding and irrelevant information, which dilutes their attention and hinders effective planning and reasoning. To mitigate this challenge, we introduce Cognitive Resource Self-Allocation (CORAL), a novel reasoning paradigm that empowers agents to proactively optimize their context. Implemented as an agent-callable working memory management toolset, CORAL allows an agent to checkpoint its task progress and verified facts, and adaptively initiate a new problem-solving episode by resetting cluttered context while preserving only the checkpointed information, enabling it to resume reasoning from a clean state. We further enhance the agent's checkpointing capabilities using a Multi-episode Agentic Reinforced Policy Optimization algorithm. On the GAIA benchmark, CORAL with an 8B model achieves 42.8% average accuracy, with particularly strong gains on the most challenging long-horizon tasks (Level 2 and Level 3). Crucially, these gains come at a substantially lower inference cost: CORAL consumes 6-7$\times$ fewer input tokens than a ReAct agent. It also generalizes to multi-hop Knowledge QA benchmarks and to the considerably harder BrowseComp, where it improves a strong proprietary agent by 8.0 points.