PaperlayerView exact revision on arXiv

Source-linked AI summary

Teaching Large Language Models to Reason with Reinforcement Learning

Alex Havrilla, Yuqing Du, Sharath Chandra Raparthy, Christoforos Nalmpantis, Jane Dwivedi-Yu, Maksym Zhuravinskyi, Eric Hambro, Sainbayar Sukhbaatar, Roberta Raileanu

arXiv:2403.04642v1cs.LG

TL;DR

The paper asks which factors most affect reinforcement-learning improvements to LLM reasoning. It compares Expert Iteration, PPO, and Return-Conditioned RL across reward schemes, model sizes, and initializations, finding similar overall performance with Expert Iteration usually best. RL improves maj@1 and pass@n together through dynamically generated diverse data, while exploration remains limited beyond solutions available from supervised fine-tuning.

  • Problem

    It remains unclear which factors among tasks, pretraining data, supervised fine-tuning data, RL algorithms, and reward sources have the biggest impact during RL fine-tuning.

  • Method

    The paper compares Expert Iteration, PPO, and Return-Conditioned RL across reward schemes, model sizes, and pretrained or SFT initializations.

  • Results

    All tested RL algorithms perform similarly on reasoning tasks, with Expert Iteration performing best in most cases and EI and PPO converging in approximately 60,000 model rollouts.

  • Takeaways & Limitations

    RL fine-tuning improves maj@1 and pass@n simultaneously by dynamically growing diverse synthetic training examples, unlike static-dataset SFT.

  • Takeaways & Limitations

    Pretraining supplies a necessary exploration prior but constrains exploration, and RL does not significantly improve pass@n beyond light supervised fine-tuning.

Abstract

from arXiv · show

Reinforcement Learning from Human Feedback (\textbf{RLHF}) has emerged as a dominant approach for aligning LLM outputs with human preferences. Inspired by the success of RLHF, we study the performance of multiple algorithms that learn from feedback (Expert Iteration, Proximal Policy Optimization (\textbf{PPO}), Return-Conditioned RL) on improving LLM reasoning capabilities. We investigate both sparse and dense rewards provided to the LLM both heuristically and via a learned reward model. We additionally start from multiple model sizes and initializations both with and without supervised fine-tuning (\textbf{SFT}) data. Overall, we find all algorithms perform comparably, with Expert Iteration performing best in most cases. Surprisingly, we find the sample complexity of Expert Iteration is similar to that of PPO, requiring at most on the order of $10^6$ samples to converge from a pretrained checkpoint. We investigate why this is the case, concluding that during RL training models fail to explore significantly beyond solutions already produced by SFT models. Additionally, we discuss a trade off between maj@1 and pass@96 metric performance during SFT training and how conversely RL training improves both simultaneously. We then conclude by discussing the implications of our findings for RLHF and the future role of RL in LLM fine-tuning.

1 Introduction

This paper studies whether feedback-driven reinforcement-learning algorithms improve LLM reasoning across reward schemes, model sizes, and initializations. It finds Expert Iteration generally performs best, with performance gains linked to diverse synthetic training data and limited exploration beyond SFT solutions.

  • Main findings: Expert Iteration performs best across all metrics for most reward setups and model initializations.The comparison includes maj@1, maj@96, rerank@96, and pass@96.
  • Main findings: EI and PPO require only a few thousand samples to converge even when initialized from a pretrained checkpoint.This makes EI nearly as sample efficient as PPO despite its simpler procedure.
  • Metric trade-offs: RL fine-tuning can improve maj@1 and pass@96 simultaneously, unlike continued supervised fine-tuning.The paper attributes this to RL generating more diverse examples during training.
  • Implications: The paper identifies limited exploration as a major factor constraining PPO’s performance and sample-complexity advantages.Models appear not to explore substantially beyond solutions already produced by SFT models.
  • Study scope: The study compares Expert Iteration, PPO, and Return-Conditioned RL for improving LLM reasoning across multiple feedback settings.It evaluates sparse, dense, heuristic, and learned rewards, using both pretrained and SFT initializations.

2 Related Work

Prior work applies RL, expert iteration, search, tools, and reward modeling to LLM reasoning and fine-tuning. The paper addresses the remaining uncertainty over which factors most influence RL-based reasoning improvement.

  • Reasoning methods: Chain-of-thought, search algorithms, and external tools have been used to improve LLM reasoning across mathematical, scientific, and code tasks.Tree of Thought combines LLM generation with breadth-first search, while code interpreters support generation and self-verification.
  • RLHF and expert iteration: RLHF commonly trains a reward model for human preferences and then performs policy improvement, most often with PPO.Recent methods also show that fine-tuning high-return responses with cross-entropy can achieve comparable performance; the paper groups these approaches as Expert Iteration.
  • Prior empirical findings: Prior expert-iteration and code-generation studies report gains that vary with model size, task, and metric.One cited GSM8K study finds larger models receive diminishing gains, while code-generation work reports substantial maj@1 improvements.
  • Open question: It remains unclear which combinations of task, data, algorithm, initialization, and reward source drive RL fine-tuning outcomes.The paper therefore analyzes these factors to identify bottlenecks and future directions for improving LLM reasoning.

3 Methods

The paper formulates LLM reasoning as deterministic-token MDPs and compares EI, PPO, and RCRL as methods for maximizing expected return. The algorithms differ in how they generate, label, and reuse solution trajectories for policy improvement.

  • Algorithm comparison: The study compares EI, PPO, and RCRL as representative algorithms for maximizing a student policy’s expected return on reasoning tasks.Each algorithm uses a different route to approximate the optimal policy.
  • PPO: PPO performs online on-policy exploration with the current student policy, then updates that policy by gradient ascent on reward.Its objective uses an estimated advantage, defined as expected return after an action minus expected return at the state.
  • Expert Iteration: Expert Iteration samples many solutions, filters incorrect or duplicate outputs, and distills the retained data into a policy with cross-entropy training.Exploration and filtering are repeated until validation performance saturates.
  • Return-Conditioned RL: Return-Conditioned RL trains policies conditioned on both the current state and a desired return.In this study, trajectories are labeled using sampled continuation correctness and balanced between good and bad solutions.

4 Experiments

Across SFT and pretrained settings, Expert Iteration generally achieves the strongest reasoning performance with sample complexity comparable to PPO. Reward-model guidance improves sample efficiency without changing final performance, while exploration diversity and initialization shape RL gains.

  • SFT initialization: Expert Iteration achieves the best performance with competitive sample complexity across model sizes and initializations.It outperforms PPO in most SFT-initialized settings, while PPO is slightly better on GSM8K without SFT initialization.
  • SFT initialization: EI reaches maj@1 accuracies of 0.485 for 7B and 0.53 for 13B models, with the best greedy accuracy after two expert iterations.Performance plateaus after the second iteration in the reported SFT-initialized experiments.
  • Reward design: ORM guidance improves sample efficiency for both EI and PPO but does not change final performance.Dense rewards provide no additional performance benefit and heuristic dense rewards can slightly harm performance by limiting solution diversity.
  • Evaluation metrics: RL primarily improves maj@1 while leaving pass@96 broadly unchanged, with most similarly sized models within 3% of the best pass@96 result.EI models show an approximately 5% rerank@96 lead, which may partly reflect ORM overfitting to EI-generated data.
  • Sample complexity: EI and PPO have comparable sample complexity from pretrained initialization, although EI can require more GSM8K samples when each prompt is oversampled.Reducing EI sampling from K = 96 to K = 4 lowers total samples while producing slightly lower final accuracy.
  • Exploration: Larger models and stronger initializations produce more diverse exploration, and models with more diverse exploration improve more from RL training.The paper links this exploration pattern to the observed performance differences across model sizes and initializations.

5 Discussion and Conclusions

Across reasoning tasks, the tested RL algorithms perform similarly, with Expert Iteration usually best and both EI and PPO converging quickly. RL improves maj@1 and pass@n together, but exploration remains constrained by pretrained and SFT behavior.

  • Expert Iteration performs best in most reasoning-task comparisons, while all tested RL algorithms perform similarly overall.
  • ∼60,000 model rollouts are sufficient for EI and PPO to converge without supervised fine-tuning.
  • Neither EI nor PPO benefits significantly from ORM guidance or denser rewards.
  • EI and PPO fine-tuning simultaneously improves maj@1 and pass@n, unlike SFT.
  • Online RL grows diverse synthetic training sets, allowing longer training without the output-diversity loss associated with static SFT data.
  • Similar performance across algorithms and rewards suggests models explore little beyond pretraining and SFT data.
  • The pretrained-model prior enables exploration in textual action spaces but also constrains early exploration, with additional SFT making this constraint worse.
  • Deterministic reasoning environments favor EI and RCRL, whereas PPO has greater advantage in highly stochastic environments.

A RCRL Label Balance

RCRL performs best when positive training data greatly outweighs negative data. Increasing negative-label proportions generally degrades performance, suggesting the 7B student learns ineffectively from negative demonstrations.

  • The experiments vary the proportions of [GOOD] and [BAD] labels to use abundant negative data that is easier to generate.
  • RCRL without SFT uses whole-solution [GOOD] and [BAD] labels after one EI iteration with K = 1 per question.
  • Increasing the proportion of negative samples generally decreases performance, with at most marginal RCRL gains.
  • Positive training data greatly outweighing negative data produces the best RCRL performance.
  • Balanced positive and negative data degrades performance, indicating the 7B student does not effectively learn from negative demonstrations.

B EI Improvement across Iterations

Expert Iteration improves accuracy over successive rounds, while solution diversity initially rises and later plateaus or declines. Larger models generate more diverse outputs, including semantic variation without always changing computational structure.

  • Maj@1 increases monotonically until convergence after at most four expert-iteration rounds.
  • SFT-initialized models converge faster than pretrained counterparts during expert iteration.
  • Exact diversity compares solution strings, whereas trace diversity compares extracted sequences of intermediate calculations.
  • Both solution-diversity measures increase during the first two rounds, after which trace diversity plateaus or slightly decreases.
  • The largest diversity increases coincide with the largest maj@1 gains during early expert-iteration rounds.
  • 13B models produce significantly more diverse outputs than 7B models throughout fine-tuning, with the gap widening over rounds.
  • After four rounds, 13B outputs can be exactly unique while trace diversity peaks earlier, indicating semantic variation without changed computational structure.

C Sample Complexities

ORM guidance can reduce sample complexity initially but does not improve converged performance and may reduce it. Ground-truth rewards remain strongest, while alternative replay techniques are compared against PPO and SFT.

  • Figures 14–17 report sample complexities for EI and PPO on GSM8K and SVAMP, including ORM-guided variants.
  • The default and ORM-guided EI comparisons are shown for GSM8K without SFT.
  • The default and ORM-guided PPO comparisons are shown for GSM8K without SFT, with ORM initially improving maj@1 before underperforming.
  • ORM guidance somewhat reduces sample complexity for EI and PPO on GSM8K but does not improve best performance.
  • Dense ORM rewards further decrease sample complexity but reduce final converged performance compared with ground-truth scoring.
  • Prioritized Level Replay and Backtracking are compared with default PPO and SFT using maj@1 scores.

D Curriculum Learning for RL

The paper tests Backtracking and Prioritized Level Replay as curriculum-learning alternatives for improving exploration during RL fine-tuning. Neither exceeds default PPO, which the authors attribute to limited exploration from pretrained and SFT starting points.

  • Backtracking: Backtracking uses high-quality supervised trajectories to improve exploration by sampling from partially complete ground-truth solutions.It removes the final intermediate step after solving from a partial prefix, with default threshold τ0 = 0.9.
  • Prioritized Level Replay: Prioritized Level Replay prioritizes problems with high learning potential, focusing exploration on tasks that are neither too easy nor too hard.Its learning-potential estimate is the average absolute advantage, and it does not require SFT data.
  • Evaluation: Figure 19 compares GSM8K maj@1 scores for PLR and Backtracking against default PPO and SFT.
  • Results: Neither Backtracking nor PLR exceeds default PPO.The authors hypothesize that limited exploration from pretrained and supervised fine-tuned models constrains both methods.

E Data augmentation

The paper augments training with synthetically generated question–answer pairs scored by how often a student recovers the intended answer. Synthetic data generally harms performance because the generator and verifier produce many unreliable or unusable examples.

  • Synthetic data construction: Synthetic pairs are generated by mapping ground-truth answers to new answers and questions, then scoring each pair by student recovery of the intended answer.The score is based on K = 20 student samples per synthetic question.
  • Data filtering: Most synthetic pairs, over 50,000, are never solved by the student and are likely unsuitable because they are impossible or too difficult.Pairs with score = 1 may also be undesirable because they are too easy, while very small positive scores may reflect incorrect solutions.
  • Data filtering: The proposed filtering targets synthetic questions whose scores lie near the middle of the score range rather than near 0 or 1.The intended neighborhood is controlled by τ.
  • Evaluation: Figure 20 reports the percentage of times the forward student model recovers the intended final answer for backwards-generated pairs.
  • Results: Introducing synthetic data degrades performance because many generated pairs are false positives.In the example, the technically correct answer is 100, but the student assigns 120 47% of the time, and the verifier repeats the same error.
  • Limitations: Larger models or separate generator and verifier models might improve synthetic-data reliability, but this remains future work.

F RCRL Step-label Generating Process

The paper contrasts process-based rewards with stepwise optimal-value estimation for identifying errors during reasoning. Stepwise ORMs approximate whether a partial solution can still lead to a correct final answer, distinguishing this from local step correctness.

  • Reward modeling: A Process Based Reward Model estimates the probability that each intermediate reasoning step is correct independently of its effect on the final answer.The paper notes that collecting human annotations for this approach would be expensive.
  • Stepwise value estimation: The optimal value function assigns 1 to an error-free solution prefix and 0 to a prefix containing a mistake that will lead to an incorrect final answer.
  • Illustrative example: The example question contains inconsistent category percentages, producing an answer of 120 from 40, 30, and 50.
  • Stepwise value estimation: Stepwise ORMs approximate the optimal value function rather than a process reward model.They differ because a prefix may contain an earlier mistake while ending with a locally valid terminal step.
Loading 2403.04642v1…