Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning

Experimental Setup

We fine-tuned base models using MRT (models that already produce traces with <think> markers). For STaR, we used two sizes of models (7B and 1.5B) on 10,000 NuminaMate samples, comparing MRT (with progress bonus) against vanilla STaR (outcome-only reward). For RL, we used 1.5B models, comparing MRT with vanilla GRPO (outcome-reward RL). We fine-tuned one model on 4,000 NuminaMath problems, while another (already fine-tuned on 40K MATH problems) was trained on 919 AIME problems. We also compared against an RL approach that penalizes token length. Evaluation used a 16K token budget, matching our fine-tuning maximum.

Pass@1 performance of RL-trained MRT models on various math reasoning benchmarks

We performed a detailed analysis of the computational costs associated with our proposed MRT method compared to classical approaches like STaR and GRPO. The analysis quantifies both forward generation and training costs using established FLOP estimation formulas.

To estimate computation costs, we employ two key formulas:

N represents the number of model parameters, Drolluts is the total number of tokens generated during inference, and Dtrain is the total number of tokens used during training.

For the STaR baseline, we sampled 200 full rollouts per problem and selected solutions that correctly solved each problem. In contrast, for MRT (STaR), we generated just 1 complete rollout per problem, then selected 10 prefixes from this rollout and sampled 20 continuations for each prefix to approximate the information gain.

When applied to the NuminaMATH dataset containing 20,000 problems, using Llama-3.1-8B-Instruct with 4,000 token completions, this approach yielded 12,000 correct solutions for training (with incorrect solutions discarded). Training proceeded for three epochs.

The total FLOP calculations are:

For the GRPO baseline and MRT (RL), we used a different sampling strategy. In MRT (RL), we first generated 1 complete rollout per problem, then selected a prefix and sampled 10 rollouts to approximate information gain. In both methods, given a prompt, we sampled 4 responses and maximized their group advantage estimations.

The resulting FLOP calculations are:

Our analysis reveals that MRT (STaR) requires only 1.01× more FLOPs (2.64192 × 10²⁰ / 2.62912 × 10²⁰) than STaR to achieve comparable performance, while using 1.7× fewer tokens during inference. Similarly, MRT (RL) uses just 1.08× more FLOPs than GRPO to achieve equivalent performance, while requiring 1.6× fewer tokens during inference.

Standard techniques for fine-tuning LLMs to use test-time compute optimize outcome reward at the end of a long trace. This does not incentivize the model to make use of intermediate tokens to make progress (i.e., probability of eventual success) and leads to 1) unnecessarily long output traces and 2) inability to make steady progress on new, hard problems as shown in (a). MRT, shown in (b), trains the LLM to minimize cumulative regret over the entire output stream (red, shaded area) by optimizing a dense reward function in addition to sparse 0/1 reward and thus alleviates both challenges in (a).