Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Qu, Yuxiao, Yang, Matthew Y. R., Setlur, Amrith, Tunstall, Lewis, Beeching, Edward Emanuel, Salakhutdinov, Ruslan, Kumar, Aviral
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866917950093524992
author Qu, Yuxiao
Yang, Matthew Y. R.
Setlur, Amrith
Tunstall, Lewis
Beeching, Edward Emanuel
Salakhutdinov, Ruslan
Kumar, Aviral
author_facet Qu, Yuxiao
Yang, Matthew Y. R.
Setlur, Amrith
Tunstall, Lewis
Beeching, Edward Emanuel
Salakhutdinov, Ruslan
Kumar, Aviral
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2503_07572
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning
Qu, Yuxiao
Yang, Matthew Y. R.
Setlur, Amrith
Tunstall, Lewis
Beeching, Edward Emanuel
Salakhutdinov, Ruslan
Kumar, Aviral
Machine Learning
Artificial Intelligence
Computation and Language
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.
title Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2503.07572