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

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

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

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

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

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

PPO-MCTS 解碼算法

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

五回合模擬結(jié)束時的搜索樹。邊上??的數(shù)量代表該邊的訪問量。

樹的構(gòu)建從一個代表當(dāng)前 prompt 的根結(jié)點(diǎn)開始。每回合的模擬包含以下四步：

1. 選擇一個未探索的節(jié)點(diǎn)。從根結(jié)點(diǎn)出發(fā)，根據(jù)以下 PUCT 公式選擇邊向下前進(jìn)，直到到達(dá)一個未探索的節(jié)點(diǎn)：

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

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

3. 評估該節(jié)點(diǎn)的價值。該步使用 PPO 的價值模型進(jìn)行推斷。該節(jié)點(diǎn)及其子邊上的變量初始化為：

4. 回溯并更新樹上的統(tǒng)計值。從新探索的節(jié)點(diǎn)開始向上回溯直至根結(jié)點(diǎn)，并更新路徑上的以下變量：

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

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

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

相比于傳統(tǒng)的蒙特卡洛樹搜索，PPO-MCTS 的創(chuàng)新之處在于：

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

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

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

文本生成實驗

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

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

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

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

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

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

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

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

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

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

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