在一项最新的研究中,来自 UW 和 Meta 的研究者提出了一种新的解码算法,将 AlphaGo 采用的蒙特卡洛树搜索算法(Monte-Carlo Tree Search, MCTS)应用到经过近端策略优化(Proximal Policy Optimization, PPO)训练的 RLHF 语言模型上,大幅提高了模型生成文本的质量。
面向大众用户发布的 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 解码算法
树的构建从一个代表当前 prompt 的根结点开始。每回合的模拟包含以下四步:
1. 选择一个未探索的节点。从根结点出发,根据以下 PUCT 公式选择边向下前进,直到到达一个未探索的节点:
该公式偏好拥有高 Q 值与低访问量的子树,因而能较好平衡 exploration 和 exploitation。
相比于传统的蒙特卡洛树搜索,PPO-MCTS 的创新之处在于:
文本生成实验
文章在四个文本生成任务上进行了实验,分别为:控制文本情绪(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 生成。