在一项最新的研究中，来自 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 算法的分析和消融实验，得出以下结论支持该算法的优势：
PPO 的价格模型比用于 PPO 训练的奖励模型（reward model）在指导搜寻方面更加有效。
对于 PPO 训练出的政策和价格模型，MCTS 是一个有效的启发式搜寻方式，其效果优于一些其它搜寻算法（如 stepwise-value decoding）。
PPO-MCTS 比其它提高奖励的方式（如运用 PPO 进行更多次迭代）具有更好的 reward-fluency tradeoff。
总结来说，本文通过将 PPO 与蒙特卡洛树搜寻（MCTS）进行结合，展示了价格模型在指导搜寻方面的有效性，并且说明了在模型部署阶段用更多步的启发式搜寻换取更高质量生成文本是一条可行之路。
更多方式和实验细节请参阅原论文。封面图片由 DALLE-3 生成。