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

隨著生成模型在自然語言處理、計算機視覺等領域持續發展，如何有效且靈活地生成離散序列資料成為研究熱點。傳統主流的生成架構如自回歸模型（Autoregressive Models, ARM）因採用固定的生成順序，雖在訓練上相對簡單，但在推論時嚴格受限於此序列，有時難以靈活應對各種下游任務和輸入變化。
近期興起的Masked Diffusion Models（MDMs）則提出了一種新穎思維：在訓練階段必須解決大量隨機遮蓋補全（infilling）問題，努力學習如何填補任意被mask掉的token，使模型訓練時更為複雜；但推論階段卻能自由選擇任何順序進行token生成，極大提升推論彈性與效率，使它成為離散序列生成中的有力候選者。

本篇2025年ICML傑出論文《Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions》由Kim等人提出，深刻分析並理論化檢視了MDMs在訓練與推論階段的「矛盾」本質，並探討如何透過「智慧安排推論順序」突破訓練時的巨大挑戰，實現效能大幅提升。

研究背景與動機

在生成式AI中，離散資料（如語言tokens、象徵符號、拼圖格）需要模型能夠對序列做合理且高品質的生成。傳統ARM假設固定序列順序，透過一對一的教師強迫（teacher forcing）學習下個token，但序列長度與決策連鎖反應導致推論時間線性成長，且生成方向無法隨意改變。

Masked Diffusion Models則透過隨機遮蓋部分token，讓模型在訓練時必須從種種重構子問題中學習，理論上這使MDMs涵蓋的潛在生成任務數量呈指數爆炸。雖然訓練難度大幅增加，MDMs推論時可採用多元生成策略（arbitrary token decoding order），使靈活度翻倍。但，這種訓練-推論「二元悖論」尚無明確理論與實證加以說明，且如何選擇推論時「最佳token生成順序」提升表現仍屬未知。

核心方法與創新

本文的首要創新，在於從理論到實驗全方位刻畫MDMs訓練上的巨量子問題挑戰與推論時彈性策略的重要性：

• 訓練階段理論分析：作者證明，在訓練MDMs時，模型實際上必須學習一個涵蓋所有可能被mask的子問題集合，這個集合的規模是指數級成長，遠超過ARM固定順序訓練的子集。也就是說，MDMs的訓練難度本質上遠高於ARM。
• 推論階段適應性順序策略：論文提出一種通過「動態決策」推斷token生成順序的方法，利用模型自身對難易程度的判斷，優先生成較容易的token，從而避開在訓練中較難學習或推理的子問題。
• 理論與實驗結合論證：結合邏輯拼圖（像是數獨）等複雜問題，展示MDMs透過適應性生成順序，能顯著提升推論表現，遠超不考慮順序的基線模型。

主要實驗結果

為顯示方法強大，作者選擇了具挑戰性的解難任務—數獨（Sudoku）謎題作為實驗舞台：

• 預訓練的MDM若不採用適應性推論策略，在數獨解題正確率不到7%。
• 透過智慧的「動態序列選擇」方案，MDM的解題率攀升至約90%，顯示適應性token生成順序在實務中有卓越效果。
• 即便面對參數量多出7倍的ARM模型，該ARM是經教師強迫訓練以學習固定正確的token生成順序，但在解題精度仍不及MDM，顯示MDM推論靈活性的優勢。

這些結果同時佐證了本文理論預測的合理性與實務可行性，也為MDMs日後應用到其他複雜序列任務奠定了信心和基礎。

對 AI 領域的深遠影響

本文從理論和應用雙重層面，系統剖析並突破了Masked Diffusion Models長期以來存在的訓練難度與推論策略挑戰，帶來多方面的影響：

1. 促進序列生成模型架構多樣化：MDMs跳脫傳統arm固定順序生成的框架，使研究者具備更彈性自由的生成策略，為序列生成問題開啟全新思路。
2. 推動順序決策策略研究：本論文強調的「根據難易程度動態決定生成順序」理念，讓未來生成模型不再被動，而能主動調整推論流程，提升性能和效率。
3. 強化離散空間問題求解能力：在數獨等邏輯問題上的優異表現，讓生成模型在求解複雜規則與限制時，更具智慧和彈性，可望推廣到自動定理證明、結構化任務等範疇。
4. 理論架構為後續研究提供藍本：作者對MDMs訓練和推理難度的量化分析，為後續模型設計和優化提供關鍵理論依據，將影響未來生成模型方法論的發展。

綜合來說，Kim等人這篇獲得ICML傑出論文獎的作品，不僅發現並解析了Masked Diffusion Models的重要本質問題，更透過「訓練迎接 worst case，推論策劃 best case」的策略，開創了一條可行且高效的離散生成新方向。對希望突破生成靈活性和推論效率瓶頸的AI工程師與學界研究生而言，這篇論文的重要見解和方法，無疑是未來深化生成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