AgentTTS: Large Language Model Agent for Test-time Compute-optimal Scaling Strategy in Complex Tasks
Authors: Fali Wang
Hui Liu
Zhenwei DAI
Jingying Zeng
Zhiwei Zhang
Zongyu Wu
Chen Luo
Zhen Li
Xianfeng Tang
Qi He
Suhang Wang
Venue: NeurIPS 2025 poster ()
Keywords: language models
scaling
test-time compute
compute-optimal inference
Primary area: deep_learning
TL;DR: We introduce AgentTTS, an LLM-agent framework that efficiently allocates test-time compute across multi-stage tasks, achieving compute-optimal performance through insight-guided adaptive search.

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.

AgentTTS: Large Language Model Agent for Test-time Compute-optimal Scaling Strategy in Complex Tasks
Fali Wang';
Hui Liu?, Zhenwei Dai Jingying Zeng', Zhiwei Zhang', Zongyu Wu', Chen Luo", Zhen Li Xianfeng Tang'. Qi He Suhang Wang'
"The Pennsylvania State University, University Park, PA, USA
2 AAmazon, Palo Alto, CA, USA
{fqw5095, zbz5349, zzw5373, szw494) &psu..ed edu (liunhu, iingngii,,nninn. zwdaii,,,,,,,,, zwdai, cheluo, xianf Camazon. com
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 sequence of heterogeneous subtasks with each subtask requires I.I.M of specific capability. Therefore, study novel problem: the test-time compute-optimal scaling in multi-stage complex tasks, aiming to select suitable models and allocate budgets per subtask maximize overall performance. TTS in multi-stage tasks introduces two fundamental challenges: (i) The combinatorial search space model and budget allocations, combined with the high cost of inference, makes brute-fort 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 L.LMs in multi-stage complex tasks. Informed by these insights, we propose AgentTTS, LLM-agent-based framework that autonomously searches for compute-optimal allocations through iterative feedback-driven interactions with the execution en- vironment. Experimental results demonstrate that AgentTTS significantly outper- forms traditional and other LLM based baselines in search efficiency, and shows improved robustness varying training set sizes and enhanced interpretability.
Introduction
Test-time scaling (TTS), which refers to allocating additional computational resources during in- ference, has shown promising results in improving the performance of large language models I9 54 26). For example, Brown et al. (2] scaled inference compute via repeated sampling of candidate solutions, using verifier to select the best prediction, enabling weak small model to outperform single-sample state-of-the-an models. Despite its effectiveness,, existing methods function, such as mathematical problem solving 245 55
(LLMs) [2] 45 55 B3
primarily designed for single-stage tasks, where the model performs a single
26 [4 or code generation 34). However, many real-world applications involve multi-stage complex tasks, for which compute alloca- tion remains underexplored. These tasks involve a sequential execution of heterogeneous subtasks accomplish complex objectives. Representative examples include retrieval-the -heattl--e-eeeeee wee-tttheeneeeeeee generation QA systems [6), waterfall-style software development (comprising requirement analysis, system design,
33
[9
Work done during an internship at Amazon.
"Code link: httt:::///thhbbbbomm//tttta/////
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

coding, and testing) [29], and multi-agent task automation (such as task decomposition, tool selection, and parameter prediction) [32 Each subtask within multi-stage workflow often requires a model
with specific capabilities 13 For example, in seerieval-heen-eeeraaiin QA task, retrieval benefits
from large models with strong long-context understanding, while generation can achieve competitive performance using smaller models with repeated sampling (see Fig. ab)).. Therefore, multi-stage tasks demand models with diverse types and levels of capabilities, challenging existing TTS methods.
To bridge this gap, we formulate and study novel problem: test-time compute-optimal scaling in multi-stage complex tasks: for a complex task composed of sequence of subtasks and each subtask has a set of candidate models to choose from, given a total compute budget, the objective to select the appropriate model and allocate the compute budget for each subtask to maximize overall task performance. This problem presents two major challenges. (i) The search space is large due to the combinatorial choices of models and budget allocations across subtasks. For instance, in software development with three subtasks and two model options each (3B and 70B), the number of configurations can reach up to 10 (see Appendix A.1 for a calculation example), with each configuration requiring time-consuming inference (often several hours), rendering brute-force search impractical. (ii) Subtasks are not independent. The compute allocation for one subtask affects the performance and budget requirements of others. As shown in Fig. lb-d), high-quality retrieval can significantly reduce the compute required by the downstream generation subtask to achieve peak performance. In contrast, poor retrieval quality necessitates increased computation for generation, either through additional sampling or the use of larger models, to compensate for degraded input quality. These challenges highlight the need for an efficient search strategy that can handle the large and interdependent search space.
To address these challenges, we first conduct preliminary experiments to characterize the behavior of LLMs under test-time scaling in multi-stage tasks. From experiments, we derive three generalizable insights across four task types: (I) Different subtasks exhibit distinct preferences between large and small models; (2) Increasing test-time compute initially improves performance, but beyond a certain point, additional compute yields diminishing no gains; (3) The compute allocated to earlier subtasks influences the scaling dynamics and compute needs of downstream subtasks. Guided by these insights, propose Agen novel LLM-agent based framework designed to efficiently search for compute-optimal budget allocations in multi-stage tasks. Recent works [20] [36] 49 (2] 23 have shown that L.LM-based agents effective in planning and searching for hyperparrmeter optimization. Given the capability of LLMS in understanding and following instructions, AgentTTS integrates observed test-time scaling insights into the LLM-agent search process. AgenTTTS consists of three key components: the Agent, Archive, and Environment. The Agent begins by generating an initial set of trials based on Insight 1, which guides model preferences for each subtask. These trials are executed by the Environment, which evaluates them on the actual task platform and returns performance feedback. The crciive stores a history of generated trials, guidelines, and feedback. In subsequent stages, the Agent produces new trials and guidelines informed by Insights 2 and 3. This iterative process continues until a predefined stopping criterion is met. By leveraging LLM-based search, AgenITTS offers two advantages: (i) interpretability. through explicit guideline generation that explains decision rationales; and (ii) robustness, navigating non-smooth search spaces commonly found in test-time scaling.
Our main contributions are: (i) We study a novel problem of test-time compute-optimal scaling for multi-stage complex tasks; (ii) We identify three key insights that uncover fundamental patterns of test-time scaling in multi-stage tasks and motivate the design of AgentTTS, an efficient LLM-agent- based framework for searching compute-optimal configurations in this problem setting. and (iii) Comprehensive evaluations on six datasets show that AgentTTS achieves strong search efficiency, transparent interpretability in generating new trials, and robustness to non-smooth search landscapes.
Related Work
Test-time Scaling and Compute-optimal Strategy. Test-time scaling (TTS) enhances LLM perfor-
mance by allocating additional compute during inference [2|4]] 39). Existing methods fall into two
categories: sequential scaling (24 6 33
26
48 which iteratively refines outputs but depends on
good initial responses, and purallel scaling 2 55
34
35
9
45, which generates multiple
outputs and aggregates them using reward- based selection (e.g., repeated sampling (2 555, Best-
of-N [35), or tree search (45).. Recent work reduces reliance on reward models by using LLMs as

---

30 3 Parallel scaling is preferred for complex tasks due to better scalability and broader solution coverage [33]. Hence, we adopt repeated sampling with fusion. Research on
fusers 13
17
lest-time compute-optimal scaling shows that small models with optimal strategies can outperform
40]. Approaches include difficulty-aware model selection [R), reward-guided voting [55. and budget-aware prompting (54). However, they focus single-stage tasks, while we extend this to multi-stage tasks, where budgets must be adaptively allocated across interdependent subtasks to maximize overall performance. A more detailed discussion is given in Appendix A.10..
larger ones 2
45
19.
54
33
LLMs for Hyperparameter Optimization. LLMs have become powerful tools for hyperparameter
optimization (HPO). surpassing traditional AutoML techniques such as Bayesian Optimization
(BO) [S B1 by leveraging contextual reasoning and prior knowledge [8] LLM-based HPO
research generally follows two directions: (1) reducing the search space and (2) directly generating
hyperparameter configurations. For the former, LLMS have been used to prune large search spaces in NAS and HPO [22
25 23 21, as seen in GPT-NAS 52 AutoM3 2 and Llambo [21]. For the
22 56 18 57 treat LLMs autonomous optimizers. Systems like AutoMMLab [49, GENIUS (62), ML MLCopilot [56., and AgentHPO 20) refine trials via feedback and experience. We extend this line of work to compute-optimal test-time scaling in multi-stage tasks. A more detailed introduction of related work is given in Appendix 10.2
latter, recent works 4
62 58
3 Preliminary Knowledge and Problem Definition
Problem Definition. We define multi-stage complex task T [Ti, II..... as comprising n simpler subtasks. Each subtask T; has a set of candidate models M; E M,,, where each M; is tailored for subtask T,. Given a fixed total computational budget B for the entire complex task T, each subtask must be assigned a portion B; such that Σ B; = B. For a subtask T; there exists a trade-off between using a larger model with fewer inference samples and a smaller model with more samples, constrained by the assigned budget B; Our research problem is defined as
Definition 1 (Test-time compute-optingt budget allocation in multi-stage tasks). Given fixed total computational budget B for a multi-stage complex task T, how can we optimally allocate the computational budget among sublasks B + {Bi, B2,...... Ba), select appropriate models M,, and effectively distribute the allocated resources to maximize overall performance?
Test-time Scaling Mode: Repeated Sampling with Fusion. We adopt the Repeated Sampling with Fusion strategy for test-time scaling, it does not rely on additional reward models or verifiers compared to Best-of-N 335
9 or tree search algorithms [5] 53 and it offers greater scalability compared to sequential scaling [53 2. Given a problem p. a language model M with parameters test-time scaling is performed by increasing the number of repeated samples k. To stimulate diverse generations, we set the temperature hyperparameter to 0.9 throughout. fusion prompt is then used to aggregate the multiple candidate solutions generated through repeated sampling:
0 fruse (S M). S = {ss S i k}, S ((s P)
(1)
where S denotes the set of sampled responses, each 8; is independently drawn from the model M conditioned on the input prompt p and model parameters 0. The fusion function fiuee integrates the sampled responses using a fusion prompt, as described in Appendix A.7 Notably, the same LLM is used for both generating solutions and performing fusion.
4 Insights of Test-time Scaling on Multi-stage Complex Tasks
Allocating computational budgets for TTS in multi-stage complex tasks has significant challenges: (i) the search space expands exponentially as the number of subtasks increases, making exhaustive search impractical; (ii) subtasks typically require models tailored to their specific characteristics, rendering uniform model assignment inefficient; (iii) the scaling strategies employed earlier subtasks directly impact subsequent stages. To address the challenges, we first conduct pilot experiments to understand fundamental patterns in I.MM behavior under TTS in complex tasks, which pave the way to design a compute-optimal scaling strategy specifically tailored for multi-stage scenarios.
3

---

4.1 Experimental Setting
In this subsection, we first briefly introduce the datasets, models, and metries used across the four multi-stage complex tasks. Then, we describe unified budget conversion approach across models and tasks, followed by an example using the inference FLOPs as the primary compute cost metric. Detailed descriptions of the datasets, models, and metrics are provided in Appendix A.4
Tasks, Datasets, and Models We conduct preliminary experiments four multi-stage tasks: (i) betrieval-based Question Answering using 2WikiMuliHopQA [1]] and HotpotQA [50] datasets, (ii) Knowledge Graph Question Answering using CWQ (36) and WebQSP [51]. (iii) Task Automation, using TaskBench [32], and (iv) Automated Software Development, using ChatDev [29]. Subtasks differ in prompt and generation lengths; retrieval uses longer prompts, while QA needs longer generations. We consider models ranging from 3B to 72B. Details of task specifications, datasets, models, and evaluation metrics are provided in Appendix A.4 and A.S
Unified Budget Conversion Across Mudels and Tasks Different models and tasks demand varying levels of computational resources. Larger models require more inference FL.OPs, and complex tasks typically involve longer input or output tokens, leading to increased compute costs. These disparities complicate fair comparisons of computational overhead across models and tasks. To address this, we propose budget normalization framework that equates the cost of fewer inference samples from larger models with that of more samples from smaller models, while explicitly accounting for task-specific compute variations.
Formally, let the total computational cost of sampling S times from model M on task T be denoted as f S, T), and the corresponding normalized compute budget as = foudger (M, T).. The cost metric may reflect user-defined preferences, such inference FLOPs, wall-clock time, monetary cost. To establish a common unit of budget, we define it as the cost of a single inference pass by the smallest model (e.g., LLaMA 3B) the lowest computationally consuming task Tiowea:
Studgal M 1, 7 Twwea) 1.
Given model My, sample count Se, and task Tr, we compute the equivalent budget B by equating its total cost to that of the smallest model executing Samallest samples on the cowest-consuming task:
B fuudger (M, S,, T?) Samalles such that Sana (Mr Sz, fcos (aoaaMMann (Mamallest Smallesss Tlowes) (3)
This formulation expresses the inference budget of any model-task pair as the number of equivalent samples that the smallest model would generate on the least computationally intensive task.
Compute Cost Metric and Budget Definition. Different stakeholders prioritize different cost metrics. Consumers are typically concerned with monetary expenses, such as the price per million input and output tokens. LLM researchers focus on computational cost, such as inference FLOPs, due constraints of limited GPU memory. Discussion of using API price as cost metric is in Appendix 6 In this work, we adopt inference FLOPs as the primary cost metric throughout. The corresponding budget conversion is stated in the theory below. The proof is Appendix(A2.
Theorem 1 (Normalized Budget Function). Given configuration (Me, St, Tt) where Mj is the model size, Se the number of samples, and Tę (Np.t, NN.e) the task specification where Np.e and Na. is the average prompt and generation lengths, the equivalent normalized budget under the base configuration (Momalesst 3B, p.lowea 128, N N.Jowsstt 64) is:
20.32
B Saudger (Me, St, Te)
+ 2(a,32
1) -1)
(4)
31
where
Mr. Manntee
31
No.e Na.e
32
Ng.e
Np.inen
Under this conversion, the budget B intuitively represents: how many passes using 3B--prrmmeerr model on the base configuration T Thoess would incur the same cost as given sampling S₁ times with model Mj on task T The base configuration corresponds to the minimal compute setting among our subtasks (see Table the Appendix) and aligns with standard defaults used in short-form QA tasks, such QA in WebQSP.
4.2 Preliminary Experimental Results and Insights
We use the Floes- budget conversion in Eq. to analyze performance variance with increasing sample count and compute budget, and to guide budget allocation in our method in the next section.

---

0.80
0.75
0.70
0.70
EM
0.60
0.75
0.73
0.00
0.70
0.40
0.6
0.00
0.55
20
40
60
20
40
60
80
20
40
80
80
20
40
60
80
Samples
Samples
Semples
Samples
0.80
0.70
0.04
0.63
0.52
0.50
0.90
0.72
LJama3-3B
Owen22 6-7B
Owen22: 5-320
0.70
Liama3-88 8B
0.40
Oownn2 5-728
Llama3-708
0.60
086
0.98
0.52
600
1003
1600
2003
2500
60
100
100
200
250
50
100
150
200
250
0
60
100
150
200
260
Budget
Budget
Budget
Budget
(8)
(b) Ret
F1
= 0.8
(c) F1 = 0.6
(d) Ret 0.35
Figure 1: Performance variance on 2WikiMultiHopQA with increasing sampling and inference
FLOPs. Top: performance by sample count; bottom: performance by log-scaled inference FLOPPs (a) Retrieval accuracy measured by Retrieval Fl (Ret_Fl) FI). (b-d) QA performance under varying retrieval quality levels measured by Gen Gen_EM, the exact match between the generated answer and ground truth
Fig. showw how subtask performance in retrieval-based question answering varies with increasing test-time sampling and computational budget (FLOPs). In Fig. [ (a), performance on the retrieval subtask exhibits a strong positive correlation with model size, while smaller models provide only marginal improvements even with increased budget. This suggests that retrieval benefits from larger models, where fewer high-capacity samples yield better results. In contrast, Fig. (b) shows that for the question answering subtask, under a limited budget (e.g.. < 10 FLOPs), LLaMA-3 3B and LLaMA-3 8B outperform LLaMA-3 70B, indicating a preference for smaller models. Similar trends are observed in Fig. (a) and across other benchmarks (see Fig. 11 in Appendix (A6). These discrepancies arise from subtask-specifi demands: retrieval emphasizes long-context understanding. favoring large models; while question answering primarily involves extracting information from retrieved content, where smaller models excel and can better exploit test-time scaling via repeated sampling. Thus, different subtasks exhibit distinct preferences between language models and sampling frequency, depending on the specific capabilities required by each subtask (Insight 1).
Fig. T (b-d) reveal a non-monotonic performance trend across models in the question answering subtask, which is sensitive to test- time compute. As the number of sampling repetitions increases, performance initially improves but often fluctuates or declines after reaching an optimal budget. For example, LLaMA-3 3B reaches peak performance 10 samples (FLOPs budget x 2 x 1023) in Fig. (b), indicating that additional compute does not always yield better results in this subtask. Similar patterns are observed Fig. I0 (b-d) and other benchmarks (see Appendix A.66). This is because, as the number of candidates grows, fusion becomes more complex and may become performance bottleneck. Smaller models, with limited capacity, struggle more under high sampling and tend to degrade, whereas larger models are more capable of handling extensive fusion and show greater tolerance. In conclusion, subtask-level test-time scaling typically exhibits an optimal budget, beyond which more budget often leads to limited or even negative returns (Insight 2).
Fig. (b-d) show how the performance of the question answering subtask varies under different levels of retrieval quality, with Fl scores approximately 0.80, 0.60, and 0.35, respectively. When high-quality retrieval is provided by larger model (Fig. (b)). budget for question answering is reached earlier. For instance, I.I. - 8B peaks at 10 samples (budget around 2 x 10 FLOPs): and (2) smaller models (3B and 8B) may outperform the larger model (70B) under the same compute budget. In contrast, when retrieval quality is low (Fig. D (d)), the peak performance is delayed. For example, with Ret_FI 0.6, LLaMA-3 8B peaks 20 samples (budget about 3 x 10 13 FLOPs), while with Ret F1 = 0.35, even 90 samples (budget near 10 FLOPs) still do not reach peak performance. Moreover, under low retrieval quality, smaller models are less likely to outperform larger ones. For instance, in Fig. (d), the peak performance of the 3B and models does not surpass even the lowest performance of the 70B model. These findings indicate that the scaling behavior of subtask is affected by the performance of its preceding subtasks. Poor retrieval increases downstream task difficulty, requiring the model to compensate for missing information, where larger models are more beneficial. Similar trends are observed in Fig. [0] (b-d) and
observe: (1) the optimal test-time
5

---

Algorithm Compute-Optimal Test-time Budget Allocation
Require: Agent A, Model set M, Subtasks T, Environment E, Total budget B.
1: Initialize experiment archive: L + 0.
2: Initialize candidate trials: C f A.initiaaliio
(Insight 1, Eq.
3: Obtain feedback: S < E.execute(C)
4: while stopping criterion not met do
5: Generate exploration guidelines: 9 Aggenerate.guieelineeeeee
(Insights 3)
6: Update experiment log: L < CU ((C,S,9))
7: Generate new candidate trials: C t A geemerate(G.M..7. M,,,,
(Eq.
8: Obtain feedback: S E.execute(C)
9:
end while
10: return Best-performing trial from L.
other benchmarks in Appendix A.6 We conclude that budget allocation for preceding subtask impacts the model selection and optimal budget of subsequent subtasks (Insight 3).
5 AgentTTS: Agent for Test-time Scaling Budget Allocation
The insights in Sec. 4 provide guidance on how to efficiently search for effective budget and model allocations in multi-stage tasks. As LLMs have shown strong capabilities in planning and reasoning. we propose AgentTTS (Agent for Test-Time compute-optimal Scaling), which integrates these insights into an L.-M-based agent to autonomously navigate the compute allocation space and capture subtask-specife model preferences and optimal budget, and inter-subtask dependencies. An illustration of AgentTTS is in Fig. Next, we elaborate on the design of AgentTTS..
Overview of AgentTTS. As shown in Fig. 2 (b-d), the framework consists of three core components: the Agent, Archive, and Environment. Initially, the Agent gener- ates a batch of candidate trials (each looks like (M1, B1, M2, .).. guided by Insight which reflects preliminary model preferences across subtasks. These trials are stored in the Archive and forwarded to the Environment for evaluation. The resulting performance feed- back is returned to the Agent and used to con- struct initial guidelines that suggest whether sub- sequent trials should prioritize smaller or larger models. Both feedback and guidelines are re- tained in the Archive for future reference. In subsequent iterations, the Agent generates new trials based the existing guidelines and adds new guidelines by Insight and Insight 3. The Archive continuously stores the evolving guidelines along with corresponding trials and performance feedback, while the Environment evaluates each batch of trials in practical execution environment and returns feedback. This loop repeats until predefined stopping criterion is met. The full procedure is outlined in Algo. I and the class diagram is shown in Fig. T6 (Appendix).
(+A)))+))
Allocation
Total Budget
Agent
Fensback
Archives
Enaluated by
LLM Pool
Environment
[a] Overall Framework AgentTTS
Candidate
candidats
Complex Task
Sub-Task
Generate
guidelines
Generote
Sub-lask (+]
condidates
Guideline Candidate
Feedbock
(b) Agent
(c) Enviromeent
(d) Archive
Figure 2: Overview LLM Agent for Test-time
Scaling Budget Allocation.
Agent component The Agent, implemented using LLM, is responsible for generating test-time budget allocation candidate trials and guidelines, as shown in Fig. (b). Although LLMs perform well in hyperparameter optimization for machine learning models, they lack knowledge of test-time scaling, a relatively new technique. To address this, we incorporate Insight into the initial search stage and Insights 2 and 3 into all stages.
Insight reveals that different subtasks prefer models with specific capabilities. Therefore, matching suitable model early steers the search toward effective configurations, reducing wasted effort on inferior options. To estimate each subtask's model preference, we perform the following during the initialization stage. Let Biaax Shotgea ( (Mi, max, 1, Ti) be the budget required for single sample from the largest available model Mimox for subtask T; and Bimiu foudga (M mins 1, T,).) the budget

---

from the smallest model. For each subtask T; we compare all candidate models that fit within B while fixing all other subtasks use their respective largest available model with one-pass inference. This ensures a fair comparison across model candidates for the target subtask Tí, Under the condition that Bmnax > B - Σj Binin, the largest model for subtask occupies too much budget to allow feasible allocations for the remaining subtasks. In this case, we progressively downsize the model until the largest available model Mam is found. This operation is repeated for each subtask.
Upon receiving initial feedback from the Envi cnmment module, the Agent summarizes the initial guidelines G (Algo. Line 5) using the prompt provided in Appendix A.7 based on performance comparisons. If the large model significantly outperforms the small model, the large one is preferred. Otherwise, the small model is preferred due to its greater flexibility in subsequent exploration. The resulting guideline specifies which model should be prioritized in the later stages of the search. The instruction to identify the preferred model is used only in the initial search stage. The overall procedure is implemented in Line of Alg. D Tand formalized in Eq. below
{( ((.., Mr, B.....) - 1 < N, M, M,, B, Bm initial stage C= c- [fcl Agenerato(G.,M,TT.))
(5)
subsequent stages
where each trial is represented as c = (M, Bi,..... Ma Br), ) M, E M, is the selected model, and B, is its assigned budget for subtask i. B denotes the total compute budget. M,,max is the largest available model for subtask i, and the configuration (..., M₁, _.... indicates that other subtasks are fixed to their respective available largest models with one-pass inference.
Insight 2 shows that increasing test-time compute initially improves performance, but beyond an optimal point, it may cause oscillation degradation within subtask. Each model thus has task-specific optimal budget: too little limits performance, while too much wastes compute and harms other subtasks. Identifying this optimal budget per subtask is key to maximizing overall performance efficiently. To support this, we put Insight 2 in the guideline generation prompt (Appendix (A.7) and ask LLM to follow Insight to: "identify the search direction for finding the optimal number of samples for each subtask." " This ensures that the agent focuses the next round of search within the right scope, enabling faster convergence optimal allocation.
Insight 3 indicates that budget allocation for a preceding subtask affects model choice and sampling needs in downstream subtasks. Under limited budgets, optimal configurations cannot be assigned to all subtasks. Allocating more resources to one may shift the optimal setup of others, making previous configurations suboptimal. This interdependence increases search complexity, as each subtask must be re-evaluated in light of others' changes. To address this, we embed an instruction into the prompt (Appendix A.7) that leverages the LLM's planning capabilities to explore allocation trade-offs across subtasks. This enables the Agent to generate search guidelines that adaptively identify critical subtasks and configurations.
The instructions derived from three insights are applied concurrently throughout the search process. Based on the updated guidelines, the Agent then generates the next round of candidate trials (Algo.
Line 7) for evaluation.
Environment & Archive Components The Environment module executes and evaluates trials in the actual runtime environment. Upon receiving trials from the Archive, it converts them into executable scripts and submits them to the task platform for execution on a small training set. After completion, performance feedhack is returned to the Agent, as shown in Fig. [c) and Lines 3 and of Algo. The Ar chi ve component stores generated guidelines and candidate trials in corresponding bases, as shown in Fig. d) and Lines 1 and 6 of Algo. [ records the search process throughout iterations and outputs the best best-performing trials upon termination.
6 Experiments
This section presents a comprehensive evaluation of AgentTTS for test-time compute-optimal budget allocation in multi-stage complex tasks. We also conduct ablation studies on integrated key insights, analyze the interpretability of budget allocations, evaluate robustness to different training sizes, and compare search effectiveness under various budget conversion methods.
Experimental Setup. We conduct experiments across six datasets covering four distinct task categories, as detailed in Appendix A.4 Baselines include two groups of hyperparameter optimization:
7

---

2WikiMultHHopOA
HolpotOA
0.85
0.9
0.45
0.80
0.80
0.8
0.75
0.00
0.7
0
0.05
0.00
0.6
0.50
0.55
0.55
0.50
0.5
0.50
10
20
30
40
10
20
40
10
20
40
WebQSP
Taskbench
Chatdev
0.00
Best
0.38
Consistency
0.78
BD
0.50
0.84
LLM-28
0.75
MLCopilot
0.45
AgentHPP
0.80
AgenITTS
0.40
0.73
0.70
Trial
Trial
Trial
Figure 3: Performance trajectories across various search methods over 50 trials. X-axis: trial count; Y-axis: best performance up to each trial. The best score is obtained from the optimal trial in prior grid search and serves as the benchmark for all methods.
traditional machine learning methods and recent LLM-based methods. The traditional baselines include Bayesian optimization (BO) [T5 31 and random search. For
approaches, consider AgentHPO [20]. MICCopoot (56]. and LLM _Z... Given that these I.LM-based methods were adapt them the test-time budget allocation. Both AgentHPO and MI.Copilot employ a two-phase pipeline: (1) feedhack-driven guideline generation and (2) guideline-informed candidate trial generation. They differ in the initialization of guidelines: AgentHPO generates initial guidelines autonomously via the LLM, while MLCopilot initializes guidelines based on performance trends observed in similar tasks. In contrast, LLM ZS directly prompts the LLM to generate candidate trials without relying on guidelines or external references. Further details of each baseline are Appendix A.9 All methods execute 50 iterations of search on training set comprising 50 samples and are subsequently evaluated on a test set containing 500 samples. We use GPT-o3-mini [27 the LLM search agent. The default budget is set as the sum of the minimum budget required for one-pass inference each subtask using the largest model.
initially proposed for hyperparameter tuning,
Main Results and Analysis. Fig- shows the search trajectories of method and the baselines for test-time compute-optimal budget allocation on the training set. The search time and test-set performance of the best trials are reported Tab. We observe: Our method outperforms baselines in both search efficiency and test-set performance in Fig. and Tab. First, ours achieves top results on most tasks and requires fewer trials, even when baselines eventually reach optimality. Second, our method reduces search time, confirming higher computational efficiency. These gains come from leveraging empirical insights to guide the search toward promising trials. Despite similar performance on the training set, our method often generalizes better, as shown in Tab. I For example, on 2WikiMultiHopAA, it exceeds the next best methods by 2% on the test This suggests that Insight 2 helps identify minimal budgets and avoid redundant sampling, improving generalization. (ii) Compared to L..-aagent-baee AgentHPO and MI.Copilot, our method converges faster, showing integrated insights improve search efficiency on compute-optimal allocation. (iii) Traditional methods such as BO often get stuck in local optima due to the non-smooth landscape, while random search is noise-tolerant but inefficient. In contrast, LLM-based methods use prior hyperparameter tuning knowledge to bypass suboptimal regions, achieving better search performance.
Table 1: Comparison of search time (in hours) and test-set performance across datasets. indicates
unavailable results or failure to find the optimal trial.
2Wiki
Hotpot
CWQ
WebQSP
Taskbench Time
ChaiDev
Method
Time
EM
Time
EM
Time
EM
Time
EM
p-Fl
Time
Cons.
Random
0.66
0.71
0.76
0.86
0.40
0.74
BO
0.60
0.71
0.76
0.85
0.52
0.75
LLM ZS
12.5
0.70
0.71
55.3
0.76
37.7
0.89
0.49
0.74
MLCopilot
12.5
0.70
46.3
0.72
48.4
0.78
20.8
0.88
0.53
0.75
AgentHPO
8.3
0.70
0.74
37.4
0.78
0.89
0.49
0.74
Ours
2.5
0.72
17.5
0.74
11.1
0.78
7.8
0.89
64.3
0.53
14.3
0.75
8

---

Training size 100
Training size 75
Training size 50
Ablation studies
0.85
0.65
0.8
0.80
0.80
0.9
0.75
0.75
0.7
0.7
0.70
0.70
BO
w/o Insight
0.85
0.66
LLM-ZS
w/o Insight
AgentHPO
0.5
wio Insight
0.80
0.60
0.5
AgentTTS
AgentTTS
0.55
0.58
0.4
20
20
Trial
Trial (b)
Trial
Trial
(a)
(c)
(d)
Figure 4: (a-c) Search trajectories under varying training sizes on 2WikiMultHHopQA. (d) Ablation study on 2WikiMuliHHoQAA with total budget of 900.
Ablation Studies. To evaluate the contribution of each insight in AgenITTS, we perform ablation studies by comparing it to variants with individual insights removed. AgentTTS-woo-nnighhtt replaces Insight I with random initialization. AgentTTS-woo-Inigh222 removes prompt components addressing diminishing returns beyond optimal budget (Insight 2) and inter-subtask dependencies (Insight 3). As shown in Fig. dd), we observe: (1) Removing Insight 1 prevents reaching optimal configurations, highlighting the role of initial model choice in guiding subsequent exploration. (2) Without Insight 2, search efficiency drops, delaying optimal trials to 29 steps, showing the need to identify per-subtask optimal budget. (3) Excluding Insight delays the optimal trial to step 38, confirming the importance of leveraging LLM planning to handle subtask dependencies.
Robustness Varying Training Sizes. A small training set might lead unstable performance, creating a non-smooth search landscape that hinders optimization. To assess the robustness of search methods under this condition, we vary the training set size (50, 75, 100 samples) and evaluate their effectiveness in finding optimal configurations. As shown in Fig. fa-c), the search efficiency of LLM- based baselines and Bayesian optimization declines with smaller training sets, whereas AgentTTS maintains strong efficiency. This indicates that our empirical insights integrated in AgentTTS help the I.LM agent effectively navigate non-smooth search spaces through valuable prior guidance.
Interpretability of Budget Allocation To illustrate the interpretability benefit from integrating prior insights into the LLM agent for test-time compute allocation, we present three studies on 2WikiMultHHoQQA in 5ig.5) The first row shows Insight I enables subtask-specific model use, favoring large models for retrieval and small ones for QA, to guide efficient trials. The second row shows Insight 2 pinpoints the optimal sampling range (5-50), narrowing the search space. The third row shows Insight 3 promotes balanced allocation, favoring one-sample niigh-capacity" retrieval and "lower-cost" QA configuration, enabling effective trade-offs for better overall performance.
Case 1: Guidelines from Insight
Model Count of Trials
Rerreval: The 72B medel delivers sniking improvenent (Roe_ [o 0.90) over the 70..9 33aund 378 toy modes This jumber explovaiiss
Retrieval Model
Count
QA Modid
Court
shosod Gian lays modil alspie sumple ccurt.
Quesion Asswesing Comparin
yielis around
sempoo,
0wen225 328
ULaM3 88
0
(2-501-traesel)
camples
Gwen2255 728
10
LLaMAB 706
coalie 38 aodld is profiratle for fuster explonation
Case 2: Guidelines from Insight
Sample Count of Trials
For Question Asswering
around saciples
QA Model
Ulam33338
-0 funther
50 yteld litll
cuts
sfould fieso pioposietty the God aasple shonshaid (o-)) iooon prytimmatoo plantese
This aforll
Sample
[s. 30, 20, 50, 9t]
Generated Triaks
Case Guidelines from Insight
Retrioval Model
Retrieval Budget
QA
QA Budget
Used Budget
Model
midel (Owen2 5-738) jax sempll highest scores.
Alloate mocz toward rctioval model quality aisect, struggr
Qwer2.5
appsans badget
crrples.
2.we2.55
B14
LLaMA3
86
500
quiliy Sohiinal proortiti nasple bolasce badger
720
88
Qwer2.5
882
LLUMAS
38
500
728
58
Figure 5: Detailed Cases of Interpreting AgentTTS Through Integrated Empirical Insights.
9

---

Performance under varying budget settings To assess the search efficiency under different budgets, we evaluate AgentTTSS under two bud-
Budget 500
Budget 2000
0.75
0.85
0.70
get settings on 2WikiMuliHHoQA: 500 (low budget: only one subtask can reach its optimum) and 2000 (high, full optima reachable but with larger search spaces). As shown in Fig. 6 we can conclude: (1) Our method finds optimal con- figurations under low budgets, confirming the benefit of integrated insights on trade-off sub- task interdependence: (2) Across low to high budget settings, our method consistently outperforms baselines in search efficiency, indicating that the insight-informed L.MM agent remains robust to the expanded search space.
0.65
Best
0.80
0.60
0.65
LLM
0.70
0.50
AgentHPO
0.65
0.45
AgentTTS
0.40
0.05
Trial
Trial
Figure 6: Comparative search results under low, medium, and high compute budget settings.
Performance under varying temperatures To assess the impact of temperature on test-time scaling, we conduct additional experiments the 2Wiki dataset. The retriever is fixed to Qwen2.5-72B (producing a single retrieval result), and we vary the number of samples and temperature values for the QA model LLaMA 3 3B. As shown in Tab. higher temperatures (e.g.. 0.9) improve performance with multiple samples, while lower temperatures (e.g., 0.1) are more effective for cases. This suggests that higher temperatures enhance output diversity and fusion in sampling- based reasoning, whereas lower temperatures promote more stable, deterministic outputs. These results align with prior findings (28), which show that cemperature-induce diversity improves sample efficiency in reasoning tasks.
(Price)
API price as the cost metric Beyond inference-compuie FLOPs, end-users are often more concerned with the mon- etary cost. Therefore, we introduce an alternative budget
07
metric: API price. Using this metric, redraw the test- time scaling curves for question answering subtask of 2WikiMultiHopQA, as shown in Fig. 7 (left). The results indicate that smaller models remain preferable un- der constrained budgets. Furthermore, Fig. (right), the proposed insights continue to improve the search efficiency and performance of AgentTTS, demon- strating strong generalization across different cost metrics.
0.70
LiameS-3B
LLM 23
BO
Llama3-88
0.5
AgantPO
Llama3-70B
AgentTT
Price ($)
Trial
illustrated in
Figure 7:
Test-time scaling and
AgentTTS search trajectories under price budget. Left: scaling curves; Right: corresponding search trajectories.
7 Conclusion
We study novel problem of test-time compute-optimal scaling in multi-stage com- plex tasks, where the search space grows ex- ponentially with the number model choices and subtasks, and budget allocation across subtasks is interdependent. Empirical anal- ysis across four task types and six datasets
Table 2: EM scores of LLaMA 3B different tem- perature values on the 2Wiki dataset. The retriever is fixed Qwen2.5-72B (single retrieval result).
Samples
Temp 0.1
Temp = 0.5
Temp 0.9
0.70
0.64
0.66
10
0.70
0.78
0.78
yields three key insights: (1) subtasks have
20
0.72
0.76
0.78
distinct preferences for small or large mod-
30
0.74
0.78
0.72
40
0.74
0.76
0.78
els: (2) test-time scaling saturates, with dimin- ishing or negative returns beyond an optimal budget; and (3) early subtask budgets influ- ence downstream scaling behavior. Informed by these insights, we propose AgentTTS, an LLM-agent- framework that iteratively searches for compute-optimal configurations through interaction with actual task platforms. Ex- periments show that AgentTTS outperforms both traditional and L.LM-based baselines in search efficiency, final test performance, and robustness to non-smooth landscapes.
50
0.74
0.72
0.78
60
0.74
0.78
0.80
70
0.76
0.72
0.74
80
0.76
0.76
0.78
90
0.74
0.72
0.76
10

---

Acknowledgment
This material based on work supported by, or in part by, the Army Research Office (ARO) under grant number W911NF-21-10019 and Cisco Faculty Research Award.
References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08772 2023.
[2] Bradley Brown, Jordan Juravsky. Ryan Ehrlich, Ronald Clark, Quoc Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint 2rXiv:2407.21177, 2024.
[3] Juntai Cao, Xiang Zhang, Raymond Li, Chuyuan Li, Shafiq Joty, and Giuseppe Carenini. Multi2: Multi-agent test-time scalable framework for multi-document processing. arXiv preprint 2rXiv:2502.20592.. 2025.
[4] Angelica Chen, David Dohan, and David So. Evoprompting: Language models for code-level neural architecture search. In seventh Conference on Neural Information Processing Systems, 2023. URL http:::t///pperrrrioo.ooo
[5] Guoxin Chen, Minpeng Liao, Chengxi Li, and Kai Fan. Alphamath almost zero: Process supervision without process. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL UtRpitttpp://ooerrvviiee http:///ppppppoeeeeeeeeeoo/aaa/a net./forumnid-Vaaaxx@0
[6] Zhibin Gou, Zhihong Shao, Yeyun Gong, yelong shen, Yujiu Yang, Nan Duan, and Weizhu Chen. CRITIC: Large language models can self-correct with tool-interactive critiquing. In The Twelfth International Conference on Learning Representations, 2024. URL httpp://
openrevies.net/Eoooee http://ww...e.oorrriio/o/oo3333
[7] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, al. The Ilama 3 herd of models. arXiv preprint 2rXiv:2407.21788 2024.
[8] Yang Gu, Hengyu You, Jian Cao, Muran Yu, Haoran Fan, and Shiyou Qian. Large language models for constructing and optimizing machine learning workflows: A survey, arXiv preprint arXiv:2411.1047 2024.
[9] Lin Gui, Cristina Garbacea, and Victor Veitch. BoNBon aligmment for large language models and the sweetness of best-of-n sampling. In The Thirty-eigtth Annual Conference on Neural Informa- tion Processing Systems, 2024. URL /ttp:stss////peeeerrr/e www.:eee/frruui haSKMro
[10] Shibo Hao, Yi Gu, Haodi Ma. Joshua Hong. Zhen Wang, Daisy Wang, and Zhiting Hu. Reason- ing with language model is planning with world model. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages
2023.
(11) Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing multi-hop qa dataset for comprehensive evaluation of reasoning steps. In Proceedings of the 28th International Conference on Computational Linguistics, pages 6625, 2020.
[12] Noah Hollmann, Samuel Müller, and Frank Hutter. Large language models for automated data science: Introducing caafe for context-aware automated feature engineering. Advances in Neural Information Processing Systems, 66:44753-44775 2023.
[13] Dongfu Jiang, Xiang Ren, and Bill Yuchen Lin. Lim-blender: Ensembling large language models with pairwise ranking and generative fusion. In Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pupers), pages 14165-14178,
2023.

---

[14] Can Jin, Hongwu Peng, Qixin Zhang, Yujin Tang. Dimitris N Metaxas, and Tong Che. Two heads are better than one: Test-time scaling of multi-agent collaborative reasoning, arXiv preprini arXiv:250400977,, 2025.
[15] Donald R Jones, Matthias Schonlau, and William Welch. Efficient global optimization of expensive black-box functions. Journal Global optimization, 13:455-492, 1998.
[16] Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et Serrivaal-uugment generation for snowledge-intensive nlp tasks. Advances in neural information processing systems, 13:9459-9474, 2020.
[17] Zichong Li, Xinyu Feng, Yuheng Cai, Zixuan Zhang, Tianyi Liu, Chen Liang, Weizhu Chen, Haoyu Wang, and Tuo Zhao. Llms can generate better answer by aggregating their own responses. arXiv preprini 2:xXi::5503.04104 2025.
[18] Minhua Lin, Hui Liu,, Xianfeng Tang, Jingying Zeng, Zhenwei Dai, Chen Luo, Zheng Li, Xiang Zhang, Qi He, and Suhang Wang. How far are Ilms from real search? comprehensive study on efficiency, completeness, and inherent capabilities. arXiv preprint arXiv:2502. 18387, 2025.
[19] Runze Liu, Junqi Gao, Jian Zhao, Kaiyan Zhang, Xiu Li, Biqing Qi, Wanli Ouyang, and Bowen Zhou. Can lb 1lm surpass 405b Ilm? rethinking compute-optimal test-time scaling. arXiv preprint arXiv:2502.000 06703, 2025.
[20] Siyi Liu, Chen Gao, and Yong Li. Large language model agent for hyper-parameter optimization. arXiv preprint 24X.022402.07881, 2024.
[21] Tennison Liu, Nicolás Astorga, Nabeel Seedat, and Mihaela van der Schaar. Large language models to enhance bayesian optimization. In The Twelfth International Conference on Learning Representations, 2024. URL httpps://openreviii ew www./eeeuuueeeee 1d=00xott000ooo
[22] Xiyu Liu, Zhengxiao Liu, Naibin Gu, Zheng Lin, Wanli Ma, Ji Xiang, and Weiping Wang. Relation also knows: Rethinking the recall and editing of factual associations in auto-regressive transformer language models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 24659 24667, 2025.
[23] Daqin Luo. Chengjian Feng, Yuxuan Nong, and Yiqing Shen. Autom31: An automated multimodal machine learning framework with large language models. In Proceedings of the 32nd ACM International Conference on Multimedia, pages 8586 8594, 2024.
[24] Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao. Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, al. Self-refine: Iterative refinement with self-feedback. Advances in Neural Information Processing Systems, 2023.
[25] Clint Morris, Michael Jurado, and Jason Zutty. Llm guided evolution-the automation of models advancing models. In Proceedings of the Genetic and Evolutionary Computation Conference, pages 377-384, 2024.
[26] Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling. arXiv preprint arXiv:2501. 2025.
[27] OpenAI. Gpt-o3-mini.
http:://AAaaafooooopppocoocoooop 2025. [Online;
accessed January 31, 2025).
[28) Nearchos Potamitis and Akhil Arora. Are retrials all you need? enhancing large language model reasoning without verbalized feedback. arXiv preprint arXiv:2504. /295/, 2025.
[29] Chen Qian, Wei Liu, Hongzhang Liu, Nuo Chen, Yufan Dang, Jiahao Li, Cheng Yang, Weize Chen, Yusheng Su, Xin Cong, et al. Chatdev: Communicative agents for software development. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 15174-15186 18, 2024.
12

---

[30] Jon Saad-Falcon, Adrian Gamarra Lafuente, Shlok Natarajan, Nahum Maru, Hristo Todorov, Etash Guha, E Kelly Buchanan, Mayee Chen, Neel Guha, Christopher Ré, al. Archon: An architecture search framework for inference-time techniques. arXiv preprint arXiv:2409. 15254, 2024.
[31] Bobak Shahriari, Kevin Swersky, Ziyu Wang, Ryan Adams, and Nando De Freitas, Taking the human out of the loop: A review of bayesian optimization. Proceedings of the IEEE, 104 (1):148-175, 2015.
[32] Yongliang Shen, Kaitao Song, Xu Tan, Wenqi Zhang, Kan Ren, Siyu Yuan, Weiming Lu, Dongsheng Li, and Yueting Zhuang. Taskbench: Benchmarking large language models for task automation. Advances in Neural Information Processing Systems, 37:4540-4455 2024.
[33] Charlie Victor Snell, Jachoon Lee, Kelvin Xu, and Aviral Kumar. Scaling LLM test- time compute optimally can be more effective than scaling parameters for reasoning. In The Thirteenth International Conference on Learning Representations, 2025. URL https //openreview/ //operreriee..nttfforuutd 4FWAwzzdd2
[34] Benedikt Stroebl, Sayash Kapoor, and Arvind Narayanan. Inference scaling f laws: The limits of Ilm resampling with imperfect verifiers. arXiv preprint arXiv:2411.1750/ 2024.
[35] Hanshi Sun, Momin Haider, Ruiqi Zhang, Huitao Yang, Jiahao Qiu, Ming Yin, Mengdi Wang, Peter Bartlett, and Andrea Zanette. Fast best-of-n decoding via speculative rejection. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems, 2024. URL cprUs
[36] Alon Talmor and Jonathan Berant. The web as knowledge-basd for answering complex questions. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long Papers), pages 641-651, 2018
[37] Ziyu Wan, Xidong Feng, Muning Wen, Stephen Marcus McAleer, Ying Wen, Weinan Zhang, and Jun Wang- Alphazero-like tree-search can guide large language model decoding and training. In Forty-first International Conference on Machine Learning, 2024. URL https //openreview. //openreviee.net/forouiee http://orormrii--CCClCeEeee
[38] Fali Wang, Runxue Bao, Suhang Wang, Wenchao Yu, Yanchi Liu, Wei Cheng, and Haifeng Chen. Infuserki: Enhancing large language models with knowledge graphs via infuser-guided knowledge integration. In Findings of the Association for Computational Linguistics: EMNLP
2024, pages 3675-3688, 2024.
[39] Fali Wang. Jihai Chen, Shuhua Yang, Ali Al-Lawati, Linli Tang, Hui Liu, and Suhang Wang. A survey on collaborating small and large language models for performance, cost-effecctieenes: cloud-edge privacy, and trustworthiness. arXiv preprint 1tfh::2510.3880 2025.
[40] Fali Wang, Minhua Lin, Yao Ma, Hui Liu, Qi He, Xianfeng Tang, Jiliang Tang, Jian Pei, and Suhang Wang. A survey on small language models in the era of large language models: Architecture, capabilities, and trustworthiness. In Proceedings of the 31 ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2, KDD '25, page 6173-6183, New York, NY, USA, 2025. Association for Computing Machinery. ISBN 97984007 14542. doi: 10.1145/371 189653771899633366666 URL pattps:///o// org/ 10.
[41] Fali Wang, Zhiwei Zhang, Xianren Zhang, Zongyu Wu, TzuHao Mo, Qiuhao Lu, Wanjing Wang, Rui Li, Junjie Xu, Xianfeng Tang, Qi He. Yao Ma, Ming Huang, and Suhang Wang. comprehensive survey of small language models in the era of large language models: Techniques, enhancements, applications. collaboration with 11ms, and trustworthiness. ACM Trans. Intell. Syst. Technol., September 2025. ISSN 2157-6904. doi: 10.1145/3768165. URL. https Just Accepted.
//doi.org/ 10. 101111111677811
[42] Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoe V Le, Ed H. Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2023. URL. https://openreviex
1PL1NIIMMM
13

---

[43] Zhepeng Wang. Runxue Bao, Yawen Wu, Jackson Taylor, Cao Xiao, Feng Zheng, Weiwen Jiang, Shangqian Gao, and Yanfu Zhang. Unlocking memorization in large language models with dynamic soft prompting. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 9782-9796, 2024.
[44] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia. Ed Chi, Quoe V Le, Denny Zhou, et al. ofaiin-of-thouhh prompting clicits reasoning in large language models. Advances in neural information processing systems, 35:24824-2483) 2022.
[45] Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, and Yiming Yang. Inference scaling laws: An empirical analysis of compute-optimal inference for Ilm problem-solving. In The Thirteenth International Conference on Learning Representations. 2025.
[46] Yuxi Xie, Kenji Kawaguchi, Yiran Zhao, James Xu Zhao, Min- Kan, Junxian He, and Michael Xie. Self-evaluation guided beam search for reasoning. Advances in Neural Information Processing Systems, 36:41618 2023.
[47] An Yang. Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2.5 technical report. arXiv preprint arXiv:2412.15115 2024.
[48] Chenxu Yang, Qingyi Si, Yongjie Duan, Zheliang Zhu, Chenyu Zhu, Zheng Lin, Cao, and Weiping Wang. Dynamic early exit in reasoning models. arXiv preprint 1585 2025.
[49] Zekang Yang, Wang Zeng, Sheng Jin, Chen Qian, Ping Luo, and Wentao Liu. Autommlab: Automatically generating deployable models from language instructions for computer vision tasks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 22056-22064, 2025.
[50] Zhilin Yang, Peng Qi. Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdinov. and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 2369-2380, 2018.
[51] Wen-tau Yih, Ming- Wei Chang, Xiaodong He, and Jianfeng Gao. Semantic parsing via staged query graph generation: Question answering with knowledge base. In Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 1321-1331, 2015.
[52] Caiyang Yu, Xianggen Liu, Yifan Wang, Yun Liu, Wentao Feng, Xiong Deng, Chenwei Tang, and Jiancheng Lv. Gpt-nas: Evolutionary neural architecture search with the generative pre- trained model. arXiv preprint arXiv:2305.0535. 2023.
[53] Fei Yu, Anningzhe Gao, and Benyou Wang. Ovm, outcome-supervised value models for plan- ning in mathematical reasoning. In Findings of the Association for Computational Linguistics: NAACL 2024, pages 858-875, 2024.
[54] Zhenrui Yue, Honglei Zhuang, Aijun Bai, Kai Hui, Rolf Jagerman, Hansi Zeng. Zhen Qin, Dong Wang, Xuanhui Wang, and Michael Bendersky. Inference scaling for long-context retrieval augmented generation. In The Thirteenth International Conference on Learning Representations, 2025. /tttss
net forum?id=FS] Fdr Ir0m1vz
[55] Zhiyuan Zeng, Qinyuan Cheng. Zhangyue Yin, Yunhua Zhou, and Xipeng Qiu. Revisiting the test-time scaling of ollkike models: Do they truly possess test-time scaling capabilities? arXiv preprint arXiv:2502. 12215, 2025.
[56] Lei Zhang, Yuge Zhang, Kan Ren, Dongsheng Li, and Yuqing Yang. Micopilot: Unleashing the power of large language models in solving machine learning tasks. In Proceedings of the 18th Conference of the European Chapter of the Association for Computational Linguisties (Volume
1: Long Papers), pages 2931-2959, 2024.
14

---

[57] Shaokun Zhang, Jieyu Zhang, Jiale Liu, Linxin Song, Chi Wang, Ranjay Krishna, and Qingyun Wu. Offline training of language model agents with functions learnable weights, In Forty-first International Conference on Machine Learning, 2024.
[58] Shujian Zhang, Chengyue Gong, Lemeng Wu, Xingchao Liu, and Mingyuan Zhou. Automl-gpt: Automatic machine learning with gpt. arXiv preprint arXiv:2305.024: 2023.
[59]
Shun Zhang, Zhenfang Chen, Yikang Shen, Mingyu Ding, Joshua B. Tenenbaum, and Chuang Gan. Planning with large language models for code generation. In The Eleventh International Conference on Learning Representations, 2023. URL
wwwprrevviieeeeefouuuuu.. net/forum?
Yd=Lr3c00EEbfL
[60] Tong Zhang, Chunyu Lei, Zongyan Zhang, Xian-Bing Meng, and CL Philip Chen. As-nas: Adaptive scalable neural architecture search with reinforced evolutionary algorithm for deep learning. IEEE Transactions on Evolutionary Computation, 25(5):830 841, 2021.
[61] Zhiwei Zhang, Fali Wang, Xiaomin Li, Zongyu Wu, Xianfeng Tang. Hui Liu, Qi He, Wenpeng Yin, and Suhang Wang. Catastrophic failure of LLM unlearning via quantization. In The Thirteenth International Conference on Learning Representations, 2025. URL https://
openreview wwt.:e//oorrrm/dddddaaaaaaaaoaa
[62] Mingkai Zheng, Xiu Su, Shan You, Fei Wang, Chen Qian, Chang Xu, and Samuel Albanie. Can gpt-4 perform neural architecture search? arXiv preprint arXiv:2304. 10970, 2023.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the main claims in the abstract and introduction align with the paper's contribution to test-time compute-optimal budget allocation in multi-stage tasks.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitation in Appendix A.II
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
16

---

Justification: We report a deviation in inference FLOPs and detail the computation procedure in Appendix A.2
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Yes, we provide sufficient experimental information for reproduction, related sections include (Experimental Setup), 4.1 A.4 A.7 Our code is available at: https
//github.ccm com htt////aaaaallaennTTTS..
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
17

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: All experiments in our paper are conducted using open-source datasets. The code is available at: https:///ithun comm/Faarralllllgaatiiiii.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.cc/pblle Sote:////uudde//eeuuuuubuii
(CodeSubiisss wwtauubmmsssooni ssiiiisiiicciyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, provide the experimental settings in Sec. 4 and Appendix A.4
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification:
Due to the high computational overhead, we do not report variance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
18

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, we provide a detailed experimental setting Appendix A.4 and time of execution in See.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http://wwwrrr.cccoooooolccccccea
Answer: [Yes]
Justification: Yes, have read and comply with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the societal impacts in Appendix A.12
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
19

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: No such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We use open-source datasets and benchmarks and have cited all sources.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
20

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[NA]
Justification:
No new datasets, models, or other assets are introduced in this work.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing was involved.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects were involved.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
21

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.MM policy (https: http:p/s/eeeeereeecc///ooorrrere for what should or should not be described.
22

---

A Technical Appendices and Supplementary Material
A.1 An Example of Search-space Size
We take automated software development as representative example, which consists of three subtasks: coding, static testing, and dynamic testing. The average prompt and generation token lengths for each subtask are summarized in Table 4 All subtasks share the same model space M (LLaMA-3 3B, LLaMA-3 70B}. Given total compute budget of B = E; B 1768, which corresponds to the cost of using the largest model (LLaMA-3 70B) for all subtasks, we provide the corresponding calculation code below. The output shows the count is about 1.8 x 10
sumpute__uudet-ee N_large, Epl, 2.3 Compute the nornalized compute budget required to generate aamples fron model M_large task with prompt length up1 and generation length Nd1.
8nsal1-3e,, Np2-128, Nd2-64)
The budget normalized with respect the unit
generating 1 sanple using the smallest model (M_small) on the lovest conputationally consuning task proapt.length-ip2, decodingg.enggth-kkk))
..
alpha M_large / N_emall Kodel size ratio
betal Mpi Nd1 a Proapt-to-generate
length ratio for the task beta2 Npi / Np2 Prompt length scaling factor relative to the reference task
beta3 Np2 Nd2
length for the reference task
# Compute normalized budget cest
return
((alphs
betal) (betal
def sand_eaaple...
prompt_led decode_len); 117 "IIrrrstillly increases S (sample count) until budget exceeded for model and task configuration.""
(nax, budget, model budd.
range 10000):
if compute (udget model_size, proupt_ decode_lea) nax_ bud_ett return
return 9999
Compute the total budget required for sanple from the largest model tasks. (udget, 70s9, 1024, 1024) bonpate_budgatt(: (1, 70c9, 1024, 512) cospute_budget 70e9, 1024, 256)
model_space - (3e9, 70e9] # Tuo model sizes: small (3B) and large (70B) sample count per
Compute
using
model
hind_semple_uppe (ppperbbounddb,, 3e8, 1024, 1024)
maz_s2 aane_senple_aper_boor
3e9, 1024, 512)
pan__sampe_upper__bunee
3e9, 1024, 256)
ipt_1ist [1024, 1024, 1024] a Prompt lengtha for each aubtask Adi_list [1024, 512, 256] Generation lengths
over configurations balid_config, count =
budget
for si in range (naz_si 1): nodell bl = compute. conpute_budget modell, Npl_list [o], Ndi_ [o])
range (nax_s2 1):
for nodel2 in model_apace.
b2 compute_ budget (=2, model2, Npi_list [i], Mdi_list (1])
if b1
break
range (max_e5
for madel3 in model_apace:
computa_budger (a3, model3, Npi_liat [2], Mdi_liat [2])
ir b2 b3 b: break salid_contig_coum +=
print ("Total valid configurations:", calid_coffig.com output should be: Total valid configuratione: 1854841
Listing I: Enumerating Valid (MI, S1, M2, S2, M3, S3) Triplets under Budget Constraint
A.2 Inference FLOPP-Based Budget Conversion
We adopt floating-pott point operations (FLOPs) as the primary computation cost metric for measuring inference cost. In repeated sampling scenarios, Transformer models can exploit batching to share
23

---

the cost of prompt encoding across multiple decoding passes. This allows for efficient reuse of the prompt representation while performing multiple, independent decoding operations.
Then, we formalize the computation of FLOPs and describe how to convert the number of decoding samples between large and small model under equivalent FLOPS budget. We use unified notation throughout in Tab.
Table 3: Notation used in FLOPs-based cost formulation.
Symbol
Description
M
Number of model parameters
L
Number of Transformer layers
D
Hidden dimension size
Np
Number of prompt tokens
Na
Number of generated (decoded) tokens per sample
S
Number of decoding samples
The calculation of budget conversion using inference FLOPss Eq. follows:
Proof. Per-token FLOPs. For a decoder-only Transformer generating tokens autoregressively. the per-token FLOPs primarily stem from components:
(i) Parameter matrix multiplications. Each non-embedding weight used exactly once per token in matrix multiplication followed by addition, resulting in approximately 2M FLOPs per token, where M is the total number of non-embedding parameters.
(ii) Attention operations. At decoding step t + 1, attention requires computing query-key dot products and applying the attention weights value vectors. Both operations involve multiply-adds contributing approximately 4 4T.DU FLOPs per token, where I, is the number of layers, D is the hidden size, and is the number of preceding tokens.
Thus, the approximate FLOPs per token is: FLOPStoken (M, D,t) 2 2M + 4LDI.
Per-phase FLOPs. The total inference cost for generation can be decomposed into two phases:
(I) Prompt Encoding:
N₂ Σ
FLOPS Sprompi (M, M,I,D,N, Np)
FLOP Stoken (M, T,, D2
(6)
t=1
2MN, + 2LDN +
(2) Decoding:
N₂
FLOPSdccode (M, Na) = FLOPs Sroken (M, L, D, +
2M + 2LNN, Na + 1)
(3) Total FLOPs: To generate decoding samples from a prompt of length Np. the total FLOPPs is: FL.OPS (S, M, J., D, Ny. Na) = FL.PPppomp (M, T., D, N₂) FL.OP (M, L, D, N,, Na) We denote this as feou FL.OPSuaaS, M, I., D, Np. Na).
(8)
Budget Conversion Across Models and Tasks. Consider converting FLOPs usage between two tasks using different models: large model (2) for Task and smaller model (s) for Task Let Me, Le, Dr and M,, L., D, be their respective parameters, and Np.1, Na,1, and N,,2, Na.2 the average prompt and decoding lengths for the two tasks.
Suppose the large model executes St decoding samples. The equivalent number of samples S, that the small model can generate under the same FLOPs budget is:
S. FLOPS (St, Me, Le, De, Np.1, Na,1) FLOP FLOPSSpampp... Ma, La, Day Np.2)
FLOPSdecdde. (M,, L,, D., Np.2, N42)
24

---

To simplify, we ignore incention-relatee terms (typically less than 1% of total FLOPs) and define: 31 N,11
,3
No11
,333
Nu.2
Mr Substituting these yields:
Na.1
N2.2
Na.2
M,
a.32 (31 + + B1
S. = 83
(10)
Normalized Budget Function. We define the unit budget B the FLOPs cost of generating single sample using the smallest model (M, = 3B) on the lowest consuming task specification, char- acterized by N₂ p,lowes 128 and Na Na,wowest = 64. This setting represents the minimal computational cost among all available subtask-model combinations.
Substituting into Eq. 10 yields:
20r32Se
S.
+ 2(0:32-
(11)
31
Thus, for a configuration defined by model size Me, sample count Sy. and task parameters Tç = (x,,))Nae) the normalized budget function is:
20.32tt
Sbudgee ( M,, Se. Tt) = S.
+ 2(a -
(12)
B1
We provide the average lengths of prompt and decoding in different subtasks in Tab. 4
We provide
the look-up tables when generating = 1, 5, 10, 45, 90 in Tab. 5
Table 4: Average prompt and generation token lengths (N,, Na) for each subtask.
Subtask
N,
N₂
2WikiMulnHHopAAARRooo Retrieval
2048
128
2WikiMultiHopQA-O0
256
2048
128
HotpoiQA-QA
256
64
CWQ-Retrieval
1024
CWQ-QA
256
SaboSP-Retrieval
1024
WebQSP-QA
128
64
Taskbench Decomposition
Tuskbench-Deccompsittee
1024
64
Taskhench- Selection
1024
256
Baakkhench-Paramett Prediction
1024
ChatDev-Cods
1024
1024
ChatDev-Staiie Test
1024
$12
ChatDev-Dynamic Test
1024
256
Table 5: Normalized compute budget when generating S = sample.
Model
Wwirgg
Ber.k
Herrgh
CWOQ-
CWO-q3
WGRF-411
FananFred
Cooe
State
Bysanic
Qn2215-720
Qvex25-18
10
38
38
45
LLaMA-L1--1
701
115
305
115
66
398
445
1518
355
460
LLaMA-31148
Gomma-22218 2.188
304
43
151
26
151
178
214
176
Gomea-22.0
100
100
13
49
5
49
142
10
55
Gotera-3-8
80
38
5
38
110
54
45
20
20
28
23
Table 6: Normalized compute budget when generating S - samples.
Modi
SWiLLR
2WHi-D4
Hos.
HoL-QA
CWQg-
CWQ-DA
WQSFF-JA
Decsti
Vnbbil
Parandhee
Cooe
Static
Dyaassic
1006
110
Unwe2.5-11
14
47
$1
11.M0--1 1-70
V15
228
1776
7105
198
208
158
161
4000
8386
4105
22200
1305
K03
11.MA--1 1411
Genmra 3-11
30
187
51
1877
102
Ganera3-40t
134
136
23
53
20
63
106
530
255
196
116
12
Ph-Jomal
34
51
25
25
221
130
50
44
25

---

Table Normalized compute budget for each subtask when generating S 10 samples.
Medell
1W11278
1ea-E
lire0h
CWO--
CWO-0s
WOSTSE
Wosr-0x
TasBal
Statie
Dysamir
Unen255-12 726
1.40
134
1335
355
162
134
6.2
214
102
1342
5563
4222
1312
1342
Qwee2.5-22 5-328
243
553
149
275
141
355
54
275
1583
1575
1002
355
ELaMA-3
325
32
157
11LMA-122005
10
12
30
13
34
12
34
8
34
54
134
04
54
G-mera-2-211
134
203
232
Genrra3331 30
33
241
261
53
119
119
30
119
230
1356
488
250
Table 8: Normalized compute budget for each subtask when generating S 45 samples.
Medid
2Wat-Rt
2Wa6--A
Bat-8
18--14
CWQ..
CWQQ-A
WOSE-R
WONF-BA
Dxcemp
ParasPrro
Code
Stadie
Drasssic
11:41
Queed 5-70
1490
455
1745
455
22 2MAA 1-700
2615
1341
18255
1141
1131
1343
1433
-157
1433
0571
13973
17171
3771
4511
135
3501
Genmaa33 2310
13300
6523
3302
1753
Gamma--00 all
166
145
254
345
181
105
181
50
161
585
1999
2355
1126
506
399
485
143
30
153
60
75
50
75
25
75
24h
1942
V00
474
246
Table 9: Normalized compute budget for each subtask when generating S = 90 samples.
CWO--
WONF- (WO-Ux
WON--0X
Tre251
ParamPid
Bysanar
Qwea25-77
400
217
-03
213
245
217
245
345
875
07255
3885
1115
675
LL-MA-3.--11
4945
1101
4545
2191
2471
2191
2471
877
2471
6771
07511
33971
171.1
$171
LLMM--3..8
2504
304
1996
944
1986
944
032
544
952
388
952
3582
26063
13502
6622
3353
Gorera-2-8
696
280
316
350
316
112
300
1100
5600
4300
2306
1136
3:38
Pli
403
317
453
203
245
317
345
345
875
0755
3385
1715
875
Phid-anii
307
117
267
187
132
132
474
3000
1862
930
474
A.3 Dollar Costs of Repeated Sampling
Table 10 is the API cost information for each model from Together AI. We do not convert it in the same manner above, the dollar serves natural unit of price.
Model Name
Parameters
Inference Cost (per 1M tokens)
Qwen2.5 72B
72B
$1.20
Qwen2.5 32B
32B
$0.80
Qwen2.5 5 7B
7B
$0.30
LLaMA-3.1-708
70B
$0.88
1 8B
8B
$0.18
I.LMA-3.2
3B
$0.06
Table 10: Inference costs per IM tokens for selected models from Together AI.
A.4 Benchmarks, Datasets, Models, Metries, and other Experimental Setup
We introduce the benchmarks, datasets, models, evaluation metrics, and other experimental setups.
Retrieval-based Question Answering uses datasets: 2WikiMultHHoPAA [II] and Hot- potQA [50], each providing 100 candidate text chunks per query. The task involves two subtasks: (I) retrieving relevant documents using LLM retriever, and (2) answering the question using another LLM with the retrieved documents as context. For both datasets, the average prompt and generation lengths are 2048 and 128 for retrieval tasks, and 256 and 64 for the QA subtasks, respectively. We employ Qwen 2.5 (7B, 32B, 72B) 47 for retrieval and LLaMA-3 (3B, 8B, 72B) 7 for answering. These model families are selected due to their strong subtask performance and wide range of model sizes, which helps mitigate the influence of pre-training differences. This setup is used consistently across the other three tasks. The evaluation metrics are retrieval F1 (Ret-FI) and exact match (EM), corresponding to the two subtasks, respectively. For the Ret-Fl metric, given the ground-truth
26

---

retrieval documents, precision is computed as the fraction of correctly retrieved documents among all predicted documents, while recall is the fraction of ground-truth documents that are successfully retrieved. The Fl score is then calculated as the harmonic mean of precision and recall. For the EM metric, an answer is considered correct only if it exactly matches the ground truth. The EM score computed as the fraction of such exactly matched answers among all predictions.
Knowledge Graph Question Answering This task uses two KGQA datasets: Complex WebQues- tions [36] and WebQSP [5]]. The objective is to retrieve relevant knowledge triplets from 100 candidates in the knowledge graph and answer the question based on the retrieved context. The task includes two subtasks: knowledge retrieval and question answering. For both datasets, the average prompt and generation lengths are 2048 and 64 for retrieval tasks, and 256 and 64 for the QA subtasks, respectively. We use Qwen 2.5 (7B, 72B) 47 for retrieval and LLaMA-3 (8B, 70B) [7] for answering. The evaluation metrics are retrieval FI (Ret-F1) and exact match (EM), corresponding to the two subtasks, respectively. The computation of Ret-Fl and EM is the same as above.
This task involves decomposing complex user instructions into subtasks and invoking external tools to execute them. We use the TaskBench benchmark [32] to evaluate LLM performance across three stages: task decomposition, tool selection, and parameter prediction. The average prompt and generation lengths are 2048 and 64 for task decomposition, 1024 and 256 for tool selection, and 1024 and 2048 for parameter prediction. LLaMA-3 (8B, 70B) ( is used for all subtasks. To evaluate LLMs' ability to understand and decompose complex tasks, assess the quality of task decomposition using three complementary ROUGE metrics, including ROUGE-L. To evaluate tool selection, we introduce Node Fí (n-F1), which measures the correctness of predicted tools by comparing them against reference set, reflecting the model's ability to identify appropriate tools for each subtask. To assess the accuracy of tool configuration, use Parameter Name & Value Fl (p-F1), which evaluates both the identification of parameter names and the correctness of their assigned values, capturing the model's ability to provide contextually appropriate configurations.
Task Automation
Automated Software Development We leverage ChatDev (29). a chat-driven framework based on the waterfall model, encompassing documentation, design, coding, static testing, and dynamic testing phases. Our focus is on three core subtasks: coding, static testing, and dynamic testing. The average prompt and generation lengths 1024 and 1024 for coding, 1024 and 512 for static testing, and 1024 and 256 for dynamic testing. We use LLaMA-3 (3B, 70B) (7] for all subtasks. Consistency measures how well the generated software code aligns with the original requirement description. is quantified as the cosine similarity between the semantic embeddings of the textual requirements and the generated code. A higher score indicates a greater degree of consistency with the specified requirements.
Other Experimental Setup To encourage diverse generations and activate LLMs' coverage during repeated sampling, we set the temperature to 0.9, while keeping all other decoding hyperparameters at their default values. Each trial is conducted using an NVIDIA H100 80GB HBM3 GPU to ensure consistent runtime performance across all datasets. For each dataset, we randomly sample 50 instances to form the training set used for search and 500 instances as the test set for final evaluation. All datasets used in our experiments are made publicly available through an anonymous link to ensure transparency and reproducibility.
A.5 Model Selection for Subtask-Specific Analysis
For each subtask, we select models from the same model family and assume that model size reflects model capacity. By comparing the performance of small and large models within the same family under test-time scaling, we can figure out whether smaller models can outperform larger ones under equivalent compute budgets. This comparison focuses solely on the relationship between test-time compute and model size. Otherwise, comparisons across different model families could introduce confounding effects due to variations in pretraining, such as differences in pre-training data or objectives. Therefore, we restrict each subtask to models within single family. During evaluation, all other non-target subtasks are fixed to use L- 70B with one-pass inference to isolate the effect of model choice in the target subtask. For each subtask, we select the model family that achieves the best performance. Fig. [ present subtask-level performance across model sizes and families. We exclude ChatDev from this analysis, as it does not provide metrics for intermediate steps.
27

---

Based on the results, we select the Qwen family for retrieval-related subtasks and the LLaMA-3 family for all others.
IWMMMuttilopOA Qceetion Anownde
HotpeROA Rotriwal
2WIcMeNiEHONN- Beirioa
Model Fussle
Modal Family
Getino
Gamma
HaMMa
LiaMAS
LlaMaa
Open
Openroder
Phil
5
CpenCoten-33
0pecidae-1..0
Openbdaadd
Cpeateke111
Quei
Owed 5323
Ower 1-708
Ower 1-30
Qw422222
0ve22-77
BM
Art Pr
(a) 2Wiki-Ret
(b) 2Wiki-QA
(c) Hotpot-Ret
Bitporo -Qustico Auowering
Qvestion
Moddi Fandy
Model Pamily
Modal Furily
LLANAS
Gemma
Gentaa
OpanCwicc
LaMMA
LuMMAS
Mai
Open
Creas
OpenCodox
4paccoded4
5
Mui
SpenCoteni..5
Owwe2.5
1ax633200
0perCoees1.00
ItD-ital
Psr onken
na
Oena
EM
(d) Hotpot- QA
(c) CWQ-Ret
(f) CWQ-QA
Decompoctiin
Mudil Fandy
Model Funely
Model Funily
Genmo
Genno
Genna
MAMMA
Grmmu-1
LAMA
Comex:00
ULMM3
Open
OpenCater
MaM
Mai
LL23A33333
Hus
122003-00
L1
Gpecuarn 58
Phd-
BOUXBLL
PPL
(g) DoskBench-Decompy
(h) CeckBench-ToolSee
(i) TaskBench-PaaamPPii
Figure& 8: Performance of various LLMs across retrieval, question answering, and task execution subtasks on three types of tasks.
A.6 More Preliminary Experimental Results
Fig. 12 L3
and 14 illustrate how subtask performance evolves in KGQA, TaskBench, and ChatDev as the test-time sampling, FLOPs, and budget increase.
28

---

0.75
0.70
Gen EM
0.00
0.50
0.5
0.50
0.80
0.30
Samples
Samples
0.80
0.70
0.64
Owen225-78
074
0.00
0.93
Oyyen 5-326
5-728
Gen
0372
Llama3-86
Llama3-3B
0.22
Liamm33708
708
0.54
0.30
10
10"
10
10
10"
10"
10"
10"
FLOPs
FLOPs
FLOPs
FLOPs
0.80
0.64
0.61
0.02
0.60
0.80
0.72
070
0.40
0.68
0.80
0.52
500
1000 1600
2000
2500
00
100
150
200
50
100
100 Budget
200
250
00
100
100
200
250
Budget
Blugge
Busget
(8)
(b) Ret
(c) Ret F1
0.6
(d) Ret 0.35
Figure 9: Performance variance on 2WikiMultiHopQA with increasing sampling, inference FLOPs, and budget. Top: performance by sample count; middle: performance by log-scaled inference
FLOPs; bottom: performance by normalized budget. (a) Retrieval accuracy measured by Retrieval FI (Ret_Ft FI). (b-d) QA performance under varying retrieval quality levels measured by Gen EM, the exact match between the generated answer and the ground truth.
0.72
0.70
0.68
0.05
070
0.00
0.8
0.64
0.62
0.50
0.60
0.4
0.45
20
40
60
20
40
60
50
20
20
40
Samples
Samples
Samples
Samples
0.00
0.7
Qwen2 5-7B
0.55
Llama3-38
0.6
Owen2 5-32B
Owen2
Gen
0.88
0.64
0.02
0.0
Llama3-B8
708
0.8
0.00
0.46
064
10
10
10
10
10
10" FLOPs
10"
FLOPs
FLOPs
FLOPs
03.9
0.70
0.8
0.22
0.68
0.70
0.08
0.88
0.64
500
1030
1500
50
100
150 Budget
200
250
50
100
150 Budgel
200
250
50
100
150
200
250
Budgel
Budget
(8)
(b) Ret
0.9
(c) Ret
0.7
(d) Ret_F1 0.4
Figure 10: Performance variance on HotpotQA with increasing sampling, inference FLOPs, and budget. Top: performance by sample count; middle: performance by log-scaled inference FLOPs: bottom: performance by normalized budget. (a) Retrieval subtask performance; (b-d) Question answering performance under different retrieval quality levels.
29

---

0.78
Gen EM
Ret F1
0.72
0.40
0.74
0.68
20
40
60
20
40
Sample lea
Samples
Samples
0.70
Qwen2.5-PP
0.78
0.74
0.00
Qwen2.5 5-728
0.75
0.74
Rel
0.90
Gen
0.73
0.70
0.40
0.72
Llama3-8B
0.71
0.68
Llama3-708
0.30
10"
10"
10"
FLOPs
FLOPs
FLOPs
0.74
0.72
0.74
Ret
0.70
0.60
100
200
300
400
500
50
100
150
200
50
100
150
200
Budget
Budget
Budget
(a)
Ret_F1
0.5
(c) Ret_FF F1
0.33
Figure 11: Performance variance on CWQ with increasing sampling, inference FLOPs, and budget. Top: performance by sample count; middle: performance by log-scaled inference FLOPs; bottom: performance by normalized budget. (a) Retrieval performance; (b-d) Question answering performance under different retrieval quality levels.
0.89
EM
Ret F1
0.80
0.80
0.40
0.80
0.35
0.82
20
40
0
20
40
Samplis lss
Samples
Samples
0.70
Qwen2.5 2-BB
Llama3-8B
0.80
Owen2.5- 5-B2B
Gen EM
0.88
0.86
Llama3-708
Ret
0.50
0.87
0.85
0.88
0.84
0.40
0.86
0.30
10
10
10
FLOPs
FLOPs
FLOPs
0.80
0.80
0.87
0.58
Rel
Gen
0.87
0.88
0.84
0.83
100
200
300
400
500
100
150
200
90
100
150
200
Budget
Budget
Budget
(a)
(b) Ret_
0.5
(c) Ret_F1 0.35
Figure 12: Performance variance on WebQSP with increasing sampling, inference FLOPs, and budget. Top: performance by sample count; middle: performance by log-scaled inference FLOPs; bottom: performance by normalized budget. (a) Retrieval performance; (b-d) Question answering performance under different retrieval quality levels.
30

---

Liama3-88
0.90
0.50
ROUGE
0.1
Liamr3377B
0.51
0.00
048
0.90
0.30
0.82
0.35
0.28
Samples
Samples
Samples
Samples
0.38
ROUGE
0.88
045
0.36
0.36
0.22
C.48
0:04
0.38
0.69
0.90
0.50
ROUGE
0.90
0.51
0.98
Q
048
0.30
0.82
0.35
0.28
9300
5000
Budget
Budget
Budget
Budget
(a)
(c) Node F1 0.8
(d) Node 0.8
Figure 13: Performance variance Taskbench-DaillyPPISSe with increasing sampling, inference FLOPs, and budget. Top: performance by sample count; middle: performance by log-scaled inference FLOPs; bottom: performance by normalized budget. (a) Task decomposition performance; (b) Tool
selection performance; (c-d) Parameter prediction under different tool-selection quality levels.
0.747
Consistency
0.745
Llama3-3B
0.742
0.742
Llamma3708
0.740
0.740
0.740
0.760
0.738
0.736
0.739
0.752
20
40
60
80
20
40
80
50
100
0
20
40
60
80
100
Samples
Samples
Samples
0.747
Consistency Consistency
0.745
0.761
0.742
0.740
0.740
0.740
0.738
0.736
0.739
0.732
0.734
0.0
0.5
1.0
1.5
9
FLOPs
1615
FLOPs
1e14
FLOPs
0.747
0.745
0.741
0.742
0.740
0.738
0.740
0.757
0.736
0.736
0.739
0.732
0.734
0.738
1000
2000
3000
4000
500
1000
1500
2000
250
500
750
1000
1250
Budgel
Budgel
Budget
(a) Coding Phase
(b) Static Test
(c) Dynamic Test
Performance variance on ChatDev with increasing sampling and inference FLOPs. Top: performance by sample count; middle: performance by log-scaled inference FLOPP; bottom: performance by normalized budget. (a-c) The consistency score during coding, static test, and dynamic test, respectively.
Figure 14:
31

---

Retrieval: The larger models shine. Quenn 5-328 (Bet_ F1 +0.5 (5.55) is much better than 78 variont (*0.34) and S-wen2.5728 pushes Ret_ Fl
("Quen2.5778" "LLaW-332-33"" 0.3997 6.54),
("Quen2.5-328 28, "1."""N633.2-38" 0.5533, 8.64),
puhog:
=q. despite mirimal sampling-so further expl should the large Qwen? 5-728.
("Quen2.5-722 1, "LLAMM-332333" 1, 0.7867, 6.86),
(*Qven2. "'uen2.5-78", 1, "LLLa..3.2-38"
66, 0.3472, 0.58),
Question Arewering: The benefit of upgrading is minimal. The smaller DA model (meta hattlllma/mmmmss 2-33) with moderate samples performs its larger counterparts order 0.02 0.06), we prefer sticking with the smaller model.
('Quen2.5: 7.7",, 1, "LLaM-3.1-88", 0.3561, 0.52),
"0 1, "LLoM-3.11708", 1, 0.3314, 0.60)
differences only
1. Optimall sample count:
each subtask, prograssheely increase samples until performancc
gains taper.
In recent QA eaperiments, performance improved from S to 15
("Owen2, 5-720",
"'Qeen2.--77 - 1, "LLaM4-3,2-38", 5, 0.7904, 0.76),
1, "LLaWW-3:2-38",
18, 0.7814, 0.74),
samples (Gen, EMt rising -0.78] but then flattine dropped beyond
("Quena, 5-720".
"LLaMM-322-3""
15, 0.7061, 0.78).
15 suggestingg 115 samples optimal.
("Qven2.5-27 5-728",
1, "LLaM4-3.2-39"
20, 8,7771, 9.78),
Bunbo
using sample already ylelds high Ret. F1,, implying additional samples offer limited benefit.
Retrizval,
"Quen2_5-728",
"0xen225-720"
1, "LLaW6-3.2-38"
1, ".
25, 0.775e, 0.74),
36, 0.7783, 8.74)
2. Busget allocation strategy:
Witth flxed, shared budget and interdependent subtasks, prioritize
on impact there, CA given main metric Gen_EM),
("Qwen2.5- 5-728"
1, "LLaVA-3.1-88", 18, 0.7854, 0.74),
Use a high-capacity, efficient retrieval model (with minimal sampline) to free budget for CA.
("Qven2. 5-728"
1, "LLaMA-3.1-88",
5, 0.7943, 9.76),
("Oven. 5-3288,
1, "LLaW6-3.2-30", 15, 0.5420, 0.66),
Experiment with different allocation mixes balance performance settings
5-128",
"LLaMM-3.2-38""
0.5549, 8.78),
assigned every
("Owen2. 5-328"
1, "LLaMa-3-1-780", 1, 0.5293, 8.60),
("Qven2.5- 18, "LLaMa-3..-300, 15, 0,3543, 0.69),
1. SSmmle counts: For GA, improvements are seen when increasing samples from to about 20, but beyond that the grins flatten dedinee for Retrieual using high-quality model (e-g. Sowen2.5-7288
("Owen2.5-22 5-7282
("Qwen2.5-785 5-78"
5, "LLaMA-3.3--8"
1,
"LLAMA-3.2-38".
30, 8.346, 8.56),
8.7867, 8.66),
sample appears optimal. Facus on finding the "inee" point where extra samples stop yielding benefts.
("Qven2.5- 5-328"
"LLaMA-3.2-30", 20, 0.5679, 8.08), 10,
('Qwen255-78"
1,
"Ll A.LaMM-3,.22 3-3-3.3-788", 1, 0.3314, 8.68)
2. Busgee allocation: Since the compute budget is shared and subtases interdepardeni prioritize subtisk that impacts Gan (AA) and experiment with reallocating budget from less sensitive subtasks (like fletiissll) to balance model cost and sample count.
Generate Guidelines
Generate New Trials
Best Trial
(a) AgenITTS 2Wiki (budget = 850).
("Qvem2.5-720" 1, "LLaMM-3.2- 1, 8,67, 0.06),
("Qwen2.5 5-328"
1.LaM--.11-006 1, 0.5293, 0.66),
+High-quality retrieval models le.s, Owen2.5 5-728 with 1 sample) vield exceilent Ret, F1 (*0.79) but consume much the budget.
("Qwen2. 70, "LLoMM-3.1-89",
90, 9.3367, 0.50),
("Qren2.5-22) 1, "L--MM-3.1-800,
90, 0.5531, 0.00),
model (e.g. Lama-3.1/708, with minimal sampling can Hlightly boost the main | EM to 0.68). Incrressing sample counts on both subtasks with lower capacity models Burio (e.g, 70-30 samples) leads to significant performance degradation.
using
("Quen2.5-77) 78, "LLaMM-3.1-70"" 1, 8,3660, 6.68)
("Qven2.5-720", "LLaMM-3.1-88", 1, 9.36,, 0.56),
sampling boost: increasing sample count (from to 10)- especially with meta-llama/lmma 3.2- 3.2-3 22-30-correlate lites with a jump in
("Qven2.5-328" 1, "LLaMM-3.2-33". 1, 0.5359, 0.66),
Gen FM (up to 74).
("Qwen2.5-78", 1, "LLaNA-3,1-708"
1, 0,3314, 0.08),
(Cawen2.5-72B)
Retrieval model sue matters: Using high-capacity retrieval model
single sample yie strong Ret_F1 without needing
10, 0,7814, 0.74), ("Qven2.S-728 1, "LLoAt--3.2-38",
("Q4en2.5-328" "LLaM6-3.1-88",
("Qwen2-5-70", "LLaMM-331--86"
1, 0.3185, 0.50),
1, 6.5553, 0.58),
Dunod
ectra samples..
3. Model trade-off: A smaller, cheaper QA model with increased sampling
("Qven2.5-320 1, "LLaM--311770",
19, 8.5719, 6.66),
outperforms larger model at sample, emphasiong the benefit of samples pure model
("Quen2.5-328 "LL.MM-32--38"
''Qwe225-778" 10, "1LaMM-3.1-88"
1, 0.5350, 0.66),
1, 6,3462, 0.56),
Budget effect gunaiuuations spending closer the budget tand to show higher Gen, EM, saggesting aggresio agressive budget use can be worthwhile if allooaate wisely.
"LLoMM-3.1-700", 1, 0.3175, 0.58) ("Qxen2.S-70", 10,
"'Qu'n2_57288', 1, "LLaMM-3,2-38", 5, 0.7984, 0.76),
The top Gen EM (0)76) was achieved using the high-capacity retrieval model (Cwen2.5 2.5 7231 with sample palred with a moderate QA model (1.mm-3.2:3B) using samples.
(*Qwenz. 20nnn2.57228 5-728", 1, "LLGM--.1-88",
18, 0.7854, 8.74),
"0uen2.5-328" 1, "LLaMk-3.1798""
10, 8.5719, 0.68),
("Qxen2.5-7e) 5-78", 39, "LL-M4-3,2-39",
50, 0.3280, 9,60),
increasing samples (e.g.. 30) for retrieval or QA often degraded
''Quen2.5.78" 78, ",1-5W-,,1888",
98, 0.3367, 8.58),
Gen suagesting that more samples don't always equate to better
('Qwen2. 5-328", 1,
"LLaM4-3,1-88",
56, 0,5531, 8.60),
Dunog
* Strong retrieval models (728 328/78) with minimal sampling outperform weaker models when lotter
performance.
"-8=" 5, "LLoM--3.17799", 25, 0,2941, 0.70),
5, "LLaMA-3.1-798" 16, 0.5566, 0.68),
samples.
(" 5-Quen2.5-78", 78, "LLaW4-3.2-38", 50, 0.3456, 8.58)
appears beneficial reletive to Uamo-33 3.1-48 when combined with the strong retrieval model.
Using Uama-3 3.333
Confi gurations consuming near max budget (820-850) with balanced parameter cho stelded the best overall
Generate Guidelines
Overloxding subtask
many samples lespecially
Generate New Trials
Inwer-capacihy models) reduces both Gen_ FM and Ret_ F1 performance.
Best Trial
(b) AgentHPO on 2Wiki (budget 850).
Figure 15: Comparison of complete trial generation and decision guidelines using TTs and AgentHPO on the 2Wiki dataset under compute budget of 850.
32

---

A.7 Prompt Design in AgentTTS for Search Guidelines and Trial Generation
We design prompts below guide the LLM agent in generating both guidelines and candidate trials. Each prompt is constructed to provide the agent with basic task information and requirements, enabling effective reasoning and planning. For guideline generation, initial and non-initial phases. The initial guideline prompt, informed by Insight 1, helps identify whether each subtask favors smaller larger models, thereby guiding early-stage search. The non-initial guideline prompt, informed by Insights 2 and 3, reflects general test-time scaling patterns and assists in balancing budget allocations across subtasks. For trial generation, the LLM agent proposes new configurations based on the existing guidelines.
distinguish between the
Fusion Prompt
You have been provided with a set of responses from various open-source models to the latest user query, which is query. Your task is synthesize these responses into single, high-quality response. It is crucial to critically evaluate the information provided in these responses, recognizing that some of it may be biased or incorrect. Your response should not simply replicate the given answers but should offer refined, accurate, and comprehensive reply to the instruction. Ensure your response well-structured, coherent, and adheres the highest standards of accuracy and reliability. Once again, the query is: query Responses from models:
LLM Prompt for Initial Guidelines
[System Prompt] You are expert in parameter optimization.
[User Prompt] am tuning test-time parameters for complex task that can be broken down into multiple simple subtasks. The parameters include the model size and the number of samples used for each subtask. Generally, increasing the number of samples improves performance. The goal to find the best parameter settings that maximize performance within a fixed budget.
I will provide the task name, description, budget, and initial search history. Please help me summarize the guidelines from the search history in the following aspect:
The search history includes initial trials that use different model sizes under the same unit budget. In these trials, smaller models fewer samples due to their higher cost. Please help me identify which model size we should explore further in each subtask. If the large model performs significantly better than the small model, we will choose the large one. Otherwise, if its performance similar or only slightly better, will prefer the smaller model.
allowed more samples, while larger models have
As the number of samples increases, performance tends to improve and then level off. Please identify the search direction for finding the optimal number of samples for each subtask, especially the point after which more samples no longer improve results.
3. The total compute budget is shared across all subtasks and they interdependent. If we could not assign the optimal model and sample count to every subtask, consider the following strategy. Leveraging your planning capability, prioritize each subtask, try different combinations of budget allocations for the other tasks, and balance the allocations across subtasks.
Inputs:
Task name: [eaas_naame
Task description: {task_desc)
Subtasks, models, prompt/generation lengths: [model_space}
(subtask.specification
Total budget: [budget}
Evaluation metrics: {metrics}
Main metric: imain_metric
Search history: (history)
Please generate your response in the most concise wording.
Response:
33

---

L.LM Prompt for Non-Initial Guidelines
[System Prompt] You are expert in parameter optimization.
[User Prompt] tuning test-time parameters for complex task that be broken down into multiple simple subtasks. The parameters include the model size and number of samples used for each subtask. Generally, increasing the number of samples improves performance. The goal is to find the best parameter settings that maximize performance within fixed budget.
I will provide the task name, description, budget, and both old and recent search history.
Please help me summarize guidelines from the recent search history in the following three aspects: 1. As the number of samples increases, performance tends to improve and then level off. Please identify the search direction for finding the optimal number of samples for cach subtask, especially the point after which more samples longer improve results.
The total compute budget shared across all subtasks, and they interdependent. If we could assign the optimal model sample count to every subtask, consider the following strategy. Leveraging your planning capability, prioritize each subtask, try different combinations of budget allocations for the other tasks, and balance the allocations across subtasks.
If you observe that the performance variance across the search history is small, please suggest that future searches focus morc exploration strategies such crossover, mutation, and random search, rather than strictly following the existing search patterns. Inputs:
Task name: (task_naadd name)
Task description: {task (task_desc} desc}
Subtasks, models, prompt/gennration lengths: fsubrask,specifccation (model_space)
Total budget: [budget}
Evaluation metries: {metrics)
Main metrie: (main_metric)
Search history: (history)
Please generate your response in the most concise wording.
Response:
LLM Prompt for Trial Generation
[System Prompt] You are
parameter optimization.
[User Prompt| am tuning test-time parameters for a complex task decomposable into multiple subtasks. Each subtask requires selecting model and the number of samples. Increusing the number of samples generally improves performance. The objective identify configurations that maximize overall performance under fixed compute hudget. provide
name, description, list sublasks, samples, the hudget calculation formula, evaluation metrics (including the main metric to optimize), search history (previous configurations tried), insights from previous experiments.
model choices available
budget, available
The budget for each subtask is computed with the following Python function:
def compute_ budgat (nun_samples_ model_aize, peonpt_length, generation_length,
N_ N_amall-3e9, Ep2-128, Nd2-64):
M_anail betal pengtt generation_ Up2
alpha model,
beta2 proupt_lenge beta3 Up2 Nd2
budgat
((alpha
beta2
hatal) (batal mun_aamples.
return budget
budget sum across subtasks. Using information provided, generate tbatch_aizal candidate configurations that follow the insights and stay within the budget. Retum exactly (batch, eime) new configurations in strict ISON format. following the schema:
"subtask__ _
'modll": "model_name" "samples":
Inputs:
Tisk nume: (tack_name)
Task description: (tas)
Subtasks, models, prompt/generation lengths:
(nodel_apace)
Total budget: (budget)
Evaluation metries: (metries)
Main metric: (uain, nettri))
Search history: (history)
Guidelines: {oxperience)
Response: Retum only (batch_size) candidates in strict ISON format.
34

---

Takk@edhEnnrromend Tunjparmeters foot
Archives
evoliste batzh(] updane Jistery))
rusjpanamalans) Bor.
Task
dduscc
&budggt faat
atsterfaces
iinterroon
modo chokoes: I56001
AgenITTS
saro beeice Ke(per)
-ogont: Agent
ChaioPT
arthie Anchiw
uift atr Say:
update, exporience)
poneratogrommptt ST
* gon now, condidatoo
gen ran, espertece))
tucgett)
Figure 16: Class Diagram of the AgentTTS Framework. A.8 Detailed Cases for Interpretability
We present a complete case study comparing Agen(TTS and AgentHPO in Fig. 15 The budget of 850 poses a challenge in balancing optimal usage across two subtasks, as the optimal retrieval configuration (Qwen2.5-72B, sample) consumes 814 units. while the optimal QA setting (LLaMA- 3-3B, 50 samples) requires 52. We observe: (1) AgentTTS identifies optimal configurations using embedded insights. whereas AgentHPO does not; (2) Insights 2, and 3 guide different stages of the search: Insight selects models (Qwen2.5-721) for retrieval, LLaMA-3-33 for QA), Insight 2 focuses sampling (reaching the optimal trial: Qwen2.5-72B, 1 LLLaMA-3-3B 15), and Insight balances budget use (yielding another optimal trial: Qwen2.5-32B, 5 LLaMA-3-3B, 20).
A.9 Baselines
We compare two baseline categories: LLM-based approaches (AgentHIP 200 ML.Copilot [56] and LLM_ZS) and traditional methods (Random Search and Bayesian Optimization (11)
AgentHPO [20] targets Hyperparameter Optimization (HPO), addressing limitations of traditional
AutoML methods such as high trial cost, complex setup, and limited interpretability by leveraging LLM--owered autonomous agents. It comprises two modules: the Creator, which converts natural language inputs (e.g., dataset, model, and goals) into initial hyperparameters, and the Executor, which trains models, logs outcomes, and returns feedback. The Creator iteratively refines HPs using this feedback, forming a closed-loop optimization. To adapt AgentHIP for our test-time scaling budget allocation setting, we provide AgentHPO with structured inputs similar to method, including task descriptions, model choices, and total budget. However, AgentHPO assumes LLMs possess sufficient ML knowledge for HPO tasks, an assumption that fails in multi-stage test-time scaling, where domain-specific knowledge is lacking. Our method addresses this gap by injecting targeted insights about test-time scaling into the agent's reasoning process, improving both search efficiency and decision interpretability.
focuses on generating machine learning solutions, addressing limitations in traditional AutoML. methods, such as Bayesian optimization (31] which often lack interpretability and struggle to incorporate human prior knowledge, e.g., regarding model architectures. Inspired by human design patterns, where individuals understand a task by recalling relevant experiences or knowledge from prior work or textbooks, MLCopilot emulates this process by leveraging LLMs to suggest solutions for new ML tasks based on historical cases. To adapt MLCopilot to our problem setting, we provide the search history of a similar task to guide the LLM in generating both prior experience and candidate trials iteratively for the target multi-stage task. However, while MLCopilot can extract surface-level patterns from related tasks, it lacks the ability to derive tailored insights necessary for efficient budget allocation in test-time scaling. This limits its effectiveness in the test- time scaling compute-optimal budget allocation problem, where fine-grained domain understanding is essential for optimal resource allocation.
MI.Copilot S6
35

---

LLM ZS LLM ZS directly prompts the LLM to generate zero-shot candidate trials across multiple iterations, using only the task description, model choices, and total budget as input. Unlike methods that leverage prior knowledge or historical search experience, L.M _S relies solely on the L.L''s pretrained knowledge to guide trial generation.
Bayesian Optimization (BO) 33 is a model- method for optimizing black-box functions and is widely used in hyperparameter tuning. For test-time scaling in multi-stage tasks, BO models the performance of candidate allocation as black-box objective and employs a surrogate model (e.g., Gaussian Process) to guide the search via acquisition functions (e.g., Expected Improvement). It balances exploration and exploitation to navigate the combinatorial space of model-budget configu- rations. However, its limited interpretability hinders its applicability. Moreover, in multi-stage tasks with complex subtask interactions and non-smooth performance landscapes, BO may converge to local optima and require a large number of evaluations.
Random Search serves simple yet widely used baseline for hyperparameter tuning. In the context of test-time scaling budget allocation for multi-stage settings, it randomly samples model- budget configurations without relying on performance history prior knowledge. Despite its simplicity, Random Search is often robust to non-smooth landscapes, which allows it to occasionally find globally optimal trials. However, its inherent randomness and lack of guidance result in significantly slower convergence compared to more informed or tailored strategics.
A.10 Detailed Related Work
A.10.1 Test-time Scaling and Compute-optimal Strategy
Inspired by the human reasoning process that tends to engage in deeper, more deliberate thinking, several studies have proposed allocating additional compute during inference to enhance task perfor- mance [4] [42. Further, other works 2 45 have observed that increasing inference-time compute exhibits a similar pattern akin to training scaling laws: additional inference compute consistently leads to improved task performance. This phenomenon is commonly referred to as test-time scal- ing (TTS). Test-time scaling technologies can be mainly classified into two categories: sequential scaling and parallel scaling. In sequential scaling, it enhances test-time computation by gener- ating progressively longer solutions along the sequence dimension. The most prevalent method is self-revision, where first generate initial response and then iteratively evaluate and refine based on self-assessment [24 33 Because sequential scaling relies on the model's ability to generate reasonably accurate initial responses, it is generally more effective for simpler tasks 33 On the other hand, parallel scaling generates multiple independent solutions in parallel and aggregates them into final answer. Common solution-level aggregation methods are Best-of-N [35] 9 which sample N complete solutions and then select the best one according to a verifier. While Tree-Search algorithms, viewed as parallel scaling at the token or step level, are against a process- based reward model to search top-K intermediate steps and further explore, including Beam Search [53] 46) and
Monte Carlo Tree Search (MCTS) [45] 33 10 37 5 55 However, they typically rely on explicit guidance signals for candidate selection. An alternative line of work directly employs an LLM as
the fuser to aggregate candidate solutions, offering greater generalization and flexibility [3 17
Zeng et al. [55] argues that parallel scaling provides better coverage and scalability than sequential scaling, as extended sequential reasoning chains may corrupt intermediate correct outputs. Similarly, 3 demonstrates that parallel scaling is more suitable for complex tasks, as it requires only the ability to generate correct final answers. Therefore, we adopt a parallel test-time scaling repeated sampling with fusion, to address complex multi-stage reasoning tasks, leveraging the strengths of small models in diverse sampling and robust aggregation.
Snell et al.
strategy-specificnn
The configuration of test-time scaling strategies plays critical role in performance improvement, giving rise to the broader problem of optimally allocating compute and budget, commonly referred to as the test-time compute-optings shown that, in certain settings, smaller models can outperform large language models under the same compute budget. These works explore how to choose between training scaling and test-time scaling, as well as how to select among different test-time scaling techniques to optimize performance. For instance, Snell et al. [K demonstrate that task difficulty strongly influences the optimal scaling strategy: for moderately difficult tasks, parallel search with small models preferred, while for simpler tasks, sequential scaling with large models is more effective. Based on this observation,
scaling strategy. Recent studies [2] 45
19
34
33
40 have
36

---

they propose a difficulty predictor to guide scaling decisions dynamically. Liu et al. [9 further extend this line of work by showing that the choice of optimal TTS strategies is highly sensitive the
design of reward functions. Yue et al. [54] propose using linear model to fit key factors influencing strategy selection within -uggaallgggeetted generation (RAG) systems. Wu et al. [45] introduce Reward Balanced Search (REBASE), novel tree search algorithm that combines weighted voting to achieve a Pareto-optimal trade-off between accuracy and inference cost. However, these approaches primarily focus single-stage task scenarios and do not systematically address the challenge of budget allocation across subtasks in multi-stage complex tasks. In this work, we extend the lest-time compute-optimal scaling framework to multi-stage complex tasks.
A.10.2
LLM for Hyperparameter Optimization
LLMs have emerged as promising tools for enhancing the efficiency and effectiveness of hy- perparameter optimization (HPO) by leveraging contextual understanding and prior knowledge 43 61 22 38), compared to traditional Automated Machine Learning (AutoML) approaches such as Bayesian Optimization (BO) (T. Existing LLM- HPO studies can be mainly categorized into two directions: (1) using LLMs to reduce the search space for traditional methods, and (2) en- abling LLMs to autonomously search optimal hyperparameters. On one hand, LLMs can significantly reduce the vast search spaces in Neural Architecture Search (NAS) and HPO
152
25
23 20 For
example, GPT-NAS 52 combines GPT models with evolutionary algorithms
60 to rapidly prune
low-quality architectures, enhancing search efficiency. Morris et al. 25 introduce "Evolution of Thought," allowing LLMs to refine architectures and hyperparameters based on feedback iteratively. AutoM3L (23] integrates external tools such as Ray.Tune for hyperparameter tuning, with LLMs recommending optimized search spaces, Llambo (20] formulates BO problems natural language, enabling LLMs to propose high-potential candidates and iteratively optimize based on historical
observations.
On the other hand, several works 4 18 configurations autonomously. AutoMM Lab mizers, iteratively refining configurations based on evaluation feedback. CAAFE 12 focuses on tabular data, where L.LMs generate new semantically meaningful features based on dataset descrip- tions. Moreover, LLM--asedd agent frameworks have emerged, leveraging feedback from machine learning platforms or historical experiments. For instance, AutoML-GPT (58) automates the pipeline from preprocessing to model selection and hyperparameter tuning, adjusting strategies based on experimental logs. MLCopilot (56) predicts hyperparameters for unseen tasks by canonicalizing prior task experiences and interactively refining solutions through LLM reasoning. AgentHPO [20] autonomously processes task descriptions, conducts experiments, and iteratively improves hyper- parameter search based on accumulated trials. Building upon these ad advancmmnnss, we extend the LLM-agent framework to the domain of test-time scaling budget allocation in multi-stage complex tasks, aiming search for compute-optima scaling configurations.
62
58
56 empower LLMs to search hyperparameter 40 and GENIUS 62 treat LLMS as black-box opti-
A.11 Limitations
Our problem focuses on multi-stage tasks with static stages, where the compute budget is allocated across fixed set of subtasks. However, in real-world applications, some tasks exhibit dynamic multi-stage behavior, where the actual runtime subtasks may vary depending on input conditions or user interaction. For example, in voice-based personal assistants mobile devices, the system may dynamically decide whether to perform document retrieval, clarification, or follow-up question generation based on the user's query and context. In such scenarios, is difficult to predefine all possible subtasks and their corresponding budget allocations. Consequently, applying our LLM- agent-based search framework these dynamic multi mtati-stage tasks presents considerable challenge.
A.12 Broader Impact
This work proposes AgentTTS, an LLM-agent-baseed framework for compute-optimaal test-time scaling in multi-stage tasks. The methodology offers significant positive societal impacts. First, by improving inference efficiency and reducing unnecessary computation, it contributes to more sus- tainable and cost-effective deployment of large language models, especially resource-constrained environments. Second, our framework promotes better alignment between model selection and
37

---

task complexity, enabling more accessible AI solutions across diverse application domains such as education, healthcare, and low-resource software development.
Nonetheless, our approach raises potential concerns. AgentTTS relies on repeated sampling from foundation L.Ms,, which amplify not only their strengths but also their limitations. For instance, hallucinations in large language models may be intensified through test-time scaling integration and efficiently propagated, increasing the risk of misinformation. Moreover, LLMs susceptible to adversarial attacks such as jailbreaks, backdoor injections, and membership inference attacks, thereby heightening the security risks associated with AgentTTS.
38


Reviews:
[{"review_id": "K7TitwpBQn", "reviewer": "Reviewer_8uj2", "summary": "The paper proposes AgentTTS, an LLM-guided agent that optimizes compute allocation across multi-stage NLP pipelines under test-time budgets. It leverages empirical insights such as diminishing returns and subtask dependencies to guide efficient search, outperforming multiple baselines on six benchmarks.", "questions": "Q1. Given that several reported gains are relatively small (e.g., 1–2 EM points), can you provide statistical significance estimates or confidence intervals for key results (perhaps through multiple seeds or bootstrap analysis)?\n\nQ2. How were the HPO baseline methods adapted to handle the multi-stage compute-aware pipeline setting? Is there evidence that they are competitive under best-case tuning?\n\nQ3. Have you considered or tested AgentTTS on pipelines involving other domains, modalities, or more complex, longer-horizon tasks? What might limit generalization to these settings?\n\nQ4. Have you considered repeating your analysis for different temperature values? For example, this [paper](https://arxiv.org/abs/2504.12951) [1] shows that increasing the temperature value influences the sample efficiency of various reasoning methods.\n\n[1] Potamitis et al. Are Retrials All You Need? Enhancing Large Language Model Reasoning Without Verbalized Feedback. https://arxiv.org/abs/2504.12951", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n\nS1. Open source code.\n\nS2. Comprehensive evaluation on six datasets.\n\nS3. Clear empirical insights are discussed in the introduction that guide the development of the method.\n\nS4. Strong empirical performance. Figure 3 shows the sample efficiency of AgentTTs while Table 1 (partly) shows its cost efficiency in terms of search time.\n\n### Weaknesses\n\nW1. Table 1: Since there are a lot of unavailable results, it’s hard to understand the table completely. For example, in the Taskbench and Chatdev case, there’s no search time included for any of the baseline search times, even though some of them achieve the same score as AgentTTs.\n\nW2. No error bars are reported, which raises a question regarding the robustness of the proposed method.\n\nW3. Domain coverage. Even though the authors benchmark their method against multiple baselines and different benchmarks, the domain that the benchmarks include is fairly limited. The used benchmarks focus on QA-style pipelines; however, I would like to see some results on different domains as well, such as math, and more complex benchmarks such as planning.", "was_revised": true, "final_justification": "I appreciate that the authors have provided detailed responses to my questions; however, as things stand, I would like to maintain my original assessment of the paper. Should the paper get accepted, I would also appreciate it if the AC urges the authors to add the results from their response in the final version of their paper.\n\nFinally, as stated in my response to the authors, I would appreciate it if the authors showcase their findings on tasks beyond \"question answering\" and from diverse domains.", "rebuttal": "We appreciate the reviewer’s valuable comments. Below, we respond to each point in detail.\n\n> W1 Table 1: Since there are a lot of unavailable results, it’s hard to understand the table completely. For example, in the Taskbench and Chatdev case, there’s no search time included for any of the baseline search times, even though some of them achieve the same score as AgentTTs.\n\nIn the orginal table, we omitted search times for methods that failed to find optimal configs on the trainng set to focus on those with strong **generalization** and **efficiency**. Reported test results reflect the generalization of training-found optimal or near-optimal configs.\n\nTo address the concern, we have included the previously missing search times in a revised version of Table below. Entries marked with an asterisk (*) indicate cases where the method did not find the best-performing configuration in the training set; in such cases, the search time corresponds to the best configuration the method could find, which often leads to suboptimal test-set performance. \n\n**Table:** Comparison of search time (in hours) and test-set performance across datasets. * indicates the method did not find the best trial in the training set.\n\n|Method|2WikiTime|2WikiEM|2WikiEM(90%CI)|HotpotTime|HotpotEM|CWQTime|CWQEM|WebQSPTime|WebQSPEM|TaskbenchTime|Taskbenchp-F1|ChatDevTime|ChatDevCons.|\n|------|---------|-------|---------------|-----------|--------|--------|------|------------|----------|---------------|---------------|-------------|---------------|\n|Random|24.3*|0.66|0.656(0.624,0.688)|16.4*|0.71|37.4*|0.76|10.0*|0.86|14.6*|0.40|18.4*|0.74|\n|BO|3.5*|0.60|0.604(0.574,0.635)|40.5*|0.71|23.8*|0.76|10.0*|0.85|14.6*|0.52|4.6*|0.75|\n|LLM_ZS|12.5|0.70|0.677(0.625,0.729)|13.1*|0.71|55.3|0.76|37.7|0.89|114.0*|0.49|2.0*|0.74|\n|MLCopilot|12.5|0.70|0.698(0.680,0.716)|46.3|0.72|48.4|0.78|20.8|0.88|23.4*|0.53|20.9*|0.75|\n|AgentHPO|8.3|0.70|0.698(0.680,0.716)|36.3|0.74|37.4|0.78|48.1|0.89|43.8*|0.49|16.9*|0.74|\n|**Ours**|**2.5**|**0.72**|**0.724(0.695,0.752)**|**17.5**|**0.74**|**11.1**|**0.78**|**7.8**|**0.89**|**64.3**|**0.53**|**14.3**|**0.75**|\n\n\n> W2 No error bars are reported, which raises a question regarding the robustness of the proposed method.\n>\n> Q1. Given that several reported gains are relatively small (e.g., 1–2 EM points), can you provide statistical significance estimates or confidence intervals for key results (perhaps through multiple seeds or bootstrap analysis)?\n\nSince the weakness and question overlap, we respond jointly. We ran 5 repeated searches on 2Wiki with different seeds and report **mean ± 90% CI** in the table above to assess robustness.\n\nThe results confirm that while there is natural variance across runs, our method consistently discovers high-quality configurations. On average, the configurations identified by our approach achieve **superior generalization performance** on the test set compared to baselines, further supporting the robustness and effectiveness of our search strategy.\n\n> W3. Domain coverage. Even though the authors benchmark their method against multiple baselines and different benchmarks, the domain that the benchmarks include is fairly limited. The used benchmarks focus on QA-style pipelines; however, I would like to see some results on different domains as well, such as math, and more complex benchmarks such as planning.\n>\n>  Q3. Have you considered or tested AgentTTS on pipelines involving other domains, modalities, or more complex, longer-horizon tasks? What might limit generalization to these settings?\n\nBoth the weakness and question concern the same point, so we address them together. While our main experiments use QA-style pipelines, **our method is domain-agnostic** and applies to any multi-stage tasks. It assumes modular we have subtasks and model choices, which generalize across domains—for example, geography VQA can be split into image parsing (e.g., maps) and reasoning (e.g., climate), each handled by specialized models.\n\n1. **Why math and planning tasks were not included initially**\n\nWe also tried mathematics or planning datasets but not included in our main evaluation because **existing benchmarks in those domains are inherently single-stage**. For example:\n\n- **Math tasks** such as *MATH*, and *MathQA* typically require long-form generation but do not involve discrete multi-stage components.\n- **Planning tasks** (e.g., ALFWorld) are often benchmarked via environment interactions rather than modular subtask execution.\n\nOur framework is designed for structured reasoning with budget allocation across dependent subtasks. Thus, without existing multi-stage decompositions, math and planning tasks fall outside the scope of the current setting.\n\n2. **New domain experiment: Multimodal reasoning on TextVQA**\n\nTo address the concern about generalization to other domains/modalities, we adapted the **TextVQA** dataset into a two-stage pipeline: **Stage 1**: Image understanding; **Stage 2**: Question-specific reasoning based on image understanding. \n\nWe then evaluated the **search efficiency** for multiple models in this multimodal setting. Results are shown below.\n\n**Table: Search Time and Performance across different methods on TextVQA with budget 1000. \\* means failure to get the optimal configuration.** \n\n|Method|SearchTime(hrs)|Performance|\n|------|----------------|-----------|\n|Random|205*|0.90|\n|Bayes|50*|0.83|\n|MLCopilot|70*|0.90|\n|AgentHPO|245*|0.90|\n|**Ours**|**65**|**0.93**|\n\nThe experimental results show our method could generalize to the multi-modal task. \n\n3. **Test-time scaling on MATH dataset**\n\nTo address the request, We **run MATH experiments** with Qwen2.5 (3B, 72B) to test how **test-time scaling** affects open-ended math reasoning. Accuracy on 500 samples is shown below with **Best-of-N aggregation**.\n\n**Table: Qwen 72B test-time scaling on MATH test set with 500 samples using Best-of-N aggregation** \n\n|RepeatCount|Budget|Accuracy|\n|-----------|------|--------|\n|1|214|0.676|\n|2|406|0.708|\n|3|598|0.724|\n|4|790|0.730|\n|5|982|0.738|\n|6|1174|0.740|\n|7|1366|0.742|\n|8|1558|0.742|\n|9|1750|0.744|\n|10|1942|0.744|\n\n\n**Table: Qwen 3B test-time scaling on MATH test set with 500 samples using Best-of-N aggregation** \n\n|RepeatCount|Budget|Accuracy|\n|-----------|------|--------|\n|1|7|0.506|\n|10|79|0.640|\n|20|159|0.670|\n|30|239|0.670|\n|40|319|0.676|\n|50|399|0.680|\n|60|479|0.684|\n|70|559|0.682|\n|80|639|0.684|\n|90|719|0.686|\n\nQwen2.5-72B costs 214 units per gen; 3B costs 7. Though 3B can beat 72B with >50 repeats, it does so at **higher total cost**. Thus, under fixed budgets, **72B is more cost-effective** for MATH.\n\n> Q2. How were the HPO baseline methods adapted to handle the multi-stage compute-aware pipeline setting? Is there evidence that they are competitive under best-case tuning?\n\n**How HPO Methods Were Adapted:**\n\nLLM-based HPO methods, such as AgentHPO, were originally developed to tune hyperparameters for machine learning tasks by having an LLM iteratively propose candidate configurations based on task descriptions and observed feedback. In our setting, **test-time scaling configuration search closely parallels HPO**, with the distinction that the parameters to be optimized are no longer traditional hyperparameters, but rather **test-time scaling configurations**—i.e., model choices and generation budgets across multiple subtasks. This conceptual alignment motivated us to adopt HPO methods as baselines for our work.\n\nWe adapted **AgentHPO** as follows:\n* Provide task description and config space to the LLM.\n* LLM generates **initial trials**, evaluated for feedback.\n* Log trials with **rationale-style experiences** per AgentHPO.\n* LLM uses past feedback to **plan next trials**.\n* Select **best training trial** for test evaluation.\n\n**Evidence of Competitiveness:**\n\nThe effectiveness of LLM-based tuning is supported by the results reported in **Table 1** and **Figure 3** in the paper. In particular:\n- **AgentHPO is among the strongest HPO baselines**, frequently identifying near-optimal configurations in training.\n- Its performance serves as a robust foundation for evaluating our method.\n\nOur method builds on this strong baseline by incorporating three key insights (as described in Section 4), which improve **efficiency** of the search. Nonetheless, the baseline AgentHPO method itself is already capable and competitive under best-case tuning, reinforcing the relevance and strength of the comparison.\n\n> Q4. Have you considered repeating your analysis for different temperature values? For example, [1] shows that increasing the temperature value influences the sample efficiency of various reasoning methods.\n\nIn our main experiments, we used a relatively **high temperature value of 0.7 or 0.9 by default**. However, we had not systematically varied the temperature values in earlier analysis.\n\nTo address this concern, we conducted additional experiments on 2Wiki to evaluate the effect of temperature on test-time scaling. Specifically, we fixed the retriever to **Qwen2.5-72B** (generating a single retrieval result) and varied the number of samples and temperature values for the **QA model LLaMA 3 3B** on the **2Wiki dataset**. The results are: \n\n**Table: EM Scores at Different Temperatures on 2Wiki**\n\n|Samples|Temp0.1|Temp0.5|Temp0.9|\n|-------|--------|--------|--------|\n|1|0.70|0.64|0.66|\n|10|0.70|0.78|0.78|\n|20|0.72|0.76|0.78|\n|30|0.74|0.78|0.72|\n|40|0.74|0.76|0.78|\n|50|0.74|0.72|0.78|\n|60|0.74|0.78|0.80|\n|70|0.76|0.72|0.74|\n|80|0.76|0.76|0.78|\n|90|0.74|0.72|0.76|\n\nKey Observations:\n* **High temps (e.g., 0.9)** boost test-time scaling by increasing output diversity, aiding fusion with more samples; **low temps (e.g., 0.1)** suit single-sample cases.\n* This aligns with [1], showing **diverse outputs are key** for sampling-based reasoning.\n\nWe hope our replies clarify your concerns. Let us know if you have further questions."}, {"review_id": "H9jHxvUCm2", "reviewer": "Reviewer_bAdt", "summary": "This paper addresses the challenge of test-time compute-optimal scaling in multi-stage complex tasks, where existing methods for single-stage tasks fall short due to the combinatorial search space and interdependent subtasks. The authors first conduct pilot experiments, deriving three key insights. Guided by these, they propose AgentTTS, an LLM-agent framework that iteratively searches for optimal compute allocations via feedback-driven interactions. AgentTTS leverages an LLM to generate candidate trials, update guidelines based on performance feedback, and adapt to subtask dependencies. Experiments show that AgentTTS outperforms traditional and LLM-based baselines in search efficiency", "questions": "See above", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths:\n1. The pilot study derives three key insights from comprehensive experiments, showing that different subtasks prefer distinct models, compute scaling has diminishing returns, and earlier budgets affect downstream tasks. These insights ground the framework in real-world LLM behavior.\n2. AgentTTS reduces search time by up to 80% compared to baselines while matching or exceeding performance, demonstrating practical utility for real-world deployment.\n3. By focusing on multi-stage task compute optimization, the paper addresses a critical gap in existing TTS research, which primarily targets single-stage tasks. This aligns with complex real-world scenarios like retrieval-then-generation QA and software development .\n\n**Weaknesses:**\n1. While Theorem 1 and the budget unit definition aim to standardize compute metrics, the experimental section rarely references these concepts, making their practical impact unclear. The complex formula may introduce unnecessary complexity without tangible benefits in results.\n2. The three core insights, subtask model preferences, optimal budget saturation, and downstream dependency, are largely intuitive and supported by common LLM observations. While validated experimentally in this paper, they may lack the novelty expected for foundational \"insights\".\n3. Inconsistent result presentation:\n   - Figure 3 shows baselines approaching 0.8 EM on 2WikiMultiHopQA, but Table 1 lists AgentTTS's peak at 0.72. Why? I thought it was because Figure 3 was reporting performance on training set, but I read that Line 325 says that Figure 3 is also the performance on test set.\n   - From Figure 3, it seems most of the methods have the same performance upper bound, why? What's the dashed line in each plot? It is also unclear for the caption potentially refers to the dashed line: \"The best score is obtained from the optimal trial in a prior grid search and serves as the benchmark for all methods\".\n   - AgentTTS appears to still be improving at 50 trials on 2WikiMultiHopQA, raising questions about its performance with more iterations. Why not present the results with more trials?", "was_revised": true, "final_justification": "My current assessment leans toward a score of 3.5. While most of my concerns have been addressed, I find the definition of the budget unit remains unclear. This aspect appears to unnecessarily complicate the paper's readability while contributing minimally to either the methodology or results. I recommend that the AC also consider the clarity of writing when making their final decision.", "rebuttal": "We thank the reviewer for the thoughtful and detailed feedback on the budget unit, the proposed insights, and others. We address the comments point-by-point below.\n\n> While Theorem 1 and the budget unit definition aim to standardize compute metrics, the experimental section rarely references these concepts, making their practical impact unclear. The complex formula may introduce unnecessary complexity without tangible benefits in results.\n\nThank you for the thoughtful comment. We appreciate the opportunity to clarify the motivation and practical role of Theorem 1 and the budget unit definition.\n\n**Purpose of Theorem 1 and Budget Unit Definition:**\n\nThe budget unit formulation in Theorem 1 is not intended to add unnecessary complexity, but rather to **standardize the notion of compute cost** across diverse tasks and models. In practice, different deployment scenarios may use varied cost metrics such as FLOPs, latency, token usage, or even monetary cost. Without a unified abstraction, comparisons between models would be inconsistent and task-specific.\n\nOur budget unit definition offers a **task- and model-invariant compute framework** that enables fair cross-task/model comparisons, supports any cost metric via normalization, and provides a consistent interface for budget-aware reasoning.\n\n**Practical Relevance in Experiments:**\n\nWhile the experimental section may not explicitly re-state Theorem 1, the concept is tightly integrated into both the **design of our method** and the **analysis of results**:\n\n1. **Insight-driven analysis (Fig. 1)** relies on plotting performance as a function of budget, enabling meaningful comparison of models under equal compute budget.\n\n2. **Method design (Eq. 5)** depends on a consistent budget unit to initialize candidate trials and perform model–sampling trade-offs during search.\n\n3. **The LLM-based search agent** uses the budget computation formula internally to assess feasibility and guide its planning strategy during the search process.\n\nThus, the theoretical framework directly informs the architecture of our search method and the interpretation of test-time scaling behavior.\n\n\n\n\n\n> W2 The three core insights, subtask model preferences, optimal budget saturation, and downstream dependency, are largely intuitive and supported by common LLM observations. While validated experimentally in this paper, they may lack the novelty expected for foundational \"insights\".\n\n\n\nWe thank the reviewer for this comment. We agree that the three insights—**subtask-specific model preference**, **optimal budget saturation**, and **downstream dependency**—are partially aligned with emerging observations in LLM literature. However, we respectfully argue that their **generalization and algorithmic application** in our work offer substantial contributions beyond prior findings.\n\n**Relation to Prior Work:**\n\nSome individual components of our insights do echo trends observed in prior studies:\n\n- **Insight 1 (Subtask-specific model preference)** aligns with findings in [1], which shows that **easier samples benefit more from small models** under test-time scaling, while **harder ones require larger models**. However, this was shown primarily in single-task, single-stage, and limited task-type settings (mainly MATH dataset), whereas we generalize this to **multi-subtask pipelines** across six diverse datasets.\n- **Insight 2 (Unimodal budget saturation)** aligns conceptually with the saturation effect in [2], where performance plateaus under majority voting. However, [2] relies on a **verifier model** for weighted aggregation, whereas our **sampling-then-fusion** approach is verifier-free, making it more general and practical.\n- **Insight 3 (Downstream dependency)**, to the best of our knowledge, is **novel**. It emerges only in **multi-stage pipelines**, where upstream outputs influence downstream task difficulty and model choice. We are the first to systematically identify and formalize this inter-stage interaction, which significantly affects budget allocation dynamics.\n\n**Contribution and Practical Role:**\n\nWhile the insights may appear intuitive in isolation, our contribution lies in:\n\n1. **Demonstrating their generality** by validating all three across **six datasets spanning four task types**: retrieval, question answering, summarization, and knowledge graph reasoning—far beyond the narrow scope of prior works.\n\n2. **Using them to inform the design of our search algorithm**, which directly incorporates these principles to improve search efficiency. It yields **substantial empirical speed-ups** in identifying optimal configurations (see Table 1).\n\n[1] *Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters*.   \n[2] *Inference Scaling Laws: An Empirical Analysis of Compute-Optimal Inference for Problem-Solving with Language Models*. \n\n> Figure 3 shows baselines approaching 0.8 EM on 2WikiMultiHopQA, but Table 1 lists AgentTTS's peak at 0.72. Why? I thought it was because Figure 3 was reporting performance on training set, but I read that Line 325 says that Figure 3 is also the performance on test set.\n\nWe appreciate the reviewer catching this inconsistency. There is a **typographical error in Line 325**, which mistakenly states that Figure 3 reports *test set* performance. \n\n**Correction**: *Fig. 3 shows the search trajectories of our method and the baselines for 325 test-time compute-optimal budget allocation on the training set.*\n\n> From Figure 3, it seems most of the methods have the same performance upper bound, why? What's the dashed line in each plot? It is also unclear for the caption potentially refers to the dashed line: \"The best score is obtained from the optimal trial in a prior grid search and serves as the benchmark for all methods\".\n\nAs to why many methods appear to approach a same upper bound: this is **by design**. The core goal of our method (and baselines) is not to exceed this bound, but to **reach it efficiently** with fewer trials. Theoretically, any method (even random search) could reach the optimum if allowed infinite time—hence they share the same asymptotic limit.\n\nThe dashed line in Figure 3 indicates the **best achievable performance** on the training set, which we use as a reference point to assess the search efficiency of all methods. Specifically:\n\n- This best score is obtained by conducting an approximated exhaustive grid search across the entire configuration space, which includes all possible combinations of model selections and sampling budgets for each subtask. Specifically, we enumerate all candidate configurations, execute them in the environment, and record their final task performance. The configuration that achieves the highest score is selected as the best. We will clarify this process more explicitly in the figure caption.\n- The purpose of this line is to serve as a **compute-optimal upper bound**—i.e., the highest training-set score obtainable with a fixed budget.\n\n\n> AgentTTS appears to still be improving at 50 trials on 2WikiMultiHopQA, raising questions about its performance with more iterations. Why not present the results with more trials?\n\nWhile Figure 3 shows a continued upward trend for AgentTTS on 2WikiMultiHopQA, we clarify that this behavior simply reflects **fine-grained fluctuations** near the optimal configuration. In practice, AgentTTS has already **found the optimal (or near-optimal) configuration** by that point. Extending the number of trials beyond 50 does not lead to meaningful gains, as the method has already converged to the compute-optimal region. \n\nWe chose 50 trials as most methods can reach the optimum within this limit. AgentHPO typically uses 10 trials, so 50 is already a generous budget that provides ample opportunity to find the best configuration. Beyond this, the search cost becomes excessive—for example, exceeding 25 hours on the 2Wiki dataset. We will make this clearer in the revised version.\n\n\nWe hope our clarifications resolve your concerns and demonstrate the robustness of our contributions. Please let us know if any further explanations are needed."}, {"review_id": "gHE1Yxgmic", "reviewer": "Reviewer_ZDCu", "summary": "This paper introduces AgentTTS, a framework designed to solve the complex and highly practical problem of test-time compute-optimal scaling for multi-stage tasks. The core challenge it addresses is how to best allocate a limited computational budget across a sequence of subtasks—like in a retrieval-then-generation system—to maximize the overall performance. The authors first establish three key empirical insights from pilot experiments: that different subtasks have distinct model preferences, that performance returns diminish with added compute, and that budget choices in early stages impact later ones. AgentTTS uses an LLM-powered agent that leverages these insights to autonomously and iteratively search for the best combination of models and budget allocations. The agent interacts with a task environment, gets feedback, and refines its strategy to efficiently find a near-optimal configuration. Results show AgentTTS outperforms baselines in efficiency and robustness.", "questions": "As mentioned in the Weakness.", "limitations": "As mentioned in the Weakness.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths：\n1. The paper effectively targets a new problem of compute-optimal scaling in multi-stage complex tasks, which is highly relevant to real-world applications where tasks involve multiple stages with diverse model requirements.\n2. The authors conduct experiments on models of various sizes （from 3B to 72B）and types (Qwen, Llama), providing a thorough evaluation of the proposed method across different model capabilities.\n3. The authors have done a great job of validating their method across a wide range of tasks, from question answering to automated software development, using six different datasets.\n\nWeaknesses：\n1. Despite the novel problem proposed by the authors, the innovation of the method itself should be further elaborated.\n2. Lacks in-depth theoretical analysis of the proposed method.\n3. The paper focuses on optimizing the inference budget, but the search process itself has a computational cost. A more direct discussion about this trade-off would be beneficial. For instance, how much deployment is needed to justify the upfront cost of the search?", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer’s careful evaluation and valuable comments on innovation, in-depth analysis of the proposed method, and the computational cost during searching. We provide detailed responses to each issue raised below.\n\n\n\n> W1 Despite the novel problem proposed by the authors, the innovation of the method itself should be further elaborated.\n\nWe appreciate the reviewer’s interest in further understanding the innovation of our proposed method.\n\nThe novelty of our method is closely tied to the three core insights presented in the paper, which directly address the unique challenges posed by **multi-stage complex tasks within a fixed budget**: (1)**Model and number of sampling across subtasks**: Each subtask may require a different model based on its reasoning complexity and domain. Our method explicitly considers how to identify the most suitable model for each subtask. (2) **Optimal sampling budget varies by subtask**: There exists an ideal number of samples for each subtask. Our framework is designed to discover this optimal per-subtask sampling rate. (3) **Global budget constraints create trade-offs**: Given a fixed compute budget, it is often impossible to allocate the optimal budget to all subtasks. Thus, we must prioritize subtasks and allocate resources in a way that **maximizes overall task performance**.\n\nTo address these intertwined challenges, we propose a **LLM-agent-based framework that autonomously searches for compute-optimal allocations** that leverages both (i) the planning capability of LLMs and (ii) our proposed insights as **informed priors**. During the search process, the agent receives performance feedback and iteratively improves its configurations based on accumulated experience, leading to **more efficient search** and **superior configurations** compared to prior baselines. Thus, the **innovation lies not only in the LLM-agent-based framework** but in how it is ***informed and guided*** by empirically validated insights that are specific to **multi-stage, budget-constrained tasks**.\n\nWe will make this connection between insights and algorithmic design clearer in the revised manuscript.\n\n\n\n\n\n> W2 Lacks in-depth theoretical analysis of the proposed method.\n\nThank you to the reviewer for highlighting the lack of in-depth theoretical analysis of our search method. In response, we provide a complete theoretical analysis below, starting from the basic notations and culminating in a justification for why an insights-informed search leads to greater efficiency.\n\n\n\n**Notation.** \nLet $\\mathcal{T} = [T_1, \\dots, T_N]$ be the subtasks. For subtask $i$, let $\\mathcal{M_i}$ be the model set and $M_i \\in \\mathcal{M_i}$  a choice. \nLet $\\mathbf{B} = (B_1, \\dots, B_N)$ be integer budgets with $\\sum_i B_i \\le B$ (total budget $B \\in \\mathbb{N}$). \nA configuration (trial) is $c = (M_1, B_1, \\dots, M_N, B_N) \\equiv (M, \\mathbf{B})$. Let $F(M, \\mathbf{B})$ denote the expected task score. AgentTTS seeks $c^\\star = \\arg\\max_{(M, \\mathbf{B})} F(M, \\mathbf{B})$.\n\n**Assumptions (Insights 1–3 formalized).**\n\n- **A1 (Subtask-specific model preference).**  For each $i$ there exists a *dominant* model $M_{i}^* \\in \\mathcal{M_i}$ and a budget interval $\\mathcal{I_i} \\subseteq \\{0, \\dots, B\\}$ such that for any fixed $ ( M_{-i}, \\mathbf{B_{-i}}) $ and all $B_{i}  \\in \\mathcal{I_{i}}$,  \n$$\nF(M_{i}^* , B_{i}; M_{-i}, \\mathbf{B_{-i}}) - \\max_{M_{i} \\ne M_{i}^*} F(M_{i}, B_{i}; M_{-i}, \\mathbf{B_{-i}}) \\ge \\Delta_{i} > 0.\n$$\n\n- **A2 (Unimodality in $B_i$).**  For fixed $M_i$ and $(M_{-i}, \\mathbf{B}_{-i})$, the function $B_i \\mapsto F(M, \\mathbf{B})$ on $\\{0, \\dots, B\\}$ is unimodal with a unique maximizer $B_i^\\star(M)$.\n  \n- **A3 (Bounded coupling).**  For fixed $M$, define the best-response map $T: \\{0, \\dots, B\\}^N \\to \\{0, \\dots, B\\}^N$ by $\n  [T(\\mathbf{B})]_i = \\arg\\max_b F(M, \\mathbf{B} \\mid B_i = b).\n  $ There exists $\\rho \\in (0, 1)$ such that $\n  \\|T(\\mathbf{B}) - T(\\mathbf{B}')\\|_2 \\le \\rho \\, \\|\\mathbf{B} - \\mathbf{B}'\\|_2 \\quad \\text{for all } \\mathbf{B}, \\mathbf{B}'.\n  $\n\n**Initialization (Model elimination; Insight 1).**  Compare models within each $\\mathcal{M}_i$ under aligned conditions (same $B_i \\in \\mathcal{I}_i$, others fixed).  By A1, the dominant model $M_i^*$ is separated by a margin $\\Delta_i$, so selecting the top model per subtask is correct.  Keeping only a small shortlist $\\widetilde{\\mathcal{M}}_i$ that contains the winner reduces the search from $\n\\prod_i |\\mathcal{M}_i| \\quad \\text{to} \\quad \\prod_i |\\widetilde{\\mathcal{M}}_i|.\n$\n\n**Budget Localization (per subtask; Insight 2).**  With $M$ fixed (after elimination), $B_i \\mapsto F(M, \\mathbf{B})$ is unimodal (A2).  A simple bracket-and-zoom (e.g., ternary or golden-section on integers) halves the budget interval each round.  After $R = \\mathcal{O}(\\log(B/\\varepsilon))$ rounds, the bracket length is $\\le \\varepsilon$, yielding $\\widehat{B}_i$ with $\n|\\widehat{B}_i - B_i^\\star(M)| \\le \\varepsilon.\n$ Hence, the number of evaluations per subtask for budget refinement satisfies $\nm_i = \\mathcal{O}(\\log(B / \\varepsilon)).\n$\n\n**Coordination and Convergence (Insights 3).** Fix $M \\in \\prod_i \\widetilde{\\mathcal{M}}_i$ and let $\\mathbf{B}^\\diamond$ be the unique fixed point of $T$ (A3).  \nIf AgentTTS updates budgets via (exact or close-to-exact) best responses, the error  $\n\\mathbf{e}^{(t)} = \\mathbf{B}^{(t)} - \\mathbf{B}^\\diamond\n$\ncontracts geometrically:  $\n\\|\\mathbf{e}^{(t)}\\|_2 \\le \\rho^t \\, \\|\\mathbf{e}^{(0)}\\|_2.\n$ The Archive retains the best-seen configuration, so the recorded best score is non-decreasing across iterations.\n\n\n\n**Why Insight 3 Accelerates Convergence:**\nA **naïve update** tweaks subtask $i$ while pretending the rest are fixed. Then move to subtask $j$ and tweak it, which **changes the conditions that made the last tweak to $i$ look good**. The result is that when come back to $i$, we often need to move it **back toward where it was**—so that we “undo” part of your previous change. This behaves like a weaker contraction with factor $\\rho_{\\text{naive}}$ (closer to 1).  \n\nIn constrast, Insight 3 lets the LLM *group and update strongly coupled subtasks together* (block best responses),  better approximating the full map $T$ and yielding a smaller factor $\\rho_{\\text{eff}} < \\rho_{\\text{naive}}$. Therefore, the insight 3 helps accelerate convergence. \n\n\n\n\n> W3 The paper focuses on optimizing the inference budget, but the search process itself has a computational cost. A more direct discussion about this trade-off would be beneficial. For instance, how much deployment is needed to justify the upfront cost of the search?\n\nThank you for raising this important point. We agree that analyzing the trade-off between the computational cost of the search process and the resulting inference-time performance is essential for assessing our method.\n\nTo address this, we now include a practical estimate illustrating how quickly the search cost can be amortized through downstream inference improvements under typical deployment scenarios. For instance, on the 2Wiki dataset, our method's search process requires approximately 2.5 GPU-hours. Given that our optimized configuration achieves a 10–20% increase in test accuracy compared to baseline methods at a fixed inference budget, this initial computational investment is offset after only a few hundred test-time executions—a frequency commonly observed in real-world batch deployments. This scenario is analogous to fine-tuning models for specific downstream tasks, where investing several GPU-hours upfront significantly enhances model performance during inference.\n\nWe believe this analysis helps clarify when the upfront search effort is worthwhile.\n\n\n\nThank you again for your thoughtful review. We welcome any additional feedback."}, {"review_id": "byrkZUbnwW", "reviewer": "Reviewer_NhkM", "summary": "This paper proposes AgentTTS,  an LLM-agent framework to efficiently allocate compute resources across subtasks in multi-stage complex tasks. This framework aims to solve test-time compute-optimal scaling in multi-stage complex tasks. By conducting empirical studies on four task types and six datasets, the authors identify three key insights—subtasks have distinct model preferences, test-time scaling yields diminishing returns, and early subtask budgets affect downstream performance—which are then integrated into AgentTTS’s iterative, feedback-driven search process. Experiments show that AgentTTS outperforms both traditional and LLM-based baselines in search efficiency, robustness, and interpretability for compute-optimal scaling in complex workflows.", "questions": "See weakness", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. The paper addresses a novel and highly relevant problem, i.e. 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.\n2. The design of AgentTTS is well-motivated by three empirical insights derived from preliminary experiments.\n3. Authors conduct experiments across four task types and six datasets, demonstrating the superiority of the proposed framework.\n\nWeakness:\n1. While the empirical insights are valuable, the paper lacks deeper theoretical understanding of why these patterns emerge. The insights appear somewhat ad-hoc without principled theoretical foundations explaining the underlying mechanisms driving these behaviors.\n2. The three insights are derived from specific task types and may not generalize to other multi-stage scenarios. The paper doesn't adequately address the boundaries of applicability for these insights.\n3. The FLOPs-based budget conversion assumes linear scaling relationships that may not hold across all model architectures and task types. Also, the choice of base configuration (3B model, 128 prompt tokens, 64 generation tokens) seems arbitrary and may bias results.\n4. It seems only GPT-o3-mini is used as the agent LLM, the paper could benefit from applying the method on more LLMs to show its generalizability.", "was_revised": true, "final_justification": "Authors's response address most of my concerns. I have updated the score accordingly.", "rebuttal": "We thank the reviewer for valuable feedback on theory and generalization. Below we address each point.\n\n> W1 ... the paper lacks deeper theoretical understanding of why these patterns emerge. The insights appear somewhat ad-hoc without principled theoretical foundations explaining the underlying mechanisms driving these behaviors.\n\nThank you to the reviewer for highlighting the lack of in-depth theoretical analysis of our proposed insights. In response, we provide a complete theoretical analysis below.\n\n**Setup.** Model $\\pi$ receives input $x$, samples a reasoning prefix $r \\sim \\pi(\\cdot \\mid x)$, then an answer $y \\sim \\pi(\\cdot \\mid x + r)$.  Let the correct answer be $Y^\\star$. The single-sample success is  \n$$\np_\\pi(x) = \\Pr_\\pi(y = Y^\\star \\mid x) = E_{r \\sim \\pi(\\cdot \\mid x)}\\big[\\pi(Y^\\star \\mid x + r)\\big].\n$$\nWith $k$ samples, we obtain $(r_{1:k}, y_{1:k})$ and an aggregator $A_k$ outputs the final prediction $y^ \\star = A_k(y_{1:k}, r_{1:k}, x)$ (abbreviated as $\\pi(y^\\star \\mid y_{1:k})$), so we have \n$$\np_\\pi^k =\\Pr_\\pi(y^ * =Y^*|x) = E_{r\\sim \\pi(\\cdot|x), y_{1:k}}[\\Pi_{i\\in[1,k]} \\pi(y_i|x+r) \\pi(Y^ * |y_{1:k})]\n$$\n\n**Insight 1: subtask-specific preference on model and on $k$**\n\n**Model preference:** Define the per-sample log-odds (against the strongest competitor $\\bar{Y}$):  \n$$\nL(r; x, \\pi) = \\log \\frac{\\pi(Y^\\star \\mid x + r)}{\\pi(\\bar{Y} \\mid x + r)}, \\qquad\n\\gamma_\\pi(x) = E_r[L(r; x, \\pi)].\n$$\n$\\gamma_\\pi(x)$ depends on both aspects:  (i) $\\pi(r \\mid x)$ producing helpful $r$, and (ii) $\\pi(y \\mid x + r)$ separating $Y^\\star$ from others. Different subtasks $i$ require different capabilities, so $\\gamma_{\\pi,i}(x)$ varies across subtasks, leading to different preferred models.\n\n**The preferred number of samples is subtask- and model-specific:** We define $k^ *$ as the optimal number of samples for a given subtask. Its value depends on two key factors: (i) the correctness of individual outputs $y_i$, and (ii) the accuracy of the aggregation function $\\pi(Y^ * \\mid y_{1:k})$. Harder subtasks generally require larger $k$ due to lower per-sample accuracy, while the effectiveness of aggregation is constrained by the model’s long-context handling ability. As a result, the optimal sampling count $k^ *$ varies across subtasks and models.\n\n**Insight 2: there is an optimal $k$**\n\nRepeated sampling helps but saturates with correlation and can worsen due to long-context aggregation.  According to Hoefffding inequality, we could get a simple composite view:  \n$$\\Pr(y^ \\star \\ne Y^ \\star) \\lesssim \\exp(-\\alpha k) + C(k), \\alpha = 2(2p_\\pi(x) - 1)^2.\n$$\n\n$C(k)$ is small for $k \\le k_0$ and then increases (e.g., due to distraction or truncation).  The sum is U-shaped, so  \n\n$$\nk^\\star = \\arg\\min_k { \\exp(-\\alpha k) + C(k) }\n$$\n\nexplains the “improve $\\to$ plateau $\\to$ degrade” pattern as $k$ increases.\n\n**Insight 3: earlier subtasks affect later model and $k$**\n\nLet $z$ be upstream output appended to the context; the downstream stage uses $(x, z)$. Then:  \n$$\n\\gamma_\\pi(x, z) \\ge \\gamma_\\pi(x), \\qquad p_\\pi(x, z) \\ge p_\\pi(x),\n$$\nso the downstream optimal number of repeats drops: $k^\\star(x, z) < k^\\star(x)$.  \n\nMoreover, informative $z$ can make a smaller downstream model adequate:  \n$$\n\\gamma_{\\pi_{\\text{small}}}(x, z) \\approx \\gamma_{\\pi_{\\text{large}}}(x),\n$$\nwith much lower cost—changing the preferred downstream model.  \n\nLet $F$ denote the expected task score. If  \n$$\n\\frac{\\partial^2 F}{\\partial B_{\\text{up}} \\partial B_{\\text{down}}} < 0,\n$$\nthen upstream and downstream budgets are substitutes:  more upstream budget reduces the downstream best-response budget and can flip the downstream argmax over models.\n\n\n> W2 The three insights are derived from specific task types and may not generalize to other multi-stage scenarios. The paper doesn't adequately address the boundaries of applicability for these insights.\n\nWe thank the reviewer for raising the important point.\n\nTo assess generalizability, we evaluated our framework and validated the proposed insights across **four diverse task types** spanning **six benchmark datasets**, introduced in Appendix A.4, including:\n\n- **Open-domain question answering** (e.g., 2WikiMultiHopQA, HotpotQA),\n- **Knowledge graph question answering** (e.g., CWQ, WebQSP),\n- **Task Automation** (e.g., Taskbench),\n- **Automated Software Development** (e.g., ChatDev).\n\nThese tasks exhibit different multi-stage structures and model behaviors.  In the **Appendix A.6**, the consistency of the observed trends across these varied domains provides strong empirical support that the proposed insights are **not specific to a single task type** but rather **reflect common patterns** in multi-stage complex tasks.\n\nTo further address the concern about generalization to other tasks, we also adapted the **TextVQA** dataset into a two-stage pipeline: **Stage 1**: Image understanding; **Stage 2**: Question-specific reasoning based on image output. \n\nWe then evaluated the **test-time scaling behavior** and **search efficiency** for multiple models in this multimodal setting. Results are shown below.\n\n**Table: Llama-3.2-11B-Vision-Instruct Summarizes Images**\n\n|Sample|1|5|10|15|20|25|30|35|40|45|50|55|60|65|70|80|90|95|\n|------|--|--|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|\n|Budget|20.0|49.33|86.0|122.67|159.33|196.0|232.67|269.33|306.0|342.67|379.33|416.0|452.67|489.33|526.0|599.33|672.67|709.33|\n|Accuracy|0.6954|0.7605|0.8419|0.8368|0.8396|0.8423|0.8421|0.8351|0.8264|0.8240|0.8289|0.8338|0.8310|0.8355|0.8400|0.8289|0.8254|0.8249|\n\n**Table: Llama-3.2-90B-Vision-Instruct Summarizes Images**\n\n|Sample|1|5|10|15|20|25|30|35|45|50|55|60|65|80|95|\n|------|-|-|--|--|--|--|--|--|--|--|--|--|--|--|--|\n|Budget|178.0|418.0|718.0|1018.0|1318.0|1618.0|1918.0|2218.0|2818.0|3118.0|3418.0|3718.0|4018.0|4918.0|5818.0|\n|Accuracy|0.6782|0.8323|0.8281|0.8237|0.8451|0.8467|0.8583|0.8667|0.8467|0.8667|0.9000|0.9000|0.9200|0.9267|0.9167|\n\n**Table: Llama-3.1-8B-Instruct Reason**\n\n|Sample|1|5|10|15|20|25|30|35|40|45|50|55|60|65|70|80|90|95|\n|--------|------|------|------|------|------|------|------|------|------|------|------|------|------|------|------|------|------|-------|\n|Budget|14.0|56.67|110.0|163.33|216.67|270.0|323.33|376.67|430.0|483.33|536.67|590.0|643.33|696.67|750.0|856.67|963.33|1016.67|\n|Accuracy|0.7741|0.7860|0.7948|0.8022|0.8140|0.8247|0.8242|0.8188|0.8145|0.8147|0.8167|0.8161|0.8156|0.8240|0.8244|0.8149|0.8120|0.8159|\n\n**Table: Llama-3.1-70B-Instruct Reason**\n\n|Sample|1|5|10|15|20|25|30|35|40|45|50|55|60|65|70|80|90|95|\n|--------|------|------|------|-------|-------|------|-------|-------|------|-------|-------|------|-------|-------|------|-------|-------|-------|\n|Budget|138.0|511.33|978.0|1444.67|1911.33|2378.0|2844.67|3311.33|3778.0|4244.67|4711.33|5178.0|5644.67|6111.33|6578.0|7511.33|8444.67|8911.33|\n|Accuracy|0.7608|0.7960|0.8285|0.8312|0.8272|0.8261|0.8346|0.8416|0.8471|0.8494|0.8476|0.8422|0.8357|0.8312|0.8294|0.8304|0.8255|0.8232|\n\n**Table: Search Time and Performance across different methods on TextVQA with budget 1000. \\* means failure to get the optimal configuration.** \n\n|Method|SearchTime(hrs)|Performance|\n|---------|-----------------|-----------|\n|Random|205*|0.90|\n|Bayes|50*|0.83|\n|MLCopilot|70*|0.90|\n|AgentHPO|245*|0.90|\n|**Ours**|**65**|**0.93**|\n\n**Key Takeaways:**\n\n* Test-time scaling boosts performance in multimodal tasks.\n* Under budget ≤ 500, **small models (e.g., 8B)** outperform **larger ones (e.g., 70B)**.\n* Our method can generalize to multimodal tasks.\n\n> W3 The FLOPs-based budget conversion assumes linear scaling relationships that may not hold across all model architectures and task types. Also, the choice of base configuration (3B model, 128 prompt tokens, 64 generation tokens) seems arbitrary and may bias results.\n\nWe justify the linear budget scaling assumption in Appendix A.2. Intuitively, FLOPs scale linearly with model size $k$ and token length $m$—i.e., $k \\times m$ increase leads to proportional compute growth.\n\nThe base config (3B, 128 prompt + 64 gen tokens) was chosen to reflect realistic deployment. The 3B size typifies SLMs, and the token lengths were empirically verified across all used benchmarks and task types to meet minimal needs. This compact yet sufficient setup serves as a standard unit for flexible budget allocation.\n\nAll methods use the same base config for fair and unbiased comparison. Though other base configs may yield different absolute budget values, our goal is to optimize allocation under a fixed total. We pick a minimal, representative unit—like choosing km vs. miles as a distance metric—valid as long as consistently applied.\n\n> W4 It seems only GPT-o3-mini is used as the agent LLM, the paper could benefit from applying the method on more LLMs to show its generalizability.\n\nWe appreciate the reviewer’s suggestion to evaluate the generalizability of our method across different agent LLMs.\n\nIn our main experiments, we used **GPT-o3-mini** as the search agent due to its strong reasoning and planning capabilities. However, we agree that it is important to assess whether our proposed method can adapt to other large language models. To address this, we conducted additional experiments using **DeepSeek-Chat** and **LLaMA 3.3-70B** as alternative agent LLMs. The results are as follows:\n\n**Table: Search Time on 2Wiki Dataset (in hours)**, -- means failure to find the best configuration \n\n|Method|GPT-o3-mini|DeepSeek-R1|Llama3.370B|\n|------|-----------|-----------|-----------|\n|LLM_ZS|12.5|12.5|--|\n|MLCopilot|12.5|12.5|--|\n|AgentTTS|8.3|16|--|\n|**Ours**|**2.5**|**3**|--|\n\nExperiments show that both our method and the baselines generalize well to DeepSeek-R1 but fail to transfer to Llama 3.3 70B. We attribute this to the search procedure’s reliance on the reasoning and planning capabilities of the LLM used for search.\n\nWe hope our clarifications resolve the reviewer's concerns. Let us know if more clarification is needed."}]

Author rebuttal:

Meta-review: The paper frames test-time scaling (TTS) for multi-stage pipelines as a compute-optimal allocation problem over (i) model choice per subtask and (ii) per-subtask sampling budgets under a fixed overall budget. Pilot studies yield three recurring patterns: (I) subtasks prefer different models, (II) returns w.r.t. repeats are unimodal / saturating, and (III) earlier-stage budget affects downstream difficulty and optimal allocations. Building on these, AgentTTS is an LLM-agent that iteratively proposes configurations, reads feedback, updates guidelines, and converges to near-optimal allocations with markedly fewer trials than baselines. Across six datasets (2Wiki, HotpotQA, CWQ, WebQSP, TaskBench, ChatDev) and models from 3B to 72B, AgentTTS matches or exceeds prior methods while reducing search time (often substantially).

Strengths:

- Compute-optimal TTS in multi-stage workflows is both under-explored and practically important.

- The three empirically backed patterns are baked into the agent’s search policy and improve sample efficiency.

- Six datasets, multiple model scales and consistent search-time reductions with comparable or better accuracy.

Weaknesses:

- The “budget unit” and FLOPs normalization remain under-explained in the main text; a train/test mismatch around Fig. 3 required rebuttal clarification.

- The three insights, while useful, are partly intuitive and the theoretical support is sketch-level and appears mainly in rebuttal.

- Core experiments skew QA-style and multimodal and math results appear only post-rebuttal and are limited.

- Success relies on a capable planner LLM and cross-agent robustness was only lightly probed (e.g., DeepSeek ok, Llama-3.3-70B struggled).

- Error bars/seed variability are not pervasive across all tables and some baselines’ search-time entries were initially missing.

Decision:

This paper is the first solid treatment of compute-optimal, multi-stage TTS with interdependent subtask choices. It reliably reaches near-optimal configs faster than strong LLM-HPO baselines. Though not exhaustive, new TextVQA and MATH results, plus cross-agent notes, show the approach is not QA-only. Authors resolved presentation issues, filled missing baselines/stats, and provided theory intuition and cost amortization.

Discussion:

- 8uj2 (score 4): Asked for error bars, baseline adaptation, tasks beyond QA, temperature effects. Authors added CI on 2Wiki (5 seeds), clarified HPO adaptations, provided TextVQA and MATH TTS studies, and a temperature grid; reviewer kept a 4.

- bAdt (score 4): Concerned budget unit/Theorem 1 relevance and result inconsistencies (Fig. 3 vs Table 1; dashed-line meaning). Authors clarified Fig. 3 is training-set trajectories, explained the upper-bound dashed line and how the budget abstraction is used throughout, plus why 50 trials suffice.

- ZDCu (score 4): Wanted clearer innovation, theory depth, and search-cost trade-off. Authors tied insights to algorithmic blocks, gave a convergence/efficiency sketch, and provided an amortization argument (e.g., ~2.5 GPU-hrs search).

- NhkM (score 4): Asked for principled theory, applicability limits, FLOPs linearity & base config choice, and agent-LLM diversity. Authors offered a formalization of the insights, defended FLOPs scaling, and added agent variants (DeepSeek ok; Llama-3.3-70B struggled).
Decision: Accept (poster)