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