Technical Summary: Reuse Your FLOPs: Scaling RL on Hard Problems by Conditioning on Very Off-Policy Prefixes

1. Problem Statement

Reinforcement learning (RL) has become the de facto method for enhancing the reasoning capabilities of Large Language Models (LLMs), particularly in mathematics and coding. However, standard on-policy RL methods face a critical bottleneck when applied to hard problems where the base model's pass rate is near zero (e.g., pass@2k ≈ 0).

In this regime, the RL policy rarely samples a correct reasoning trace. Consequently:

• Vanishing Learning Signals: The policy gradient estimates become zero or near-zero because all sampled trajectories receive zero reward.
• Compute Waste: The system expends massive amounts of FLOPs on sampling without receiving any learning signal, causing training to stall and rewards to plateau.
• Inefficiency of Existing Solutions:
  • Supervised Fine-Tuning (SFT) on off-policy traces: While this can bootstrap the model, it often leads to entropy collapse and memorization, hindering subsequent exploration.
  • Direct Off-Policy RL: Using off-policy traces directly via importance weighting suffers from high-variance gradient estimates and instability.

The core challenge is how to reuse the "wasted" compute from prior inference or RL runs (which may contain rare correct traces) to guide the online RL policy without introducing the instabilities associated with direct off-policy supervision.

2. Methodology: PrefixRL

The authors propose PrefixRL, a framework that conditions the on-policy RL process on prefixes of successful off-policy traces rather than using those traces as direct supervision targets.

Core Mechanism

1. Data Collection: A dataset of off-policy traces ($D_{off}$) is collected, typically via large-scale rejection sampling on the base model until a correct trace is found for each problem.
2. Problem Augmentation: For each original problem $x$, the system extracts a prefix (the first $h$ tokens) from a correct off-policy trace $y_x$. This prefix is appended to the original problem prompt to create a prefixed problem: $\text{concat}(x, y_{x, 1:h})$.
3. Training Objective: The model is trained using on-policy RL on a mixture of:
   • Original (no-prefix) problems.
   • Prefixed problems.
   • Crucially: Gradients are masked on the off-policy prefix tokens. The model is only updated on the tokens it generates after the prefix.

Theoretical Guarantees

The paper provides theoretical proofs regarding the PrefixRL objective:

• Consistency: If the off-policy prefixes are correct and realizable by a policy within the model class, any maximizer of the PrefixRL objective is also a maximizer of the standard RL objective. The method does not bias the optimal solution.
• Sample Efficiency: PrefixRL reduces the suboptimality gap compared to standard RL. By placing the policy in higher-rewarding states (states where a correct solution is already partially formed), it reduces the exploration burden. Theoretical bounds suggest a reduction in the number of on-policy samples required, scaling favorably with the context length.

Key Phenomenon: Back-Generalization

Empirically, the authors discover back-generalization: training only on prefixed problems significantly boosts performance on the original, unprefixed problems, even though the model was never explicitly trained on the unprefixed versions during that phase.

• Mechanism: The model does not simply imitate the prefix. Instead, it learns strategies that generalize to the unprefixed state.
• Strategy Discovery: The model can unlearn suboptimal strategies present in the prefix and discover new, better strategies. It alters the next-token distribution on unseen states, effectively "backtracking" from the high-reward states provided by the prefix to find better paths in the unprefixed space.
• Relatedness: Back-generalization is most effective when the prefix (or in-context problem) is semantically related to the target problem.

3. Key Contributions

1. PrefixRL Framework: A novel method to recycle off-policy FLOPs by conditioning on correct prefixes, avoiding the instabilities of direct off-policy RL or SFT-induced entropy collapse.
2. Theoretical Analysis: Proof that PrefixRL is consistent with standard RL objectives and offers provable sample efficiency gains by reducing the exploration burden.
3. Discovery of Back-Generalization: Identification of a phenomenon where training on augmented (prefixed) problems transfers performance to original problems, often outperforming standard generalization across related tasks.
4. Self-Improvement Loop: Demonstration of a practical pipeline where off-policy traces are generated via rejection sampling from the base model itself, creating a self-improving cycle without external teacher models.

4. Experimental Results

The authors evaluated PrefixRL on hard math reasoning benchmarks (DAPO, OMNI-MATH, AIME '25, HMMT '25) using Llama-3.1-8B and Qwen-3-4B models.

• Compute Efficiency: PrefixRL achieves the same training reward 2× faster than the strongest baseline (SFT on off-policy data followed by RL), even after accounting for the compute cost of the initial rejection sampling.
• Accuracy Gains:
  • On hard training problems, PrefixRL increases final accuracy by >45% (over 3× relative) compared to the baseline.
  • On held-out benchmarks like AIME '25, PrefixRL improves pass@1 by 12% over the strongest baseline in a compute-matched comparison.
• Robustness to Source: The method remains effective even when off-policy prefixes are sourced from a different model family (e.g., using Qwen prefixes to train a Llama policy), demonstrating flexibility in practical settings.
• Training Dynamics:
  • Stability: Unlike SFT+RL, PrefixRL preserves token-level entropy, preventing exploration collapse.
  • Gradient Quality: PrefixRL exhibits higher gradient norms and lower variance compared to standard off-policy RL (which suffers from high-variance importance weights).
  • Efficiency: The model generates shorter responses on average, indicating it reaches decisive steps earlier and reduces "unproductive wandering."

5. Significance and Claims

The paper claims that PrefixRL fundamentally changes the approach to scaling RL on hard problems. Instead of viewing off-policy data as a target for imitation (which leads to instability or collapse) or discarding it as noise, PrefixRL treats it as a conditioning mechanism to guide exploration.

The significance lies in:

• Reusing FLOPs: It transforms previously wasted compute (failed attempts that eventually yielded a correct trace via rejection sampling) into a direct signal for learning.
• Solving the "Hard Problem" Regime: It enables RL training on problems where on-policy rollouts almost never succeed, a regime where standard RL stalls.
• Back-Generalization: The discovery that conditioning on prefixes can lead to superior performance on unprefixed tasks suggests a powerful mechanism for leveraging prior compute that goes beyond simple generalization, potentially reshaping how future RL algorithms incorporate historical data.

The authors conclude that harnessing back-generalization is a crucial ingredient for bootstrapping prior compute and enabling continuous self-improvement in reasoning LLMs.