PaperlayerView exact revision on arXiv

Source-linked AI summary

VinePPO: Refining Credit Assignment in RL Training of LLMs

Amirhossein Kazemnejad, Milad Aghajohari, Eva Portelance, Alessandro Sordoni, Siva Reddy, Aaron Courville, Nicolas Le Roux

arXiv:2410.01679v2cs.LGcs.CL

TL;DR

LLM reasoning produces delayed rewards, so assigning credit to individual steps is difficult, while PPO’s value network often estimates those steps poorly. VinePPO uses Monte Carlo rollouts from intermediate states to estimate values without a separate critic. Across mathematical reasoning tasks, it outperforms PPO and other baselines, with less wall-clock time and better generalization at a given training accuracy.

  • Problem

    In reasoning tasks with delayed rewards, it is unclear how much each step contributes, and PPO’s value network often estimates returns poorly and barely exceeds random chance in ranking candidate steps.

    Source: §1 · pp. 1–2

  • Method

    VinePPO replaces the separate critic with Monte Carlo value estimates from auxiliary rollouts that restart from intermediate states by refeeding partial context.

    Source: §1 · p. 2

  • Results

    VinePPO consistently outperforms PPO and other credit-assignment-free baselines across mathematical reasoning tasks, reaching baseline peak performance with less wall-clock time and achieving higher test accuracy at a given training accuracy.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results underscore the importance of accurate credit assignment in RL training of LLMs and identify VinePPO as an alternative to value-network approaches.

    Source: §1 · p. 2

    The value derivation relies on the language environment being deterministic.

    Source: §4 · p. 4

Abstract

from arXiv · show

Large language models (LLMs) are increasingly applied to complex reasoning tasks that require executing several complex steps before receiving any reward. Properly assigning credit to these steps is essential for enhancing model performance. Proximal Policy Optimization (PPO), a common reinforcement learning (RL) algorithm used for LLM finetuning, employs value networks to tackle credit assignment. However, recent approaches achieve strong results without it, raising questions about the efficacy of value networks in practice. In this work, we systematically evaluate the efficacy of value networks and reveal their significant shortcomings in reasoning-heavy LLM tasks, showing that they often produce poor estimate of expected return and barely outperform a random baseline when comparing alternative steps. This motivates our key question: Can improved credit assignment enhance RL training for LLMs? To address this, we propose VinePPO, a straightforward approach that leverages the flexibility of language environments to compute unbiased Monte Carlo-based estimates. Our method consistently outperforms PPO and other baselines across MATH and GSM8K datasets in less wall-clock time (up to 3.0x). Crucially, it achieves higher test accuracy for a given training accuracy, capturing more generalization signal per sample. These results emphasize the importance of accurate credit assignment in RL training of LLM.

1. Introduction

LLM reasoning involves delayed rewards, making it difficult to identify which steps help or hinder a solution. The paper finds shortcomings in PPO’s value network and introduces VinePPO, which uses Monte Carlo value estimates to improve credit assignment.

1. Introduction

  • Delayed effects make credit assignment difficult because some reasoning steps matter while others are irrelevant or detrimental.Figure 1 illustrates a case where only step s2 provides a key insight, while most steps show little or no advantage over the preceding step.Source: §1 · p. 1, Figure 1 · p. 2
  • PPO uses a value network to estimate expected cumulative rewards for intermediate actions and uses those estimates to compute advantages for policy updates.This mechanism is retained in RL-based LLM finetuning.Source: §1 · p. 1
  • VinePPO estimates intermediate-state values with Monte Carlo rollouts, enabled by resetting language environments through refeeding partial context, rather than using a separate critic.The rollouts provide value estimates while preserving PPO’s overall framework.Source: §1 · p. 2
  • VinePPO consistently outperforms PPO and credit-assignment-free baselines across mathematical reasoning tasks, reaching baseline peak performance with less wall-clock time and higher test accuracy at a given training accuracy.The paper also reports better KL-divergence trade-offs and a better generalization slope.Source: §1 · p. 2
  • PPO’s value network often misestimates intermediate values and barely outperforms random chance when ranking candidate steps.On chain-of-thought reasoning tasks, it can fail to distinguish promising from unproductive steps.Source: §1 · p. 2

2. Related Work

Prior approaches to credit assignment in LLM training range from treating whole responses or intermediate steps uniformly to using a learned value network. VinePPO instead uses Monte Carlo samples for value estimation within PPO, distinct from inference-time tree search.

2. Related Work

  • DPO treats a response as one bandit action, rejection sampling trains on high-reward responses, and RLOO and GRPO use average sampled reward as a policy-gradient baseline.These approaches simplify or discard PPO’s value-network credit assignment; DPO (Rafailov et al., 2023), RestEM, RLOO, and GRPO (Shao et al., 2024).Source: §2 · p. 2
  • Earlier Monte Carlo methods used in DPO detect key reasoning errors, while this work aims to address credit assignment within RL training itself.Hwang et al. (2024) and Setlur et al. (2024) introduce those DPO methods.Source: §2 · p. 2
  • RLOO and GRPO treat intermediate states equally using group-average returns, whereas PPO predicts state values and VinePPO estimates them from auxiliary rollouts.GRPO additionally normalizes returns to unit variance; RLOO’s baseline can be viewed as a Monte Carlo value estimate only for the initial state.Source: Figure 2 · p. 3
  • VinePPO uses Monte Carlo samples solely to estimate values during PPO training, rather than using Monte Carlo tree search to select inference-time moves.The approach is orthogonal to methods that build verifiers or reward models, since it optimizes a given reward within PPO rather than designing new ones.Source: §2 · p. 3

3. Background

The background formulates LLM generation as a finite-horizon RL problem with sequence-level rewards and a KL constraint. PPO updates its policy using advantages estimated from values predicted by a critic.

3. Background

  • LLM finetuning maximizes expected undiscounted return while constraining KL divergence from the reference policy, with the policy initialized from that reference.The objective uses a prompt dataset, sequence-level reward, and a coefficient controlling the KL penalty.Source: §3 · p. 4
  • Language generation is modeled as a token-level MDP: states append generated tokens, transitions are deterministic, and reward is applied only when the response ends.Intermediate actions receive zero reward, creating delayed feedback across the trajectory.Source: §3 · p. 4
  • Policy-gradient updates use the advantage to increase the probability of better-than-average actions and decrease it for worse-than-average actions.PPO constrains these policy updates through clipping for stability.Source: §3 · p. 4
  • The value function is the expected cumulative reward from a state under the current policy, and the advantage compares action value with state value.PPO uses advantage-weighted policy gradients.Source: §3 · p. 4
  • PPO estimates advantages using a learned value network trained to minimize mean squared error against empirical returns.In LLM RL tuning, the value network is often initialized from the initial SFT policy or a reward model, with a scalar prediction head.Source: §3 · p. 4

4. Accurate Credit Assignment with VinePPO

VinePPO replaces PPO’s value network with Monte Carlo value estimates by replaying partial contexts, enabling unbiased estimation in the language environment. These estimates feed PPO’s advantage computation, with rollout count trading variance against sampling cost.

4. Accurate Credit Assignment with VinePPO

  • VinePPO replays partial contexts to sample continuations from intermediate states, averaging their returns into Monte Carlo value estimates for advantage computation.The language environment permits resetting to any intermediate state by refeeding its token context.Source: §4 · p. 5
  • Auxiliary rollouts estimate values only; they do not directly contribute to policy-gradient updates because credit assignment is unavailable for them.VinePPO changes only PPO’s advantage computation, isolating the effect of improved credit assignment.Source: §4 · p. 5
  • For any K ≥1, VinePPO’s advantage estimator yields an unbiased policy-gradient estimate of expected return.The PPO framework then uses that estimate to update the policy on the original trajectory.Source: §4 · p. 5
  • Increasing K reduces Monte Carlo estimator variance but requires more sampling and computational effort.The method groups states within a reasoning step and assigns one advantage to all tokens in that step, trading granularity for efficiency.Source: §4 · p. 5

5. Experimental Setup

The experiments compare full-parameter finetuning on mathematical reasoning datasets using shared SFT starting checkpoints. Training and checkpoint selection are controlled to compare VinePPO with PPO and alternatives that omit explicit credit assignment.

5. Experimental Setup

  • Experiments use base DeepSeekMath 7B and RhoMath 1.1B models, finetuned on MATH competition problems and GSM8K grade-school math problems.Each task’s training set produces its initial SFT policy, and model names refer to those SFT checkpoints.Source: §5 · p. 5
  • Comparators include standard PPO, RLOO, GRPO, RestEM, and DPO+; all methods start from the same SFT checkpoint and are compared by test Pass@1.Except VinePPO and PPO, the baselines omit explicit credit assignment by assigning equal weight to response tokens.Source: §5 · p. 5
  • All methods consume the same episode budget, while VinePPO inherits PPO’s hyperparameters and changes only advantage estimation.PPO receives an extensive hyperparameter search; other baselines follow their implementations with additional tuning for RLOO and GRPO.Source: §5 · pp. 5–6

6. Results

VinePPO achieves the highest test accuracy across the evaluated models and datasets, reaches PPO’s peak performance faster, and shows stronger generalization at matched training accuracy. Ablations also report improved accuracy and faster overall convergence with more auxiliary rollouts.

6.1. Task Performance

  • VinePPO achieves the highest test accuracy on both models and datasets, consistently outperforming standard PPO and other baselines; its performance gap widens on the more challenging MATH dataset.Source: §6.1 · pp. 6–7

6.2. Computational Efficiency

  • 3.0x less time for RhoMath 1.1B and 1.51x less time for DeepSeekMath 7B let VinePPO match PPO’s peak accuracy with fewer steps.VinePPO iterations take up to 5x longer for RhoMath 1.1B and 2x longer for DeepSeekMath 7B, but require 9x and 2.8x fewer steps, respectively.Source: §6.2 · p. 7, Figure 6 · p. 7

6.3. Generalization Slope

  • VinePPO has the strongest generalization gains at a given training accuracy, while RestEM overfits near the end.Figure 5 describes VinePPO as having the steepest generalization slope at matched training fit.Source: §6.3 · p. 7

6.4. Effect of K

  • Increasing K reduces Monte Carlo variance and improves accuracy; although iterations take longer, overall convergence is faster.The ablation varies K from 1 to 3 and 9 on RhoMath 1.1B, and reports the same direction of accuracy improvement as K increases.Source: §6.4 · pp. 7–8, Figure 7 · p. 7

7. Why and How Value Networks Fail

VinePPO’s value estimates track ground-truth values more accurately than PPO’s, while PPO often misclassifies states and performs near chance when ranking candidate actions.

Value prediction analysis

  • The analysis establishes ground-truth values by averaging returns from 256 Monte Carlo samples per reasoning step.It compares value predictions against this reference for DeepSeekMath 7B on MATH.Source: §7 · p. 8

Accuracy

  • VinePPO achieves 70-90% value-prediction accuracy throughout training, while PPO improves to 65% under the within-0.05 correctness criterion.VinePPO estimates are unbiased, whereas PPO's value network has high bias and often labels bad states as good and vice versa.Source: §7 · p. 8

Top Action Identification

  • PPO’s value network ranks the highest-ground-truth next step near chance for most training, while VinePPO identifies it with high accuracy throughout.The evaluation samples five possible next steps from a shared initial state.Source: §7 · p. 8

Error Per Reasoning Step

  • PPO’s value-estimation error worsens as reasoning progresses, whereas VinePPO’s error decreases as later steps become more deterministic.The proposed explanation is that PPO relies on memorization for early steps and struggles to generalize to more diverse later steps; longer context makes later steps more deterministic for VinePPO.Source: §7 · p. 8

8. Discussion

The discussion frames VinePPO as evidence that better credit assignment improves LLM RL training and identifies compute scaling and algorithm-design assumptions as research directions.

Discussion

  • VinePPO is presented as a stepping stone for identifying and fixing PPO’s broken credit assignment, with better credit assignment improving LLM RL training.The discussion also calls VinePPO the first RL post-training algorithm to scale generalization slope by scaling post-training compute.Source: §8 · p. 8
  • Algorithms with better generalization trends are valuable when truly challenging, verifiable reasoning tasks are limited.The discussion connects this direction to scaling post-training compute.Source: §8 · p. 8
  • For already capable LLMs, the discussion favors careful compute-intensive weight steering over spending compute mainly on gathering environment samples.It contrasts this with Deep RL's random-policy starting point, where quickly improving performance makes additional environment samples more useful.Source: §8 · p. 9

Impact Statement

The appendix states the paper’s impact scope, reviews PPO’s language-generation setup and critic, and describes dataset-specific reasoning-step segmentation.

Impact Statement

  • The paper targets improved LLM reasoning, while noting possible misuse for sophisticated misinformation and calling for appropriate safeguards.It says societal impact depends on responsible deployment and continued ethical consideration.Source: Impact Statement · p. 9

A. Reviewing PPO

  • PPO models each response as a token-level episode, with zero intermediate rewards and the sequence-level reward applied at termination.The state concatenates the prompt and generated tokens; each action generates the next token, and the language-environment transition is deterministic.Source: A. Reviewing PPO · p. 13

A.1. Value Network

  • PPO trains a separate value network, typically initialized from the initial policy or reward model, and uses GAE for advantage estimates.The appendix reports λ = 1 works best in its experiments, simplifying GAE when γ = 1.Source: §A.1 · pp. 13–14

B. Reasoning Step Separation Examples

  • MATH solutions are split at natural boundaries without breaking formulas, with long steps further divided at logical breakpoints; GSM8K uses newlines alone.The MATH procedure further segments steps exceeding 100 characters when possible.Source: B. Reasoning Step Separation Examples · p. 14

C.1. Datasets

The experiments use MATH and GSM8K mathematical-reasoning datasets, with separate training, validation, and test splits and illustrated step-separated solutions.

MATH

  • MATH has 12,500 training problems, split into 11,500 training and 500 validation problems, plus 500 test problems.Generated answers are checked against the ground truth using the final answer marked by \boxed{} in solutions.Source: §C.1 · p. 14

GSM8K

  • GSM8K has 7,473 training problems, split into 7,100 training and 373 validation problems, plus 1,319 test problems.Final answers are typically integers marked by ####, making response verification straightforward.Source: §C.1 · p. 14

Dataset examples

  • Figures B.1 and B.2 show MATH and GSM8K solutions separated into reasoning steps, highlighted in distinct colors with † marking boundaries.The experiments focus on mathematical reasoning datasets with step-by-step solutions.Source: §C.1 · p. 14, Figure B.1 · p. 15, Figure B.2 · p. 15

C.2. PPO Implementation

The appendices document PPO and baseline implementations, evaluation and hyperparameter choices, compute, and software. They describe measures taken to support reproducible comparisons across VinePPO and alternative methods.

C.2. PPO Implementation

  • PPO normalizes advantages across batches and data-parallel ranks, uses binary terminal rewards, handles truncated responses at the final token, disables dropout, and fixes the KL coefficient.Advantage normalization is shared by PPO and VinePPO; assigning truncated responses a negative reward did not improve performance.Source: §C.2 · pp. 15–16

C.3. SFT Models

  • Four SFT starting models cover DeepSeekMath 7B and RhoMath 1.1B on MATH and GSM8K, with checkpoints selected by validation accuracy after three training epochs.Learning rates were swept for each model before selecting its starting checkpoint.Source: §C.3 · pp. 16–17

C.4. Evaluation

  • Evaluation uses test-set Pass@1, averaged over 16 seeded rounds, with one response sampled per prompt at temperature 0.35 and a 1024-token limit.Correctness is determined by whether the final answer matches the ground truth.Source: §C.4 · p. 17

C.5. Other Baselines

  • RestEM trains on self-generated correct responses without action-specific advantages, and the authors report that it overfits training data.GRPO uses a group-return baseline, while RLOO excludes each trajectory's own return; DPO-Positive adds a term that discourages lowering correct-response probability below the reference model's.Source: §C.5 · pp. 17–19

C.6. Hyperparameters

  • PPO hyperparameters were swept on RhoMath 1.1B across MATH and GSM8K, then reused for other experiments and VinePPO, with additional tuning for RLOO and GRPO.RestEM and DPO-Positive were also configured to match PPO's training-sample counts.Source: §C.6 · pp. 19–20

C.7. Compute

  • RhoMath 1.1B PPO and VinePPO runs used four A100 GPUs, while DeepSeekMath 7B runs used eight H100 GPUs; Table 5 reports average MATH training-step times.DeepSeekMath 7B RestEM used four A100 GPUs.Source: §C.7 · p. 20, Table 5 · p. 20

C.8. Software Stack

  • The software stack combines Hugging Face, DeepSpeed, and vLLM, which also generates VinePPO's Monte Carlo samples; DeepSeekMath 7B inference reached up to 30K tokens per second on eight H100 GPUs.A separate vLLM engine runs on each GPU rank.Source: §C.8 · p. 20

D.1. Training Plots

Additional training plots compare accuracy, value prediction, and accuracy relative to KL divergence. They also examine temperature tolerance, value-prediction errors, and examples of how VinePPO assigns advantages across reasoning steps.

D.1. Training Plots

  • VinePPO maintains higher test accuracy than PPO on GSM8K, with a narrower gap for DeepSeekMath 7B, lower KL divergence, and faster per-iteration time.The authors attribute the faster GSM8K iterations to its typically shorter responses.Source: §D.1 · p. 21, Figure D.6 · p. 25

D.2. Explained Variance and Mean Absolute Error (MAE) of Value Prediction During Training

  • VinePPO achieves the highest explained variance and lowest mean absolute error among the compared methods, while PPO's improving explained variance and decreasing error indicate stable training.RLOO and GRPO are included to illustrate how their baselines deviate from ground-truth values.Source: §D.2 · p. 21, Figure D.5 · p. 24

D.3. KL Divergence

  • VinePPO achieves higher accuracy at the same KL divergence, which the authors attribute to more precise credit assignment and avoiding updates on zero-advantage steps.The plots compare accuracy against KL divergence during training.Source: §D.3 · p. 21, Figure D.7 · p. 25, Figure D.8 · p. 25

D.4. Temperature Tolerance

  • Across temperatures 0.6, 0.8, and 1.0, VinePPO benefits from higher temperatures and converges faster, whereas PPO fails to benefit and diverges at 1.0.The comparison covers the initial third of training steps.Source: §D.4 · p. 21, Figure D.10 · p. 26

D.5. Value Prediction Analysis

  • Additional figures plot predicted-value distributions against ground truth across MATH and GSM8K, and show per-step MAE for both model sizes and datasets.The per-step plots cover DeepSeekMath 7B and RhoMath 1.1B on both datasets.Source: §D.5 · p. 21, Figure D.21 · p. 31, Figure D.22 · p. 32
  • RhoMath 1.1B GSM8K wall-clock comparisons are not direct because those runs used different hardware.This caveat applies to the GSM8K runs, not all timing comparisons.Source: §D.5 · p. 21

E. More Examples of Advantages in VinePPO

  • The examples distinguish consequential reasoning steps: s2 can provide a key insight or contain an error that prevents solving, while s3 and s4 may affect success without being insightful.The appendix supplies further examples of VinePPO advantage computation.Source: Figure E.25 · p. 33, Figure E.26 · p. 34, Figure E.27 · p. 34
Loading 2410.01679v2…