Optimisation du raisonnement en temps réel par la politique d'optimisation relative au budget

Optimizing Anytime Reasoning via Budget Relative Policy Optimization

May 19, 2025

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

cs.AI

Résumé

L'optimisation du calcul au moment du test est cruciale pour améliorer les capacités de raisonnement des grands modèles de langage (LLM). Les approches existantes utilisent généralement l'apprentissage par renforcement (RL) pour maximiser une récompense vérifiable obtenue à la fin des traces de raisonnement. Cependant, ces méthodes optimisent uniquement la performance finale sous un budget de jetons fixe et important, ce qui nuit à l'efficacité tant en formation qu'en déploiement. Dans ce travail, nous présentons un nouveau cadre, AnytimeReasoner, pour optimiser la performance de raisonnement à tout moment, visant à améliorer l'efficacité des jetons et la flexibilité du raisonnement sous des contraintes de budget de jetons variables. Pour y parvenir, nous tronquons le processus de réflexion complet pour qu'il s'insère dans des budgets de jetons échantillonnés à partir d'une distribution a priori, obligeant le modèle à résumer la réponse optimale pour chaque réflexion tronquée afin de vérification. Cela introduit des récompenses denses vérifiables dans le processus de raisonnement, facilitant une attribution de crédit plus efficace dans l'optimisation RL. Nous optimisons ensuite les politiques de réflexion et de résumé de manière découplée pour maximiser la récompense cumulative. De plus, nous introduisons une nouvelle technique de réduction de variance, l'Optimisation Relative des Politiques par Budget (BRPO), pour améliorer la robustesse et l'efficacité du processus d'apprentissage lors du renforcement de la politique de réflexion. Les résultats empiriques dans des tâches de raisonnement mathématique démontrent que notre méthode surpasse systématiquement GRPO pour tous les budgets de réflexion sous diverses distributions a priori, améliorant à la fois la formation et l'efficacité des jetons.

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.