今天小編分享的科學經驗:清華&英偉達揭露ICML最佳論文隐藏數值問題,簡單掩碼模型是等價更好選擇,歡迎閱讀。
掩碼離散擴散模型,可能并沒有看上去那麼厲害。
這是清華及英偉達研究人員最新提出的觀點。
他們發現,作為離散擴散模型中性能最強的類别,掩碼擴散模型可能有點 " 被包裝過度 " 了。為啥呢?
第一,這類模型所宣稱的超高性能,其實是由于一個技術上的小瑕疵,用 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
論文鏈接:
https://arxiv.org/abs/2409.02908
論文博客:
https://zhengkw18.github.io/blog/2024/mdm/
一鍵三連「點贊」「轉發」「小心心」
歡迎在評論區留下你的想法!
— 完 —
學術投稿請于工作日發郵件到:
标題注明【投稿】,告訴我們:
你是誰,從哪來,投稿内容
附上論文 / 項目主頁鏈接,以及聯系方式哦
我們會(盡量)及時回復你
點亮星标
科技前沿進展每日見