Otimização de Raciocínio em Tempo Real via Política de Otimização Relativa ao Orçamento

Optimizing Anytime Reasoning via Budget Relative Policy Optimization

May 19, 2025

Autores: Penghui Qi, Zichen Liu, Tianyu Pang, Chao Du, Wee Sun Lee, Min Lin

cs.AI

Resumo

A escalabilidade do cálculo em tempo de teste é crucial para aprimorar as capacidades de raciocínio de modelos de linguagem de grande escala (LLMs). As abordagens existentes geralmente empregam aprendizado por reforço (RL) para maximizar uma recompensa verificável obtida ao final de trajetórias de raciocínio. No entanto, tais métodos otimizam apenas o desempenho final sob um orçamento grande e fixo de tokens, o que prejudica a eficiência tanto no treinamento quanto na implantação. Neste trabalho, apresentamos uma nova estrutura, o AnytimeReasoner, para otimizar o desempenho de raciocínio em qualquer momento, visando melhorar a eficiência de tokens e a flexibilidade do raciocínio sob restrições variáveis de orçamento de tokens. Para alcançar isso, truncamos o processo completo de pensamento para se ajustar a orçamentos de tokens amostrados de uma distribuição prévia, compelindo o modelo a resumir a resposta ideal para cada pensamento truncado para verificação. Isso introduz recompensas densas verificáveis no processo de raciocínio, facilitando uma atribuição de crédito mais eficaz na otimização de RL. Em seguida, otimizamos as políticas de pensamento e resumo de maneira desacoplada para maximizar a recompensa cumulativa. Além disso, introduzimos uma nova técnica de redução de variância, a Otimização de Política Relativa ao Orçamento (BRPO), para aumentar a robustez e eficiência do processo de aprendizado ao reforçar a política de pensamento. Resultados empíricos em tarefas de raciocínio matemático demonstram que nosso método supera consistentemente o GRPO em todos os orçamentos de pensamento sob várias distribuições prévias, aprimorando tanto o treinamento quanto a eficiência de tokens.

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.