PaperlayerView exact revision on arXiv

Source-linked AI summary

Reuse your FLOPs: Scaling RL on Hard Problems by Conditioning on Very Off-Policy Prefixes

Amrith Setlur, Zijian Wang, Andrew Cohen, Paria Rashidinejad, Sang Michael Xie

arXiv:2601.18795v2cs.LGcs.AIcs.CL

TL;DR

Hard reasoning problems often yield too few correct on-policy traces for standard RL to learn efficiently. PrefixRL reuses successful off-policy prefixes as conditioning context while keeping completion updates on-policy, achieving faster and higher-reward training with transfer to held-out benchmarks.

  • Problem

    Standard on-policy RL for LLM reasoning can stall on hard problems because correct traces are rarely sampled, wasting sampling FLOPs without a learning signal.

  • Method

    PrefixRL conditions on correct off-policy trace prefixes and applies on-policy RL to complete them, using off-policy data to guide exploration without imitating it directly.

  • Results

    PrefixRL reaches the same training reward 2× faster than the strongest mid-training baseline, improves training accuracy by >45% and over 3× relative, and transfers gains to held-out benchmarks.

  • Takeaways & Limitations

    Back-generalization lets training on prefixed problems improve unprefixed performance, supporting reuse of prior compute for continuous self-improvement.

  • Takeaways & Limitations

    Off-policy RL can require much more off-policy data to address instability, increasing upfront collection cost, while some potentially useful traces may exceed the model architecture's representational capacity.

Abstract

from arXiv · show

Typical reinforcement learning (RL) methods for LLM reasoning waste compute on hard problems, where correct on-policy traces are rare, policy gradients vanish, and learning stalls. To bootstrap more efficient RL, we consider reusing old sampling FLOPs (from prior inference or RL training) in the form of off-policy traces. Standard off-policy methods supervise against off-policy data, causing instabilities during RL optimization. We introduce PrefixRL, where we condition on the prefix of successful off-policy traces and run on-policy RL to complete them, side-stepping off-policy instabilities. PrefixRL boosts the learning signal on hard problems by modulating the difficulty of the problem through the off-policy prefix length. We prove that the PrefixRL objective is not only consistent with the standard RL objective but also more sample efficient. Empirically, we discover back-generalization: training only on prefixed problems generalizes to out-of-distribution unprefixed performance, with learned strategies often differing from those in the prefix. In our experiments, we source the off-policy traces by rejection sampling with the base model, creating a self-improvement loop. On hard reasoning problems, PrefixRL reaches the same training reward 2x faster than the strongest baseline (SFT on off-policy data then RL), even after accounting for the compute spent on the initial rejection sampling, and increases the final reward by 3x. The gains transfer to held-out benchmarks, and PrefixRL is still effective when off-policy traces are derived from a different model family, validating its flexibility in practical settings.

1 Introduction

PrefixRL reuses correct off-policy traces as prefixes while keeping RL updates on-policy, addressing stalled learning on hard problems. It improves compute efficiency and final performance, with gains transferring from prefixed to original problems through back-generalization.

  • On hard problems with pass@2k ≈0, on-policy RL rarely samples correct traces, so learning stalls after spending substantial sampling compute.
  • PrefixRL appends prefixes from correct off-policy traces to problems, then runs on-policy RL on both prefixed and original no-prefix problems.Gradients are masked on the off-policy prefix, while the prefix places the policy in higher-rewarding states and strengthens the learning signal.
  • PrefixRL is theoretically consistent with standard RL and reduces the suboptimality gap with fewer samples by a factor of context length when prefixes are correct and realizable.
  • Back-generalization transfers improvements from prefixed training to original no-prefix problems, and learned strategies need not imitate the supplied prefix.The phenomenon changes next-token distributions on prefixes not trained directly and can exceed standard RL generalization across related problems or environments.
  • 2× higher compute efficiency and >45% higher final training accuracy were achieved versus SFT on rejection-sampled data followed by RL, including rejection-sampling cost.The >45% gain is over 3× relative on no-prefix training problems, and gains transferred to AIME ’25 and other held-out evaluations.
  • PrefixRL remained effective when prefixes came from Qwen3-4B-instruct while the RL policy was distilled Llama-3.1-8B-instruct.The setting provided similar compute and accuracy gains, indicating flexibility in the off-policy data source and model size.

2 Preliminaries

The paper formalizes hard-problem RL using verifiable rewards and previously collected correct traces, primarily obtained through rejection sampling. It characterizes the sampling-gradient failure mode and the assumptions governing off-policy trace realizability.

  • A problem is hard for the base policy when its pass@512 is approximately zero, and the training objective maximizes expected verifiable reward over the problem distribution.Correct responses receive reward 1 and incorrect responses receive reward 0.
  • Previously spent inference compute is represented as a dataset of correct off-policy traces from the base policy or models fine-tuned from it.
  • Rejection sampling collects one correct trace per problem, requiring 1/p_x samples in expectation when the base policy's pass@1 is p_x.The resulting off-policy dataset has one trace for each problem.
  • The theoretical setup assumes the empirical off-policy-trace distribution can be perfectly realized by some policy in the model class.In practice, traces may instead come from inference procedures using extra compute or oracle access and may not be representable by the architecture's compute depth.
  • REINFORCE estimates policy gradients from n independent traces and a group-baseline advantage, but on hard problems all samples typically fail, making the advantage and gradient usually zero.

3 PrefixRL: On-Policy RL on Very Off-policy Prefixes

PrefixRL reuses correct off-policy traces by conditioning on their prefixes while applying on-policy RL to complete both prefixed and original problems. The method is designed to avoid off-policy instability, preserve the standard RL optimum under stated assumptions, and improve sample efficiency.

  • PrefixRL framework: Prefixing can place the policy in strategy-revealing states where the probability of future success rises sharply after a key state is visited.Figure 3 varies prefix length across five problems and measures accuracy relative to the full off-policy prefix length.
  • PrefixRL framework: PrefixRL creates prefixed problems by appending prefixes from correct off-policy traces, then runs on-policy RL while masking gradients on those prefixes.Training also includes the original no-prefix problems, and prefix lengths are selected to expose strategy-revealing states.
  • PrefixRL framework: Unlike direct SFT or off-policy RL on correct traces, PrefixRL conditions on off-policy prefixes without training the policy directly on those tokens.This avoids the exploration degradation, entropy collapse, and high-variance importance-weight problems associated with direct use of off-policy data.
  • Sample-efficiency guarantees: PrefixRL converts logged-prefix information into sample-complexity advantages over standard RL, with a smaller suboptimality bound and fewer on-policy samples needed for a given reward.The theoretical analysis uses a natural policy-gradient variant and includes an optimization term plus a critic-approximation term.
  • Objective consistency: Under realizability and correctness of the off-policy traces, maximizing the PrefixRL objective also maximizes the standard no-prefix RL objective J(π).The consistency result applies when correct traces are realized by an optimal policy in the policy class.
  • Sample-efficiency guarantees: In a worst-case construction, standard RL is bottlenecked by exponentially rare rewarding traces, while PrefixRL needs only polynomially many samples to learn optimal actions in reverse.PrefixRL outsources exploration on longer prefixes to the off-policy prefixes before exposing the policy to shorter prefixes.

4 Back-Generalization Boosts the Learning Signal in PrefixRL

Back-generalization is the performance improvement on no-prefix problems when RL trains only on prefixed counterparts. PrefixRL transfers shared and sometimes newly discovered strategies beyond the provided prefix, with stronger transfer when prefix and suffix problems are related.

  • 4.1 PrefixRL Improves No-Prefix Performance Even When Training Only on Prefixed-Problems: Back-generalization improves no-prefix performance even though RL trains only on prefixed problems.Training on shorter or moderately long prefixes progressively improves shorter-prefix and eventually no-prefix performance; very long prefixes slow this transfer.
  • 4.1 PrefixRL Improves No-Prefix Performance Even When Training Only on Prefixed-Problems: Long-prefix training creates a severe train/test mismatch, delaying but not eliminating transfer to no-prefix problems.The figure reports measurable no-prefix gains after longer training despite the mismatch.
  • 4.2 PrefixRL can Discover New Strategies Beyond What is Present in the Prefixed Problem: Strategy usage remains tightly coupled between prefixed and no-prefix responses, consistent with shared internal representations across trained and unseen states.The coupling occurs even though no-prefix problems are never trained on, supporting a function-approximation account rather than simple solution stitching.
  • 4.2 PrefixRL can Discover New Strategies Beyond What is Present in the Prefixed Problem: PrefixRL can suppress strategies present in off-policy prefixes and upweight rare alternatives, transferring these changes to no-prefix behavior.The paper gives Erdős–Gallai suppression and Dirichlet Theorem upweighting as examples.
  • 4.3 Which Prefixes Back-Generalize the Most in PrefixRL? Analysis via In-Context Learning: Related prefix and suffix problems enable stronger transfer than unrelated problems, with PrefixRL improving both the trained and in-context problems.For related P2 and P3, PrefixRL reaches 63% pass@4 on P2 and 60% pass@4 on untrained P3, whereas unrelated pairs show no comparable gain.

5 Experiments and Results

PrefixRL is evaluated on hard math problems against on-policy RL and off-policy baselines, using rejection-sampled prefixes and compute-matched comparisons. It improves efficiency, accuracy, held-out pass@k, problem coverage, exploration, and gradient behavior.

  • Experimental setup: PrefixRL trains on three prefixed variants and one no-prefix variant per problem, using prefixes cut from 40%–80% of rejection-sampled correct traces.The experiments use 1k hard problems and compare against RL, SFT+RL, importance-weighted off-policy RL, and LUFFY.
  • Compute efficiency and training accuracy: 2× compute-efficiency improvement over the strongest SFT+RL baseline remains after accounting for initial rejection-sampling cost.PrefixRL also achieves 45% greater training accuracy for Llama-3.1-8B and 30% greater accuracy for Qwen3-4B in same-compute comparisons.
  • Held-out evaluation: Over 10% absolute pass@k gains across AIME’25, HMMT’25, and IMO-AnswerBench hold-out evaluations extend from k=1 through k≤64.On AIME’25, the gap reaches +18 points at k=8 and +28 points at k=64.
  • Problem coverage: PrefixRL steadily expands the set of solvable training problems at compute-matched pass@32, while competing baselines largely saturate.Its pass@1 improvements are also more uniform across problems, unlike vanilla RL’s concentration on a small subset.
  • Off-policy data sources: Prefixes from a different model family remain similarly effective on train and test problems, supporting robustness to off-policy data sources.The Llama-generated prefixes required more rejection-sampling compute, accounting for most curve differences.
  • Training dynamics: PrefixRL reduces all-negative batches and eventually samples fewer tokens per no-prefix trace while achieving higher reward rates.Shorter sampled traces lower gradient variance; prefixes also produce higher expected gradient norms than direct off-policy weighting.

6 Related Work and Discussion

The discussion positions PrefixRL among methods that learn from off-policy rollouts, condition on hints, or reset to off-policy states. Its distinction is using off-policy prefixes as conditioning context while keeping updates on-policy and accounting for data-collection compute.

  • Off-policy learning and conditioning: Unlike methods that imitate entire off-policy trajectories or partial targets, PrefixRL conditions on off-policy prefixes and completes them with on-policy RL.This avoids the instabilities associated with using off-policy data as direct training targets.
  • Resetting to off-policy states: PrefixRL adapts the resetting idea to reasoning LLMs while emphasizing self-improvement from rejection-sampled traces rather than human-labeled data.The paper reports effective resets from a relatively small dataset of correct traces on low-pass-rate problems.
  • Compute accounting and scope: PrefixRL’s compute advantage remains after including the inference cost of collecting off-policy traces, unlike the comparison identified for prior resetting work.The paper also reports gains on standardized evaluations and robustness to prefixes from different model families.

7 Conclusion

The conclusion presents back-generalization as the mechanism that lets PrefixRL incorporate off-policy feedback without directly imitating it. It frames this capability as useful for bootstrapping prior compute and continuous self-improvement.

  • Conclusion: PrefixRL differs from imitation-based methods by using back-generalization to incorporate off-policy data as conditioning context while keeping updates on-policy.The paper describes the design space as broad because conditioning is flexible and the mechanisms of back-generalization remain incompletely understood.
  • Conclusion: Back-generalization leverages the model’s own capabilities to incorporate off-policy feedback, supporting bootstrapping of prior compute and continuous self-improvement.

B Using Off-Policy Traces as Supervision Targets

Directly supervising on off-policy traces can reduce exploration or destabilize RL. The experiments contrast these failure modes with the motivation for PrefixRL’s prefix-conditioning approach.

  • SFT supervision: SFT on correct off-policy traces can cause entropy collapse, reducing exploration and eventually lowering pass@64 during RL.The observed effect is attributed to the model sharpening responses already available after SFT.
  • Importance-weighted off-policy RL: Importance-weighted off-policy RL can suffer large gradient variance or biased gradients from clipping and token-level weighting.The associated training behavior includes gradient-norm spikes and instability.

C Omitted Proofs

The omitted proofs establish PrefixRL’s consistency with standard RL, analyze its suboptimality using natural policy gradients, and provide a worst-case performance-gap lower bound.

  • Theorem 3.2 proves that any maximizer of the PrefixRL objective also maximizes the standard RL objective.
  • Theorem 3.3 bounds the suboptimality gap for natural policy gradient optimization of PrefixRL.
  • Proposition 3.4 lower-bounds the worst-case performance gap between PrefixRL and standard RL.

C.1 Proof of Theorem 3.2

The proof shows that PrefixRL’s objective preserves standard-RL optimality under correct, realizable off-policy traces, while its analysis uses explicit assumptions about the policy, critic, and data.

  • PrefixRL constructs prefixed problems from correct off-policy traces and defines an objective over their full-transcript rewards.The analysis compares this objective with the no-prefix RL objective.
  • For any PrefixRL policy, a no-prefix policy can reproduce its behavior by emitting the same prefix before sampling the suffix, yielding a uniform objective upper bound.
  • Under realizability, a policy that deterministically reproduces the correct off-policy traces attains the optimal standard-RL value.
  • Therefore, any PrefixRL objective maximizer also maximizes the standard RL objective.
  • The practical implementation uses REINFORCE, whereas the theoretical algorithm analyzes natural policy gradients with fitted Q functions and mirror-ascent updates.
  • The guarantees assume an expressive Q-function class, correct and realizable off-policy data, and finite KL divergence between the fitting policy and the base LLM.
  • The proof’s final bound introduces no extra factor of H.

C.3 Proof of Proposition 3.4

The separation proof constructs a hidden-string problem where standard on-policy RL needs exponentially many samples, while PrefixRL uses the successful off-policy trajectory to obtain a non-exponential guarantee.

  • Without off-policy data, standard on-policy RL samples only from its current policy and observes terminal rewards.
  • The constructed MDP hides a binary string of length H, and exactly one length-H action sequence earns reward 1.
  • Any such algorithm has an instance with expected suboptimality gap at least 1−(TN +1)2−H after TN episodes.
  • The exponential lower bound follows because each rollout succeeds with probability 2−H and zero-reward rollouts eliminate at most one candidate string.
  • PrefixRL samples states along the optimal trajectory from Doff, forcing visitation of rewarding states and yielding a non-exponential sample-complexity bound.
  • For uniform initialization, the relevant KL divergence is log2 under the state-averaged convention or Hlog2 under the summed convention, neither exponential in H.
  • Yao’s minimax principle converts the average-case separation into a fixed-instance separation.
  • A stitching explanation proposes that improved completion from off-policy intermediate states can improve no-prefix performance when those states are later sampled naturally.

D.2 Hard problems used in the in-context back-generalization experiment

The in-context back-generalization experiment compares transfer from structurally similar versus unrelated hard problems, using P2→P3 as the similar pair and P1→P3 as the dissimilar pair.

  • The experiment trains on an in-context solved hard problem followed by a target hard problem and compares structurally similar and unrelated pairings.
  • P1 is a geometry and conic-projection problem with base-model Pass@16 of 0.119, whereas P2 is a graph-feasibility problem with Pass@16 of 0.074.
  • P3 is a graph-feasibility problem with base-model Pass@16 of 0.063.
  • P2 and P3 share graph-theoretic reasoning about global constraints from local degree requirements, while P1 does not share that scaffold with P3.
  • The controlled comparison reports substantially stronger transfer for the compatible P2→P3 pairing than for the incompatible P1→P3 pairing.

D.3 Back-generalization under model-family mismatch

PrefixRL transfers across model families, but back-generalization depends on prefix-length coverage and compatibility with the target model. Improvements can reach no-prefix problems without reproducing the specific off-policy trace.

  • Cross-family setup: PrefixRL trains Llama3.1-8B-Instruct on correct prefixes sourced from Qwen3-4B-Instruct, evaluating performance across prefix lengths including no-prefix problems.Training prefixes occupy a restricted length band, while evaluation spans the full prefix range.
  • Back-generalization: Cross-family prefixes improve near the trained region and can propagate to shorter prefixes and no-prefix performance when prefix coverage is moderately broad.The transfer is less robust than in the same-family setting.
  • Prefix distribution: Long-prefix-skewed training weakens transfer to shorter prefixes and no-prefix problems, remaining limited after 800 iterations.The injected states may improve only a narrow, late-prefix slice rather than earlier states that dominate no-prefix rollouts.
  • Mechanism: Back-generalization does not require memorizing the in-context off-policy trace: its negative log-likelihood barely decreases, while the final policy uses a different correct sequence.The findings suggest similarity between prefixed and no-prefix solutions but weak similarity to the specific injected trace.
  • Mechanism: The proposed interpretation is that distinct long-chain-of-thought token sequences can induce similar internal states, enabling transfer despite backtracking and self-correction.This interpretation explains why the model can transfer performance without cloning the prefix.

E.2 Evaluation Protocols and FLOPs accounting.

The evaluation protocol measures original no-prefix performance and cumulative compute, including rejection-sampling overhead. Reversing cross-family transfer shows asymmetric gains: Llama prefixes help Qwen less, with higher upfront harvesting cost.

  • Evaluation protocol: All Section 5 evaluations use original no-prefix problems, with pass@k estimated from 256 samples per problem and 95% confidence intervals where possible.The protocol also reports compute-matched comparisons using cumulative FLOPs.
  • FLOPs accounting: Cumulative compute includes rollout sampling, gradient-update tokens, and upfront rejection-sampling cost when constructing off-policy traces.If R attempts are needed for one correct trace, the added upfront cost is 2RND.
  • Cross-family transfer: Reversing transfer by training Qwen with Llama-sourced prefixes produces smaller gains in both no-prefix training accuracy and AIME 2025 performance.This establishes asymmetric cross-family transfer rather than uniformly flexible prefix reuse.
  • Cross-family transfer: The asymmetry is attributed to source-prefix quality and compatibility with the target model’s representations and solution strategies.Prefixes from a stronger source can encode higher-quality intermediate reasoning states, while weaker-source rejection sampling also costs more upfront.
  • Training diagnostics: PrefixRL’s signal-to-noise evaluation tracks expected-gradient norm and sampled-gradient standard deviation throughout training.The standard-deviation estimate uses coordinate-wise first and second moments and corresponds to the trace of the gradient covariance.
Loading 2601.18795v2…