Generalizing Test-time Compute-optimal Scaling as an Optimizable Graph
Test-Time Scaling (TTS) improves large language models (LLMs) by allocating
additional computation during inference, typically through parallel,
sequential, or hybrid scaling. However, prior studies often assume fixed
collaboration architectures (e.g., topologies) and single-model usage,
overlooking that optimal architectures and model combinations can vary across
tasks. Therefore, we study the novel problem of searching for compute-optimal
model combinations and architectures in TTS under a fixed budget. We formalize
it as a multi-LLM collaboration graph, where nodes encode roles and LLM model
assignments, and edges capture information flow. This problem is challenging
because (i) the combinatorial search space is prohibitively large, and (ii)
task-specific requirements demand tailored designs. To address these, we
reformulate the problem as probabilistic graph optimization and, through pilot
experiments, derive three empirical insights into TTS collaboration graphs.
Guided by these insights, we propose Agent-REINFORCE, an LLM-agent-augmented
framework that mirrors the REINFORCE pipeline by mapping
sampling-gradient-update to sampling-feedback-update, where feedback serves as
a textual gradient to update the probabilistic graph and efficiently search for
optimal multi-LLM collaboration graphs. Experiments show that Agent-REINFORCE
outperforms both traditional and LLM-based baselines in sample efficiency and
search performance, and effectively identifies optimal graphs under joint
objectives of accuracy and inference latency.
69096c6c812eca10f9cc6190
Fali Wang
69096c6c812eca10f9cc6191
Jihai Chen
69096c6c812eca10f9cc6192
Shuhua Yang
69096c6c812eca10f9cc6193
Runxue Bao
69096c6c812eca10f9cc6194
Tianxiang Zhao
69096c6c812eca10f9cc6195
Zhiwei Zhang
69096c6c812eca10f9cc6196
Xianfeng Tang
69096c6c812eca10f9cc6197
Hui Liu
69096c6c812eca10f9cc6198
Qi He
69096c6c812eca10f9cc6199
Suhang Wang
Generalizing Test-time Compute-optimal Scaling as an Optimizable Graph
Test-Time Scaling (TTS) improves large language models (LLMs) by allocating
additional computation during inference, typically through parallel,
sequential, or hybrid scaling. However, prior studies often assume fixed
collaboration architectures (e.g., topologies) and single-model usage,
overlooking that optimal architectures and model combinations can vary across
tasks. Therefore, we study the novel problem of searching for compute-optimal
model combinations and architectures in TTS under a fixed budget. We formalize
it as a multi-LLM collaboration graph, where nodes encode roles and LLM model
assignments, and edges capture information flow. This problem is challenging
because (i) the combinatorial search space is prohibitively large, and (ii)
task-specific requirements demand tailored designs. To address these, we
reformulate the problem as probabilistic graph optimization and, through pilot
experiments, derive three empirical insights into TTS collaboration graphs.
Guided by these insights, we propose Agent-REINFORCE, an LLM-agent-augmented
framework that mirrors the REINFORCE pipeline by mapping
sampling-gradient-update to sampling-feedback-update, where feedback serves as
a textual gradient to update the probabilistic graph and efficiently search for
optimal multi-LLM collaboration graphs. Experiments show that Agent-REINFORCE
outperforms both traditional and LLM-based baselines in sample efficiency and
search performance, and effectively identifies optimal graphs under joint
objectives of accuracy and inference latency.
Agent-REINFORCE optimizes multi-LLM collaboration graphs for test-time scaling, improving sample efficiency and search performance under accuracy and latency constraints.
Test-Time Scaling large language models parallel scaling sequential scaling hybrid scaling collaboration architectures multi-LLM collaboration graph probabilistic graph optimization Agent-REINFORCE REINFORCE pipeline textual gradient