擴散模型（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 規劃鋪平道路。」