Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions

近年來，生成模型在離散領域的研究持續蓬勃發展，尤其在文字、組合問題等領域，如何有效地建模離散資料序列是一大挑戰。傳統的自回歸模型（Autoregressive Models, ARMs）憑藉其將問題拆解成序列化的單步預測而大放異彩，然而同時也面臨推理階段低並行度及固有的序列依賴困境。本文由 Kim 等人於 ICML 2025 發表的獲獎論文《Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions》深入探討一種新興架構——遮蔽擴散模型（Masked Diffusion Models, MDMs），試圖揭示其訓練與推理過程中不同的困難與優勢，並提出有效策略突破現有限制。

研究背景與動機

生成模型通常面對的重要項目是如何在訓練與推理中達成有效平衡。自回歸模型因為直接學習序列的條件機率分布，方便訓練及理論分析，但其嚴格的序列解碼順序導致推理時無法並行，速度受限且模式易受「錯誤累積」影響。相反地，擴散模型近年在連續資料（如影像）生成中展現卓越成效，透過學習從噪聲反推資料分布，並允許更加靈活的生成步驟安排。

對於離散資料，傳統擴散模型較不適用，因為其基於連續空間設計。Masked Diffusion Models（MDMs）應運而生，結合遮蔽隨機遮蔽部分tokens並重建的機制，將生成任務視為多重遮蔽填補（infilling）問題。此架構在訓練時需要解決大量潛在的遮蔽組合，形成極為龐大的問題集合，理論上是計算不可行的，但MDMs能成功訓練並展現良好生成能力。

本研究即聚焦於深入理解MDMs在訓練與推理中所面臨的「穩健性挑戰」，以及不同token解碼順序對模型性能的重大影響。作者提出「Train for the Worst, Plan for the Best」的觀點：在訓練中，MDMs被迫面對最困難的遮蔽重建子問題（worst-case），而推理階段則可策略性挑選最易解的token順序（best-case），藉此極大提升推理效能與生成質量。

核心方法與創新

本文核心分析將MDMs與ARMs的訓練與推理上的計算複雜度進行深入比較。在理論層面，作者證明MDMs訓練時面臨的子問題數量隨序列長度呈指數爆炸，遠超自回歸模型。這種「最壞情境」訓練策略雖乍看難以承受，卻培養模型擁有泛化多種遮蔽修復任務的能力。

然而，重點在推理階段的「最佳計畫」策略。與自回歸模型被綁定於固定解碼順序不同，MDMs可根據當下生成狀態動態選擇下個token填補的順序，極大提升靈活度。論文提出一套適應性（adaptive）token排序方法，透過評估生成過程中的不確定度或難度，優先生成最「容易」解的tokens，繞過高難度遮蔽問題的直接挑戰。

此外，作者在演算法設計上，結合探索（exploration）與利用（exploitation）策略平衡，且採用對MASK策略的多樣化遮蔽方案，促使模型在訓練中涵蓋更廣泛的子問題分布，強化對推理階段多樣解碼路徑的適應性。

主要實驗結果

為驗證理論分析與方法有效性，論文選擇離散邏輯謎題代表——數獨（Sudoku）作為測試基準，凸顯模型對組合推理及順序決策的考驗。實驗結果令人驚豔：

• 預訓練的MDM在沒有調整解碼順序時，解答正確率不到7%，顯示單純隨機或固定順序下，困難子問題大幅限制性能。
• 採用論文提出的適應性token解碼順序後，解答正確率大幅攀升至約90%，展現驚人的解決困難任務能力，並優於參考自回歸模型。
• 即使與擁有7倍參數、自回歸且經過教師強制訓練（teacher forcing）以學習固定正確順序的強型基準相比，MDM在正確率上仍有顯著優勢。

同樣的策略也在語言模型和其他離散生成任務中展示強化效果，表示方法具有較好普適性和適用性。

對 AI 領域的深遠影響

此論文對離散生成模型的發展具有劃時代意義。首先，提出了全新視角審視生成模型中「訓練最壞案例與推理最佳策略」間的平衡問題，豐富我們對模型能力的理解，指出不必將固定解碼順序視為必然限制。

其次，MDMs在離散領域展現彈性，使得模型能透過靈活的搜尋策略，避開計算瓶頸及高難度子問題，提高推理效率與成功率。這在實際應用中，尤其是組合規劃、邏輯推理與自然語言生成等場景，有著顯著的實用價值。

此外，此工作也為擴散模型邁向離散領域提供了理論與實驗基礎，有望引領更多結合擴散機制與遮蔽式生成的後續研究，促進模型在多模態、結構化數據與複雜決策問題上的突破。

總結來說，《Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions》不僅從理論上探討MDMs的複雜性，更透過創新的推理策略實現其潛力，為離散生成模型設計提供嶄新思路及實證基礎。這項研究大幅推動了擴散模型在離散域的應用前沿，對未來生成模型的靈活性與效率提升具深遠貢獻。

---

論文資訊
📄 Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions
👥 Kim, Shah, Kontonis, Kakade, Chen
🏆 ICML 2025 · Outstanding Paper
🔗 arxiv.org/abs/2502.06768