Making PPO even better: Value-Guided Monte-Carlo Tree Search decoding

Thanks for the insightful feedback! We really appreciate your recognition of our paper’s clarity and the effectiveness of our method.

Here is our response to your concerns: (we have updated our paper submission to reflect the appropriate changes, changes are rendered in blue font color)

1. Inference-time efficiency: The significant inference-time overhead is indeed a weakness of our method, which we have acknowledged in the “Limitations” section. We argue that PPO-MCTS may be found useful in several scenarios:
   • When the quality of generated text is critical, PPO-MCTS offers a way to improve text quality out of existing trained models, at some expense of additional inference-time compute;
   • Our method shows the promise of using MCTS to explore and rollout better trajectories during PPO training (improving upon the commonly-used temperature-sampling rollout), which we briefly mentioned in the “Conclusions and Future Work” section. This has the potential to improve PPO training itself, and will not incur additional inference-time latency (assuming not using MCTS for inference). To achieve this, our proposed PPO-MCTS decoding method would be an indispensable component.
2. Turning off reward whitening: We acknowledge that our PPO training settings are slightly different from the original in Lu et al. (2022). However, the baseline we report and compare with is the PPO model trained with our settings, so our comparisons are fair. We added a note in Tables 1&2 to clarify this.
3. Higher baseline for knowledge introspection: Our experimental results on the PPO baseline are slightly lower than Liu et al. (2022), possibly due to discrepancy in compute environments and random variation. For consistency, we compared PPO-MCTS with the baseline results produced by our experiments.
4. PPO with even more steps: Thanks for your suggestion, we added more experimental results to Table 6 (rows “PPO(9x more steps)”, “PPO(4x more steps)+MCTS”, and “PPO(9x more steps)+MCTS”). PPO converged at “9x more steps”, and yet its goal satisfaction rate (86.34) is still lower than PPO+MCTS (86.72). MCTS is able to improve the goal satisfaction rate by at least 10 points across all these PPO models.
5. MCTS with fewer simulations: We added more experimental results to Table 5, with the number of simulations $S=\{1,2,5\}$ (in addition to existing results where $S=\{10,20,50\}$). (Note that $S=1$ reduces MCTS to greedy decoding, and thus has very low diversity metrics.) We do observe that the goal satisfaction rate steadily increases as $S$ increases.
6. Applying temperatures to priors: To get diversity, we can apply a temperature $\tau_e$ when computing the prior $p(a|s)$ by following: Let $l(a|s)$ be the pre-softmax logits of $p(a|s)$. Normally you would do $p(a|s) = \exp{(l(a|s))} / \sum_{a’} \exp{(l(a’|s))}$. When applying a temperature, we do $p(a|s) = \exp{(l(a|s) / \tau_e)} / \sum_{a’} \exp{(l(a’|s) / \tau_e)}$. When using a $\tau_e > 1.0$, the distribution becomes less spiky.
7. The “r” in Eq.4: We apologize for the confusion. We improved the notation in Equations 4-6, where the previous $r$ is actually $r(s, \tilde{a})$, the step-level reward for taking a hypothetical action $\tilde{a}$ at state $s$. Since this is an unrealized, hypothetical decoding state living deep down in the search tree, we can’t use $r_t$ to denote this reward, but the definition is the same as $r_t$ in Eq.1. We also added an algorithm (Algorithm 1, in Appendix A.3) to clarify the backup process of MCTS.

Thanks for your valuable feedback! We really appreciate your recognition of the competitiveness of our approach.

Below is our response to your concerns: (we have updated our paper submission to reflect the appropriate changes, changes are rendered in blue font color)

1. Prior success of MCTS on off-policy RL: Thanks for noting that prior work has reported success of MCTS in combination with off-policy RL algorithms. We’d like to note that PPO is an on-policy RL algorithm, and therefore our findings may be seen as an extension to existing knowledge of what type of models MCTS can be useful with.
2. Novelty:
   • Replacing V with Q: Thanks for pointing out the original source of this idea. Other reviewers (GmAu and ZQMA) also agree with this. We added a citation to Kocsis and Szepesvari (2006), and removed this point from our novelty claim.
   • The KL term: We appreciate your perspective, but we respectfully hold a different viewpoint on this. The goal of search is not to make the KL term vanish – the KL term is conventionally included in the PPO training objective, and in MCTS it is a part of the value/Q functions that guides the search. Also, the KL term steers decoding towards the reference policy, while the prior in Eq 3 makes the decision informed by the PPO policy (which is regularized by the KL term to stay close to the reference policy during PPO training).
   • Initializing Q with V: Based on our investigation of several MCTS implementations as listed below, this technique is not adopted, so we concluded that it is not standard for MCTS. If you can share evidence that shows otherwise, we’d greatly appreciate it. [1, 2, 3]
   • Forbidding exploration after terminal states: Thanks for pointing out that this is standard practice. Reviewer (GmAu) also agrees with this. We removed this point from our novelty claim.
3. Backup algorithm: We apologize for the confusion. We improved the notation in Equations 4-6, and added an algorithm (Algorithm 1, in Appendix A.3) to clarify the backup process of MCTS. To answer your question, a NaN result cannot happen, because to get a zero denominator you need $N(s’) = 0$ for all $s’$ (where $s’$ is a child node of state $s$, where the child edge (a.k.a., the action taken) is $a$). At least one of the $N(s’)$ must be non-zero, because $s$ is an intermediate node and the “expand” stage of the current simulation round has explored one of its subtrees.
4. Clarity: Hopefully our response point #3 above addresses your comment that our description of MCTS “is somewhat difficult to understand”, since aside from the modifications we carefully described, the MCTS we use is standard as in prior literature. If there’s anything else we can make it easier to follow, please don’t hesitate to let us know ;)

Thanks for the insightful feedback! We appreciate your recognition of the potential significance of our proposed method.

Here is our response to your concerns: (we have updated our paper submission to reflect the appropriate changes, changes are rendered in blue font color)

1. Title and method name: We agree that our proposed method is an inference-time algorithm that complements PPO, rather than improving PPO itself. We will change our title to Generating more preferable text with Value-Guided Monte-Carlo Tree Search decoding. We will also change the method name to PPO+MCTS to better reflect that we are combining two existing approaches.
2. Novelty:
   • Replacing V with Q: Thanks for pointing out the original source of this idea. Other reviewers (Svcj and ZQMA) also agree with this. We added a citation to Kocsis and Szepesvari (2006), and removed this point from our novelty claim.
   • The scale of Q and c_puct: The unknown scaling factor in Q cannot be mitigated by tuning $c_\text{puct}$, because the unknown scaling factor is instance-specific (Appendix A.5). Empirically, we added experiments with different values of $c_\text{puct}$ while turning off the “initialize Q with V” technique: while increasing $c_\text{puct}$ does mitigate the exploration degeneracy, it causes degradation of task performance. For details please see additions to Table 5 (Appendix D).
   • Forbidding exploration after terminal states: Thanks for pointing out that this is standard practice. Reviewer (Svcj) also agrees with this. We removed this point from our novelty claim.
3. Structural: Thanks for the suggestion. We believe it is better to keep the introduction of MCTS algorithm in the Method section for sake of fluency in narration, because there are several parts tailored to PPO models and language generation tasks.
4. Using value function in place of Monte-Carlo rollout: Indeed we can confirm that AlphaGo-Zero already uses this technique, and we added a citation there (in addition to Chaffin et al. (2021)).

Thanks for your detailed and thoughtful feedback! We really appreciate your recognition of the diversity of tasks and comprehensiveness in our experiments.

Here is our response to your concerns: (we have updated our paper submission to reflect the appropriate changes, changes are rendered in blue font color)

1. Efficiency as evaluation axis: To investigate the efficiency-performance tradeoff, we added more experimental results to Table 5, and now it reports performance with the number of simulations $S=\{1,2,5,10,20,50\}$, representing increasing amounts of compute. We do observe that the goal satisfaction rate steadily increases with more inference-time compute.
2. More experiments on best-of-n (with reward model): Thanks for your suggestion, we added more experimental results to Table 6 (rows “PPO + best-of-$n$[R]” and “no PPO + best-of-$n$[R]”). PPO + best-of-$n$[R] has much lower performance than our PPO-MCTS (72% vs 87%). From these results, we can see that best-of-$n$ is not as good as MCTS, and thus MCTS is needed.
3. Comparison with beam search: Since diversity is an integral part of the sentiment steering task, beam search is not applicable there. We reported a baseline method similar to beam search, referred to as “stepwise-value” in Table 6 and also the 3rd paragraph of Sec 5.1. This method evaluates candidate tokens using the value function and keeps more promising tokens (as would beam search do), and supports diversity by sampling from the (softmax of) value function. As reported in Table 6, MCTS is substantially better than this stepwise decoding method.
4. Novelty of “Replacing V with Q”: Thanks for letting us know that this was adopted by Leblond et al. As other reviewers (GmAu and Svcj) pointed out, this technique was also adopted in the original formulation of MCTS. Therefore, we removed this point from our novelty claim.
5. Q-function and discounting factor: We acknowledge that most RLHF implementations set the discounting factor $\gamma = 1$. However, our reason for using the Q-function is still valid, because it captures the step-level reward $r(s, a)$ (Eq.4) which acts as a regularization term.
6. KL in the HH-RLHF task: Due to an accidental deletion of model checkpoints, we cannot retrieve the KL numbers on this task, and we apologize for this. In Table 4, we added the perplexity of generated text against the PPO reference policy (which was in the original experiment results we recorded). PPO-MCTS has similar perplexity as vanilla decoding from the PPO policy, which indicates that our method maintains good fluency and regularization on this task.
7. Nucleus sampling in HH-RLHF: We use nucleus sampling for the HH-RLHF task because, we experimentally found that using greedy decoding results in very verbose and repetitive outputs and has slightly lower reward, and temperature sampling can be unrobost due to the long-tail distribution issue (Holtzman et al. 2020).
8. PPO+best-of-n, value model: Yes, it means decoding $n$ samples and then choosing the best sample according to the PPO value model. We use the value model instead of reward model because, in some practical use cases, the reward model may be unavailable (e.g., in the toxicity reduction task, maybe PerspectiveAPI is offline or shut down). PPO-MCTS does not rely on the reward model, when using approximations (Appendix A.4).
9. Fluency numbers: Our fluency numbers look different than Quark’s because Quark used an incorrect perplexity formula (confirmed via their code and by Quark’s author). Their PPL is the inverse probability of the full generated sequence, whereby the correct PPL should have a length-normalization factor. Also, single-digit PPLs on text generated by an LM (as in our results) are more plausible than double-digit numbers.
10. Removing the adaptive KL: We removed the adaptive KL coefficient because many recent PPO implementations have removed it by default (ref: AlpacaFarm, Fine-Grained RLHF, Rainier), likely because it does not hurt performance in language tasks. It is definitely viable to keep track of the final value of the adaptive KL coefficient and use it for MCTS decoding; we kept the KL coef constant for simplicity.

If there’s any additional experiments we can add to better support our proposed method, please don’t hesitate to let us know. We look forward to your further review, and happy to discuss in more depth.