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

隨著生成模型的快速發展，文本、生物序列等離散資料的建模也引起了廣泛關注。傳統上，自回歸模型（Autoregressive Models, ARMs）因其簡單直觀的序列生成方式，成為離散序列生成的主流框架。然而，自回歸模型推理採用固定序列解碼，導致生成效率與靈活性受限。近年來，掩碼擴散模型（Masked Diffusion Models, MDMs）逐漸崛起，提供了一種在訓練與推理階段分別取捨的全新視角，也帶來了潛在的運算挑戰與解碼策略問題。Kim 等人在其於 ICML 2025 獲獎論文《Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions》中，深入分析了 MDMs 在訓練與推理階段的核心困難，並提出有效的策略突破瓶頸，對推動離散生成模型的理論與實踐均有重要貢獻。

研究背景與動機

離散序列生成問題長期以來以自回歸架構為主，因其透過「逐步預測下一個標記（token）」的方式，確保序列依賴性。然而，這樣的架構在推理時間存在著嚴重的瓶頸：生成必須嚴格遵守固定的時間步驟，無法隨機存取或並行化，造成速度瓶頸。此外，自回歸模型極其依賴預先設定的解碼順序，即使存在更有效率或靈活的解碼路徑，模型也難以利用。

為了克服這些限制，掩碼擴散模型提出了一種基於「隨機遮蔽與重建」的訓練過程，藉由不斷從部分遮蔽的輸入恢復出完整序列，模型學會解決大量不同的「補全（infilling）」任務。在推理階段，這使得解碼順序極具彈性，理論上可以根據當前資訊選擇最合適的標記順序，潛力巨大。

然而，這樣的彈性也帶來挑戰：從訓練角度看，MDMs 需面對指數級多樣的補全問題，每個問題均可能極其困難，模型需在各種困難程度的子問題中平衡學習。而從推理角度看，若能找到優化的解碼順序策略，將能大幅提升生成效能與準確率，否則任意解碼順序可能導致生成品質不佳。

核心方法與創新

本論文的首要創新在於系統性剖析 MDM 訓練與推理中「token ordering（標記序）角色」的雙面向影響。具體分為：

1. 理論與實驗證明 MDM 訓練的嚴酷性：論文從理論角度證明 MDM 在訓練時須解決的大量補全子問題中，存在計算不可行的高難度子問題，相較自回歸模型的漸次預測，MDM遭遇的子問題複雜度呈指數增長。也就是說，讓模型學會解決所有可能的掩碼組合，其難度遠超單一路徑的自回歸訓練。
2. 提出並驗證動態解碼策略以提升推理性能：此處為論文創新的關鍵——作者設計了一套「adaptive decoding order（自適應解碼順序）」方法，推理過程中根據模型當前信心與剩餘未解碼標記的難易度，自主決定下一個解碼的標記。此策略讓模型能避開最困難的補全子問題，聚焦於最有把握的位置，逐步累積準確生成。此舉在推理階段極大地提升了 MDM 的性能。

此外，為驗證方法可靠性與廣泛性，作者針對邏輯推理型的難題（如數獨）進行深入實驗，突顯該方法在極端案例中卓越的實用價值：

• 在預訓練 MDM 基礎上應用自適應解碼順序，數獨解題正確率從不足 7% 飆升到約 90%。
• 這一表現在模型效率方面遠超參考的自回歸基線：後者雖經過教師強迫（teacher forcing）以學習正確解碼順序且參數規模是 MDM 的七倍，仍不及前者。

主要實驗結果

論文的實驗設計嚴密，包含理論分析支持與多種離散序列數據集上的實徵驗證。重點包含：

• 計算難度評估：透過數學證明與模擬，清楚描繪 MDM 所遇子問題的指數複雜性，揭示標準訓練可能受阻於困難子問題，造成學習效率下降。
• 自適應解碼策略效果驗證：在多種推理場景下，展現依靠模型信度反覆調整解碼順序，對提升生成品質與準確率有關鍵促進作用。
• 邏輯謎題（數獨）專案：相較於傳統自回歸模型和非自適應 MDM，作者提出的 adaptive MDM 在解題成功率、參數效率、推理速度均有明顯優勢。

整體實驗鏈條完整，充分證實了論文提出的方法不僅理論合理，更在實踐中具備強大競爭力。

對 AI 領域的深遠影響

本論文在 AI 生成模型領域牽動多個重要趨勢與問題：

1. 重新定義離散生成的訓練與推理范式：傳統自回歸模型束縛於固定序列解碼，自適應掩碼擴散模型在解碼順序的自由度開創了新的思路，有望於語言、編碼、規劃等多種場景打破速度與靈活性的瓶頸。
2. 突顯訓練與推理中子問題多樣性與難度的平衡挑戰：該發現警示研究者，設計能針對子問題難度自我調節學習策略的模型是未來方向，單純追求泛化覆蓋所有掩碼配置不具實務可行性。
3. 推動 Adaptive Decoding 策略的進一步研究：自適應解碼的方法證明在邏輯推理等結構嚴謹問題上的優勢，鼓勵擴展至更多複雜生成任務，如程式碼生成、知識問答、多模態生成等，提升模型推理智慧與效率。
4. 跨領域理論與應用結合的典範：本論文同時兼顧嚴謹的理論證明與生動的應用驗證，展現現代 AI 研究應追求學術深度與實務突破雙贏的典範，啟發後續相關生成模型設計。

總結而言，Kim 等人的工作不僅深化了我們對掩碼擴散模型內部機制與潛力的理解，也為離散序列生成領域提出了切實可行的新解，具備推動整體生成模型發展的廣泛意義。未來，隨著硬體與算法進一步演進與完善，結合自適應解碼的掩碼擴散模型將有望在語言生成、邏輯推理、自動編程等多個 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