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

Optimizing Anytime Reasoning via Budget Relative Policy Optimization

May 19, 2025

Авторы: Penghui Qi, Zichen Liu, Tianyu Pang, Chao Du, Wee Sun Lee, Min Lin

cs.AI

Аннотация

Масштабирование вычислительных ресурсов на этапе тестирования имеет решающее значение для улучшения способности к рассуждению у крупных языковых моделей (LLM). Существующие подходы обычно используют обучение с подкреплением (RL) для максимизации проверяемой награды, получаемой в конце цепочек рассуждений. Однако такие методы оптимизируют только итоговую производительность при большом и фиксированном бюджете токенов, что снижает эффективность как в обучении, так и в эксплуатации. В данной работе мы представляем новую структуру, AnytimeReasoner, для оптимизации производительности рассуждений в режиме реального времени, которая направлена на повышение эффективности использования токенов и гибкости рассуждений при различных ограничениях на бюджет токенов. Для достижения этого мы сокращаем полный процесс мышления, чтобы он укладывался в выборочные бюджеты токенов из априорного распределения, заставляя модель суммировать оптимальный ответ для каждого сокращенного процесса мышления для проверки. Это вводит проверяемые плотные награды в процесс рассуждений, способствуя более эффективному распределению заслуг в оптимизации RL. Затем мы оптимизируем политики мышления и суммирования раздельно, чтобы максимизировать совокупную награду. Кроме того, мы вводим новую технику снижения дисперсии, Budget Relative Policy Optimization (BRPO), для повышения устойчивости и эффективности процесса обучения при усилении политики мышления. Эмпирические результаты в задачах математического рассуждения показывают, что наш метод стабильно превосходит GRPO при всех бюджетах мышления в различных априорных распределениях, улучшая как обучение, так и эффективность использования токенов.

English

Scaling test-time compute is crucial for enhancing the reasoning capabilities of large language models (LLMs). Existing approaches typically employ reinforcement learning (RL) to maximize a verifiable reward obtained at the end of reasoning traces. However, such methods optimize only the final performance under a large and fixed token budget, which hinders efficiency in both training and deployment. In this work, we present a novel framework, AnytimeReasoner, to optimize anytime reasoning performance, which aims to improve token efficiency and the flexibility of reasoning under varying token budget constraints. To achieve this, we truncate the complete thinking process to fit within sampled token budgets from a prior distribution, compelling the model to summarize the optimal answer for each truncated thinking for verification. This introduces verifiable dense rewards into the reasoning process, facilitating more effective credit assignment in RL optimization. We then optimize the thinking and summary policies in a decoupled manner to maximize the cumulative reward. Additionally, we introduce a novel variance reduction technique, Budget Relative Policy Optimization (BRPO), to enhance the robustness and efficiency of the learning process when reinforcing the thinking policy. Empirical results in mathematical reasoning tasks demonstrate that our method consistently outperforms GRPO across all thinking budgets under various prior distributions, enhancing both training and token efficiency.