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

隨著生成模型在自然語言處理與序列建模領域的重要性日益增長，研究者不斷探索更高效且靈活的模型架構。在離散序列生成領域中，自回歸模型（Autoregressive Models, ARM）長期以來是主流，這類模型透過固定的序列順序一步步預測下一個標記。然而，ARMs 在訓練與推論過程中都受限於此固定序列順序，造成推論彈性不足，且在序列較長時推論速度瓶頸顯著。相較之下，掩碼擴散模型（Masked Diffusion Models, MDMs）因其能在推論時靈活地以任意順序生成標記，近期受到關注，展現出可與 ARMs 媲美甚至超越的潛力。

本篇由 Kim 等人在 ICML 2025 發表並榮獲 Outstanding Paper 獎的論文《Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions》，聚焦於深入剖析 MDMs 中標記生成的排序策略，並提供理論、實證與實務三方面的洞見，以解決當前 MDM 面臨的關鍵挑戰。

研究背景與動機

自回歸模型雖然在序列生成任務上表現優異，但受限於其固定的生成順序，導致推論時間難以加速且無法利用更靈活的生成策略。掩碼擴散模型作為新興的生成架構，透過在訓練階段學習無數種「填空」(infilling) 問題，理論上可在推論時任意決定生成標記的順序，進而提升推論效率與靈活度。然而，MDMs 的訓練複雜度極高，因為他們必須同時學會大量的子問題，這背後的組合爆炸，使得推論時選擇最優順序成為必要。

更重要的是，實際應用中如何選擇推論順序直接影響生成品質。若推論策略不當，模型將面臨難解的子問題，導致生成結果品質低落。論文作者因此提出問題：MDMs 如何在「訓練時學習最壞（worst-case）」的問題、「推論時規劃最佳（best-case）」的生成順序間取得平衡？同時，此問題的理論與實踐意義將觸及序列生成模型設計的根本。

核心方法與創新

本研究從理論與實驗雙重視角入手。首先，作者分析並證明 MDMs 在訓練時面臨的子問題複雜性遠高於 ARMs：ARM 固定的生成順序將決策空間壓縮，而 MDMs 必須學習任意掩碼組合下的預測，面臨指數級的子問題。這解釋了為何 MDM 訓練效果有時不穩定，甚至面臨訓練效率瓶頸。

其次，為了克服此瓶頸，論文關鍵貢獻在於提出了一套「自適應變換推論順序」策略。該策略在推論階段根據模型對各標記生成難度的估計，動態選擇生成順序，優先生成「容易」的標記，避免陷入高度不確定或資訊不足的子問題。此方法形象地被作者比喻為「為最壞情況訓練，為最好情況規劃」，即訓練時讓模型面對最嚴苛的挑戰，推論時則聰明地避開困難，從而實現效能提升。

具體而言，該方法包括：1）評估每個尚未生成標記的生成難度或不確定性；2）根據此測度動態調整生成順序；3）跨多次迭代回饋調整策略，確保生成流程收斂穩定。這種機制本質上將 MDM 的靈活性與策略性適應能力結合，極大提升了推論效率與準確率。

主要實驗結果

作者將所提出的自適應推論策略在多個離散序列生成任務中驗證，包括經典的邏輯謎題如數獨（Sudoku）求解。數獨問題具有明確的邏輯結構且推理複雜，非常適合作為生成策略評估基準。

實驗結果令人印象深刻。對於在固定順序下預訓練的 MDMs，傳統推論方式成功率不過 7% 以下，但透過自適應推論策略，成功率飆升至約 90%，大幅提升近 13 倍。此外，該策略下的 MDM 不僅成功率超過同樣任務規模下的 ARMs，且其模型參數量僅為 ARMs 的七分之一，跳脫了參數量與效能的正相關限制。

這特別值得注意的是，ARMs 常透過教師強化（teacher forcing）明確學習固定順序，而 MDMs 則不必嚴格依賴順序標的，展現出更高的泛化與策略彈性。更進一步地，該策略在其他複雜離散生成任務中也展現出一致的優勢，證明其方法的廣泛性與實用性。

對 AI 領域的深遠影響

這份研究成果為生成模型社群提出了全新的視角：在複雜離散領域，訓練階段需要面對多樣且困難的子問題，而推論階段則可透過策略性選擇生成順序來避開難題，兩者的配合才能發揮模型真正潛力。

首先，本研究拓展了生成模型的訓練與推理理論，明確指出固定生成序列並非唯一途徑，更靈活的推論調度策略是解決序列生成效率與品質問題的關鍵。此理論框架可影響未來模型設計，鼓勵研究者擺脫單一生成序列的限制。

其次，對於實務場景，尤其是需要解決高難度推理問題或多樣化生成任務，MDM 搭配自適應推論策略提供了高效能且節省參數的解決方案。從文字生成、抽象推理、遊戲策略到程式碼生成等多領域均可受益。

最後，此論文促使社群重新思考深度生成模型的推理順序問題，啟發了後續相關研究在如何融合學習與推理策略、自適應調度等方向上的探索，將推論靈活性作為核心研究課題之一。

總結而言，《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