Малый RL-контроллер, большая языковая модель: адаптивная выборка под управлением RL для масштабирования на этапе тестирования

Аннотация

Масштабирование во время тестирования улучшает производительность рассуждений больших языковых моделей, но влечет за собой значительные затраты как на общие вычислительные ресурсы, так и на задержку. Существующие методы адаптивной выборки частично смягчают эту проблему, динамически принимая решение о прекращении выборки, однако они обычно опираются на эвристические правила или основаны на предположениях о распределении. В данной работе мы формулируем адаптивную выборку как марковский процесс принятия решений (MDP). Мы обучаем легковесный контроллер выборки с помощью обучения с подкреплением (RL) для совместного обеспечения баланса между правильностью ответов, задержкой и вычислительными затратами. На каждом раунде контроллер принимает решение о прекращении выборки или о получении дополнительных образцов. Наш метод является легковесным, поскольку он полагается только на статистику окончательных ответов, и может быть обучен и развернут на CPU. Мы также показываем, что полученная структура допускает интерпретацию как лагранжева релаксация задачи условной оптимизации с явными ограничениями на бюджет. Эксперименты с сильными базовыми методами, такими как ASC и ESC, показывают, что наш метод достигает улучшенных компромиссов между правильностью ответов, количеством раундов выборки и общим количеством требуемых образцов.

English

Test-time scaling improves the reasoning performance of large language models but incurs substantial cost in both total computation and latency. Existing adaptive sampling methods partially mitigate this issue by dynamically deciding when to stop sampling, yet they typically rely on heuristic rules or rely on distribution assumptions. In this work, we formulate adaptive sampling as a Markov decision process (MDP). We train a lightweight sampling controller with reinforcement learning (RL) to jointly balance answer correctness, latency, and computation cost. At each round, the controller decides to stop sampling or to acquire additional samples. Our method is lightweight which only relies on statistics of final answers, and can be trained and deployed on CPU. We further show that the resulting framework admits an interpretation as the Lagrangian relaxation of a constrained optimization problem with explicit budget constraints. Experiments against strong baselines such as ASC and ESC show that our method achieves improved trade-offs among answer correctness, sampling rounds, and total samples required.