OpenAI o1背后的技术：LLM的快思考与慢思考路线之MCTS
OpenAI o1背后的技术：LLM的快思考与慢思考路线之MCTS
在上一篇文章[1]中，我们初步探索了基于EBM-MCTS的方法，并在多个数学数据集上完成实验验证。相比使用更多的SFT语料等方法，EBM-MCTS能够极大提升已有的SFT模型的数学能力[1][7]。而且，EBM不依赖更强的LLM、不需要LLM自己提出问题并评估中间步骤答案的好坏等。能量函数还能通过无监督方法训练（也不需要更大的尺寸即可提升SFT模型的效果）。近期，笔者看到了更多类似的文章，大部分在ai-agent的task-planning中探索。我们将基于这些文章，进一步探讨这些方法的'first-principle'以及未来的方向。
MCTS
Monte-Carlo-Tree-Search[2] 主要用于解决复杂序列决策问题，基于policy-model和evaluation-function，找到reward最大的决策序列。evaluation-function评估基于policy-model生成的partial sequences，实现更好的exploitation和exploration的平衡。该方法在棋类游戏中应用广泛。其核心包含四个步骤: selection, expansion, simulation和backpropagation。
Selection: 根据一定准则在叶子节点中选择一个节点。为了保证exploitation(visit-count/reward)和exploration(未被访问的节点)，通常使用UCT[2]作为叶子节点:
近期，文献[3]提出了使用PUCT(引入了policy的prior-probability):
Expansion: 上一步中选择的节点被标记为已探索并与其父节点建立链接关系。
Simulation: 从root节点到被expanded的节点作为prefix-sequence，应用policy-model生成后续的序列。通常基于一定的终止条件终止simulation，如生成结束符号、达到指定的rollout长度等。
Backpropagation: 在simulation中产生的sequence，应用evaluation-function评估reward（如质量分、输赢）等，并将reward从结束节点反传回根节点，更新路径节点的访问次数和节点reward。
从MCTS的基本原理，我们可以看到，MCTS可以看成某种慢思考，即通过policy-model快速生成路径，并通过MCTS不断评估和更新中间过程，最终产生更准确的决策输出。
LLM-MCTS
随着ai-agent的兴起，LLM带来了更多的可能性，不仅限于传统NLP/NLG任务以及作为核心系统的子模块，其本身可以作为复杂系统的决策中枢，改变传统软件、复杂系统的工作模式。当LLM作为ai-agent的决策中枢时，核心问题是在用户给定的指令下，自动将指令分解（复杂问题分解为若干子问题）、指令泛化（子问题的关联问题泛化）、复杂子问题的工具调用以及最终的解决方案生成。
类比人类解决问题，通常在某一领域有丰富经验的专家，对于领域问题具备一定的一步到位解决能力，俗称‘肌肉记忆’。通过大量的刻意练习，将中间思考步骤逐渐固化为本能，不需要刻意思考也能快速提出最终的解决方案。这种方式可以粗略理解为“greedy-decoding”，即通过大量刻意练习（包括但不限于SFT(mentor手把手教)和RLHF(外部反馈)等），将技能固化为本能即模型的快思考。
当面对新的复杂问题时，快思考往往难以一步到位给出好的解决方案，还可能局限于已有的领域知识，给出错误但confident的解决方案---陷入快思考陷阱或者定势思维。这个时候，慢思考往往能带来新的思路，帮助我们不断评估中间过程和结果，最终从‘first-principle’的角度解决新的复杂问题。这一范式可以对应到MCTS，即不断提出问题、中间解决方案和评估，最终，根据中间过程的“价值”(访问次数/reward)输出最优的解决方案。
当我们将LLM与MCTS结合时，赋予LLM慢思考的能力，一般需要解决以下问题：
高效的rollout
evaluation-function
node-representation and edge-representation
下面，我们主要分析三篇文章[3][4][5]，均是将LLM与MCTS结合，提升模型的生成、planning能力。
MAKING PPO EVEN BETTER: VALUE-GUIDED MONTE-CARLO TREE SEARCH DECODING
PPO-MCTS巧妙地利用了PPO训练过程的产出物：value-function： 估计当前 (state, action) 下的 reward，而不需要完整地 rollout （rollout 的计算复杂度较高）。为了估计 ，PPO-MCTS 作了如下近似作为 的初始化：使用 state-value-function 近似 state-action-value-function。
即在当前的状态 （MCTS 的 prefix-sequence）下，select 节点 的累计期望 reward 即 仅与状态有关，而与 select 的节点 无关（相当于 action 的分布为均匀分布）。
当进行 backpropagation 时，通过以下方式更新 state-action-value-function（state-action-value-function 的标准更新方法）：
这里，的计算方法为：
然而，实际应用时，存在几个约束：
reward-model无法获取。此时，我们可以使用value-function近似reward-function，即用 近似 。
reference-model无法获取。此时，我们无法根据step-level-reward估计state-value-function，所以，我们只能粗略的近似 。这种情况相当于去掉了policy-model和reference-model的kl-divergence，可能存在退化以及reward-hacking。
由于PPO优化过程存在诸多trick，如Reward normalization，Reward whitening，KL term clamping 以及 Adaptive KL coefficient。当使用这些trick时，step-level-reward的计算不准确，所以，我们需要使用 。（当然，如果有消融实验，对比 step-level-reward 和 的结果，对于实践会更有帮助）。
Large Language Models as Commonsense Knowledge for Large-Scale Task Planning
该方法的核心思路如上图所示。该文章主要处理language-instructed object rearrangement tasks。核心在于:
state-representation。该文章主要解决object rearrangement，其状态包括 object、robot-position、moveable-items和moveable-containers。
action-representation。object picking, object placing, moving, opening a container, and closing a container。
MCTS的node和edge。将objects建模为node，将物体之间的关系建模为edge。
该文章将MCTS与ChatGPT结合，MCTS的使用较为常规。但在MCTS的加持下，ChatGPT具备了更强的任务规划能力，相比ChatGPT-policy，提升了10个点以上。具体细节可参考对应文章。
ALPHAZERO-LIKE TREE-SEARCH CAN GUIDE LARGE LANGUAGE MODEL DECODING AND TRAINING
该文章探讨了我们提到的三个基本问题：MCTS中的node、edge表示、evaluation-function以及rollout。
node表示：
将word作为节点。适用于短输出。
将句子作为节点。适用于长文。（EBM-MCTS中，也是类似的考虑，将句子作为node）
edge表示：当我们有语言模型后，node之间的转移可以通过LLM计算。某种程度上，可以看成一种model-based-RL。所以，edge表示了node之间的转移概率且可以通过LLM计算。
evaluation-function：将value-function作为MCTS的evaluation-function。value-function的损失函数为（这部分内容也在PPO中有类似计算过程，详见[3]）：
其中， 可以通过MC-estimate或者 估计。
此外，这篇文章介绍了三个MCTS的变种：
MCTS-：该方法主要应用于alphazero。在当前state下，做N次simulations（select, expand, backup）。node的value使用value-function估计。然后根据node访问次数，选择一个node并进入下一个state。直到达到terminal-node。由于有value-function，不必rollout也能估计节点reward并backup。
MCTS[2]：相比MCTS-，经典MCTS需要select、expand、rollout、backup。而且，只能从terminal-node的reward回溯更新路径上的中间节点reward。并且，总是从initial state node出发。在空间较大的时候计算复杂度更高。
MCTS-Rollout：结合了以上两个方法的优点。能够基于value-function估计中间节点的reward并进行backup（而不需要rollout到terminal-node）。并且，每次从root-node开始模拟。
这里，是否需要rollout取决于evaluation-function。如果evaluation-function能够估计 或者 ，则可以使用 MCTS-，反之，则需要rollout并基于terminal-node的reward（只能利用最后输赢的reward更新节点reward）。
未来
LLM-based的MCTS高效实现，在复杂任务或者OOD上取得更理想的效果。
Rollout一般需要更多的计算量。为了避免rollout，我们需要估计一个较好的value-function，用于评估当前state下的reward。一般而言，需要通过RL估计 或者。[3]使用了PPO的value-function(critic-model)做一些近似或者直接训练一个value-function[5]。也能借鉴Q-learning如Implicit-language-q-learning[6]等，使用offline-method等提升样本效率。
当SFT进行到一定程度后，可以借鉴multi-agent-RL的方法，进行多智能体的训练，提升LLM解决复杂问题的能力，实现更好的alignment。
参考文献
皓天：EBM-based Global Rank+MCTS for COT?
Monte Carlo Tree Search: A Review of Recent Modifications and Applications
MAKING PPO EVEN BETTER: VALUE-GUIDED MONTE-CARLO TREE SEARCH DECODING
Large Language Models as Commonsense Knowledge for Large-Scale Task Planning
ALPHAZERO-LIKE TREE-SEARCH CAN GUIDE LARGE LANGUAGE MODEL DECODING AND TRAINING
Offline RL for Natural Language Generation with Implicit Language Q Learning
No Train Still Gain. Unleash Mathematical Reasoning of Large Language Models with Monte Carlo Tree Search Guided by Energy Function