在建模復(fù)雜的數(shù)據(jù)分布方面，擴(kuò)散生成模型表現(xiàn)出色，不過它的成果大體上與表征學(xué)習(xí)（representation learning）領(lǐng)域關(guān)聯(lián)不大。

通常來說，擴(kuò)散模型的訓(xùn)練目標(biāo)包含一個(gè)專注于重構(gòu)（例如去噪）的回歸項(xiàng)，但缺乏為生成學(xué)習(xí)到的表征的顯式正則化項(xiàng)。這種圖像生成范式與圖像識(shí)別范式差異明顯 —— 過去十年來，圖像識(shí)別領(lǐng)域的核心主題和驅(qū)動(dòng)力一直是表征學(xué)習(xí)。

在表征學(xué)習(xí)領(lǐng)域，自監(jiān)督學(xué)習(xí)常被用于學(xué)習(xí)適用于各種下游任務(wù)的通用表征。在這些方法中，對(duì)比學(xué)習(xí)提供了一個(gè)概念簡單但有效的框架，可從樣本對(duì)中學(xué)習(xí)表征。

直觀地講，這些方法會(huì)鼓勵(lì)相似的樣本對(duì)（正例對(duì)）之間相互吸引，而相異的樣本對(duì)（負(fù)例對(duì)）之間相互排斥。研究已經(jīng)證明，通過對(duì)比學(xué)習(xí)進(jìn)行表征學(xué)習(xí)，可以有效地解決多種識(shí)別任務(wù)，包括分類、檢測(cè)和分割。然而，還沒有人探索過這些學(xué)習(xí)范式在生成模型中的有效性。

鑒于表征學(xué)習(xí)在生成模型中的潛力，謝賽寧團(tuán)隊(duì)提出了表征對(duì)齊 (REPA) 。該方法可以利用預(yù)訓(xùn)練得到的現(xiàn)成表征模型的能力。在訓(xùn)練生成模型的同時(shí)，該方法會(huì)鼓勵(lì)其內(nèi)部表征與外部預(yù)訓(xùn)練表征之間對(duì)齊。

REPA 這項(xiàng)開創(chuàng)性的成果揭示了表征學(xué)習(xí)在生成模型中的重要性；然而，它的已有實(shí)例依賴于額外的預(yù)訓(xùn)練、額外的模型參數(shù)以及對(duì)外部數(shù)據(jù)的訪問。

簡而言之，REPA 比較麻煩，要真正讓基于表征的生成模型實(shí)用，必需一種獨(dú)立且極簡的方法。

這一次，MIT 本科生 Runqian Wang 與超 70 萬引用的何愷明出手了。他們共同提出了 Dispersive Loss，可譯為「分散損失」。這是一種靈活且通用的即插即用正則化器，可將自監(jiān)督學(xué)習(xí)集成到基于擴(kuò)散的生成模型中。

• 論文標(biāo)題：Diffuse and Disperse: Image Generation with Representation Regularization
• 論文鏈接：https://arxiv.org/abs/2506.09027v1

分散損失的核心思想其實(shí)很簡單：除了模型輸出的標(biāo)準(zhǔn)回歸損失之外，再引入了一個(gè)用于正則化模型的內(nèi)部表征的目標(biāo)（圖 1）。

直覺上看，分散損失會(huì)鼓勵(lì)內(nèi)部表征在隱藏空間中散開，類似于對(duì)比學(xué)習(xí)中的排斥效應(yīng)。同時(shí)，原始的回歸損失（去噪）則自然地充當(dāng)了對(duì)齊機(jī)制，從而無需像對(duì)比學(xué)習(xí)那樣手動(dòng)定義正例對(duì)。

一言以蔽之：分散損失的行為類似于「沒有正例對(duì)的對(duì)比損失」。

因此，與對(duì)比學(xué)習(xí)不同，它既不需要雙視圖采樣、專門的數(shù)據(jù)增強(qiáng)，也不需要額外的編碼器。訓(xùn)練流程完全可以遵循基于擴(kuò)散的模型（及基于流的對(duì)應(yīng)模型）中使用的標(biāo)準(zhǔn)做法，唯一的區(qū)別在于增加了一個(gè)開銷可忽略不計(jì)的正則化損失。

與 REPA 機(jī)制相比，這種新方法無需預(yù)訓(xùn)練、無需額外的模型參數(shù)，也無需外部數(shù)據(jù)。憑借其獨(dú)立且極簡的設(shè)計(jì)，該方法清晰地證明：表征學(xué)習(xí)無需依賴外部信息源也可助益生成式建模。

帶點(diǎn)數(shù)學(xué)的方法詳解

分散損失

新方法的核心是通過鼓勵(lì)生成模型的內(nèi)部表征在隱藏空間中的分散來對(duì)其進(jìn)行正則化。這里，將基于擴(kuò)散的模型中的原始回歸損失稱為擴(kuò)散損失（diffusion loss），將新引入的正則化項(xiàng)稱為分散損失（Dispersive Loss）。

如果令 X = {x_i} 為有噪聲圖像 x_i 構(gòu)成的一批數(shù)據(jù)，則該數(shù)據(jù)批次的目標(biāo)函數(shù)為：

其中，L_Diff (x_i) 是一個(gè)樣本的標(biāo)準(zhǔn)擴(kuò)散損失，L_Disp (X) 則是依賴于整個(gè)批次的分散損失項(xiàng)，λ 是其加權(quán)項(xiàng)。

在實(shí)踐中，該團(tuán)隊(duì)沒有應(yīng)用任何額外的層（如，投射頭），而是直接將分散損失應(yīng)用于中間表示，不增加額外的可學(xué)習(xí)參數(shù)。

該方法是自成一體且極簡的。具體而言，它不會(huì)改變?cè)?L_Diff 項(xiàng)的實(shí)現(xiàn)：它不引入額外的采樣視圖，也不引入額外的數(shù)據(jù)增強(qiáng)，并且當(dāng) λ 為零時(shí)，它剛好就能約簡為基線擴(kuò)散模型。

這種設(shè)計(jì)之所以可行，是因?yàn)橐氲姆稚p失 L_Disp (X) 僅依賴于同一輸入批次中已經(jīng)計(jì)算出的中間表示。這不同于標(biāo)準(zhǔn)對(duì)比學(xué)習(xí) —— 在標(biāo)準(zhǔn)對(duì)比學(xué)習(xí)中，額外的增強(qiáng)和視圖可能會(huì)干擾每個(gè)樣本的回歸目標(biāo)。

前面也說過，分散損失的行為類似于「沒有正例對(duì)的對(duì)比損失」。在生成模型的背景下，這個(gè)公式是合理的，因?yàn)榛貧w項(xiàng)提供了預(yù)先定義的訓(xùn)練目標(biāo)，從而無需使用「正例對(duì)」。這與先前關(guān)于自監(jiān)督學(xué)習(xí)的研究《Understanding contrastive representation learning through alignment and uniformity on the hypersphere》一致，其中正例項(xiàng)被解釋為對(duì)齊目標(biāo)，而負(fù)例項(xiàng)則被解釋為正則化的形式。通過消除對(duì)正例對(duì)的需求，損失項(xiàng)可以定義在任何標(biāo)準(zhǔn)批次的（獨(dú)立）圖像上。

從概念上講，可以通過適當(dāng)移除正例項(xiàng)，從任何現(xiàn)有的對(duì)比損失中推導(dǎo)出分散損失。就此而言，「分散損失」一詞并非指特定的實(shí)現(xiàn)，而是指一類鼓勵(lì)實(shí)現(xiàn)分散的通用目標(biāo)。下文將介紹分散損失函數(shù)的幾種變體。

基于 InfoNCE 的分散損失變體

在自監(jiān)督學(xué)習(xí)中，InfoNCE 是被廣泛使用且有效的對(duì)比損失變體。作為案例研究，該團(tuán)隊(duì)提出了與 InfoNCE 損失相對(duì)應(yīng)的分散損失。

數(shù)學(xué)形式上，令 z_i = f (x_i) 表示輸入樣本 x_i 的生成模型的中間表示，其中 f 表示用于計(jì)算中間表示的層的子集。原始 InfoNCE 損失可以被解讀為分類交叉熵目標(biāo)，它會(huì)鼓勵(lì)讓正例對(duì)之間具有高相似度，而負(fù)樣本對(duì)之間具有低相似度：

其中， 表示一對(duì)正例（例如，通過對(duì)同一幅圖像進(jìn)行數(shù)據(jù)增強(qiáng)獲得的數(shù)據(jù)），(z_i,z_j) 表示包含正例對(duì)和所有負(fù)例對(duì)（即 i ≠ j）的任意一對(duì)樣本。D 表示相異度函數(shù)（例如，距離），τ 是一個(gè)稱為溫度的超參數(shù)。 D 的一個(gè)常用形式是負(fù)余弦相似度：

在 (2) 式的對(duì)數(shù)中，分子僅涉及正例對(duì)，而分母包含批次中的所有樣本對(duì)。根據(jù)之前的一些研究，可以將公式 (2) 等效地重寫為：

其中，第一項(xiàng)類似于回歸目標(biāo)，它最小化 z_i 與其目標(biāo)  之間的距離。另一方面，第二項(xiàng)則會(huì)鼓勵(lì)任何一對(duì) (z_i,z_j) 盡可能距離拉遠(yuǎn)。

為了構(gòu)造對(duì)應(yīng)的分散損失，這里只保留第二項(xiàng)：

該公式也可以被視為一種對(duì)比損失（公式 (3)），其中每個(gè)正例對(duì)由兩個(gè)相同的視圖  組成，使得  為一個(gè)常數(shù)。等式 (4) 就等價(jià)于

只差一個(gè)常數(shù)項(xiàng) log(batch size)，而這個(gè)常數(shù)項(xiàng)不會(huì)影響優(yōu)化過程。 從概念上講，此損失定義基于參考樣本 z_i。為了得到定義在一批樣本 Z = {z_i} 上的形式，這里按照之前的研究可將其重新定義為：

此損失函數(shù)對(duì)于批次內(nèi)的所有樣本具有相同的值，并且每個(gè)批次僅計(jì)算一次。在該團(tuán)隊(duì)的實(shí)驗(yàn)中，除了余弦相異度之外，我們還研究了平方?? 距離：。使用這種 ?? 形式時(shí)，只需幾行代碼即可輕松計(jì)算出分散損失，如算法 1 所示。

等式 (6) 中定義的基于 InfoNCE 的分散損失類似于前述先前關(guān)于自監(jiān)督學(xué)習(xí)的論文中的均勻性損失（盡管這里沒有對(duì)表示進(jìn)行 ?? 正則化）。在那篇論文中的對(duì)比表示學(xué)習(xí)，均勻性損失被應(yīng)用于輸出表示，并且必須與對(duì)齊損失（即正則項(xiàng)）配對(duì)。而這里的新公式則更進(jìn)一步，移除了中間表示上的對(duì)齊項(xiàng)，從而僅關(guān)注正則化視角。

該團(tuán)隊(duì)注意到，當(dāng) j = i 時(shí)，就不需要明確排除項(xiàng) D (z_i,z_j)。由于不會(huì)在一個(gè)批次中使用同一圖像的多個(gè)視圖，因此該項(xiàng)始終對(duì)應(yīng)于一個(gè)恒定且最小的差異度，例如在?? 的情況下為 0，在余弦情況下為 -1。因此，當(dāng)批次大小足夠大時(shí)，這個(gè)項(xiàng)在那個(gè)對(duì)數(shù)中的作用是充當(dāng)一個(gè)常數(shù)偏差，其貢獻(xiàn)會(huì)變小。在實(shí)踐中，無需排除該項(xiàng)，這也簡化了實(shí)現(xiàn)。

分散損失的其他變體

分散損失的概念可以自然延伸到 InfoNCE 之外的一類對(duì)比損失函數(shù)。

任何鼓勵(lì)排斥負(fù)例的目標(biāo)都可以被視為分散目標(biāo)，并實(shí)例化為分散損失的一種變體。基于其他類型的對(duì)比損失函數(shù)，該團(tuán)隊(duì)構(gòu)建了另外兩種變體。表 1 總結(jié)了所有三種變體，并比較了對(duì)比損失函數(shù)和分散損失函數(shù)。

鉸鏈損失（Hinge Loss）

在對(duì)比學(xué)習(xí)的經(jīng)典公式中，損失函數(shù)定義為獨(dú)立損失項(xiàng)之和，每個(gè)損失項(xiàng)對(duì)應(yīng)一個(gè)正例對(duì)或負(fù)例對(duì)。正例對(duì)的損失項(xiàng)為；負(fù)例對(duì)的損失項(xiàng)公式化為平方鉸鏈損失，即 ，其中 ε>0 為邊界值。為了構(gòu)造分散損失函數(shù)，只需舍棄正例對(duì)的損失項(xiàng)，僅計(jì)算負(fù)例對(duì)的損失項(xiàng)即可。見表 1 第 2 行。

協(xié)方差損失（Covariance Loss）

另一類（廣義）對(duì)比損失函數(shù)作用于表征的互協(xié)方差矩陣。這類損失函數(shù)可鼓勵(lì)互協(xié)方差矩陣接近單位矩陣。

舉個(gè)例子，對(duì)于論文《Barlow twins: Self-supervised learning via redundancy reduction》中定義的損失（它計(jì)算一個(gè)批次中兩個(gè)增強(qiáng)視圖的歸一化表征之間的互協(xié)方差矩陣），將 D×D 互協(xié)方差記為 Cov，其元素以 (m,n) 為索引。則該損失函數(shù)會(huì)使用損失項(xiàng) (1 ? Cov_mm)2 鼓勵(lì)對(duì)角線元素 Cov_mm 為 1，使用損失項(xiàng) 鼓勵(lì)非對(duì)角線元素 Cov_mn (?m≠n) 為 0，，其中 w 為權(quán)重。

在這里的分散損失中，該團(tuán)隊(duì)只考慮了非對(duì)角線元素 Cov_mn。由于不使用增強(qiáng)視圖，因此互協(xié)方差就簡化為基于單視圖批次計(jì)算的協(xié)方差矩陣。在這種情況下，當(dāng)表征經(jīng)過??正則化后，對(duì)角線元素 Cov_mm 自動(dòng)等于 1，因此無需在損失函數(shù)中顯式地處理。最終的分散損失為 。見表 1 第 3 行。

使用分散損失的擴(kuò)散模型

如表 1 所示，所有分散損失的變體都比其對(duì)應(yīng)的分散損失更簡潔。更重要的是，所有分散損失函數(shù)都適用于單視圖批次，這樣就無需進(jìn)行多視圖數(shù)據(jù)增強(qiáng)。因此，分散損失可以在現(xiàn)有的生成模型中充當(dāng)即插即用的正則化器，而無需修改回歸損失的實(shí)現(xiàn)。

在實(shí)踐中，引入分散損失只需進(jìn)行少量調(diào)整：

1. 指定應(yīng)用正則化器的中間層；
2. 計(jì)算該層的分散損失并將其添加到原始擴(kuò)散損失中。

算法 2 給出了訓(xùn)練偽代碼，其中包含算法 1 中定義的分散損失的具體形式。

該團(tuán)隊(duì)表示：「我們相信，這種簡化可極大地促進(jìn)我們方法的實(shí)際應(yīng)用，使其能夠應(yīng)用于各種生成模型。」

分散損失的實(shí)際表現(xiàn)如何？

表 2 比較了分散損失的不同變體及相應(yīng)的對(duì)比損失。

可以看到，在使用獨(dú)立噪聲時(shí)，對(duì)比損失在所有研究案例中均未能提高生成質(zhì)量。該團(tuán)隊(duì)猜想對(duì)齊兩個(gè)噪聲水平差異很大的視圖會(huì)損害學(xué)習(xí)效果。

而分散損失的表現(xiàn)總是比相應(yīng)的對(duì)比損失好，而前者還避免了雙視圖采樣帶來的復(fù)雜性。

而在不同的變體中，采用 ?? 距離的 InfoNCE 表現(xiàn)最佳。因此，在其它實(shí)驗(yàn)中，該團(tuán)隊(duì)默認(rèn)使用基于?? 的 InfoNCE。

另外，該團(tuán)隊(duì)還研究了不同模塊選擇以及不同 λ（控制正則化強(qiáng)度）和 τ（InfoNCE 中的溫度）值的影響。詳見原論文。

另外，不管是在 DiT（Diffusion Transformer）還是 SiT（Scalable Interpolant Transformers）上，分散損失在所有場(chǎng)景下都比基線方法更好。有趣的是，他們還觀察到，當(dāng)基線性能更強(qiáng)時(shí)，相對(duì)改進(jìn)甚至絕對(duì)改進(jìn)往往還會(huì)更大。

總體而言，這種趨勢(shì)有力地證明了分散損失的主要作用在于正則化。由于規(guī)模更大、性能更強(qiáng)的模型更容易過擬合，因此有效的正則化往往會(huì)使它們受益更多。

圖 5 展示了 SiT-XL/2 模型生成的一些示例圖像。

當(dāng)然，該團(tuán)隊(duì)也將新方法與 REPA 進(jìn)行了比較。新方法的正則化器直接作用于模型的內(nèi)部表示，而 REPA 會(huì)將其與外部模型的表示對(duì)齊。因此，為了公平起見，應(yīng)同時(shí)考慮額外的計(jì)算開銷和外部信息源，如表 6 所示。

REPA 依賴于一個(gè)預(yù)訓(xùn)練的 DINOv2 模型，該模型本身是從已在 1.42 億張精選圖像上訓(xùn)練過的 11B 參數(shù)主干網(wǎng)絡(luò)中蒸餾出來的。

相比之下，新提出的方法完全不需要這些：無需預(yù)訓(xùn)練、外部數(shù)據(jù)和額外的模型參數(shù)。新方法在將訓(xùn)練擴(kuò)展到更大的模型和數(shù)據(jù)集時(shí)非常適用，并且該團(tuán)隊(duì)預(yù)計(jì)在這種情況下正則化效果會(huì)非常好。

最后，新提出的方法可以直接泛化用于基于一步式擴(kuò)散的生成模型。

在表 7（左）中，該團(tuán)隊(duì)將分散損失應(yīng)用于最新的 MeanFlow 模型，然后觀察到了穩(wěn)定持續(xù)的改進(jìn)。表 7（右）將這些結(jié)果與最新的一步擴(kuò)散 / 基于流的模型進(jìn)行了比較，表明新方法可增強(qiáng) MeanFlow 的性能并達(dá)到了新的 SOTA。