Optimisation du calcul au moment du test via un réglage fin par méta-apprentissage par renforcement

Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning

March 10, 2025

Auteurs: Yuxiao Qu, Matthew Y. R. Yang, Amrith Setlur, Lewis Tunstall, Edward Emanuel Beeching, Ruslan Salakhutdinov, Aviral Kumar

cs.AI

Résumé

L'entraînement des modèles à utiliser efficacement les ressources de calcul au moment du test est crucial pour améliorer les performances de raisonnement des LLM (Large Language Models). Les méthodes actuelles y parviennent principalement par le biais d'un ajustement fin sur des traces de recherche ou en exécutant un apprentissage par renforcement (RL) avec une récompense binaire (0/1), mais ces approches exploitent-elles efficacement les ressources de calcul au moment du test ? Continueraient-elles à s'adapter à mesure que le budget augmente ? Dans cet article, nous tentons de répondre à ces questions. Nous formalisons le problème de l'optimisation des ressources de calcul au moment du test comme un problème de méta-apprentissage par renforcement (meta-RL), offrant ainsi une perspective théorique sur l'allocation de ces ressources. Cette perspective nous permet de considérer le flux de sortie long du LLM comme une série d'épisodes exécutés au moment du test et nous conduit à utiliser une notion de regret cumulé sur les tokens de sortie comme mesure de l'efficacité des ressources de calcul. De la même manière que les algorithmes de RL parviennent à équilibrer exploration et exploitation pendant l'entraînement, minimiser le regret cumulé permettrait également d'atteindre le meilleur équilibre entre exploration et exploitation dans le flux de tokens. Bien que nous montrions que les modèles de pointe ne minimisent pas le regret, il est possible d'y parvenir en maximisant une récompense dense supplémentaire, combinée à la récompense binaire (0/1) du RL. Cette récompense supplémentaire correspond aux « progrès » réalisés par chaque bloc suivant dans le flux de sortie, quantifiés par le changement dans la probabilité de succès final. En nous appuyant sur ces insights, nous développons le Meta Reinforcement Fine-Tuning (MRT), une nouvelle classe de méthodes d'ajustement fin pour optimiser les ressources de calcul au moment du test. MRT permet d'obtenir un gain relatif de 2 à 3 fois en performance et environ 1,5 fois en efficacité token pour le raisonnement mathématique par rapport au RL basé sur la récompense binaire.

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.