Optimización de Razonamiento en Tiempo Real mediante Optimización de Políticas Relativas al Presupuesto

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

Resumen

La escalabilidad del cómputo en tiempo de prueba es crucial para mejorar las capacidades de razonamiento de los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés). Los enfoques existentes suelen emplear aprendizaje por refuerzo (RL, por sus siglas en inglés) para maximizar una recompensa verificable obtenida al final de las trazas de razonamiento. Sin embargo, estos métodos optimizan únicamente el rendimiento final bajo un presupuesto de tokens grande y fijo, lo que dificulta la eficiencia tanto en el entrenamiento como en la implementación. En este trabajo, presentamos un marco novedoso, AnytimeReasoner, para optimizar el rendimiento de razonamiento en cualquier momento, con el objetivo de mejorar la eficiencia de los tokens y la flexibilidad del razonamiento bajo restricciones variables de presupuesto de tokens. Para lograrlo, truncamos el proceso de pensamiento completo para que se ajuste a presupuestos de tokens muestreados de una distribución previa, obligando al modelo a resumir la respuesta óptima para cada pensamiento truncado con fines de verificación. Esto introduce recompensas densas verificables en el proceso de razonamiento, facilitando una asignación de crédito más efectiva en la optimización de RL. Luego, optimizamos las políticas de pensamiento y resumen de manera desacoplada para maximizar la recompensa acumulada. Además, introducimos una técnica novedosa de reducción de varianza, la Optimización de Políticas Relativas al Presupuesto (BRPO, por sus siglas en inglés), para mejorar la robustez y eficiencia del proceso de aprendizaje al reforzar la política de pensamiento. Los resultados empíricos en tareas de razonamiento matemático demuestran que nuestro método supera consistentemente a GRPO en todos los presupuestos de pensamiento bajo diversas distribuciones previas, mejorando tanto la eficiencia del entrenamiento como la de los 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.