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

近年來，隨著生成模型的需求日益多元與複雜，離散領域上的生成任務（如文字、邏輯謎題、組合優化等）逐漸成為研究熱點。傳統上，自回歸模型（Autoregressive Models, ARM）透過順序預測下一個元素取得成功，但在推理階段存在固有的順序限制，導致推理速度以及彈性受限。而蒙版擴散模型（Masked Diffusion Models, MDMs）作為一種新興的生成方法，探索出在訓練與推理階段不同的策略平衡，提供了更具彈性的生成流程。ICML 2025 年獲得傑出論文獎的 "Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions" 一文，由 Kim 等人深入揭示並理論化 MDMs 訓練和推理過程中的核心挑戰，並提出創新性的適應性解碼順序策略，極大提升模型效能和實用價值。

研究背景與動機

傳統的自回歸模型擅長逐一預測序列中下一個 token，藉此完成文本生成、序列決策等任務。其優點是訓練過程簡單且自然地學習序列依賴關係，但缺點是推理時必須嚴格按序解碼，造成推理速度慢且不易調整生成順序。

相比之下，擴散模型（Diffusion Models）原則上是在連續空間的生成模型領域取得重大突破，但當應用到離散資料（如文字、棋盤狀態、拼圖組合）時，設計直接的拓展方法相對困難。Masked Diffusion Models 引入蒙版機制，將輸入序列部分遮蔽，訓練模型學習在被遮蔽部分「填補」正確 token 的能力。這種方法在推理時允許任意順序的解碼，大大提升了推理靈活性與效率。

然而，MDMs 的強大也伴隨著代價：訓練過程中，模型要學習從任意遮蔽模式下正確還原序列，等同於要解決大量、甚至是指數級數量的補齊子問題，造成訓練計算複雜度大幅攀升。更進一步，推理階段若不加策略隨機或固定解碼順序，隨機填補容易陷入困難模棱兩可的決策環境，導致性能不佳。

因此，論文核心動機是：全面剖析 MDMs 於訓練與推理間平衡的本質，理解解碼順序如何影響模型表現，並探討如何設計有效策略改善 MDMs 的實際應用效果，特別是在具有高度結構性、邏輯推理性的任務中表現。

核心方法與創新

理論分析：訓練困難度與子問題規模 論文首先從理論角度分析 MDMs 與 ARMs 在訓練階段所面對的挑戰。發現 MDM 模型實際上需要同時學習解決大量隱藏的補齊問題，數量成指數成長，遠超過 ARM 逐步學習序列的複雜度。因此 MDM 的「訓練在最壞情況下」的這一特性，成為模型設計和優化的關鍵制約。

推理策略：自適應解碼順序 核心貢獻在於推理階段，提出了一種「Plan for the Best」的策略：透過自適應地挑選解碼的 token 順序，動態避開最棘手、最不確定的子問題，讓模型在更易解決的子問題上先下手為強。這種策略類似於人類在做邏輯謎題時的思考方式，會先處理容易推論，後處理難度較高的部分。

具體實現上，論文提出基於模型內部不確定度評估，調整下一個解碼 token 的選取策略。不同於傳統編碼順序固定、或隨機抽樣，這樣自適應機制能有效提升生成質量與推理效率。

主要實驗結果

論文在多個代表離散生成任務上展示創新推理策略的威力，其中以邏輯性強的 Sudoku 謎題為例，最具說服力。以下是實驗亮點：

• 預訓練 MDMs 在隨機解碼順序下的解謎準確率不足 7%。
• 採用自適應解碼順序後，解謎準確率大幅攀升至約 90%，表現躍升近 13 倍。
• 相比之下，擁有七倍參數且透過教師強制（teacher forcing）明確學習解碼順序的自回歸模型，準確率卻低於本論文提出的 MDM 搭配自適應解碼策略。
• 實驗中同時驗證了訓練過程的計算負擔，說明 MDM 在理論訓練複雜度上較 ARM 高，但透過推理階段靈活順序安排彌補此劣勢。

額外實驗還驗證了在其他類型的離散組合生成任務，這種「訓練最壞、推理最好」的策略同樣有效，展現模型的通用性與實用價值。

對 AI 領域的深遠影響

本論文的發現具有多方面深遠意義：

• 理論理解層面：首次系統揭露 Masked Diffusion Models 訓練與推理中的本質困難，提供從理論層面理解模型行為的框架，推動離散生成模型理論的深化。
• 方法論創新：開創性提出自適應解碼順序策略，突破傳統固定序列生成的框架，展現了動態規劃和不確定度導向決策在推理階段的潛力。
• 實務應用價值：提升了 MDMs 在對結構性、邏輯性要求高的任務上的可用性，引發未來在自然語言處理、符號推理、組合優化等多領域的廣泛應用可能。
• 推動生成模型多樣性：相較於依賴序列化解碼的自回歸模型，MDM 強調推理階段的靈活性與多樣性，豐富了生成模型的設計空間。

總結來說，Kim 等人這篇論文透過深刻解析與創新方法，不僅優化了Masked Diffusion Models的性能，更引領研究者重新思考順序在離散生成模型中的角色和策略。未來隨著算力與理論的發展，此類模型有望成為離散結構生成與邏輯推理的重要基石，助力人工智慧更智慧地處理複雜的現實問題。

---

論文資訊
📄 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