VinePPO: Refining Credit Assignment in RL Training of LLMs

We thank the reviewer for their time and effort in reviewing the paper. We now address the concerns raised in the review.

Novelty

MC estimation is one of the classic methods for value function estimation in RL. This paper did not introduce additional novel contributions. In fact, in PPO, GAE with $\lambda =1$ is also a standard form of MC estimation. Given that VinePPO primarily replaces GAE with its special case, the novelty is limited.

VinePPO is fundamentally different from setting the λ = 1 in GAE within PPO and does not replace GAE. In fact, both PPO and VinePPO use GAE. Notably, GAE takes value estimates as inputs to compute advantages. The key difference lies in how these value estimates are computed: VinePPO use auxiliary rollouts while PPO trains a critic network. In both methods, GAE’s λ is still a hyperparameter. Our hyperparameter search (see section C.6.) found λ = 1 for PPO and λ = 0 for VinePPO.

Estimating the values of intermediate states via auxiliary MC was known to be theoretically possible but never deemed practical in DeepRL [1]. This is because auxiliary rollouts require resetting the environment to intermediate states: an uncommon assumption in Deep RL. However, language environments inherently support this capability. VinePPO is, to our knowledge and as noted by Reviewer p3iT, the first to exploit this property, turning an idea once considered impractical in DeepRL into an effective method in RL for LLMs. Moreover, this work goes beyond merely demonstrating better performance; it also provides thorough experiments to explain why.

[1] Schulman, John et al. “Trust Region Policy Optimization.” ArXiv abs/1502.05477 (2015)

Variance and efficiency of MC estimation

While we agree with the reviewer that in theory variance of MC estimation might necessitate higher values of K, we empirically establish that even values as small as K=9 demonstrate superior performance (Fig 4), efficiency (Figs 6-7), and stability (Fig D.10) over all baselines. Our extensive study of compute efficiency in Sec. 6.2 and 6.4 shows that while each iteration takes longer with higher K, the improved credit assignment leads to less wall-clock time for reaching a target test accuracy.

Thus it is intuitive to use more samples (e.g., k=256 in Section 7) to validate the effectiveness of MC estimation. However, it cannot support the effectiveness of VinePPO (k=9) in the main results.

Please note that in Fig 3, VinePPO values estimates are computed with K=9 and they’re 3x more accurate than the next best baseline in terms of MAE (VinePPO(K=9): 0.03, PPO: 0.11).

Questions

Given that estimating the expectation of E(R) from a single episode is intuitively unreliable, why does VinePPO (K=1) still outperform PPO?

This is a great question. First, we clarify that K refers to the auxiliary rollouts for value estimation. Regarding the question, both variance and bias can hinder performance. The value network in PPO exhibits high bias in its value estimates (see Fig 3 and Fig D.15). VinePPO even with K=1 provides an unbiased estimate of the value. Although the variance with K=1 might be high, the high bias of the PPO’s value estimates empirically result in worse performance. Note that such subpar performance of value network has been reported in previous work [2].

[2] Ahmadian et al, "Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs" (2024)

Efficiency was a key focus of our work, and we conducted a thorough analysis of VinePPO’s efficiency in the paper. Across all experiments VinePPO consistently demonstrated superior efficiency and performance compared to baselines. We hope this clarification address your concerns and encourage a fresh evaluation of our work.

We are thrilled that the reviewer liked our work and happy to hear that our results were successfully replicated. We now address the question.

Questions

How would vinePPO perform when training a reasoning model (e.g., o1/R1) with extremely long COT? How can we decide when to perform MC samples when the COT is really long?

This is a great question. While it is best answered empirically, we believe that credit assignment becomes even more critical in long-CoT models (such as O1/R1). As trajectories grow longer, the distance between the final reward and each individual action increases, making credit assignment more challenging. We hypothesize that in such settings, methods like VinePPO, which enable fine-grained credit assignment, could be more impactful. Moreover, VinePPO could be further optimized by adaptively deciding when to perform additional MC rollouts, for example, based on computed confidence intervals (also known as stopping criteria), which we leave for future work.

We thank the reviewer for their time and their focus on details in their feedback. We now address the raised concerns:

The proposed method is not new, albeit it's not commonly seen for solving CA in RL training LLMs

While estimating intermediate-states values via auxiliary MC was known to be theoretically possible, it was never deemed practical in DeepRL [1]. This is because auxiliary rollouts require resetting the environment to intermediate states of an already sampled trajectory: a rare assumption in DeepRL. In fact, even today, none of standard DeepRL frameworks implement or offer any feature resembling the "Vine" approach. However, reset to intermediate states come naturally in language environments. VinePPO is, to our knowledge and as noted by Reviewer p3iT, the first to exploit this property, turning an idea once considered impractical in DeepRL into an effective method in RL for LLMs. Furthermore, this work goes well beyond demonstrating improved performance and provides thorough analysis to explain the underlying reasons of such improvement.

[1] Schulman, John et al. “Trust Region Policy Optimization.” ArXiv abs/1502.05477 (2015)

Questions

In figure 3, the plots for GRPO and RLOO, are they depicting the values for the initial state only or for intermediates too?

Figure 3 includes all steps including the intermediate steps to assess the fine-grained credit assignment throughout trajectories.

Why does GRPO and RLOO perform so much worse? Is it a lack of hyperparameter tuning or samples per prompt?

This is an interesting, yet not surprising result. GRPO/RLOO lack any fine-grained credit assignments. In other words, they assign the same advantage to all the tokens in a response. This is known to be problematic in RL: actions with negative ground-truth advantage may be boosted if their sampled trajectory succeeds, while highly advantageous tokens may be penalized if their trajectory fails. All PPO-based variants, such as VinePPO, GRPO, and RLOO inherit PPO’s hyperparameters. However, the lack of fine-grained CA in GRPO/RLOO (see Fig. 3 for their MAE) led to less stable training, requiring even an additional HP tuning compared to VinePPO (see Section C.6). Please note that the same sample per prompt is used for all experiments. In other words, the number of training trajectories, responses, is equal between all algorithms.

We hope we have adequately addressed all of your concerns and would be happy to clarify further if needed. With this in mind, we hope the reviewer increases their rating of our paper.

We appreciate the reviewer’s attention to detail and positive feedback. Below, we address the questions asked:

Questions

What about evaluation on more OOD benchmark datasets like AIME

We considered benchmarks like AIME too difficult for the base LLMs used in our experiments. However, this is an interesting question that could be explored with newer and stronger models.

Is VinePPO still effective for larger-scale models?

While this question is ultimately empirical (and unfortunately beyond our current resources), we hypothesize that in larger models, which often produce longer responses, the credit assignment problem could be even more critical. In very long chains-of-thoughts, where reward is only received at the end, attributing credit to individual actions becomes increasingly challenging. Therefore, we believe methods like VinePPO, which offer better credit assignment, should be more relevant in such settings.