【Openai o1背后技术】Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning

Monte Carlo Tree Search Boosts Reasoning via Iterative Preference Learning

Github：https://github.com/YuxiXie/MCTS-DPO

一、动机

大语言模型在偏好对齐环节可以提高模型的性能。目前有诸多工作尝试将偏好对齐通过迭代的形式进行改进:

It involves a cycle that begins with the current policy, progresses through the collection and analysis of data to generate new preference data, and uses this data to update the policy

目前有一些工作尝试这种迭代式对齐：

• Reinforced self-training (rest) for language modeling
• Large language models can self-improve
• Self-rewarding language models

在强化学习生态中，一个典型的工作AlphaZero就是采用这种迭代式的训练，其通过结合神经网络、强化学习以及蒙特卡洛树搜索（MCTS）实现这个迭代式过程。

然而在强化学习中，MCTS是一个N-step自举法，即对整个trajectory进行多步Reward计算。然而在大模型中，如何使用这种N-step自举法？目前很多preference pair都是站在样本（instance-lever）层面进行打标的，这会导致损失一些细节信息，对于MCTS这种需要以step层面进行学习的方式会不友好。

Conventionally, preference data is collected at the instance level. The instance-level approach employs sparse supervision, which can lose important information and may not optimally leverage the potential of MCTS in improving the LLMs

二、方法

本文提出一种迭代式的DPO算法，通过MCTS算法来抽取偏好数据，并用迭代式地训练Policy模型。整个流程大致如下所示：

• 首先初始化一个policy模型${\pi }_{{\theta }^{(0)}}$，以及一个prompt数据集${\mathcal{D}}_{\mathcal{P}}$；
• 在第$i$次迭代时，先采样一组prompt，并使用上一轮的policy模型${\pi }_{{\theta }^{(i-1)}}$为每个prompt生成若干个response；
• 使用一个不断进化的Reward标准来抽取偏好数据${\mathcal{D}}_i$；
• 基于这个新的偏好数据，训练新的policy模型${\pi }_{{\theta }^{(i)}}$；

这一过程比较类似于online版本的DPO偏好训练。

在抽取偏好数据时，采用MCTS算法，将instance-level的偏好转换为step-wise。

2.1 MCTS获得 Step-wise偏好数据

假设$x$为prompt，$s_t$表示大模型生成推理过程中reasoning chain中的前$t$步，$a$表示从$s_t$进入下一个时刻$s_{t+1}$的动作，换句话说$a$表示当前时刻要执行的推理步骤（动作）。
为此，当前所有可能的推理步骤（动作空间）可以表示为${\pi }_{\theta }(a|x,s_t)$。MCTS会根据当前已知的状态预测未来N步骤之后的奖励情况，例如预测下一最佳推理状态可表示为：

其中$Q(s_t,a)$表示基于当前已有的reasoning chain，完成当前的推理$a$后会得到的Reward奖励值。$\frac{\sqrt{N(s_t)}}{1+N(s_{t+1})}$则用于平衡探索（exploring）与利用（exploiting）之间的关系，。

• 探索：更多地试探其他可能的推理路径；
• 利用：取奖励最大的动作$a$作为下一步的推理。

为了确保在树搜索过程中，在搜索过程中，采用Self-evaluation。evaluation的模板${\text{prompt}}_{\text{eval}}$如下所示：

基于这个evaluation prompt${\text{prompt}}_{\text{eval}}$，让当前的额policy模型充当一个evaluator，对当前前$t$步骤的推理结果$s_t$进行预测，得到一个score：

$$\mathcal{C}(s_t)={\pi }_{\theta }(\text{A}|{\text{prompt}}_{eval},x,s_t)$$

Self-evaluation相关工作：Decomposition enhances reasoning via self-evaluation guided decoding.

另外，如果大模型生成的$s_t$在格式上完成了推理（即整个生成已经完成）且正确，那么记作$\mathcal{O}(s_t)=1$，若未完成则为0，若推理结果错误则为-1。
为此，可以得到一个Reward打分函数：

$$R(s_t)=\mathcal{O}(s_t)+\mathcal{C}(s_t)$$

当整个树搜索完成扩张（Expand）之后，后面需要进行回溯（Backup），更新公式如下：

其中$N(s_t)$是一个计数器。
假设整棵树的深度为$T$，即reasoning chain最多有$T$个步骤。在每个步骤时$t\in [1,T]$，都将会构建一个pair，其中正样本为具有最高$Q$值的路径，负样本则为具有最低$Q$值的路径。因此，最终可以获得$T$个pair。

2.2 迭代式DPO

考虑到偏好数据中可能会存在噪声，此时采用conservation version DPO

参考文献：A note on dpo with noisy preferences & relationship to ipo.
conservation version是指对于一些可能是噪声的偏好数据，将$(y_w>y_l)$逆转为$(y_l>y_w)$。

定义一个标签平滑系数：

DPO中两个偏好样本之间的Reward差为：

$$h_{{\pi }_{\theta }}^{y_w,y_l}=\log \frac{{\pi }_{\theta }(y_w|x)}{{\pi }_{ref}(y_w|x)}-\log \frac{{\pi }_{\theta }(y_l|x)}{{\pi }_{ref}(y_l|x)}$$

那么，loss定义为：

相当于一部分$(y_w,y_l)$会被反转。

三、实验

基座模型：Mistral-7B，
基座模型进行SFT训练，训练数据为：https://huggingface.co/datasets/akjindal53244/Arithmo-Data
训练设备：4台A100（40G）
训练细节：

数据集：

• GSM8K、MATH
• ARC、CSQA、OpenBookQA、AI2Science

实验结果：