在一項最新的研究中，來自 UW 和 Meta 的研究者提出了一種新的解碼算法，將 AlphaGo 采用的蒙特卡洛樹搜索算法（Monte-Carlo Tree Search, MCTS）應用到經過近端策略優化（Proximal Policy Optimization, PPO）訓練的 RLHF 語言模型上，大幅提高了模型生成文本的質量。

PPO-MCTS 算法通過探索與評估若干條候選序列，搜索到更優的解碼策略。通過 PPO-MCTS 生成的文本能更好滿足任務要求。

論文鏈接：https://arxiv.org/pdf/2309.15028.pdf

面向大眾用戶發布的 LLM，如 GPT-4/Claude/LLaMA-2-chat，通常使用 RLHF 以向用戶的偏好對齊。PPO 已經成為上述模型進行 RLHF 的首選算法，然而在模型部署時，人們往往采用簡單的解碼算法（例如 top-p 采樣）從這些模型生成文本。

本文的作者提出采用一種蒙特卡洛樹搜索算法（MCTS）的變體從 PPO 模型中進行解碼，并將該方法命名為 PPO-MCTS。該方法依賴于一個價值模型（value model）來指導最優序列的搜索。因為 PPO 本身即是一種演員 - 評論家算法（actor-critic），故而會在訓練中產生一個價值模型作為其副產品。

PPO-MCTS 提出利用這個價值模型指導 MCTS 搜索，并通過理論和實驗的角度驗證了其效用。作者呼吁使用 RLHF 訓練模型的研究者和工程人員保存并開源他們的價值模型。

PPO-MCTS 解碼算法

為生成一個 token，PPO-MCTS 會執行若干回合的模擬，并逐步構建一棵搜索樹。樹的節點代表已生成的文本前綴（包括原 prompt），樹的邊代表新生成的 token。PPO-MCTS 維護一系列樹上的統計值：對于每個節點 s，維護一個訪問量和一個平均價值；對于每條邊，維護一個 Q 值。

五回合模擬結束時的搜索樹。邊上??的數量代表該邊的訪問量。

樹的構建從一個代表當前 prompt 的根結點開始。每回合的模擬包含以下四步：

1. 選擇一個未探索的節點。從根結點出發，根據以下 PUCT 公式選擇邊向下前進，直到到達一個未探索的節點：

該公式偏好擁有高 Q 值與低訪問量的子樹，因而能較好平衡 exploration 和 exploitation。

2. 展開上一步中選擇的節點，并通過 PPO 的策略模型（policy model）計算下一個 token 的先驗概率。

3. 評估該節點的價值。該步使用 PPO 的價值模型進行推斷。該節點及其子邊上的變量初始化為：

4. 回溯并更新樹上的統計值。從新探索的節點開始向上回溯直至根結點，并更新路徑上的以下變量：

每回合模擬的四個步驟：選擇、展開、評估、回溯。右下為第 1 回合模擬結束后的搜索樹。

若干回合的模擬結束后，使用根結點子邊的訪問量決定下一個 token，訪問量高的 token 被生成的概率更高（這里可以加入溫度參數來控制文本多樣性）。加入了新 token 的 prompt 作為下一階段搜索樹的根結點。重復這一過程直至生成結束。

第 2、3、4、5 回合模擬結束后的搜索樹。

相比于傳統的蒙特卡洛樹搜索，PPO-MCTS 的創新之處在于：

1. 在選擇步驟的 PUCT 中，使用 Q 值替代了原版本中的平均價值。這是因為 PPO 在每個 token 的獎勵中含有一個 action-specific 的 KL 正則化項，使策略模型的參數保持在信任區間內。使用 Q 值能夠在解碼時正確考慮這個正則化項：

2. 在評估步驟中，將新探索節點子邊的 Q 值初始化為該節點的評估價值（而非原版本 MCTS 中的零初始化）。該更改解決了 PPO-MCTS 退化成完全 exploitation 的問題。

3. 禁止探索 [EOS] token 子樹中的節點，以避免未定義的模型行為。

文本生成實驗

文章在四個文本生成任務上進行了實驗，分別為：控制文本情緒（sentiment steering）、降低文本毒性（toxicity reduction）、用于問答的知識自省（knowledge introspection）、以及通用的人類偏好對齊（helpful and harmless chatbots）。

文章主要將 PPO-MCTS 與以下基線方法進行比較：（1）從 PPO 策略模型采用 top-p 采樣生成文本（圖中的「PPO」）；（2）在 1 的基礎上加入 best-of-n 采樣（圖中的「PPO + best-of-n」）。

文章評測了各方法在每個任務上的目標完成率（goal satisfaction rate）以及文本流暢度（fluency）。

左：控制文本情緒；右：降低文本毒性。

在控制文本情緒中，PPO-MCTS 在不損害文本流暢度的情況下，目標完成率比 PPO 基線高出 30 個百分點，在手動評測中的勝率也高出 20 個百分點。在降低文本毒性中，該方法的生成文本的平均毒性比 PPO 基線低 34%，在手動評測中的勝率也高出 30%。同時注意到，在兩個任務中，運用 best-of-n 采樣并不能有效提高文本質量。

左：用于問答的知識自省；右：通用的人類偏好對齊。

在用于問答的知識自省中，PPO-MCTS 生成的知識之效用比 PPO 基線高出 12%。在通用的人類偏好對齊中，文章使用 HH-RLHF 數據集構建有用且無害的對話模型，在手動評測中勝率高出 PPO 基線 5 個百分點。

最后，文章通過對 PPO-MCTS 算法的分析和消融實驗，得出以下結論支持該算法的優勢：

1. PPO 的價值模型比用于 PPO 訓練的獎勵模型（reward model）在指導搜索方面更加有效。
2. 對于 PPO 訓練出的策略和價值模型，MCTS 是一個有效的啟發式搜索方法，其效果優于一些其它搜索算法（如 stepwise-value decoding）。
3. PPO-MCTS 比其它提高獎勵的方法（如使用 PPO 進行更多次迭代）具有更好的 reward-fluency tradeoff。

總結來說，本文通過將 PPO 與蒙特卡洛樹搜索（MCTS）進行結合，展示了價值模型在指導搜索方面的有效性，并且說明了在模型部署階段用更多步的啟發式搜索換取更高質量生成文本是一條可行之路。

更多方法和實驗細節請參閱原論文。封面圖片由 DALLE-3 生成。