隨著生成模型的蓬勃發展,尤其在離散資料領域(如自然語言、符號序列等)中,模型如何高效且靈活地進行生成成為研究熱點。過去以自回歸模型(Autoregressive Models, ARMs)為主流,因其訓練與推論流程相對直觀,但推論時只能嚴格按照固定的序列順序逐步解碼,靈活度與速度受限。近年來,掩碼擴散模型(Masked Diffusion Models, MDMs)作為一種新興生成架構,成功引入可自由調整解碼順序的能力,使推論變得更靈活、高效,但訓練過程則須解決龐大的填充(infilling)問題,計算複雜度大增。Kim 等人於 ICML 2025 所發表的《Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions》一文,針對這兩種模型在訓練與推論中的「順序策略」問題,提出深入的理論分析與實證驗證,並因其原創性與實用性榮獲該年度的 Outstanding Paper 獎。
研究背景與動機
自回歸模型(如 GPT 類模型)在離散序列生成上展現卓越成果,其依序生成下一個 token 的形式,使訓練和推論均相對簡單。然而,這種一詞接一詞的生成策略在推論時有明顯瓶頸:生成過程無法並行,並且無法根據內容或上下文彈性決定生成順序,導致對長序列的生成效率及結果品質受限。
掩碼擴散模型(MDMs)則開創性地借鑒了圖像擴散模型的成功架構,將生成問題表述成一系列「去噪」或「填空」任務:模型學習在大量可能被遮蔽的token位置中填入正確內容。這使得推論時可以不受先前生成序列位置所限,按照任意順序解碼任意 token,大幅增加解碼自由度和潛在效率。可是,這種靈活性的代價是,訓練時模型要同時學習無數子問題,每個填空設置對應一種「子序列推理」任務,導致計算複雜度呈指數爆炸。
因此,本論文的動機在於:如何理論與實務層面解析訓練時的「最難子問題」和推論時的「最優解碼策略」,期望透過「針對最壞情況訓練,而在推論時規劃最佳順序」的策略,讓 MDMs 在解決複雜離散序列生成任務時發揮最大效能。
核心方法與創新
本研究首先從理論層面分析了 MDMs 與 ARMs 在訓練上的差異。作者以形式化方法證明,MDMs 在訓練時必須應對指數級增長的子任務空間,這些子任務即各種可能的遮蔽組合填充問題,遠較 ARMs 僅需學習依序生成的單一路徑複雜得多。換言之,MDMs 天生面臨訓練計算複雜度的瓶頸。
接著,論文重點探討推論階段。由於 MDMs 允許按任意順序生成 token,如何選擇最佳的解碼順序成為關鍵。作者提出一套「自適應解碼順序策略」,根據當前已生成的部分信息動態決定下一個要解碼的 token,藉此避開「難解的子問題」,有效減少探索空間與錯誤累積。這種策略類似問題導向的「貪婪搜索改良」,但特別適配離散掩碼擴散的架構,顯著提升了推論質量與速度。
技術上,論文結合理論證明與優化算法設計,並提出一個啟發式方法來估算當下選擇 token 的困難度,指導解碼順序調整。此外,作者針對 Sudoku 等邏輯謎題這類結構嚴謹、需要高階推理能力的離散任務,進行嚴格的驗證與對比實驗,證明該策略的實用價值。
主要實驗結果
實驗部分以 Sudoku 解題任務為核心測試場景。此類問題代表了離散序列生成中高難度且結構性強的挑戰,適合驗證模型的推理能力和解碼策略的有效性。透過比較,作者發現:
- 預訓練的 MDM 透過傳統固定解碼順序,解題正確率僅約 7% 以下,表現不甚理想。
- 引入自適應解碼順序策略後,解題正確率大幅提升至約 90%。
- 同時,相較於訓練參數量是 MDM 7 倍的自回歸模型(並且經過教師強化(training by teacher forcing)以學習理想解碼順序),MDM 不只匹敵更具參數規模和順序資訊的自回歸模型,甚至在準確率上優於對手。
此外,作者進一步分析了自適應策略擇序的效果,證明動態序列規劃使模型能集中資源攻克「最難子問題」,從而避免無謂的錯誤累積,強化推論過程的穩定性,這在許多其他離散生成任務中同樣適用且具潛力。
對 AI 領域的深遠影響
本論文最核心的學術貢獻,是首次全面系統地揭示離散擴散模型訓練與推論順序設計的雙重挑戰,並創新提出「針對最壞情況訓練、推論時規劃最佳策略」的通用思維框架。這一想法顛覆了以往自回歸模型固定序列生成的觀念,為離散領域生成模型的效率與效果提升開創新路徑。
從應用層面看,自適應解碼順序策略意味著未來生成模型不必受限於單一路徑解碼,將能靈活適配多變且複雜的推理任務,提升在自然語言處理、程式碼生成、邏輯推理甚至科學計算等領域的表現。該研究的方法與理論框架亦可啟發多模態生成和更廣泛的結構化生成問題。
此外,本論文對如何在龐大複雜性與實際效能間找到平衡的洞見,對擴散模型乃至其他類型生成模型的設計策略具有指導意義。在擴散模型快速發展背景下,其提出的動態解碼策略和訓練分析可望成為未來改良擴散模型、提高推論效率及生成品質的重要參考。
總結而言,Kim 等人的工作不僅深化了對離散擴散模型本質難題的理解,更提出極具實用價值的解決方案。這項研究代表了生成模型領域的一大突破,有望推動新一代更強大且靈活的生成架構誕生,影響深遠且持續。
論文資訊
📄 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