Generalización de la Escalación Óptima de Cómputo en Tiempo de Prueba como un Grafo Optimizable

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

October 29, 2025

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

cs.AI

Resumen

El Escalado en Tiempo de Prueba (TTS) mejora los modelos de lenguaje grande (LLM) mediante la asignación de cómputo adicional durante la inferencia, típicamente a través de escalado paralelo, secuencial o híbrido. Sin embargo, estudios previos a menudo asumen arquitecturas de colaboración fijas (por ejemplo, topologías) y uso de un único modelo, pasando por alto que las arquitecturas y combinaciones de modelos óptimas pueden variar entre tareas. Por lo tanto, estudiamos el novedoso problema de buscar combinaciones de modelos y arquitecturas óptimas en cómputo dentro de TTS bajo un presupuesto fijo. Lo formalizamos como un grafo de colaboración multi-LLM, donde los nodos codifican roles y asignaciones de modelos LLM, y las aristas capturan el flujo de información. Este problema es desafiante porque (i) el espacio de búsqueda combinatorio es prohibitivamente grande, y (ii) los requisitos específicos de la tarea exigen diseños a medida. Para abordarlos, reformulamos el problema como una optimización de grafos probabilística y, mediante experimentos piloto, derivamos tres ideas empíricas sobre los grafos de colaboración TTS. Guiados por estas ideas, proponemos Agent-REINFORCE, un marco aumentado con agentes LLM que refleja la pipeline REINFORCE al mapear muestreo-gradiente-actualización a muestreo-retroalimentación-actualización, donde la retroalimentación sirve como un gradiente textual para actualizar el grafo probabilístico y buscar eficientemente grafos de colaboración multi-LLM óptimos. Los experimentos muestran que Agent-REINFORCE supera a las líneas base tradicionales y basadas en LLM en eficiencia de muestreo y rendimiento de búsqueda, e identifica efectivamente grafos óptimos bajo objetivos conjuntos de precisión y latencia de inferencia.

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.