最近的研究強調了擴散模型與表征學習之間的相互作用。擴散模型的中間表征可用于下游視覺任務，同時視覺模型表征能夠提升擴散模型的收斂速度和生成質量。然而，由于輸入不匹配和 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 總結了實驗結果。其性能明顯低于預訓練，進一步驗證了對該范式的分析。