何愷明改進了謝賽寧的REPA：極大簡化但性能依舊強悍

2025-06-13 08:53:00

謝賽寧團隊提出了表征對齊 (REPA)?。該方法可以利用預訓練得到的現成表征模型的能力。

在建模復雜的數據分布方面，擴散生成模型表現出色，不過它的成果大體上與表征學習（representation learning）領域關聯不大。

通常來說，擴散模型的訓練目標包含一個專注于重構（例如去噪）的回歸項，但缺乏為生成學習到的表征的顯式正則化項。這種圖像生成范式與圖像識別范式差異明顯 —— 過去十年來，圖像識別領域的核心主題和驅動力一直是表征學習。

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

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

鑒于表征學習在生成模型中的潛力，謝賽寧團隊提出了表征對齊 (REPA) 。該方法可以利用預訓練得到的現成表征模型的能力。在訓練生成模型的同時，該方法會鼓勵其內部表征與外部預訓練表征之間對齊。

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

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

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

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

分散損失的核心思想其實很簡單：除了模型輸出的標準回歸損失之外，再引入了一個用于正則化模型的內部表征的目標（圖 1）。

直覺上看，分散損失會鼓勵內部表征在隱藏空間中散開，類似于對比學習中的排斥效應。同時，原始的回歸損失（去噪）則自然地充當了對齊機制，從而無需像對比學習那樣手動定義正例對。

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

因此，與對比學習不同，它既不需要雙視圖采樣、專門的數據增強，也不需要額外的編碼器。訓練流程完全可以遵循基于擴散的模型（及基于流的對應模型）中使用的標準做法，唯一的區(qū)別在于增加了一個開銷可忽略不計的正則化損失。

與 REPA 機制相比，這種新方法無需預訓練、無需額外的模型參數，也無需外部數據。憑借其獨立且極簡的設計，該方法清晰地證明：表征學習無需依賴外部信息源也可助益生成式建模。

帶點數學的方法詳解

分散損失

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

如果令 X = {x_i} 為有噪聲圖像 x_i 構成的一批數據，則該數據批次的目標函數為：

其中，L_Diff (x_i) 是一個樣本的標準擴散損失，L_Disp (X) 則是依賴于整個批次的分散損失項，λ 是其加權項。

在實踐中，該團隊沒有應用任何額外的層（如，投射頭），而是直接將分散損失應用于中間表示，不增加額外的可學習參數。

該方法是自成一體且極簡的。具體而言，它不會改變原始 L_Diff 項的實現：它不引入額外的采樣視圖，也不引入額外的數據增強，并且當 λ 為零時，它剛好就能約簡為基線擴散模型。

這種設計之所以可行，是因為引入的分散損失 L_Disp (X) 僅依賴于同一輸入批次中已經計算出的中間表示。這不同于標準對比學習 —— 在標準對比學習中，額外的增強和視圖可能會干擾每個樣本的回歸目標。

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

從概念上講，可以通過適當移除正例項，從任何現有的對比損失中推導出分散損失。就此而言，「分散損失」一詞并非指特定的實現，而是指一類鼓勵實現分散的通用目標。下文將介紹分散損失函數的幾種變體。

基于 InfoNCE 的分散損失變體

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

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

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

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

其中，第一項類似于回歸目標，它最小化 z_i 與其目標之間的距離。另一方面，第二項則會鼓勵任何一對 (z_i,z_j) 盡可能距離拉遠。

為了構造對應的分散損失，這里只保留第二項：

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

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

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

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

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

分散損失的其他變體

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

任何鼓勵排斥負例的目標都可以被視為分散目標，并實例化為分散損失的一種變體?；谄渌愋偷膶Ρ葥p失函數，該團隊構建了另外兩種變體。表 1 總結了所有三種變體，并比較了對比損失函數和分散損失函數。

鉸鏈損失（Hinge Loss）

在對比學習的經典公式中，損失函數定義為獨立損失項之和，每個損失項對應一個正例對或負例對。正例對的損失項為；負例對的損失項公式化為平方鉸鏈損失，即，其中 ε>0 為邊界值。為了構造分散損失函數，只需舍棄正例對的損失項，僅計算負例對的損失項即可。見表 1 第 2 行。

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

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

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

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