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上的高吞吐量并行執行。

稀疏MCTD：減少Rollout長度

MCTD雖然將軌跡分段為子規劃，但每個子規劃內部仍然是相對密集的軌跡。Fast-MCTD引入了軌跡粗化，從根本上縮短了有效規劃時域。通過軌跡粗化 (trajectory coarsening) 在更高的抽象層次上進行rollouts，從而減少rollout的長度和總計算成本。

具體來說，在訓練擴散模型之前，通過每隔H步進行下采樣，構建粗粒度軌跡數據集.使用在這些壓縮表示上訓練的專用稀疏擴散規劃器來建模粗粒度軌跡。由此，涉及規劃的子軌跡數量大大減少，降低總體搜索復雜度及去噪成本。

在迷宮導航測試中，Fast-MCTD相對于標準MCTD實現了約80-110倍的顯著加速，而性能損失極小。

而在機械臂操作中，Fast-MCTD在保持MCTD性能的同時，顯著提升了效率。

在視覺迷宮中，Fast-MCTD表現出顯著的效率提升，比 MCTD 快 25-60 倍，而在更大的迷宮中甚至超越了MCTD。

可以說，Fast-MCTD 在保持或提升規劃性能的同時，實現了數量級的速度提升（最高100倍），成為了更實用和可擴展的解決方案 。

ONE MORE THING

這兩篇論文的一作均來自韓國科學技術院（KAIST）的博士生尹在植（Jaesik Yoon）。

本文的另一位作者則是尹在植的指導老師安成鎮（Sungjin Ahn），安成鎮教授是韓國科學技術院和紐約大學的聯聘教授。

他的研究方向包括：可擴展貝葉斯推理、深度學習以及人工智能與認知科學的交叉領域，并多次擔任NeurIPS、ICM、ICLR等頂會AC。

他于加州大學歐文分校獲得博士學位，在Max Welling教授指導下專注于近似貝葉斯推理研究。隨后在蒙特利爾大學的MILA實驗室進行博士后研究，師從深度學習先驅、圖靈獎得主Yoshua Bengio教授。

論文鏈接：
[1]https://arxiv.org/pdf/2502.07202
[2]https://arxiv.org/pdf/2506.09498
學術主頁：
[1]https://jaesikyoon.com/
[2]https://mlml.kaist.ac.kr/sungjinahn