Source-linked AI summary
Tapered Off-Policy REINFORCE: Stable and efficient reinforcement learning for LLMs
Nicolas Le Roux, Marc G. Bellemare, Jonathan Lebensold, Arnaud Bergeron, Joshua Greaves, Alex Fréchette, Carolyne Pelletier, Eric Thibodeau-Laufer, Sándor Toth, Sam Work
TL;DR
Off-policy REINFORCE struggles with negative examples, data reuse, and instability when training data differs from the current policy. TOPR uses an asymmetric, tapered importance-sampling approach without explicit KL regularization, and experiments show stable improvements across reasoning tasks while leveraging both positive and negative examples. The paper also identifies dataset balance and the baseline parameter as important to off-policy performance.
Problem
Off-policy REINFORCE has limited data reuse and unstable behavior with negative rewards when the training distribution differs from the current policy.
Method
TOPR uses an asymmetric, tapered policy gradient with truncated importance sampling and no explicit KL penalty.
Results
TOPR improves reasoning performance while leveraging positive and negative examples, including when the current policy differs substantially from the data distribution.
Takeaways & Limitations
Using negative examples avoids wasted inference, while dataset balance and REINFORCE’s baseline parameter are critical to off-policy performance.
Takeaways & Limitations
Truncated importance sampling remains sensitive to dataset composition and reward-baseline choice and can slow optimization.
Abstract
from arXiv · showhide
We propose a new algorithm for fine-tuning large language models using reinforcement learning. Tapered Off-Policy REINFORCE (TOPR) uses an asymmetric, tapered variant of importance sampling to speed up learning while maintaining stable learning dynamics, even without the use of KL regularization. TOPR can be applied in a fully offline fashion, allows the handling of positive and negative examples in a unified framework, and benefits from the implementational simplicity that is typical of Monte Carlo algorithms. We demonstrate the effectiveness of our approach with a series of experiments on the GSM8K and MATH reasoning benchmarks, finding performance gains for training both a model for solution generation and as a generative verifier. We show that properly leveraging positive and negative examples alike in the off-policy regime simultaneously increases test-time accuracy and training data efficiency, all the while avoiding the ``wasted inference'' that comes with discarding negative examples. We find that this advantage persists over multiple iterations of training and can be amplified by dataset curation techniques, enabling us to match 70B-parameter model performance with 8B language models. As a corollary to this work, we find that REINFORCE's baseline parameter plays an important and unexpected role in defining dataset composition in the presence of negative examples, and is consequently critical in driving off-policy performance.
1 Introduction
TOPR addresses instability and limited data reuse in off-policy REINFORCE by using an asymmetric policy gradient that remains stable without explicit KL regularization. Experiments on GSM8K and MATH examine positive and negative examples, dataset balance, and the role of REINFORCE’s baseline.
- Motivation and contribution: TOPR remains stable when the model differs substantially from the data distribution while fully leveraging positive and negative examples.This addresses off-policy REINFORCE’s difficulty reusing past data when negative rewards are present.
- Method: TOPR uses an asymmetric policy gradient based on trajectories sampled from a data-generating policy µ and their associated rewards R(τ).The formulation also uses a usual clipping function.
- Method: TOPR requires no explicit KL penalty for stable behavior, making it simpler to implement and computationally more efficient.KL regularization can mitigate instability but slows learning and requires additional hyperparameter tuning.
- Empirical findings: TOPR continues improving reasoning performance after π differs substantially from µ, unlike comparisons involving PPO, DPO, and naive REINFORCE.The comparison is reported in Figure 1.
- Empirical findings: GSM8K and MATH experiments identify positive/negative dataset balance and REINFORCE’s baseline parameter as important determinants of off-policy performance.The baseline, commonly used for variance reduction, also balances the dataset, and its best value depends on more than mean return.
2 Off-policy policy optimization
Off-policy optimization reuses trajectories from a reference distribution, but naive REINFORCE becomes unstable with negative examples and distribution shift. Existing remedies trade stability for discarded data, slower learning, or limited progress, motivating truncated importance sampling and related objective modifications.
- Off-policy policy optimization: Off-policy optimization reuses trajectories sampled from a reference distribution that may differ from the current language-model policy.This supports asynchronous generation, trajectory reuse, and expert data, especially when inference is computationally costly.
- The problem with naive REINFORCE: Negative trajectories make naive off-policy REINFORCE unbounded above, encouraging probabilities of supported trajectories toward zero and potentially causing degenerate behavior.The negative-reward term can drive logits toward negative infinity unless safeguards are applied.
- The problem with naive REINFORCE: Early stopping, baselines, and KL regularization can mitigate instability, but they effectively ignore some negative trajectories and limit off-policy learning.KL regularization also slows learning and adds implementation complexity.
- Supervised fine-tuning: Removing negative examples yields a bounded, stable cross-entropy objective that quickly mimics the positive-example distribution.This approach resembles reward-weighted supervised fine-tuning.
- Supervised fine-tuning: Discarding negative examples wastes inference and constrains progress because difficult tasks may provide few positives and the policy remains closer to the data distribution.Finding positive examples may require additional grading or filtering machinery.
- Truncated importance sampling: Importance sampling recovers an unbiased on-policy gradient but suffers high variance, especially for sequences whose importance ratios multiply across many steps.Large ratios can have destructive effects on negative trajectories, while small reference probabilities can slow positive-example learning.
3 TOPR: Tapered off-policy REINFORCE
TOPR combines asymmetric truncation with importance sampling to use positive and negative trajectories off-policy while controlling instability. Its canonical design accelerates learning from positive examples, attenuates negative examples, and links the baseline to dataset composition and off-policy behavior.
- TOPR uses two sets of truncation limits to downweight unlikely negative trajectories while allowing positive trajectories to be upweighted.The framework uses a+ ≤ b+ and a− ≤ b− to control the two reward classes asymmetrically.
- Setting a− = 0 progressively reduces negative trajectories’ gradient contribution, whereas any a− > 0 can cause model degeneracy unless all rewards are nonnegative.
- Choosing a+ > 0 preserves a minimum learning rate for positive trajectories and accelerates learning when their policy probability is low.This is intended to avoid the quasi-local-minima problem associated with REINFORCE in high-dimensional action spaces.
- Upper truncation limits control gradient variance, especially for large negative-example ratios early in training and large positive-example ratios late in training.
- The canonical TOPR update combines an SFT update for positive examples with a TIS update for negative examples, enabling acceleration and less brittle negative-example handling.
- The taper function equals the importance ratio on [a, b] and provides a lower bound on it; with a+ = b+ = 1 and nonnegative rewards, TOPR optimizes a lower bound on J(π).
- A baseline changes the effective positive-example proportion to p̃ = p(1 − c) / (1 + (1 − 2p)c), so it modulates dataset composition without discarding examples.Discarding negative examples is equivalent to using a baseline of −1.
4 Results
Across GSM8K and MATH, TOPR improves reasoning-model training in fully offline settings by using positive and negative examples, while remaining stable as training becomes off-policy. Its benefits include fewer invalid generations, higher accuracy and inference efficiency, robustness to dataset composition and clipping, and continued gains across iterations.
- Cross-method comparison: TOPR rapidly improves on the base model and remains effective when training becomes substantially off-policy.Compared with existing approaches, TOPR continues improving after the policy diverges from the data-generating policy.
- Failure modes: TOPR minimizes invalid generations and avoids the degeneration observed for naive REINFORCE during training.On GSM8K, TOPR teaches the model to avoid incorrect formatting, supporting solution generation without additional formatting tools.
- Negative examples: Using negative examples substantially improves performance, inference efficiency, and the distribution of correct GSM8K generations.TOPR reduces questions with no or few correct solutions, increases majority-vote efficiency, and nearly doubles MATH pass@1 accuracy over the base model.
- Dataset composition: Performance is maximal around 10-20% effective positive examples, regardless of whether the dataset contains 10% or 50% positive samples.Performance drops when the effective positive proportion is too small or too large, especially above 50%.
- Stability: TOPR remains more robust than alternatives when positive examples are scarce and gradient clipping is increased.Ratio truncation stabilizes training under off-policy gradient growth, while standard importance sampling can harm the model on negatively skewed data with gradient clipping set to 100.0.
- Iterated training and verification: TOPR continues improving across iterations and supports efficient training of both reasoning models and generative verifiers.Iterated TOPR training improves GSM8K and MATH performance, while a TOPR-trained 8B verifier becomes substantially more effective on MATH generations.
5 Conclusion and future work
The paper argues that a principled change to REINFORCE enables stable and efficient off-policy training while using both positive and negative examples. It identifies offline deployment with different data-generating policies and applications beyond LLMs as future directions.
- Keeping all generated data, avoiding KL regularization, and using negative examples improve efficiency during dataset generation and training.The paper also links negative-example use to fewer solutions needing generation at test time.
- Future work: The experiments are limited to self-taught iterations where the data-generating policy is the model at the beginning of each iteration.
- Future work: Future work includes applying TOPR offline with a different data-generating policy, despite anticipated numerical challenges.
- Future work: The study trains large language models, while the authors suggest investigating TOPR in other reinforcement-learning domains such as games and robotics.
A Proof of Proposition 3.1
The proof separates positive and negative trajectories and decomposes a reward-weighted distribution using the divergence between the data-generating and current policies. It then breaks the resulting expression into positive and negative components.
- Positive trajectories satisfy R(τ) ≥ 0, whereas negative trajectories satisfy R(τ) < 0.
- The proof defines a reward-weighted distribution before decomposing the objective into positive and negative contributions.
- The derivation includes KL(µ ∥π), the divergence from the data-generating policy µ to the current policy π, and a constant C0.
- The assembled expression includes a constant C independent of π.
B The impact of including negative examples
This section recalls the gradient of the TOPR objective and sets up the trajectory-logit notation used to analyze negative examples.
- The TOPR objective’s gradient is recalled as the starting point for analyzing its update behavior.
- The analysis represents positive- and negative-return trajectories with logits z+ and z−, respectively.
B.1 Exact gradient
The exact-gradient analysis compares gradients assigned to positive and negative trajectories through their logits. It explains how including negative examples changes the gradient difference and why SFT learns less in this formulation.
- Exact gradient: The exact-gradient analysis computes total gradients for the positive and negative trajectory logits.
- Exact gradient: Because logits are shift-invariant, the analysis studies their gradient difference to characterize the effect of negative examples.
- Exact gradient: Setting SFT’s negative coefficients a− and b− to zero makes the gradient difference larger by removing the negative term.
- Exact gradient: For b− > 0, the quantity a−R(τ−) is positive before being set to zero; removing it makes gradients for positive and negative trajectories more similar and reduces learning.
B.2 Stochastic estimate of the gradient
TOPR’s gradient analysis separates sampled and unsampled trajectories to explain how truncation controls stochastic variance. Including positive and negative terms for unseen trajectories localizes policy updates toward sampled trajectories.
- TOPR introduces unsampled positive and negative trajectories alongside sampled trajectories to analyze their respective gradient contributions.The unsampled trajectories are ν+ and ν−, with logits s+ and s− and positive and negative returns.
- The gradient terms for unsampled trajectories have opposite signs, so including both generally reduces their combined gradient norm.
- Including both unseen-trajectory terms increases the gradient of sampled positive trajectories while decreasing that of sampled negative trajectories.
- A positive negative-example baseline shifts gradient mass from unseen to seen trajectories, making policy changes more localized.This analysis concerns the effect of setting b− > 0.
C The impact of the baseline on KL regularization in TOPR
In off-policy optimization, the baseline affects more than gradient variance: under clipping, its sign changes how the policy relates to the behavior policy. Negative baselines restrain deviations, whereas positive baselines encourage determinism.
- Unlike standard REINFORCE, adding a baseline in off-policy policy optimization changes the gradient behavior rather than only its variance.
- A negative baseline discourages the policy from exceeding the behavior policy across trajectories, acting as a softer, deviation-dependent KL regularizer.
- A positive baseline encourages larger policy probabilities and therefore makes the policy more deterministic.
- These baseline effects arise solely from clipping and oppose the stochasticity effect reported by Chung et al. [2021].
D From reinforcement learning to supervised learning
The reinforcement-learning objective combines supervised learning over examples with reinforcement learning over sampled completions, making the sampling distribution over questions an important design choice. The analysis motivates emphasizing difficult questions, and Fig. 9 shows that increasing the baseline shifts relative weight toward them.
- The objective combines a supervised-learning sum over examples with a reinforcement-learning expectation over completions conditioned on each example.
- More completions should be generated for difficult questions because the base model produces fewer correct completions for them.This follows the motivation that more uniform loss across training examples can improve generalization.
- Datasets with fewer effective positive examples assign more weight to difficult questions.
- Allocating completions mainly to easy questions can widen the performance gap among examples because easy questions have more positive completions.
- As the baseline c increases and the effective positive-example proportion decreases, relative weight shifts toward difficult questions.The experiment generated 128 solutions per training question and grouped questions into ten buckets by expected base-model accuracy.
- The authors report a correlation between TOPR performance and the curve shape showing how question difficulty is weighted.