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

在生成模型迅速發展的今天，如何平衡訓練效率與推論靈活性成為一大挑戰。《Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions》一文由Kim等學者於ICML 2025發表並獲得Outstanding Paper獎項，針對近期備受關注的Masked Diffusion Models（MDMs，遮罩擴散模型）提出了深入而系統的理論分析與實驗驗證。此篇論文不僅深化了我們對MDMs訓練與推論之間微妙權衡的理解，更提出了一套創新的推論策略，顯著提升MDMs在離散序列生成上的性能，為該領域帶來重要突破與啟發。

一、研究背景與動機

傳統序列生成任務中，自回歸模型（Autoregressive Models, ARMs）是主流架構。它們透過依序預測下一個token，學習生成序列，推理簡單且訓練效率較高。然而ARMs在推論時必須依序生成，限制了其靈活性且無法並行，加重實務應用上的延時負擔。

近年興起的Masked Diffusion Models（MDMs）則透過隨機遮蔽序列中的多個token，訓練模型學習重構缺失部分，鼓勵模型能在推論過程中以

任意順序生成token

。這類模型在推論階段具有高度彈性，能快速平行解碼，理論上擁有更優的推論速度與靈活性。

然而MDMs也有挑戰：在訓練階段它必須學習處理數以指數增長的多重填補子問題（infilling subproblems），遠比ARMs同等規模模型面臨更複雜的問題。先前研究多半對此缺乏系統性分析，對於如何從理論與實務面平衡這兩者仍舊模糊不清。同時，MDMs雖然推論靈活，但如何選擇最佳或優良的token生成順序以避開困難子任務，仍是一大難題。

二、核心方法與創新

本文結合理論證明與實驗探索，從兩個核心角度展開：

1. 訓練難題的理論刻畫：作者首先從計算複雜性視角出發，數理分析証實MDMs在訓練需面對的填補子問題複雜度遠高於ARMs。ARMs透過固定的序列順序，解決順序生成任務；相對地，MDMs必須同時學習解決指數級的多種遮罩組合子問題，令訓練過程在計算量與優化難度上大幅增加。此發現幫助社群正確認識MDMs的訓練瓶頸與挑戰，不再僅靠直覺判斷。
2. 適應性token生成順序策略：面對複雜的训练任務與推論彈性，作者提出創新的「基於場景自適應決策的token解碼順序演算法」。該方法根據當前解碼狀態與難易度評估，動態選擇接下來要生成的token，刻意避開難以預測的子問題，逐步減少不確定度。此策略結合啟發式與深度模型驅動，使MDMs得以「計劃最好的推論路徑」，儘管模型訓練時「針對最壞的（最難的）子問題」進行優化。

此方法打破了以往MDMs固定或隨機生成順序的框架，首次明確將推論中的解碼順序問題作為策略優化目標，顯著提升模型效能與實用性。

三、主要實驗結果

為驗證理論與方法，論文設計多項實驗，涵蓋合成邏輯謎題（例如數獨、Sudoku）及語言生成任務：

• 計算複雜性實驗：實證結果支持理論分析，驗證MDMs在解決複雜填補問題時，模型訓練中確實存在更高的計算負擔與難易度。
• 推論順序策略驗證：在數獨問題上，預訓練的MDMs採用適應性解碼策略，解題正確率從原本不足7%飆升至約90%，大幅超越固定順序的ARMs。更令人驚艷的是，即使是ARMs模型規模擴大7倍，且特別透過teacher forcing技術學習「最佳解碼順序」，仍無法超越此MDMs策略。
• 多域驗證：論文還展示此順序選擇策略對多種不同類型的文字與序列生成任務均具廣泛適用性與強健效果，說明其方法具備良好的泛化能力。

四、對 AI 領域的深遠影響

此篇論文不僅填補了Masked Diffusion Models訓練與推論間理論與實務的鴻溝，更提出了一條全新的觀點：

• MDMs的設計哲學應是「訓練搭配最壞情境，但推論以最優策略達致最佳成效」，體現了機器學習中對抗性思考與靈活性調節的完美結合。
• 論文所提出的適應性token生成順序演算法，開啟了序列生成任務中解碼順序策略設計的新方向。未來Generative AI不再是被動順序式生成，而能結合決策與策劃機制，智能選擇生成路徑以突破傳統限制。
• 在具體應用層面，MDMs的新推論策略有望推動邏輯推理、編碼生成、文本完成等領域快速進步，尤其在面對復雜結構或多解序列任務時展現獨特優勢。
• 此外，本文的計算複雜性分析有助於後續研究者優化MDMs的訓練過程，設計更高效的資料擴散與遮蔽策略，提升擴散模型在離散領域的實用性與規模化能力。

總結來說，Kim等人的這篇研究不僅深化了學界對Masked Diffusion Models的理解，促使我們重新審視生成模型中訓練與推論的權衡問題，也展示了策略性解碼在提升模型推論能力上的巨大潛力。這些貢獻對生成式模型設計、優化與應用均具有深遠影響，值得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