PaperlayerView exact revision on arXiv

Source-linked AI summary

Don't throw away your value model! Generating more preferable text with Value-Guided Monte-Carlo Tree Search decoding

Jiacheng Liu, Andrew Cohen, Ramakanth Pasunuru, Yejin Choi, Hannaneh Hajishirzi, Asli Celikyilmaz

arXiv:2309.15028v3cs.CLcs.AIcs.LG

TL;DR

Prior guided-decoding methods may evaluate incomplete text with heuristics or classifiers not trained for partial sequences. PPO-MCTS uses PPO’s value model to guide MCTS decoding, improving preferability across four generation tasks; the approach also incurs inference-time compute overhead.

  • Problem

    Guided decoding often relies on heuristics or classifiers not trained to evaluate incomplete sequences, while PPO’s value model for partial outputs is commonly discarded.

    Source: §1 · p. 2, §2.3 · p. 4

  • Method

    PPO-MCTS combines a PPO policy and its associated value model with MCTS, using the value model to evaluate partial sequences during decoding.

    Source: §1 · p. 2

  • Results

    Across four tasks, PPO-MCTS improves preferability; sentiment-steering success is 30% (absolute) higher than direct sampling from the same PPO policy.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support retaining PPO value models for inference-time search, since they can effectively guide MCTS decoding.

    Source: §1 · p. 3, §7 · p. 9

    PPO-MCTS adds significant inference-time compute overhead, estimated naively at 2S times direct sampling when policy and value models share architecture and size.

    Source: §5.1 · p. 9, Limitations and Ethics Statement · p. 10

Abstract

from arXiv · show

Inference-time search algorithms such as Monte-Carlo Tree Search (MCTS) may seem unnecessary when generating natural language text based on state-of-the-art reinforcement learning such as Proximal Policy Optimization (PPO). In this paper, we demonstrate that it is possible to get extra mileage out of PPO by integrating MCTS on top. The key idea is not to throw out the value network, a byproduct of PPO training for evaluating partial output sequences, when decoding text out of the policy network. More concretely, we present a novel value-guided decoding algorithm called PPO-MCTS, which can integrate the value network from PPO to work closely with the policy network during inference-time generation. Compared to prior approaches based on MCTS for controlled text generation, the key strength of our approach is to reduce the fundamental mismatch of the scoring mechanisms of the partial outputs between training and test. Evaluation on four text generation tasks demonstrate that PPO-MCTS greatly improves the preferability of generated text compared to the standard practice of using only the PPO policy. Our results demonstrate the promise of search algorithms even on top of the aligned language models from PPO, and the under-explored benefit of the value network.

1 Introduction

The introduction argues that PPO’s value model is suited to evaluating partial sequences, yet is often discarded, and proposes using it to guide MCTS decoding. Across four tasks, PPO-MCTS improves generated-text preferability over standard PPO decoding.

1 Introduction

  • PPO’s value model estimates expected return for partial outputs under its policy, making it a natural guided-decoding evaluator that prior work often lacks.Prior approaches may use heuristics or classifiers not trained to evaluate incomplete sequences.Source: §1 · p. 2
  • PPO-MCTS applies MCTS to PPO policy/value pairs, using the value model to evaluate partial sequences and search for higher-reward outputs.In the example, look-ahead shifts the next-token choice from “bad” to “person,” yielding a positive-sentiment continuation.Source: §1 · p. 2
  • The method initializes children’s Q values from their parent’s V value to encourage exploration when combining PPO models with MCTS.Source: §1 · p. 2
  • PPO-MCTS improves preferability across four tasks; sentiment-steering success rises 30% absolute, toxicity falls 34% relative, introspected knowledge is 12% more useful, and chatbot win rate rises 5% absolute.Human evaluation favors PPO-MCTS by a 20% absolute margin on sentiment steering and a 30% relative margin on toxicity reduction.Source: §1 · pp. 2–3
  • PPO-MCTS also outperforms longer PPO training and best-of-n decoding, both of which directly optimize for rewards.Source: §1 · p. 3

2 Preliminaries

The preliminaries frame text generation as sequential decisions, introduce guided decoding and PPO, and explain why PPO’s often-discarded value model can evaluate partial outputs.

2.1 Notation

  • Text generation can be modeled as sequential decisions: the state is the prompt plus generated response so far, and the action is the next token.Source: §2.1 · p. 3

2.2 Guided Decoding

  • Guided decoding evaluates partial-output states against a goal and combines that evaluation with a language-model policy, often using search for look-ahead.Source: §2.2 · p. 3

2.3 Proximal Policy Optimization (PPO)

  • PPO trains policy and value models together; the value model estimates returns for computing the policy’s advantage and is trained against empirical returns.Source: §2.3 · p. 4
  • PPO deployments commonly discard or do not release the value model, although it is a suitable evaluator for guided decoding.Source: §2.3 · p. 4

3 Method

PPO-MCTS uses the PPO policy to propose actions and its value model to score partial sequences during tree search. After simulations, it decodes from root-child visit counts.

3 Method

  • PPO-MCTS balances exploitation and exploration by using tree search to evaluate partial sequences proposed by the policy.The policy proposes promising actions, while the evaluation function scores partial outputs.Source: §3 · p. 4
  • Each token is decoded after S simulations over a tree whose nodes are states and edges are actions.The search tracks node visit counts and mean values, and edge Q-values.Source: §3 · p. 4
  • PUCT selects actions using edge Q-values and PPO policy priors while accounting for node visit counts.The Q-values derive from the PPO value model, and selection continues from the root to an unexplored node.Source: §3 · p. 5
  • At expansion, PPO policy priors link the top-k actions to child nodes, and the PPO value model evaluates the selected node.PPO-MCTS does not use Monte-Carlo rollouts for efficiency; it initializes each new child-edge Q-value from the parent’s value to encourage exploration.Source: §3 · p. 5
  • Backup updates visit counts, mean values, and edge Q-values along the path from the evaluated node to the root.After simulations, decoding uses normalized root-child visit counts, with temperature controlling sampling or greedy selection.Source: §3 · p. 5
  • Node expansion stops at [EOS], but simulations reaching a terminal node proceed directly to backup to preserve visit counts.Source: §3 · p. 5

4 Experimental Setup

PPO-MCTS is evaluated on four text generation tasks: sentiment steering, toxicity reduction, knowledge introspection, and helpful and harmless chatbots.

4 Experimental Setup

  • Sentiment steering tests whether MCTS improves success in generating continuations with the target sentiment for opposite-sentiment prompts.The experiments use OpenWebText prompts and PPO-trained models.Source: §4 · p. 6
  • Toxicity reduction tests whether MCTS further reduces toxicity while generating fluent and diverse continuations.The task uses prompts from RealToxicityPrompts and PPO-trained models.Source: §4 · p. 6
  • Knowledge introspection tests whether MCTS increases the utility of generated commonsense knowledge for downstream question answering.The introspection models are trained with PPO to maximize the generated knowledge’s utility to the QA model.Source: §4 · p. 6
  • The helpful and harmless chatbot task investigates PPO-MCTS for aligning language models to human preferences using Anthropic’s HH-RLHF data.Source: §4 · p. 6

5 Results

PPO-MCTS outperforms direct decoding from the PPO policy across four tasks. The analyses compare value-guided search with alternative scorers, decoding strategies, training lengths, and hyperparameter settings, and discuss inference-time overhead.

5 Results

  • Across all four tasks, PPO-MCTS outperforms direct PPO-policy decoding; on sentiment steering it also beats best-of-n decoding and longer PPO training.The paper describes MCTS decoding as a more regularized way to obtain higher-reward generations.Source: §5 · p. 6
  • Sentiment steering success rises by 34% absolute for positive and 26% absolute for negative sentiment, with comparable fluency and diversity.Human evaluation favors PPO-MCTS by 22% on positive sentiment and 18% on negative sentiment.Source: §5 · pp. 6–7, Table 1 · p. 7
  • Maximum toxicity falls by 34% relative versus direct PPO-policy sampling, while fluency and diversity remain comparable.Human evaluation also favors PPO-MCTS for toxicity reduction by a 30% relative margin, with comparable fluency and topicality.Source: §5 · p. 7, Table 2 · p. 7
  • Knowledge introspection improves downstream QA accuracy by 0.42% absolute and knowledge usefulness by 12% relative; chatbot average reward rises by 0.04 and human win rate by 5% absolute.The reported chatbot comparison is against the PPO policy baseline.Source: §5 · pp. 7–8, Table 3 · p. 7, Table 4 · p. 7

5.1 Analyses and Ablations

  • Replacing the value model with a reward model lowers both goal satisfaction and fluency; the value model handles partial sequences and is tailored to the policy.The reward model predicts full-sequence reward, whereas the value model estimates expected return under the current policy.Source: §5.1 · p. 8
  • Stepwise value decoding yields much lower goal satisfaction and fluency than MCTS, while four-times-longer PPO training also has inferior goal satisfaction on positive sentiment.Longer training has significantly worse perplexity on positive sentiment and comparable perplexity but lower goal satisfaction on negative sentiment.Source: §5.1 · p. 8
  • Higher expansion temperature can improve goal satisfaction and diversity at a fluency cost; higher decoding temperature increases diversity with relatively small penalties.Increasing simulations and branching factor generally improves goal satisfaction and diversity at the expense of fluency; the experiments use S = k = 50 and τd = τe = 2.0.Source: §5.1 · p. 8
  • Naively, tree construction makes PPO-MCTS 2S times slower than direct sampling, but reusing the subtree under the decoded token mitigates the overhead.At least ⌈S/k⌉ tree nodes need not be recomputed; in practice, usually many more can be reused.Source: §5.1 · p. 9

6 Related Work

Prior guided decoding uses token-level value functions or step-level verifiers, often paired with search, while MCTS has also been applied to several language tasks. PPO-MCTS differs by using the jointly trained PPO value model alongside its policy.

6 Related Work

  • Guided decoding steers a generative LM with an auxiliary function that evaluates partial outputs, using token-level value functions or step-level verifiers and often search.Prior value functions commonly use hand-crafted rules or separate discriminative models, whereas PPO-MCTS uses a value model tailored to the generative LM and target metric.Source: §6 · p. 9
  • Prior MCTS text-generation work includes planning and reasoning, plug-and-play control, and machine translation, with differing choices of policy and value models.PPL-MCTS uses sequence classifiers not designed for partial sequences; Leblond et al. (2021) train a value model to fit an already-trained policy.Source: §6 · p. 9
  • PPO practitioners have generally used the trained policy at inference while discarding its value model; the authors claim PPO-MCTS is the first to combine jointly trained PPO policy and value models for improved generation.Quark (Lu et al., 2022), DPO (Rafailov et al., 2023), and SLiC do not train an accompanying value model.Source: §6 · p. 9

7 Conclusion and Future Work

The paper demonstrates MCTS decoding on PPO-trained policy and value models, showing the usefulness of value models trained as a byproduct of aligning language models to human preference. Future work could explore MCTS for policy optimization, while addressing its efficiency and language action-space challenges.

7 Conclusion and Future Work

  • PPO-MCTS applies MCTS decoding to PPO-trained policy and value models, demonstrating the usefulness of value models produced during alignment to human preference.Source: §7 · p. 9
  • Future work could use MCTS as a policy-optimization operator, but constructing its search tree is less efficient than PPO policy rollout.Source: §7 · p. 9
  • MCTS policy updates toward visit counts may lack sufficient significance in language because of its large action space.Source: §7 · p. 9

Limitations and Ethics Statement

PPO-MCTS improves generated-text preferability but introduces significant inference-time compute overhead. Its altered generation behavior may also create a risk when the value model is adversarially tampered with.

Limitations and Ethics Statement

A.1 A note on implementation

The appendix describes PPO-MCTS implementation, PPO's combined policy and value objectives, and the MCTS backup procedure. It also explains approximations used when inference-time reward or reference-policy information is unavailable, and how PPO implementation choices can require those approximations.

A.1 A note on implementation

  • PPO-MCTS is implemented as a PyTorch plugin that can replace model.generate() and works with HuggingFace Transformers and Accelerate.Source: §A.1 · p. 15

A.2 The PPO Training Objectives

  • PPO training combines policy and value objectives: the policy uses an estimated advantage and clipped surrogate, while the value model matches empirical returns with an MSE objective.The final objective is a linear combination of the policy and value objectives.Source: §A.2 · p. 15

A.3 The MCTS Backup Algorithm

  • The MCTS backup stage takes the root, selected node, step-level reward, and tree statistics as input, then updates the tree statistics by backing up reward plus discounted child value.Source: §A.3 · p. 15

A.4 Approximations

  • When the reward model is unavailable at inference, PPO-MCTS approximates terminal reward with the value estimate, although that estimate was not trained to match terminal reward.The text describes this as a reasonable approximation, especially where value-model parameters are initialized from a trained reward model.Source: §A.4 · pp. 15–16
  • When the reference policy is unavailable, PPO-MCTS approximates Q with the mean value estimate, dropping KL regularization and bringing more risk of reward hacking.Source: §A.4 · p. 16

A.5 Implementation Choices in PPO and their Implications on MCTS

  • Reward normalization, reward whitening, unknown KL-term clamping, or an unknown final adaptive KL coefficient can require approximating terminal reward with V or Q with V.Reward whitening makes the value estimate’s batch-specific scale unknown, preventing its direct addition to raw step-level rewards.Source: §A.5 · p. 16
  • When PPO uses an adaptive KL coefficient whose final value is unknown, Q must be approximated with V.Source: §A.5 · p. 16

B.1 Task 1: Sentiment Steering

The appendix specifies model sources and decoding configurations for four tasks: sentiment steering, toxicity reduction, knowledge introspection, and helpful and harmless chatbots. The setups vary in model family, search budget, decoding choice, and whether Q or terminal reward is approximated.

B.1 Task 1: Sentiment Steering

  • Sentiment steering uses GPT2-large PPO models trained for 100 steps, with up to 20 generated tokens and 50 MCTS simulations per token at branching factor 50.The setup approximates terminal reward with the value model because the reward model is unavailable at decoding time.Source: §B.1 · pp. 16–17

B.2 Task 2: Toxicity Reduction

  • Toxicity reduction uses GPT2-large PPO models trained for 500 steps, with up to 20 generated tokens and 20 MCTS simulations per token at branching factor 20.As in sentiment steering, terminal reward is approximated with the value model at decoding time.Source: §B.2 · p. 17

B.3 Task 3: Knowledge Introspection

  • Knowledge introspection uses Rainier PPO policy and value models based on T5-large, with 10 MCTS simulations per token at branching factor 10.The setup approximates terminal reward with the value model because computing it requires the question’s ground-truth answer.Source: §B.3 · p. 17

B.4 Task 4: Helpful and Harmless Chatbots

  • Helpful and harmless chatbots use LLaMA-7b-based models, with a reward-model-initialized value model and 20 MCTS simulations per token at branching factor 20.Greedy decoding is used at finalization, and Q is approximated with the mean value because reward whitening and adaptive KL were used in PPO training.Source: §B.4 · p. 17

C.1 Sentiment Steering

The sentiment and toxicity human evaluations compare generations on their task-specific criterion, as well as fluency and topicality.

C.1 Sentiment Steering

  • Sentiment steering evaluations used 400 comparisons, with three workers per comparison, judging sentiment, fluency, and topicality.The 400 comparisons came from 100 positive and 100 negative prompts, with two generations sampled per decoding method for each prompt.Source: §C.1 · p. 18, Figure 3 · p. 18

C.2 Toxicity Reduction

  • Toxicity reduction evaluations used 200 comparisons, with three workers per comparison, judging toxicity, fluency, and topicality.The comparisons used 100 prompts, with two generations sampled per decoding method for each prompt.Source: §C.2 · p. 18

D Additional Analyses

Additional analyses examine PPO-MCTS hyperparameters and extend the sentiment-steering comparisons; removing Q-value initialization creates a diversity and goal-satisfaction trade-off.

Additional analyses

  • Table 6 continues the sentiment-steering results with additional baselines and comparisons, including a PPO replica trained under a slightly different setting.The replica is reported to have similar performance to the referenced PPO model.Source: D Additional Analyses · p. 18, Table 6 · p. 20
  • Table 5 varies Q←V initialization, cpuct, simulation count, branching factor, and temperatures for visit counts and expansion priors.The caption identifies these as PPO-MCTS hyperparameters analyzed on sentiment steering.Source: Table 5 · p. 19
  • Without “initialize Q with V,” diversity falls; raising cpuct may mitigate this, but large cpuct lowers goal satisfaction, and the best rate is 7 points below standard PPO-MCTS.These experiments vary cpuct while disabling Q initialization from V.Source: D Additional Analyses · p. 19
Loading 2309.15028v3…