清華&英偉達揭露ICML最佳論文隱藏數值問題，簡單掩碼模型是等價更好選擇 | ICLR25
掩碼離散擴散模型，可能并沒有看上去那么厲害。
這是清華及英偉達研究人員最新提出的觀點。
他們發現，作為離散擴散模型中性能最強的類別，掩碼擴散模型可能有點“被包裝過度”了。為啥呢？
第一，這類模型所宣稱的超高性能，其實是由于一個技術上的小瑕疵，用32位計算時，模型會產生一種“降溫”效果，使模型看起來表現很好，但實際上只是多樣性被降低了。用更精確的64位計算，就會發現它們的表現并不如宣稱的那么好。
第二，這些模型引入了“時間”的概念，看起來很高級，但研究發現這完全沒必要。
第三，這些模型其實與已有的簡單掩碼模型完全等價，只要正確設置簡單模型的參數，就能達到相同效果。
目前，這篇研究已入選ICLR 2025。
具體說了啥？一起來看。
背景
隨著SEDD獲得ICML 2024最佳論文獎，起源于D3PM的離散擴散模型迎來了復興并成為自回歸范式的有力競爭者，在文本、蛋白質等離散序列生成任務上掀起了研究熱潮。
作為離散擴散模型中性能最強的類別，掩碼式離散擴散模型（簡稱掩碼擴散模型）在后續工作中被進一步簡化，從而在理論形式上與連續空間擴散模型對齊。
掩碼擴散模型通過引入一個連續的“時間”或“噪聲水平”的概念，定義了一個從原始數據逐漸“加噪”（掩碼）到完全掩碼狀態的前向過程，以及一個學習從掩碼狀態逐步“去噪”（預測被掩碼部分）恢復數據的反向（生成）過程。
在使用生成式困惑度（Gen PPL）作為衡量文本生成質量的指標時，掩碼擴散模型在先前工作中均顯示出了隨采樣步數增加的性能提升，并在足夠多步數下超越自回歸模型。
這種對比是否公平？同時，作為離散空間中的“擴散”模型，是否意味著其可以借鑒標準擴散模型相關算法來增強性能？
論文從訓練和采樣兩個方面對掩碼擴散模型進行解構。
掩碼擴散模型與掩碼模型的差異
盡管掩碼擴散模型借鑒了擴散模型的框架，但其核心操作與經典的掩碼語言模型（如 BERT，Mask-Predict） 或掩碼圖像生成模型（如 MaskGIT）有著相似之處：都是對部分內容進行掩碼（masking），然后預測被掩碼的內容。
BERT在訓練時只會掩碼一小部分token，適用任務為表征學習、文本理解而非生成，而Mask-Predict與MaskGIT擴大了掩碼比例的范圍并可用于文本、圖像生成。
相比于掩碼模型，掩碼擴散模型引入了一個關鍵的復雜性：時間步（time step）。其訓練和采樣都嚴格依賴于一個預先定義的、隨時間變化的掩碼（噪聲）調度。
模型需要根據當前的時間步 t 來預測原始數據。
具體而言，它和掩碼模型的差異體現在：
在訓練中，掩碼模型被掩碼的token數量及不同掩碼比例對應的損失權重可以隨意設定；掩碼擴散模型同一時間對應被掩碼token的數量是不確定的，不同時間的分布及權重需要特殊設置使得損失構成模型似然（likelihood）的證據下界（ELBO）。
在采樣中，掩碼模型按照token為粒度，逐token解碼；掩碼擴散模型以時間為粒度進行離散化，從時間t轉移到更小的時間s時，每個token被解碼與否通過概率采樣決定，被解碼token的數量是不確定的。
掩碼擴散模型的采樣存在隱性數值問題
先前評估掩碼擴散模型性能的關鍵指標之一是Gen PPL，其通過計算參考模型（如GPT-2）對模型生成內容的“驚訝程度”來衡量生成質量。
然而，Gen PPL 指標對采樣過程中的超參數（如采樣溫度）極為敏感，并且可以通過調整這些參數輕易地“刷低”數值，但這并不代表模型本身的生成能力有實質提升。
本研究首次揭露，掩碼擴散模型特有的采樣過程存在隱藏的數值問題，即使在常用的32位浮點數精度下也會帶來類似于降低溫度的效果。
為了顯示這一點，論文額外測試了生成句子的熵（entropy）來衡量生成多樣性。
隨著采樣步數的增加，Gen PPL不斷下降并超過自回歸模型（左圖），然而熵也在持續降低（右圖）。
當采樣過程以64位精度進行時，熵穩定在與自回歸模型類似的水平，而Gen PPL則顯著升高并遠遠落后于自回歸模型。
論文通過進一步的數學推導，從理論上解釋了這一溫度降低效果的根源。
具體而言，在[0,1)區間上均勻采樣的浮點數實際范圍為[0,1-ε]，其中ε是一個接近0的小數，這會導致基于Gumbel-max技巧的類別采樣（categorical sampling）存在截斷問題。
最終的分布不服從原先的類別概率π，而會偏移到
，其中
這一偏移會加強原先概率已經比較大的類別，從而達到類似降低溫度的效果。
通過對類別采樣部分的代碼做對照試驗，文章驗證了64位采樣+手動截斷確實可以復現32位采樣的效果。
同時，上述數值問題對于單個token并不顯著，逐token解碼的模型（如自回歸模型、掩碼模型）在32位下基本不受影響。
然而，此問題會在掩碼擴散模型中額外影響所有token之間的交互，導致某些token被優先解碼，進一步降低生成多樣性。
可以說，這是掩碼擴散模型采樣過程獨有的數值問題。
掩碼模型與掩碼擴散模型的等價性
先前工作從最優網絡的角度證明了掩碼擴散模型中的時間并不必要，本論文進一步在訓練和采樣兩方面證明掩碼擴散模型和掩碼模型的等價性。
具體而言：
在訓練損失函數上，掩碼擴散模型與時間有關的似然下界等價于掩碼模型的以token為粒度的損失函數，只要滿足：（1）被掩碼token的個數n在1和L之間均勻采樣，其中L是序列的總長度（2）預測損失對n取均值，即施加“似然權重”1/n來實現最大似然訓練。
需要注意的是，對不同時間/掩碼比例施加的權重并不影響網絡在無限容量下的最優值，而決定了網絡訓練過程中的重點優化區域。
文本生成的自回歸范式采取了最大似然訓練，而在圖像上，最大似然訓練往往會帶來生成質量的下降。
在采樣過程上，掩碼擴散模型逆時間的采樣過程可以通過論文提出的首達采樣器（first-hitting sampler），轉化為與掩碼模型相同的逐token采樣，最多需要L步便可達到沒有離散化誤差的精確采樣，而掩碼擴散模型原有采樣過程需要對時間無限細分才能完全精確。同時，采用逐token解碼可以避免上文所述的隱藏數值問題。
結語
掩碼擴散模型引入的“時間”概念可能不僅無益，反而有害（導致數值問題和不必要的復雜性）。
同時，其雖然帶有“擴散”兩字，但與連續空間上的擴散模型及其相關算法關系不大，如在論文中，作者仿照擴散模型為掩碼擴散模型開發了高階采樣算法，其并不如連續空間中的加速效果顯著。
在實踐中，使用掩碼擴散模型、引入連續時間相關的訓練/采樣過程或許并不必要，簡單的掩碼模型（如 MaskGIT 及其變種）在概念上更簡潔，實現上更穩定，并且在理論上具有同等的潛力。
掩碼模型作為自回歸模型使用隨機token順序和雙向注意力機制的變種，同樣是基于似然的模型，可以作為建模離散數據生成的另一種選擇。
由于雙向注意力和KV cache機制不兼容，掩碼模型在長上下文的推理速度上相較自回歸模型存在瓶頸。
近期工作通過在雙向注意力和因果注意力機制之間插值、使用隨機順序的自回歸模型等方法使模型保持雙向感知能力的同時，推理速度向自回歸模型靠近。
也有工作探究非掩碼類型的離散擴散模型與連續空間擴散模型的理論聯系，其擴散機制更加屬實，而非如掩碼擴散模型一樣是可有可無的噱頭。
論文第一作者鄭凱文為清華大學計算機系三年級博士生，在ICML、NeurIPS、ICLR發表擴散模型相關一作5篇。文章通訊作者為朱軍教授，合作者張欽圣、陳永昕、毛含子為英偉達研究員，劉洺堉為英偉達副總裁與Deep Imagination研究組主管。
論文標題：Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling
