PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning

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

arXiv:2503.07572v1cs.LGcs.AIcs.CL

TL;DR

The paper asks whether existing methods use test-time compute efficiently and continue discovering solutions as budgets grow. It formalizes compute allocation as meta-RL, models reasoning as episodic adaptation, and introduces MRT with dense progress rewards. MRT improves final performance, regret, and token efficiency relative to outcome-reward approaches in the reported math-reasoning settings.

  • Problem

    Existing search-trace fine-tuning and outcome-reward RL may not efficiently use test-time compute or scale discovery to harder problems.

  • Method

    MRT formulates test-time compute optimization as minimizing cumulative regret over episodic output streams, using a dense reward based on each episode’s improvement in success probability.

  • Results

    MRT demonstrates improved final performance, lower regret, and better extrapolation to higher test-time budgets than outcome-reward approaches.

  • Takeaways & Limitations

    Progress-sensitive rewards provide a way to train LLMs to allocate test-time tokens while balancing exploration and exploitation.

  • Takeaways & Limitations

    The paper leaves open how to instantiate an iterative training procedure that maximizes progress more directly.

Abstract

from arXiv · show

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.

1. Introduction

The paper asks whether current methods use test-time compute efficiently and scale to harder problems, then frames the issue as balancing exploration and exploitation across output episodes. MRT minimizes cumulative regret with dense progress rewards alongside outcome rewards, improving performance and token efficiency.

  • Motivation: Current methods use search-trace fine-tuning or outcome-reward RL to produce long reasoning traces, but their test-time compute use remains an open question.These approaches include reflection, planning, linearized search, and 0/1 correctness rewards.
  • Method: MRT treats the output stream as multiple episodes and learns a meta-RL strategy that balances exploration with exploitation.Episodes can try different approaches, verify or revise prior attempts, and eventually produce an answer.
  • Motivation: Outcome rewards at the end do not incentivize intermediate progress, producing unnecessarily long traces and unstable progress on hard problems.Standard fine-tuning can reinforce traces that make subpar progress but eventually succeed.
  • Method: Cumulative regret measures the gap between the model’s success likelihood and an oracle comparator over the token budget, providing a target for efficient test-time computation.Minimizing regret encourages steady progress while avoiding premature commitment and excessive exploration.
  • Findings: State-of-the-art outcome-reward RL models fail to make steady progress as episodes increase, especially on harder unseen problems.FLOPs-matched evaluation finds that fewer episodes with majority voting can sometimes be more effective.
  • Findings: 1.6-1.7x token efficiency improvement is achieved by MRT over STaR and GRPO when fine-tuning Llama3.1 models for backtracking.In another setting, MRT consistently outperforms outcome-reward RL across aggregated math benchmarks at the 1.5B scale.
  • Findings: MRT lowers regret and improves steady progress even when extrapolating to 2x larger token budgets than used in training.It reduces output length while boosting accuracy, unlike typical length-constraining methods.

2. Related Work

Related work scales test-time compute through search, long chains of thought, verifiers, and distillation, while highlighting limitations of unfamiliar traces and length-based control. MRT instead connects compute allocation to regret minimization and dense progress rewards.

  • Scaling test-time compute: Search-trace fine-tuning teaches LLMs to simulate in-context search, but unfamiliar traces can cause memorization and limit gains.Earlier approaches also use verifiers for best-of-N or beam search.
  • Reasoning with long chains of thought: Outcome-reward RL supports long chains of thought for search, planning, introspection, and correction, but long outputs may use too many tokens.Length penalties have been explored, although length may not directly correlate with accuracy.
  • Reward design: MRT’s dense reward is inspired by regret minimization and does not require an LLM judge, unlike some length- or judge-based reward designs.Trace distillation through supervised fine-tuning is described as orthogonal to MRT’s focus on improving RL directly.
  • Meta RL: Meta-RL work motivates treating new problems as tasks requiring adaptation through initial test-time episodes.The paper applies this perspective to how LLMs use test-time compute.

3. Preliminaries and Background

The paper models an LLM as an RL policy trained on problem–solution data and evaluated by whether its generated token stream earns a task reward. Meta-RL differs by training across tasks and allowing test-time adaptation episodes.

  • Problem setup: The setup assumes a reward function r(x, ·) that evaluates any token stream, such as checking whether a mathematical answer is correct.The training data consists of problems paired with oracle solution traces ending in correct answers.
  • Problem setup: The LLM is modeled as an RL policy π trained on problem–solution pairs to produce token streams that achieve high reward on test problems.The objective is evaluated on problems sampled from a test distribution.
  • Meta RL primer: Meta-RL trains across tasks with different rewards or dynamics and evaluates adaptation through episodes executed on the test task.This differs from evaluating only zero-shot performance.

4. Problem Formulation: Optimizing Test-Time Compute as Meta RL

The paper formulates test-time compute allocation as a budget-constrained meta-RL problem: output streams are segmented into episodes, and effective policies balance exploration and exploitation while making steady progress. Cumulative regret formalizes this progress relative to an optimal budget-agnostic comparator.

  • 4.1. Optimizing Test-Time Compute: The optimization objective maximizes expected test reward subject to an average output-length budget C_0 for every test problem.The evaluation budget may exceed the typical length of a correct response.
  • 4.1. Optimizing Test-Time Compute: Extra tokens can support indirect discovery operations such as backtracking, but the training budget need not equal the deployment budget.Backtracking may help the model identify steps to avoid before reaching a correct answer.
  • 4.1. Optimizing Test-Time Compute: Fixed-budget outcome-reward training can encourage redundant responses, fail under smaller deployment budgets, and provide no incentive for intermediate progress.Forcing early termination can instead cause premature commitment on complex problems.
  • 4.1. Optimizing Test-Time Compute: A budget-agnostic strategy must trade off exploration and exploitation so that behavior remains effective across multiple sufficiently large test-time budgets.The proposed formulation aims to avoid commitment to a particular training budget.
  • 4.2. Characterizing Optimal Use of Test-Time Compute: The output stream is segmented into meaningful contiguous blocks or episodes, turning test-time generation into an adaptation procedure and a meta-RL problem.Episodes may represent attempts, verification and generation cycles, or search paths separated by backtracking.
  • 4.2. Characterizing Optimal Use of Test-Time Compute: The framework requires steady progress while balancing verification, alternative strategies, and other exploratory actions against extending a committed approach.A fixed episode count is assumed for exposition, with excess natural episodes merged when necessary.
  • 4.2. Characterizing Optimal Use of Test-Time Compute: Cumulative regret measures the accumulated difference between an optimal comparator’s expected reward and a meta-prover’s reward after conditioning on prior episodes.The meta-prover may be the same policy or a separate policy that estimates the best possible answer.
  • 4.2. Characterizing Optimal Use of Test-Time Compute: Lower regret growth indicates more meaningful progress as the test-time budget increases.Increasing or large regret indicates that episodes did not make meaningful progress.

5. Case Study: Analyzing SoTA DeepSeek-R1

The DeepSeek-R1 case study evaluates whether additional reasoning episodes steadily improve solution accuracy and reduce cumulative regret. Results show that progress is reliable with few episodes but can stall or degrade when problems require many episodes.

  • Evaluation: [maj@p]j truncates reasoning after j episodes, samples p immediate solutions, and majority-votes them into one answer.The metric tests whether additional episodes provide meaningful progress before final-answer generation.
  • Evaluation: Figure 3 compares direct pass@k and [maj@p]j while grouping solutions with total episode counts of 6–10, 26–30, and 41–45.Blue points combine five episodes, and accuracy is plotted against test-time tokens and episode index.
  • Findings: With 6–10 episodes, cumulative regret is low and each new episode continuously reduces regret, while [maj@p]j and the direct baseline improve more slowly.The reported comparison concerns the top row of Figure 3.
  • Findings: For 41–45 episodes and other high-episode cases, accuracy does not consistently increase and sometimes degrades with subsequent episodes.Sequential reasoning can lose to both majority-voting and direct baselines when solutions contain more episodes.

6. The Meta Reinforcement Finetuning (MRT) Paradigm

MRT frames test-time reasoning as meta-RL over sequential episodes and uses progress as a dense reward to approximate cumulative-regret minimization. The design aims to balance exploration and exploitation while improving compute efficiency across budgets.

  • Core paradigm: MRT fine-tunes LLMs to directly optimize a surrogate for cumulative regret over test-time episodes.The framework treats steady progress as the target for budget-agnostic reasoning.
  • Motivation: Outcome-reward optimization treats intermediate episodes equally when eventual success occurs, failing to account for each episode’s contribution toward solving the problem.This can waste tokens on easy problems and leave insufficient budget for discovering solutions to hard problems.
  • Surrogate objective: MRT uses episode-level progress in a meta-RL surrogate because cumulative regret cannot be directly optimized without the unknown optimal comparator policy.The surrogate follows the idea that increasing the meta-prover’s estimated success likelihood with more episodes indicates lower regret.
  • Dense reward: The progress bonus measures the change in the meta-prover’s estimated success likelihood after adding an episode.It is defined relative to the prior context and is used as a dense reward component.
  • Novelty: Unlike action-level supervision, MRT evaluates progress across complete reasoning episodes rather than individual actions within an episode.This aligns the training signal with the meta-step structure of test-time reasoning.

7. Practical Instantiations: Dense Rewards for Optimizing Test-Time Compute

MRT is instantiated with open-ended thought blocks and backtracking search, using either STaR or online RL to select and optimize traces with dense progress rewards. Both variants evaluate intermediate prefixes by terminating reasoning and estimating solution accuracy.

  • Open-ended parameterization: Open-ended parameterization encloses unconstrained reasoning episodes between <think> and </think> markers.The objective is optimized with both STaR and RL.
  • Backtracking search: Backtracking search alternates complete solution attempts with error detection and revision until reaching a correct answer.The model identifies errors at specific steps and returns to earlier steps to correct them.
  • Training variants: MRT has STaR and RL variants that maximize dense rewards based on progress.The STaR variant uses self-generated traces for supervised fine-tuning, while the RL variant optimizes progress-enhanced rewards online.
  • STaR: The STaR variant segments generated traces into episodes and retains traces that achieve maximum progress and eventually produce correct solutions.It repeats self-training over filtered, high-quality traces.
  • RL: The RL variant estimates rewards for thought prefixes with a meta-prover, then samples rollouts that either terminate after the prefix or continue reasoning.Training optimizes the progress-augmented reward rather than only the binary outcome reward.

8. Experimental Evaluation

MRT improves both mathematical reasoning performance and token efficiency over outcome-reward training, while better progress management supports extrapolation to larger test-time budgets. Length penalties can improve efficiency but sacrifice peak accuracy, whereas dense progress rewards and budget curricula offer a better trade-off.

  • MRT consistently outperforms outcome-reward training, achieving state-of-the-art results for its model size.
  • 2-3x larger relative performance gains are obtained with MRT than with outcome-reward GRPO.
  • MRT preserves out-of-distribution AMC 2023 performance after narrow AIME fine-tuning, unlike the comparison outcome-reward approach.
  • Simple length penalties improve token efficiency but reduce pass@1 accuracy, while MRT’s dense rewards increase performance and slightly reduce response length.
  • 5x fewer tokens on AIME 2024 and around 4x fewer tokens on MATH500 are required by MRT to match the base model’s performance.
  • MRT improves token efficiency over outcome-reward training by 1.2-1.6x, with broader summaries reporting approximately 1.5x scaling over GRPO.
  • MRT maintains lower normalized regret than base and outcome-reward models within and beyond the 16K training budget, indicating stronger progress at larger budgets.
  • An 8K budget phase yields lower regret and greater progress than later 16K training, supporting iterative budget curricula alongside dense rewards.

9. Discussion, Future Work, and Conclusion

The paper frames test-time compute optimization as cumulative-regret minimization and proposes MRT to reward progress during generation. It reports improved performance and extrapolation, while identifying open questions about the meta-prover, base-model behaviors, rollout implementation, and compute tradeoffs.

  • Findings: Outcome-reward RL often fails to optimize regret or answer novel questions within the token budget because it cannot reward intermediate progress.The paper attributes this to insufficient discrimination among trajectories under fixed-budget binary rewards.
  • Conclusion: MRT formulates test-time compute optimization as minimizing cumulative regret through a surrogate objective that rewards progress at intermediate episodes.The dense bonus measures improvement in success probability under a best-guess policy conditioned on the current generation.
  • Findings: MRT empirically improves final performance, lowers regret, and extrapolates better to higher test-time budgets.These outcomes are presented as evidence that progress-sensitive rewards improve test-time compute efficacy.
  • Future work: The choice of meta-prover remains unresolved, including whether alternative μ-free dense-reward parameterizations could perform better.The current work uses a policy that produces the best-guess solution conditioned on the thinking trace.
  • Future work: MRT’s effectiveness may depend on broader base-model strategy coverage, computationally efficient branched rollouts, and a better-understood train-time versus test-time compute tradeoff.The paper notes that current base models produce relatively narrow behaviors and that its implementation increases policy-gradient variance.

Appendices

The appendices extend MRT to backtracking search, where episodes alternate between solving and error correction. They describe warmstart SFT requirements, progress-based filtering and rewards, and supporting visualizations of training and rollout behavior.

  • Backtracking search: Backtracking search alternates solution attempts with error detection and selection of an earlier point from which to continue.This parameterization explicitly encourages error-correction and backtracking behavior.
  • Backtracking search: MRT traverses two paths sharing a prefix, using backtracking to generate a corrected attempt rather than stitching unrelated traces.The accompanying construction compares MRT-style traversal with RISE-style approaches.
  • Training: The STaR variant filters on-policy traces using correctness of the corrected attempt and high progress backtracks, while RL directly optimizes progress-adjusted rewards.Several backtracking rounds can be run iteratively even though training optimizes one backtrack.
  • Warmstart SFT: Warmstart SFT is required because pretrained base models rarely sample meaningful backtracking operations, which would otherwise inhibit self-generated-rollout training.Warmstart traces are constructed without human supervision from beam-search trees and traversal heuristics.
  • Warmstart SFT: Backtracking to high-success-rate prefix nodes produced an SFT dataset easier to fit without memorization than arbitrary stitched or direct-answer traces at the same token budget.Warmstart SFT was not needed for open-ended parameterizations from R1-distilled checkpoints.

A.3. Progress Made by MRT Compared to Outcome-Reward Training

The appendix evaluates MRT’s progress signal against outcome-reward training and reports higher progress on backtracking episodes. It also documents implementation settings and pass@k and maj@k evaluations across AIME and MATH500.

  • Progress comparison: MRT exhibits net positive and higher progress over the backtracking episode than RISE and outcome-reward RL.Progress histograms are computed on evaluation rollouts and use the reward values prescribed by MRT.
  • Implementation: The RL procedure fine-tunes the model using progress and correctness rewards, with progress estimated from additional terminating rollouts.The implementation samples partial rollouts and computes progress using Definition 6.1.
  • Implementation: MRT experiments include STaR and RL variants evaluated with R1-Distill-Qwen and DeepScaleR checkpoints, using customized progress- or information-gain-weighted losses.The appendix lists hyperparameter tables for the STaR and RL configurations.
  • Evaluation: MRT pass@k and maj@k evaluations compare performance across AIME and MATH500 as the number of sampled tokens or episodes increases.The figures cover R1-Distill-Qwen-1.5B and 7B models with RL or STaR training.

D. Full Analysis of DeepSeek-R1

The full DeepSeek-R1 analysis studies how performance changes across explicitly segmented reasoning episodes. It uses selected Omni-MATH and AIME evaluations, best-guess answer extraction, and comparisons between episode scaling, pass@k, and majority-vote metrics.

  • Experimental setup: The analysis focuses on 40 Omni-MATH problems spanning difficulty levels 4, 4.5, 5, and 5.5 to assess progress on challenging unseen problems.It samples four DeepSeek-R1-Distill-Qwen-32B responses per problem at temperature 0.7 with an 8192-token maximum.
  • Episode construction: Episodes are segmented using explicit phrases marking disruptions in logical flow, with each episode constrained to at least three steps.The phrase list is provided in Figure 21.
  • Best-guess evaluation: For episode prefixes, the model is prompted to terminate reasoning, summarize its work, and provide an answer to approximate the best-guess policy μ.Figure 22 supplies the termination prompt used for this extraction.
  • Scaling comparisons: The analysis compares episode-based [maj@p]j with direct pass@k for k = 1, 2, 8, 16, 32 across Omni-MATH and AIME 2024.Episode groups use multiples of five for Omni-MATH and three for AIME, producing ranges in total episodes.

E. Additional regret analysis of MRT models

MRT exhibits lower normalized regret than STaR, base, and GRPO, indicating greater progress across sequential reasoning episodes. Additional analyses examine scaling, extrapolation, and trajectory behavior.

  • Scaling analysis: MRT STaR is evaluated against STaR and a direct base-model baseline while scaling compute across episodes on an Omni-MATH subset.The analysis compares episode-based scaling with pass@k and majority-vote baselines.
  • Normalized regret: MRT achieves the lowest normalized regret in both STaR and RL settings, indicating more progress across sequential episodes than majority voting with fewer episodes.The comparison covers MRT STaR against STaR and Base, and MRT RL against Base and GRPO.
  • Extrapolation: MRT STaR extrapolates better than the other approaches under 2x, 4x, and 6x budget forcing, but performance dips at 8x.STaR’s performance decreases throughout, while MRT RL starts more accurately and extrapolates similarly to the other approaches.
  • Trajectory examples: Backtracking trajectories train the model to recognize an earlier mistake, return to the prior step, and redo the subsequent step with a correction.The example distinguishes the mistake, correction, and detected backtracking step.
  • Trajectory examples: In the open-ended setting, the model changes its logical approach several times before being forced to stop thinking and produce a solution.The trajectory illustrates unconstrained intermediate reasoning followed by termination.
Loading 2503.07572v1…