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

隨著生成模型在自然語言處理、圖像合成以及結構化數據等領域的快速發展，如何高效且靈活地建構具備強大生成能力的模型，一直是研究熱點。在離散域（例如文本生成）中，自回歸模型（Autoregressive Models, ARMs）長久以來被視為主流，憑藉其簡明的左至右生成順序，能有效捕捉序列依賴關係。然而，ARMs 在推理階段必須嚴格按照既定的序列進行逐步解碼，限制了生成的靈活性與效率。近年來，掩碼擴散模型（Masked Diffusion Models, MDMs）因能在推理時以幾乎任意的順序解碼，而成為面對離散域生成的新興方案，帶來了訓練複雜度與推理靈活性之間的新挑戰。

研究背景與動機

傳統 ARMs 透過 teacher forcing 於訓練階段學習固定且嚴格的解碼順序，這使得模型生成過程清晰且易於優化，然而這種順序限制同時也使推理效率較低，且難以利用平行或非順序的解碼策略。與之相對，MDMs 在訓練階段不僅要學習填補「掩碼位置」的多種可能狀況，這種多樣子問題的總數呈指數級增長，因而訓練極具挑戰；但它們的優勢是推理時不必拘泥於固定解碼順序，理論上可以根據具體問題靈活調整令生成更有效率。

本論文由 Kim 等人提出，目的在於深入理解 MDMs 訓練時碰到的「最壞情況學習」問題，與推理時「最佳順序選擇」的潛力。他們試圖透過理論分析與實驗驗證，揭露 MDMs 在離散生成任務中訓練與推理策略間的權衡，並提出一套能在推理階段智能選擇解碼順序的方法，大幅提升模型效能並降低推理難度。

核心方法與創新

首先，論文在理論層面證明 MDM 面臨的核心困難是：訓練階段必須同時解決一組「指數級」的補全子問題，而其中多數子問題是計算不可行的，與 ARMs 僅需解決固定序列子問題明顯不同。也就是說，MDMs「被迫訓練於最壞情況」（Train for the Worst），使得訓練成本與挑戰大幅增加。

另一方面，在推理階段，MDMs 能夠不受限於預先設定的固定順序，其解碼令牌的順序可以動態調整。作者提出一種基於模型自身信心水準的「適應性解碼策略」，該策略根據當前已生成資訊動態決定下一個要解碼的令牌位置，藉此「規避較難的子問題」，將推理任務拆成多個相對容易的子問題逐一解決（Plan for the Best）。此策略的關鍵點在於充分利用模型內在不確定性評估，實時調整解碼順序，提升生成的準確性與效率。

此外，研究者進一步針對經典邏輯推理問題 Sudoku 進行實驗。與傳統 ARMs 需要繁複設計解碼順序不同，MDMs 在採用他們提出的適應性推理策略後，展現了顯著突破。

主要實驗結果

論文先在 Sudoku 題庫上對比實驗發現：

• 採用固定順序推理的預訓練 MDMs，其準確率極低，甚至低於 7%。此結果印證此類模型若不利用順序靈活性，將因為難解子問題而表現不佳。
• 引入作者提出的適應性解碼策略後，MDMs 的解答準確率躍升至約 90%，大幅超越固定順序的限制。
• 相較於 ARMs，即使後者的參數量大約是 MDMs 的七倍，且透過 teacher forcing 專門學習最適合的解碼順序，MDMs 依然在準確率上取得更優表現。

額外分析也彰顯該適應性推理策略減少了推理過程中「最難令牌優先解碼」的概率，顯著降低了模型陷入最壞子問題的機會，提升了推理效率和效能。此外，作者也驗證了該策略在其他離散生成任務中的泛用性，展示出良好的擴展潛力。

對 AI 領域的深遠影響

此篇論文的貢獻，不僅是在理論與實驗上確立了 MDMs 在離散生成領域中的新地位，更提出了切實可行的解決方案來突破訓練和推理的困境。對於 AI 研究社群而言，他們擺脫了傳統 ARMs 固有的編碼順序限制，為未來設計更加靈活且高效的生成模型提供了嶄新的思路。

具體而言，適應性解碼策略引入了動態規劃與不確定性估計的概念，這對設計可針對複雜、結構化離散問題（如邏輯推理、程式碼生成、結構化文本等）有深遠影響。此方法可望激發後續研究聚焦於推理階段的動態策略優化，推動生成模型往更高效且泛化更強方向發展。

此外，該文揭露了 MDMs 在「訓練面臨指數級挑戰」的本質，強調了強化訓練策略與模型架構設計的重要性。未來研究可能會從減少訓練複雜度或引入更聰明的訓練樣本挑選機制著手，進一步降低資源需求，提升可用性。

總體而言，As Kim 等人這篇在 ICML 2025 被評為 Outstanding Paper 的研究，不僅深化了我們對離散擴散模型本質的理解，也為 AI 生成技術的靈活推理策略帶來革命性的啟示，預期將在離散生成模型及跨領域應用中持續發酵並激發新一波技術革新。

---

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