Don't Lose the Thread: Empowering Long-Horizon LLM Agents with Cognitive Resource Self-Allocation
Zheng Zhu, Xiaozhuan Liang, Xi Chen, Bei Yu, Jiaya Jia
AI 速览:提出CORAL推理范式,让LLM智能体通过可调用的工作记忆管理工具主动检查点化任务进度并重置杂乱上下文,以缓解长程任务中的认知过载。
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.
Shuao Jia, Zichao Ling, Chen Bai, Bei Yu, Kang Zhao, Jianwang Zhai
Efficiently deploying Deep Neural Networks (DNNs) - particularly Large Language Models (LLMs) - on tensor accelerators is essential for achieving high throughput and energy efficiency at the edge. However, the combinatorial explosion of the joint design space spanning intra-layer mapping and inter-layer fusion renders existing approaches inadequate. In this work, we present FADiff, a gradient-based optimization framework that automatically discovers optimal mapping and fusion strategies to accelerate DNN inference. We first formulate a unified, differentiable analytical cost model that accurately estimates the energy and latency of individual layer mappings as well as diverse fusion configurations. We then combine CP-SAT-based pre-enumeration of feasible discrete candidates with augmented-Lagrangian penalties over continuous relaxations, enabling a constrained gradient-based search that efficiently navigates the vast design space to identify optimal joint strategies. Experimental results show that FADiff consistently outperforms prior methods, delivering superior optimization quality with significantly faster convergence in both energy and latency.
Zehua Pei, Hui-Ling Zhen, Lancheng Zou The Chinese University of Hong Kong, Xianzhi Yu Huawei Technologies Ltd., Wulong Liu, Sinno Jialin Pan, Mingxuan Yuan, Bei Yu Department of Computer Science and Engineering, The Chinese University of Hong Kong
Scaling large language models (LLMs) improves performance but significantly increases inference costs, with feed-forward networks (FFNs) consuming the majority of computational resources. While Mixture-of-Experts (MoE) architectures can reduce this cost through sparse activation, restructuring existing dense models into MoEs typically requires extensive retraining on hundreds of billions of tokens.We propose an analytical post-training framework that rapidly restructures FFNs into sparse MoE architectures using only a small calibration dataset. The method analyzes neuron activation patterns to partition neurons into always-active shared experts and conditionally activated routed experts, then constructs a router analytically from representative neuron statistics, enabling immediate deployment or optional lightweight fine-tuning. This approach applies both to dense models and recursively to existing MoE models for hierarchical sparsity.Experiments demonstrate up to 1.17\times speedup in compute-bound scenarios with only minutes of processing and 2k-sample fine-tuning, outperforming methods requiring orders of magnitude more resources.
尚无 AI 解读,展示原文摘要
cs.CL
Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) · 2026 · 正式发表
Sitong Wu The Chinese University of Hong Kong, Haoru Tan, Xichen Zhang Hong Kong University of Science and Technology, Bin Xia, Wenhu Zhang, Xiaojuan Qi, Bei Yu Department of Computer Science and Engineering, The Chinese University of Hong Kong, Jiaya Jia Department of Computer Science and Engineering, Hong Kong University of Science and Technology
Reinforcement Learning (RL) with sparse outcome rewards suffers from inefficient credit assignment in complex LLM reasoning tasks. While utilizing stronger LLMs as teachers to derive dense token-level supervision offers a cost-effective alternative to proprietary reward models, it relies on the flawed assumption that teachers are perfect oracles. In reality, teacher models exhibit capability limitations and uncertainty, producing noisy signals that make student policies susceptible to reward hacking. To address this, we propose Teacher Reward Adaptive Calibration (TRAC), a robust framework that filters noisy supervision by dynamically modulating teacher influence via a multi-granularity calibration mechanism. TRAC evaluates teacher reliability across three principled dimensions: problem-level expertise, trajectory-level discrimination, and token-level confidence. Furthermore, we integrate TRAC with Group Relative Policy Optimization (GRPO), formulating as TRAC-GRPO, which treats calibrated teacher-derived reward as an additive advantage reshaping term to ensure fair advantage estimation. Extensive experiments demonstrate that TRAC effectively mitigates teacher noise, significantly enhancing the reasoning capabilities and training stability of LLMs compared to standard baselines. The code will be available at: https://github.com/JIA-Lab-research/TRAC.
尚无 AI 解读,展示原文摘要
cs.CL
Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) · 2026 · 正式发表