AgentTTS: Large Language Model Agent for Test-time Compute-optimal Scaling Strategy in Complex Tasks

2025-07-26

AI Summary

This paper addresses the problem of optimal test-time computational resource scaling for multi-stage complex tasks. We propose AgentTTS, a framework that jointly optimizes model selection and budget allocation across subtasks. AgentTTS employs a large language model (LLM) as an intelligent agent that performs environment-feedback-driven iterative search, integrating reinforcement learning–inspired interaction mechanisms with empirically grounded heuristics to achieve dynamic and interpretable resource scheduling. Compared to conventional static allocation and existing LLM-based baselines, AgentTTS achieves significant improvements in search efficiency, robustness to variations in training data scale, and decision interpretability. To our knowledge, it is the first work to formulate multi-stage computational scaling as an end-to-end learnable combinatorial optimization problem. Extensive experiments on diverse complex reasoning tasks demonstrate its effectiveness and strong generalization capability.

Technology Category

• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Learning to Search
• Multiagent Systems: Agent Communication

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Agentic search
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

Abstract

Test-time scaling (TTS) enhances the performance of large language models (LLMs) by allocating additional compute resources during inference. However, existing research primarily investigates TTS in single-stage tasks; while many real-world problems are multi-stage complex tasks, composed of a sequence of heterogeneous subtasks with each subtask requires LLM of specific capability. Therefore, we study a novel problem: the test-time compute-optimal scaling in multi-stage complex tasks, aiming to select suitable models and allocate budgets per subtask to maximize overall performance. TTS in multi-stage tasks introduces two fundamental challenges: (i) The combinatorial search space of model and budget allocations, combined with the high cost of inference, makes brute-force search impractical. (ii) The optimal model and budget allocations across subtasks are interdependent, increasing the complexity of the compute-optimal search. To address this gap, we conduct extensive pilot experiments on four tasks across six datasets, deriving three empirical insights characterizing the behavior of LLMs in multi-stage complex tasks. Informed by these insights, we propose AgentTTS, an LLM-agent-based framework that autonomously searches for compute-optimal allocations through iterative feedback-driven interactions with the execution environment. Experimental results demonstrate that AgentTTS significantly outperforms traditional and other LLM-based baselines in search efficiency, and shows improved robustness to varying training set sizes and enhanced interpretability.

Problem

• Optimize compute resources for multi-stage LLM tasks
• Address interdependent model and budget allocation challenges
• Autonomously search compute-optimal allocations via feedback-driven interactions

Innovation

• Autonomous LLM-agent for compute-optimal allocation
• Iterative feedback-driven search strategy
• Multi-stage complex task optimization