Bengio参与，扩散模型+蒙特卡洛树搜索实现System 2规划

f1BPwX

扩散模型（Diffusion Model）通过利用大规模离线数据对轨迹分布进行建模，能够生成复杂的轨迹。与传统的自回归规划方法不同，基于扩散的规划器通过一系列去噪步骤可以整体生成完整轨迹，无需依赖前向动力学模型，有效解决了前向模型的关键局限性，特别适用于具有长周期或稀疏奖励的规划任务。
尽管扩散模型具有这些优势，但如何通过利用额外的测试时间计算（TTC）来有效提高规划精度仍然是一个悬而未决的问题。一种潜在的方法是增加去噪步骤的数量，或者增加采样次数。然而，已知增加去噪步骤带来的性能提升会迅速趋于平缓，而通过多个样本进行独立随机搜索的效率非常低，因为它们无法利用其他样本的信息。
另一方面，蒙特卡洛树搜索（MCTS）则具有强大的 TTC 可扩展性。通过利用迭代模拟，MCTS 可以根据探索性反馈改进决策并进行调整，使其可以随着计算量的增加而有效地提升规划准确度。
这种能力使 MCTS 成为了许多 System 2 推理任务的基石，例如求解数学问题和生成程序。
然而，与基于扩散的规划器不同，传统的 MCTS 依赖于前向模型来执行树 rollout。这也就继承了其局限性，包括失去全局一致性。
除了局限于离散动作空间之外，生成的搜索树在深度和宽度上都可能变得过大。这会导致计算需求变得非常大，特别是当任务场景涉及到长远规划和大动作空间时。
那么，关键的问题来了：为了克服扩散模型和 MCTS 各自的缺陷，同时提升基于扩散的规划的 TTC 可扩展性，可以将扩散模型与 MCTS 组合起来吗？又该怎么去组合它们？
近日，Yoshua Bengio 和 Sungjin Ahn 领导的一个团队为上述问题提供了一个答案，提出了蒙特卡洛树扩散（MCTD）。这是一种将扩散模型的生成力量与 MCTS 的自适应搜索功能相结合的全新框架。该团队表示：「我们的方法将去噪（denoising）重新概念化为一个树结构过程，允许对部分去噪的规划进行迭代评估、修剪和微调。」

• 论文标题：Monte Carlo Tree Diffusion for System 2 Planning
  
• 论文地址：https://arxiv.org/pdf/2502.07202v1
  

蒙特卡洛树扩散（MCTD）

简单来说，MCTD = 扩散模型 + MCTS。该框架整合了基于扩散的轨迹生成以及 MCTS 的迭代搜索能力，可实现更加高效和可扩展的规划。
具体方法上，MCTD 有三项创新。

• 第一，MCTD 将去噪（denoising）过程重构成了一种基于树的 rollout 过程，于是便能在维持轨迹连贯性的同时实现半自回归的因果规划。
  
• 第二，其引入了引导层级作为元动作（meta-action），从而可实现「探索」与「利用」的动态平衡，进而确保在扩散框架内实现自适应和可扩展的轨迹优化。
  
• 第三，其采用的模拟机制是快速跳跃去噪（fast jumpy denoising）。从名字也能看出来，该机制的效率肯定很高 —— 不使用成本高昂的前向模型 rollout 即可有效估计轨迹质量。
  

基于这些创新，便可以在扩散过程中实现 MCTS 的四大步骤，即选择、扩展、模拟和反向传播，从而有效地将结构化搜索与生成式建模组合到了一起。

蒙特卡罗树扩散（MCTD）的两种视角。
上图的（a）为 MCTS 视角：展示了 MCTD 一轮的四个步骤 —— 选择、扩展、模拟和反向传播 —— 在一个部分去噪树上的过程。每个节点对应一个部分去噪的子轨迹，边标记为二元引导级别（0 = 无引导，1 = 有引导）。在新节点扩展后，执行「跳跃」去噪以快速估计其值，然后沿着树中的路径反向传播。
上图的（b）为扩散视角：同一过程被视为在去噪深度（纵轴）和规划范围（横轴）上的部分去噪。每个彩色块表示在特定噪声水平下的部分去噪规划，颜色越深表示噪声越高。不同的扩展（0 或 1）在规划方向上创建分支，代表替代的轨迹优化。值得注意的是，整行同时去噪，但去噪水平不同。
MCTD 框架将这两种视角统一了起来。整体的算法过程如下所示：

MCTD 的效果得到了实验的验证
该团队也通过实验验证了蒙特卡洛树扩散的效果。他们采用的评估任务套件是 Offline Goal-conditioned RL Benchmark（OGBench），其中涉及包括迷宫导航在内的多种任务以及多种机器人形态和机器臂操作。
下表 1 展示了质点和机器蚁在中、大、巨型迷宫中的成功率，可以看到 MCTD 的表现远超其它方法。

下图展示了三个规划器的规划结果以及实际的  rollout。

三种规划器 ——Diffuser、Diffusion Forcing 和 MCTD 生成的规划与实际展开的对比。虽然 Diffuser 和 Diffusion Forcing 未能生成成功的轨迹规划，但 MCTD 通过自适应优化其规划取得了成功。

在点阵迷宫中等任务中，使用二元引导集 {无引导，引导} 的 MCTD 树搜索过程可视化展示。每个节点对应一个部分去噪的轨迹，其中左图显示带噪声的部分规划，右图显示快速去噪后的规划。搜索通过选择无引导或引导来扩展子节点，评估每个新生成的规划，并最终收敛到高亮的叶节点作为解决方案。
下面两个表格则展示了不同方法在机器臂方块操作任务以及视觉点迷宫任务上的结果。

总体而言，MCTD 在长期任务上的表现优于现有方法，可实现卓越的可扩展性，并得到高质量的解决方案。
 该团队表示：「未来还将探索自适应计算分配、基于学习的元动作选择和奖励塑造，以进一步提高性能，为更具可扩展性和灵活性的 System 2 规划铺平道路。」