图灵奖得主加持，蒙特卡洛树搜索×扩散模型杀回规划赛道｜ICML 2025 Spotlight

henry 发自凹非寺
量子位 | 公众号 QbitAI

一个不起眼的迷宫导航任务，却能让一众模型“原形毕露”。

Diffuser和Diffusion Forcing双双翻车，通关率低得可怜。

唯独有一个模型，通关率高达 100%。

而它就来自图灵奖得主Yoshua Bengio与其团队提出的全新方法：蒙特卡洛树扩散（Monte Carlo Tree Diffusion, MCTD）。

这个方法将“上古时代”的蒙特卡洛树搜索，和当下热门的扩散模型结合在一起，突破了扩散模型在长程任务推理阶段缺乏可扩展性的瓶颈，并成功入选ICML 2025的Spotlight。

Make MCTS Great Again？

如何在探索（Explore）未知可能性以寻找更优解和利用（Exploit）当前已知最佳方案之间取得平衡，一直是复杂决策和长程规划任务的核心挑战之一。

一个过于强调探索的系统可能效率低下，在大量平庸选项中徘徊；而一个过于强调利用的系统则可能过早收敛到局部最优，错过全局最佳解。

对于扩散模型来说，它虽然能够通过去噪过程实现高质量、全局一致的序列生成（对数据分布的“利用”），但缺乏在不确定性下主动探索不同未来路径的能力。

而MCTS恰恰具备通过树形搜索结构进行高效探索和局部优化的能力，因而能够系统地权衡探索与利用，并在多个决策点进行智能选择。

由此，MCTD将扩散模型的全局一致性生成优势与MCTS的局部探索决策能力相结合。通过将轨迹划分为多个子规划来作为MCTS节点，并对不同子规划实施差异化的去噪调度，实现了在长程规划中探索与利用的平衡，提高了规划的效率和质量。

通过MCTS实现异步控制

在传统的扩散模型中，尤其是在生成轨迹时，模型通常将整条轨迹视为一个整体进行去噪，并用N来表示n个时间步的状态。

与之相反，MCTD并不把整个N个时间步的轨迹作为一个整体去噪，而是将它划分为S个时间段。在论文中，MCTD则先将完整轨迹X划分5个没有重叠的子轨迹。

例如，如果 N=500,S=5，那么一个可能的划分是：

这些子规划在每个时间段（如1-100）是独立的，它们之间没有共享的时间步。

由此，每个子规划的结束可以看作是完整轨迹的一个切片。当模型对一个特定的子规划进行去噪时，这段子规划内部的所有时间步都会同时参与到去噪过程中，遵循针对该子规划设定的统一的去噪调度。

而在不同的子规划之间，MCTS决定了不同子规划的去噪进度和深度，并通过其四个阶段表现出来。

Selection：从已有的MCTS中，使用UCB（Upper Confidence Bound，在树中选择最有前景的节点）策略选择一个表示部分去噪轨迹片段（即子规划）的节点。
Expansion：从选定的子规划节点的末端状态出发，根据不同的动作或决策，生成一个或多个新的子规划节点。这些新节点代表了从当前状态开始的下一段未探索的轨迹片段。同时，为了进行更智能的规划，MCTS还通过元动作引导级别（Guidance Levels as Meta-Action）来确定子节点的引导级别。高引导级别意味着更精细地、明确地去噪（利用），而低引导级别则意味着去噪过程可以更加自由，允许尝试新的路径（探索）。
Simulation：从新扩展的子规划节点开始，MCTD 会利用扩散模型进行“跳跃式去噪”（Jumpy Denoising），快速生成一个从该子规划开始到轨迹末端的粗略但完整的未来轨迹序列。然后，通过一个奖励函数评估这个粗略轨迹的价值。
Backpropagation：将“模拟”阶段得到的整个轨迹的奖励值，从模拟开始的子规划节点，沿着树向上，更新其所有祖先子规划节点的访问次数和累积奖励。这些更新将指导未来的Selection阶段，使得 MCTS 能够更好地利用高回报路径，探索未知的路径。

由此，模型一方面将传统的“状态”和“动作”的粒度提升到了“子规划”和“子规划间连接”的粒度；另一方面，则通过MCTS过程，控制前后子规划的降噪进度，实现异步控制，从而能够更高效地处理长序列生成和规划问题。

实验结果

MCTD在包括迷宫导航、机械臂操作、视觉迷宫（部分可观测环境）等任务上进行了测试。

在迷宫导航任务中，MCTD在所有地图尺寸（medium/large/giant）上均接近100%成功率，显著优于基线方法。

在机械臂立方体操作中，所有方法在单方块任务上性能相当。而MCTD-Replanning在多方块场景中表现出显著的性能优势，将双方块任务的成功率从22%提升至50%。

在视觉迷宫中，MCTD优于所有基线，表明其在高维感知空间中的鲁棒性。

最后，随着推理计算预算的增加（如增加最大去噪步数），MCTD成功率持续提升，而Diffuser/Diffuser-Random Search收益有限，验证了MCTD的推理可扩展性。

总体而言，尽管MCTD通过将基于搜索的规划与扩散模型结合，在推理阶段的可扩展性上取得了提升，但由于其类似“系统二”的深度推理方式，计算成本仍然较高。

此外，由于MCTD在大规模搜索空间中效率较低——即使采用了低维的元动作（meta-actions），评估多个轨迹假设的计算开销依然很高。

因此，如何提升整体效率成为了MCTD进一步改进的目标。

Fast-MCTD：加速100倍！

为了解决了MCTD计算开销大，推理时间长的缺点，研究团队又进一步推出了快速蒙特卡洛树扩散框架（Fast Monte Carlo Tree Diffusion，Fast-MCTD，相比前作MCTD，在特定任务上的推理速度提升了100倍。

在原始的MCTD框架中，主要有两个效率瓶颈：

首先，MCTS算法在设计上是顺序的，每次迭代（一次完整的Selection、Expansion、Simulation、Backpropagation）完成后，才会更新搜索树的统计信息（如节点访问次数和价值估计）。这种串行更新机制限制了算法的并行执行能力。

其次，扩散模型在生成轨迹时需要执行多次迭代去噪操作。当处理长轨迹时，每一次去噪都是一个计算密集型任务，导致整体计算开销巨大。

因此，为了降低树搜索和迭代去噪的计算开销，同时保留 MCTD 强大的规划能力，Fast-MCTD集成了两种关键的优化技术：并行MCTD和稀疏MCTD。

并行MCTD：提升并行性

MCTD沿用了MCTS的顺序性，即每次模拟完成后才更新树。并行MCTD引入了并行处理，这是Fast-MCTD与前作最显著的区别。

并行MCTD允许K个并发的rollouts。每个rollout在共享的、固定快照（fixed snapshot）的搜索树上独立进行。

树的更新（价值估计和访问计数）只有在整个批次的所有rollouts完成后才统一应用。不过，当批处理量增大时，树的统计信息会变得过时，降低选择的准确性，从而影响规划性能。

为了解决上述问题，并行MCTD引入冗余感知选择 (Redundancy-Aware Selection - RAS):它在每个并行搜索阶段临时引入一个辅助访问计数变量，顺序跟踪当前批次中的选择，并在延迟树更新后重置。

这修改了标准UCT的选择准则，通过一个超参数惩罚当前批次中已被选中的节点，鼓励其他rollouts探索树的不同部分。

由于扩散模型去噪操作昂贵，并行MCTD提出了统一的批处理策略，在扩展和模拟阶段同时处理多个由 RAS 选择的子规划。它通过调度噪声级别和同步DDIM（Denoising Diffusion Implicit Models）更新来批处理去噪步骤。为了处理子规划和不同引导级别，子规划被填充并打包成统一形状的张量，以实现GPU上的高吞吐量并行执行。