Cognitive resource management helps AI agents handle long tasks better
Don't Lose the Thread: Empowering Long-Horizon LLM Agents with Cognitive Resource Self-Allocation
Computation and Language
Summary
Large language model agents can get confused when working on long tasks because their memory gets crowded with too much information. The authors introduced a technique called CORAL that helps these agents organize what they remember by saving important facts and clearing out cluttered info. This lets the agents start fresh when needed but still remember key details. Their method improved performance on tough long-step tasks while using much less computing effort.
What this means in practice
- •For ai product developers: Enhance AI assistants’ ability to solve complex multi-step problems by managing memory more efficiently with checkpointing.$Commercial implications: Enables building advanced AI software that solves longer, more complex tasks accurately while reducing compute costs.
- •For knowledge management teams: Improve automated question answering systems to handle multi-hop queries by maintaining verified facts and resetting cluttered context.
Authors
Zheng Zhu, Xiaozhuan Liang, Xi Chen, Bei Yu, Jiaya Jia
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.