最近的研究強調了擴散模型與表徵學習之間的相互作用。擴散模型的中間表徵可用於下游視覺任務，同時視覺模型表徵能夠提升擴散模型的收斂速度和生成質量。然而，由於輸入不匹配和 VAE 潛在空間的使用，將視覺模型的預訓練權重遷移到擴散模型中仍然具有挑戰性。
為了解決這些問題，來自高德地圖的研究者提出了統一自監督預訓練（USP, Unified Self-Supervised Pretraining），該方法透過在變分自編碼器（VAE）的潛在空間中進行潛在掩碼建模（Masked Latent Modeling）預訓練，預訓練得到的 ViT 編碼器等權重可以無縫遷移到下游任務，包括影像分類、語義分割以及基於擴散模型的影像生成。
論文名稱：USP: Unified Self-Supervised Pretraining for Image Generation and Understanding 論文地址：https://arxiv.org/pdf/2503.06132 程式碼地址：https://github.com/cxxgtxy/USP
USP 在理解任務上實現有競爭力的表現；在生成任務中，USP 顯著加速 DiT 和 SiT 模型收斂速度，比從頭訓練快 11.7 倍（DiT-XL）和 46.6 倍（SiT-XL）。
研究背景
在過去十年中，預訓練 - 微調（pretraining-finetuning）正規化在影像識別任務中取得了顯著成功。但在影像生成領域，該正規化的探索仍然有限。
DDAE 等近期的研究表明，生成模型不僅可以合成高質量影像，還可以學習到優越的視覺表徵，影像理解和影像生成之間存在著深層次的聯絡。
例如，iGPT（Image GPT）探索了基於畫素空間的自迴歸預訓練，但該方法在大規模資料集和模型上擴充套件時面臨著嚴重的計算成本問題。此外，這種方法與擴散模型不相容。
為了彌合這一差距，REPA 提出透過對齊擴散模型與預訓練的視覺模型（如 DINOv2）的表徵，可以讓擴散模型更高效地學習判別特徵，從而提高訓練效率和生成質量。然而，REPA 存在幾個主要問題：
1. 高昂的計算成本：REPA 依賴於 DINOv2 這樣的大規模預訓練視覺模型，其預訓練需要超過 22,000 GPU 小時（A100），計算資源需求極高。
2. 額外的教師網路（Teacher Network）：使用 DINOv2 作為教師模型會增加 GPU 視訊記憶體消耗，同時降低擴散模型的訓練速度。
儘管已有研究揭示了生成和理解任務間的聯絡，但仍然有一些關鍵問題尚未解決：
1. 預訓練是否對擴散模型的訓練是可行且必要的？
2. 是否可以找到一種同時適用於生成和理解任務的預訓練方法？
3. 現有的 “預訓練 - 微調” 正規化是否能成功應用於生成模型？
該論文旨在提出一種簡單而有效的方法來解決這些問題。
方法設計
實現能夠同時適用於影像理解和生成的統一預訓練 - 微調正規化面臨諸多挑戰：
C1: 輸入不匹配：影像理解模型通常接收乾淨的影像作為輸入，而擴散模型接受的是新增了噪聲的影像。 C2: 結構不匹配：生成模型多為基於 VAE 的潛空間擴散模型，而大多數影像理解任務並不使用 VAE。此外， ViT 結構在影像生成任務中通常會進行修改。 C3: 損失函式和標籤格式不同：影像理解任務和影像生成任務通常採用不同的最佳化目標，這使得直接共享預訓練模型變得困難。
儘管面臨這些挑戰，研究者也觀察到了一些有利的現象：
P1: 神經網路對噪聲具有魯棒性：預訓練的視覺模型在噪聲或資料增強下仍然可以保持較高的分類精度，例如在 ImageNet-C 資料集上測試。說明即使擴散模型處理的是加噪影像，預訓練模型仍可以學習到有效的特徵。 P2: 擴散模型可以學習到判別性特徵：擴散模型能學習到用於影像分類等任務的判別性特徵。如果能夠有效地對齊這些表徵，擴散模型的收斂速度和最終效能都可以得到顯著提升。 P3: ViT 結構具有較強的適應性：儘管 ViT 在應用到擴散模型時經歷了一定修改（如 AdaLN-Zero 層歸一化和額外的條件輸入）。但如果設計得當，這些修改仍然可以與 ViT 的預訓練權重相容。 P4: VAE 具有強大的壓縮和重建能力：擴散模型中使用的 VAE（如 SD-VAE）能夠有效地保留原始影像的重要資訊。即使在 VAE 的潛空間中進行訓練，仍然可以獲得高質量的視覺特徵。
基於以上觀察，本文的研究者提出了統一的自監督預訓練架構，見下圖 1：
USP 架構基於一個簡單的自編碼器（Autoencoder），但在 VAE 潛空間中進行操作，而非畫素空間。輸入影像首先經過 VAE 編碼到潛空間，並透過 PatchConv 進行圖片分塊。部分塊按照設定 mask 比例被隨機掩碼，未掩碼的塊輸入到 ViT 編碼器，而解碼器負責重建掩碼塊，損失函式僅使用簡單的 MSE loss。在預訓練階段，VAE 引數被凍結，僅訓練 ViT 編碼器。預訓練完成後，ViT 編碼器的權重可用於初始化下游任務，如分類、分割和生成。
在將預訓練模型權重適配到下游理解和生成模型時，針對影像分類任務，ViT 編碼器的預訓練權重可直接繼承，無需額外調整，且仍然使用 Class Token 作為最終表徵。
對於生成任務，由於 DiT 和 SiT 的結構與 ViT 略有不同，對初始化策略進行了最佳化。首先，在 AdaLN-Zero 層歸一化中，恢復可訓練的偏置（β）和縮放因子（γ），使其與預訓練的 ViT 權重對齊。其次，由於預訓練是在 224×224 進行，而 ImageNet 生成任務通常在 256×256 進行，因此本文采用 Bicubic Interpolation 擴充套件位置編碼。最後，由於生成任務不需要 class token，在 DiT/SiT 中直接將其移除。這種初始化策略確保了 ViT 預訓練權重能夠無縫適配到下游分類和生成任務，而不引入額外計算開銷或儲存需求。
實驗設定
本論文涵蓋三種模型規模，見表 1。預訓練階段所有模型共享相同的解碼器 —— 由 8 個 Transformer 塊組成。
在額外實驗中，將預訓練時長擴充套件到 1600 輪，以證明 USP 在更長預訓練時間上的可擴充套件性。為了與 MAE 進行公平比較，本文在 224×224 解析度上進行預訓練，儘管消融實驗表明更高的解析度可以帶來更好的效能。
影像生成實驗
本文在兩種基於 Transformer 的擴散模型，DiT 和 SiT 上驗證了 USP。評估均在 ImageNet 256×256 上進行，使用 50,000 個樣本，不使用 CFG。
表 2 顯示了在不同規模的 DiT 模型上的對比結果。USP 在所有模型規模上均顯著提升了生成質量，且隨著訓練時間延長，生成質量不斷提高。相比最近的 DiT 變體在 2.5M 步的 FID，USP 僅在 400K 內就能達到更好的效果。
表 3 顯示了在不同規模的 SiT 模型上的對比結果。USP 表現出和 DiT 一致的提升效果。同時，表 3 與近期利用表徵對齊來提升 DiT/SiT 效能的方法進行了比較，USP 在所有模型設定下均優於其他方法。
影像理解實驗
論文在 ImageNet-1k 資料集上進行了線性探測（Linear Probe）和微調（Fine-tuning）的影像分類評估。線上性探測（LP）任務上，USP 的效能優於 MAE；在微調（SFT）任務上，USP 表現與 MAE 相當，表 5 總結了分類結果：
進一步，論文在 ADE20 資料集上進行了分割效能評估。表 6 顯示了 USP 在單尺度 mIoU 指標上的表現，相比 MAE 提升了 0.5%。
消融實驗
研究者進行了全面的消融實驗，以探討 USP 設計中不同元件的影響。例如 VAE、輸入解析度、掩蔽率等。更多的消融實驗說明見原論文。
討論
VAE 在影像理解任務中的作用
在影像分類任務中應用 VAE 並不是最理想的選擇。因為 VAE 的核心目標是在保證重建能力的同時儘可能壓縮資訊，而原始影像本身是無損的，因此直接在原始影像上進行分類可能更高效。然而，我們的實驗表明，如果使用高質量的 VAE 進行編碼，影像分類任務的效能至少可以達到與標準方法相當的水平。
研究者認為，這種現象的主要原因是：
VAE 的潛空間編碼仍然能保留足夠的判別資訊，即使經過壓縮，仍能支援良好的分類表現。 VAE 提供了一種對抗噪聲的方式，透過潛空間中的資訊提取，模型可能學習到更魯棒的特徵。
工作機制（對比 REPA）
為了更深入地理解 USP，研究者在 DiT-XL/2 訓練過程中，對不同層的線性探測效能進行了分析（見圖 4）。比較了以下幾種情況：
1.DiT-XL/2 預訓練後的線性探測效能（“Pre.”）
2.DiT-XL/2 生成微調後的線性探測效能（“Ft.”）
3.SiT-XL/2 模型的線性探測效能
4.SiT-XL/2 在應用 REPA 方法後的線性探測效能
主要發現：
與 REPA 不同，USP 不依賴額外的對齊損失，而是透過精心設計的初始化策略，讓模型自動找到最適合線性分類的層。 經過 40 萬步訓練後，USP 的第 20 層成為最優的線性分類層，這表明 USP 能夠自適應地最佳化表徵學習，在分類和生成任務之間找到平衡點。 REPA 透過人為設計的表徵對齊方式來增強判別能力，但這種方法可能會限制生成模型的潛力。
這些實驗表明，USP 的初始化策略比基於表徵對齊的方法更高效，並且更適用於同時提升分類和生成任務的統一框架。
對基礎正規化的另一種視角
論文中使用經過 800 輪預訓練且掩碼比例為 0.75 的 ViT-Large 模型研究了影像修復任務。如圖 5 所示，USP 在影像修復方面顯著優於 MAE，凸顯了強表徵能力對有效修復的重要性。
這一結果與下面基於擴散的生成框架相契合，表明過度使用監督標籤微調編碼器以增強判別能力，並不會顯著提升影像生成效果。
為了實證這一假設，論文中採用了一個監督微調模型，模型在 ImageNet 驗證集上的準確率達 82.6%，並使用其初始化 DiT-B/2 進行 400K 步訓練。表 8 總結了實驗結果。其效能明顯低於預訓練，進一步驗證了對該正規化的分析。
