Обобщение масштабирования с оптимальными вычислительными затратами на этапе тестирования как оптимизируемого графа

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

October 29, 2025

Авторы: Fali Wang, Jihai Chen, Shuhua Yang, Runxue Bao, Tianxiang Zhao, Zhiwei Zhang, Xianfeng Tang, Hui Liu, Qi He, Suhang Wang

cs.AI

Аннотация

Метод масштабирования на этапе тестирования (Test-Time Scaling, TTS) улучшает работу больших языковых моделей (LLM) за счёт распределения дополнительных вычислительных ресурсов в процессе вывода, обычно с помощью параллельного, последовательного или гибридного масштабирования. Однако предыдущие исследования часто предполагают фиксированные архитектуры коллаборации (например, топологии) и использование единой модели, упуская из виду, что оптимальные архитектуры и комбинации моделей могут различаться в зависимости от задачи. Поэтому мы исследуем новую проблему поиска вычислительно-оптимальных комбинаций моделей и архитектур в TTS при фиксированном бюджете. Мы формализуем её как граф коллаборации множества LLM, где узлы кодируют роли и назначения моделей LLM, а рёбра захватывают поток информации. Эта проблема является сложной, поскольку (i) комбинаторное пространство поиска непомерно велико и (ii) требования, специфичные для задачи, требуют индивидуальных проектов. Для решения этих проблем мы переформулируем задачу как вероятностную оптимизацию графа и, с помощью пилотных экспериментов, выводим три эмпирических инсайта о графах коллаборации TTS. Руководствуясь этими инсайтами, мы предлагаем Agent-REINFORCE — фреймворк, дополненный LLM-агентами, который воспроизводит конвейер REINFORCE, отображая последовательность "сэмплирование-градиент-обновление" в "сэмплирование-фидбэк-обновление", где фидбэк служит текстовым градиентом для обновления вероятностного графа и эффективного поиска оптимальных графов коллаборации множества LLM. Эксперименты показывают, что Agent-REINFORCE превосходит как традиционные, так и LLM-базлайны по эффективности использования сэмплов и производительности поиска, а также эффективно находит оптимальные графы при совместных целях точности и задержки вывода.

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.