Generalizando a Escala Ótima de Computação em Tempo de Teste como um Grafo Otimizável

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

Resumo

O Escalonamento no Momento do Teste (TTS) melhora os modelos de linguagem de grande escala (LLMs) ao alocar computação adicional durante a inferência, tipicamente através de escalonamento paralelo, sequencial ou híbrido. No entanto, estudos anteriores frequentemente assumem arquiteturas de colaboração fixas (por exemplo, topologias) e uso de modelo único, negligenciando o facto de que as arquiteturas e combinações de modelos ótimas podem variar consoante a tarefa. Portanto, estudamos o novo problema de procurar combinações de modelos e arquiteturas computacionalmente ótimas no TTS sob um orçamento fixo. Formalizamo-lo como um grafo de colaboração multi-LLM, onde os nós codificam funções e atribuições de modelos LLM, e as arestas capturam o fluxo de informação. Este problema é desafiador porque (i) o espaço de busca combinatória é proibitivamente grande, e (ii) os requisitos específicos da tarefa exigem projetos personalizados. Para resolver estes desafios, reformulamos o problema como uma otimização probabilística de grafos e, através de experiências piloto, derivamos três perspetivas empíricas sobre os grafos de colaboração TTS. Guiados por estas perspetivas, propomos o Agent-REINFORCE, uma estrutura aumentada por agente-LLM que espelha o *pipeline* REINFORCE ao mapear amostragem-gradiente-atualização para amostragem-*feedback*-atualização, onde o *feedback* serve como um gradiente textual para atualizar o grafo probabilístico e procurar eficientemente grafos de colaboração multi-LLM ótimos. Experiências mostram que o Agent-REINFORCE supera as linhas de base tradicionais e baseadas em LLM em eficiência de amostragem e desempenho de busca, e identifica eficazmente grafos ótimos sob objetivos conjuntos de precisão e latência de inferência.

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.