Generalizzare il ridimensionamento ottimale del calcolo al tempo di test come grafo ottimizzabile

Generalizing Test-time Compute-optimal Scaling as an Optimizable Graph

October 29, 2025

Autori: Fali Wang, Jihai Chen, Shuhua Yang, Runxue Bao, Tianxiang Zhao, Zhiwei Zhang, Xianfeng Tang, Hui Liu, Qi He, Suhang Wang

cs.AI

Abstract

Il Test-Time Scaling (TTS) migliora i grandi modelli linguistici (LLM) allocando calcolo aggiuntivo durante l'inferenza, tipicamente attraverso scaling parallelo, sequenziale o ibrido. Tuttavia, gli studi precedenti spesso assumono architetture di collaborazione fisse (ad esempio, topologie) e un utilizzo a modello singolo, trascurando il fatto che le architetture ottimali e le combinazioni di modelli possano variare a seconda dei task. Pertanto, studiamo il nuovo problema della ricerca di combinazioni di modelli e architetture computazionalmente ottimali nel TTS sotto un budget fisso. Lo formalizziamo come un grafo di collaborazione multi-LLM, in cui i nodi codificano ruoli e assegnazioni di modelli LLM, e gli archi catturano il flusso informativo. Questo problema è impegnativo perché (i) lo spazio combinatorio di ricerca è proibitivamente grande e (ii) i requisiti specifici del task richiedono progettazioni su misura. Per affrontarli, riformuliamo il problema come un'ottimizzazione probabilistica di grafi e, attraverso esperimenti pilota, ricaviamo tre intuizioni empiriche sui grafi di collaborazione TTS. Guidati da queste intuizioni, proponiamo Agent-REINFORCE, un framework potenziato da agenti LLM che rispecchia la pipeline REINFORCE mappando sampling-gradiente-aggiornamento in sampling-feedback-aggiornamento, dove il feedback funge da gradiente testuale per aggiornare il grafo probabilistico e cercare efficientemente grafi di collaborazione multi-LLM ottimali. Gli esperimenti mostrano che Agent-REINFORCE supera le baseline tradizionali e basate su LLM sia nell'efficienza campionaria che nelle prestazioni di ricerca, e identifica efficacemente grafi ottimali sotto obiettivi congiunti di accuratezza e latenza d'inferenza.

English

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.