Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions 深度解讀

在生成模型領域，尤其是處理離散資料的場景中，近年出現了一種名為 Masked Diffusion Models（MDMs） 的新興架構。這種方法被認為是自回歸模型（Autoregressive Models, ARMs）的一個有力替代方案，因其在訓練與推理階段展現出截然不同的能力與挑戰。Kim 等人在 ICML 2025 以論文《Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions》榮獲 Outstanding Paper，本文將深入解析這篇獲獎論文的研究動機、方法創新、實驗成果與對 AI 領域的啟示。

研究背景與動機

傳統上，離散生成任務（如自然語言處理中的語言模型、結構化問題求解等）多數依賴自回歸模型。自回歸架構會按照固定的順序逐步生成每個輸出元素（Token），透過前一步的輸出條件化下一步。然而此種方法在推理時必須遵循嚴格的順序，限制了推理的彈性與效率，且訓練階段需要額外技巧（如教師強迫 teacher forcing）以穩定模型學習和避免暴露偏差（exposure bias）。

為了克服此限制，研究者轉向擴散模型（Diffusion Models）架構中針對離散資料的變體，即 Masked Diffusion Models（MDMs）。MDMs 在訓練時將任務設計成學習從隨機遮蔽狀態恢復完整序列，即要學會「填空（infill）」各種可能的遮蔽組合。這讓模型能在推理階段自由選擇生成順序，具有高度的解碼靈活性與多樣策略空間。

然而，這也帶來新的挑戰：

• 在訓練階段，MDMs需要學習解決指數級多樣的遮蔽組合子問題，理論與實務上的計算複雜度遠高於自回歸模型。
• 在推理階段，選擇適合的生成順序並非易事，不同的令牌解碼次序將直接影響生成質量與成功率。

因此，Kim 等人決定系統性探索「訓練中面對的困難問題」與「推理中選擇合適編碼策略」兩方面的平衡與機制，以挖掘 MDMs 潛力並解決其固有瓶頸。

核心方法與創新

本論文的核心貢獻可分為兩大部分：

1. 理論與實證揭示 MDMs 訓練中的困難問題

作者從理論上分析，MDMs 在訓練階段要應對的填空任務實質上是解決一個組合性極強的填補問題集合。這些子問題中包含許多計算複雜度極高甚至無法有效求解的「最壞」情況，他們描述為“train for the worst”。透過數理分析和實驗驗證，證明這種計算困難遠超自回歸模型中一步一步預測下一個 token 的簡單任務。這提醒了研究者，僅依靠端到端訓練是不夠的，必須有策略調整或輔助技術應對此難題。

2. 自適應推理策略以突破推理瓶頸

另一方面，論文提出靈活且高效的自適應生成順序選擇策略（adaptive token decoding order）。此策略根據每一步模型對尚未生成令牌的預測困難度動態調整解碼次序，優先生成相對容易、確定性的 token，從而推遲或繞過解決難度過高的子問題。譬如在解決須邏輯推理的問題時（如 Sudoku, 結構謎題等），動態順序選擇大幅提升解決成功率。

透過這種「先為最壞打基礎（train for the worst），而計劃時迎向最好（plan for the best）」的策略，MDMs 在推理時展現出強大且靈活的生成能力，可動態調整生成流程，打破傳統固定序列的限制。

主要實驗結果

為了驗證理論與算法主張，作者針對多種離散生成任務設計了廣泛實驗：

• 理論驗證：在訓練難度分析上，實驗結果與理論模型契合，證明設計的多樣填空子問題確實涵蓋計算複雜度高的案例。
• 邏輯謎題解決：在 Sudoku 等經典邏輯推理任務中，論文展示：對預訓練好的 MDM，透過自適應解碼順序推理，模型解謎成功率由原本不足 7% 提升到約 90%。這一驚人成果不僅大幅超越基本隨機或固定順序解碼策略，甚至勝過參數量是其七倍、且利用教師強迫訓練得知“正確解碼序”的先前自回歸模型。
• 靈活順序優勢：實驗還揭示，不同生成任務中，自適應排序能因應具體難度分布及上下文信息靈活調整，有效避開訓練期帶來的部分困難子問題。

這些成果充分證明，MDMs 在結合合理訓練策略與推理策略時，具備潛力成為更為強大且通用的離散生成模型。

對 AI 領域的深遠影響

這篇論文的發表不僅豐富了離散生成領域的理論基礎，更在技術實踐層面具備多項深遠啟示：

1. 打破固定序列生成的限制：傳統自回歸模型固守嚴格的生成序列，限制了推理彈性和效率。MDMs 提供了一條可變解碼順序的新路徑，開創解碼策略自適應的全新範式。
2. 理論視角促進訓練設計精進：研究揭露了 MDMs 訓練中存在的最壞子問題，促使後續工作思考如何降低樣本困難度、設計更有效的遮蔽策略或輔助機制，提高模型學習效率與泛化能力。
3. 強化結構與邏輯推理能力：透過自適應解碼策略，MDMs 不僅在語言生成，對結構化邏輯問題的解決也有突破，提醒我們生成模型未來可向更廣泛的結構化人工智能方向邁進。
4. 跨領域影響力：該方法同時對語言、程式碼生成、組合優化問題等離散生成任務具有潛力，引導了融合擴散模型與自回歸模型優點的混合策略探索。

總結而言，Kim 等人提出的「Train for the Worst, Plan for the Best」理念，不僅為 Masked Diffusion Models 提供了清晰的理論基石與實務策略，同時也展示出 ADAPTIVE 解碼策略激發潛力的強大威力。這項研究突破了傳統生成模型的局限，為未來離散生成系統設計樹立典範，並將持續推動 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