Оптимизация вычислительных затрат во время тестирования с помощью мета-обучения с подкреплением для тонкой настройки

Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning

March 10, 2025

Авторы: Yuxiao Qu, Matthew Y. R. Yang, Amrith Setlur, Lewis Tunstall, Edward Emanuel Beeching, Ruslan Salakhutdinov, Aviral Kumar

cs.AI

Аннотация

Обучение моделей эффективному использованию вычислительных ресурсов во время тестирования имеет решающее значение для улучшения способности крупных языковых моделей (LLM) к рассуждению. Современные методы в основном достигают этого путем тонкой настройки на основе траекторий поиска или применения обучения с подкреплением (RL) с бинарной наградой (0/1), но эффективно ли эти подходы используют вычислительные ресурсы во время тестирования? Будут ли они масштабироваться с увеличением доступных ресурсов? В данной работе мы пытаемся ответить на эти вопросы. Мы формализуем задачу оптимизации вычислительных ресурсов во время тестирования как задачу мета-обучения с подкреплением (meta-RL), что дает принципиальный подход к распределению этих ресурсов. Такой подход позволяет рассматривать длинный поток вывода LLM как состоящий из нескольких эпизодов, выполняемых во время тестирования, и приводит нас к использованию понятия кумулятивного сожаления (cumulative regret) по токенам вывода как способа измерения эффективности использования вычислительных ресурсов. Подобно тому, как алгоритмы RL находят оптимальный баланс между исследованием и эксплуатацией во время обучения, минимизация кумулятивного сожаления также обеспечивает наилучший баланс между исследованием и эксплуатацией в потоке токенов. Хотя мы показываем, что современные модели не минимизируют сожаление, этого можно достичь, максимизируя плотную бонусную награду в сочетании с бинарной наградой RL (0/1). Этот бонус представляет собой «прогресс», достигнутый каждым последующим блоком в потоке вывода, измеряемый изменением вероятности конечного успеха. Используя эти идеи, мы разрабатываем Meta Reinforcement Fine-Tuning (MRT) — новый класс методов тонкой настройки для оптимизации вычислительных ресурсов во время тестирования. MRT приводит к относительному улучшению производительности в 2–3 раза и примерно к 1,5-кратному увеличению эффективности использования токенов в задачах математического рассуждения по сравнению с RL, основанным на бинарной награде.

English

Training models to effectively use test-time compute is crucial for improving the reasoning performance of LLMs. Current methods mostly do so via fine-tuning on search traces or running RL with 0/1 outcome reward, but do these approaches efficiently utilize test-time compute? Would these approaches continue to scale as the budget improves? In this paper, we try to answer these questions. We formalize the problem of optimizing test-time compute as a meta-reinforcement learning (RL) problem, which provides a principled perspective on spending test-time compute. This perspective enables us to view the long output stream from the LLM as consisting of several episodes run at test time and leads us to use a notion of cumulative regret over output tokens as a way to measure the efficacy of test-time compute. Akin to how RL algorithms can best tradeoff exploration and exploitation over training, minimizing cumulative regret would also provide the best balance between exploration and exploitation in the token stream. While we show that state-of-the-art models do not minimize regret, one can do so by maximizing a dense reward bonus in conjunction with the outcome 0/1 reward RL. This bonus is the ''progress'' made by each subsequent block in the output stream, quantified by the change in the likelihood of eventual success. Using these insights, we develop Meta Reinforcement Fine-Tuning, or MRT, a new class of fine-tuning methods for optimizing test-time compute. MRT leads to a 2-3x relative gain in performance and roughly a 1.5x gain in token efficiency for math reasoning compared to outcome-reward RL.