最近的研究強調(diào)了擴散模型與表征學習之間的相互作用。擴散模型的中間表征可用于下游視覺任務，同時視覺模型表征能夠提升擴散模型的收斂速度和生成質(zhì)量。然而，由于輸入不匹配和 VAE 潛在空間的使用，將視覺模型的預訓練權(quán)重遷移到擴散模型中仍然具有挑戰(zhàn)性。

為了解決這些問題，來自高德地圖的研究者提出了統(tǒng)一自監(jiān)督預訓練（USP, Unified Self-Supervised Pretraining），該方法通過在變分自編碼器（VAE）的潛在空間中進行潛在掩碼建模（Masked Latent Modeling）預訓練，預訓練得到的 ViT 編碼器等權(quán)重可以無縫遷移到下游任務，包括圖像分類、語義分割以及基于擴散模型的圖像生成。

• 論文名稱：USP: Unified Self-Supervised Pretraining for Image Generation and Understanding
• 論文地址：https://arxiv.org/pdf/2503.06132
• 代碼地址：https://github.com/cxxgtxy/USP

USP 在理解任務上實現(xiàn)有競爭力的表現(xiàn)；在生成任務中，USP 顯著加速 DiT 和 SiT 模型收斂速度，比從頭訓練快 11.7 倍（DiT-XL）和 46.6 倍（SiT-XL）。

研究背景

在過去十年中，預訓練 - 微調(diào)（pretraining-finetuning）范式在圖像識別任務中取得了顯著成功。但在圖像生成領(lǐng)域，該范式的探索仍然有限。

DDAE 等近期的研究表明，生成模型不僅可以合成高質(zhì)量圖像，還可以學習到優(yōu)越的視覺表征，圖像理解和圖像生成之間存在著深層次的聯(lián)系。

例如，iGPT（Image GPT）探索了基于像素空間的自回歸預訓練，但該方法在大規(guī)模數(shù)據(jù)集和模型上擴展時面臨著嚴重的計算成本問題。此外，這種方法與擴散模型不兼容。

為了彌合這一差距，REPA 提出通過對齊擴散模型與預訓練的視覺模型（如 DINOv2）的表征，可以讓擴散模型更高效地學習判別特征，從而提高訓練效率和生成質(zhì)量。然而，REPA 存在幾個主要問題：

1. 高昂的計算成本：REPA 依賴于 DINOv2 這樣的大規(guī)模預訓練視覺模型，其預訓練需要超過 22,000 GPU 小時（A100），計算資源需求極高。

2. 額外的教師網(wǎng)絡（Teacher Network）：使用 DINOv2 作為教師模型會增加 GPU 顯存消耗，同時降低擴散模型的訓練速度。

盡管已有研究揭示了生成和理解任務間的聯(lián)系，但仍然有一些關(guān)鍵問題尚未解決：

1. 預訓練是否對擴散模型的訓練是可行且必要的？

2. 是否可以找到一種同時適用于生成和理解任務的預訓練方法？

3. 現(xiàn)有的 “預訓練 - 微調(diào)” 范式是否能成功應用于生成模型？

該論文旨在提出一種簡單而有效的方法來解決這些問題。

方法設計

實現(xiàn)能夠同時適用于圖像理解和生成的統(tǒng)一預訓練 - 微調(diào)范式面臨諸多挑戰(zhàn)：

• C1: 輸入不匹配：圖像理解模型通常接收干凈的圖像作為輸入，而擴散模型接受的是添加了噪聲的圖像。
• C2: 結(jié)構(gòu)不匹配：生成模型多為基于 VAE 的潛空間擴散模型，而大多數(shù)圖像理解任務并不使用 VAE。此外， ViT 結(jié)構(gòu)在圖像生成任務中通常會進行修改。
• C3: 損失函數(shù)和標簽格式不同：圖像理解任務和圖像生成任務通常采用不同的優(yōu)化目標，這使得直接共享預訓練模型變得困難。

盡管面臨這些挑戰(zhàn)，研究者也觀察到了一些有利的現(xiàn)象：

• P1: 神經(jīng)網(wǎng)絡對噪聲具有魯棒性：預訓練的視覺模型在噪聲或數(shù)據(jù)增強下仍然可以保持較高的分類精度，例如在 ImageNet-C 數(shù)據(jù)集上測試。說明即使擴散模型處理的是加噪圖像，預訓練模型仍可以學習到有效的特征。
• P2: 擴散模型可以學習到判別性特征：擴散模型能學習到用于圖像分類等任務的判別性特征。如果能夠有效地對齊這些表征，擴散模型的收斂速度和最終性能都可以得到顯著提升。
• P3: ViT 結(jié)構(gòu)具有較強的適應性：盡管 ViT 在應用到擴散模型時經(jīng)歷了一定修改（如 AdaLN-Zero 層歸一化和額外的條件輸入）。但如果設計得當，這些修改仍然可以與 ViT 的預訓練權(quán)重兼容。
• P4: VAE 具有強大的壓縮和重建能力：擴散模型中使用的 VAE（如 SD-VAE）能夠有效地保留原始圖像的重要信息。即使在 VAE 的潛空間中進行訓練，仍然可以獲得高質(zhì)量的視覺特征。

基于以上觀察，本文的研究者提出了統(tǒng)一的自監(jiān)督預訓練架構(gòu)，見下圖 1：

USP 架構(gòu)基于一個簡單的自編碼器（Autoencoder），但在 VAE 潛空間中進行操作，而非像素空間。輸入圖像首先經(jīng)過 VAE 編碼到潛空間，并通過 PatchConv 進行圖片分塊。部分塊按照設定 mask 比例被隨機掩碼，未掩碼的塊輸入到 ViT 編碼器，而解碼器負責重建掩碼塊，損失函數(shù)僅使用簡單的 MSE loss。在預訓練階段，VAE 參數(shù)被凍結(jié)，僅訓練 ViT 編碼器。預訓練完成后，ViT 編碼器的權(quán)重可用于初始化下游任務，如分類、分割和生成。

在將預訓練模型權(quán)重適配到下游理解和生成模型時，針對圖像分類任務，ViT 編碼器的預訓練權(quán)重可直接繼承，無需額外調(diào)整，且仍然使用 Class Token 作為最終表征。

對于生成任務，由于 DiT 和 SiT 的結(jié)構(gòu)與 ViT 略有不同，對初始化策略進行了優(yōu)化。首先，在 AdaLN-Zero 層歸一化中，恢復可訓練的偏置（β）和縮放因子（γ），使其與預訓練的 ViT 權(quán)重對齊。其次，由于預訓練是在 224×224 進行，而 ImageNet 生成任務通常在 256×256 進行，因此本文采用 Bicubic Interpolation 擴展位置編碼。最后，由于生成任務不需要 class token，在 DiT/SiT 中直接將其移除。這種初始化策略確保了 ViT 預訓練權(quán)重能夠無縫適配到下游分類和生成任務，而不引入額外計算開銷或存儲需求。

實驗設置

本論文涵蓋三種模型規(guī)模，見表 1。預訓練階段所有模型共享相同的解碼器 —— 由 8 個 Transformer 塊組成。

在額外實驗中，將預訓練時長擴展到 1600 輪，以證明 USP 在更長預訓練時間上的可擴展性。為了與 MAE 進行公平比較，本文在 224×224 分辨率上進行預訓練，盡管消融實驗表明更高的分辨率可以帶來更好的性能。

圖像生成實驗

本文在兩種基于 Transformer 的擴散模型，DiT 和 SiT 上驗證了 USP。評估均在 ImageNet 256×256 上進行，使用 50,000 個樣本，不使用 CFG。

表 2 顯示了在不同規(guī)模的 DiT 模型上的對比結(jié)果。USP 在所有模型規(guī)模上均顯著提升了生成質(zhì)量，且隨著訓練時間延長，生成質(zhì)量不斷提高。相比最近的 DiT 變體在 2.5M 步的 FID，USP 僅在 400K 內(nèi)就能達到更好的效果。

表 3 顯示了在不同規(guī)模的 SiT 模型上的對比結(jié)果。USP 表現(xiàn)出和 DiT 一致的提升效果。同時，表 3 與近期利用表征對齊來提升 DiT/SiT 性能的方法進行了比較，USP 在所有模型設置下均優(yōu)于其他方法。

圖像理解實驗

論文在 ImageNet-1k 數(shù)據(jù)集上進行了線性探測（Linear Probe）和微調(diào)（Fine-tuning）的圖像分類評估。在線性探測（LP）任務上，USP 的性能優(yōu)于 MAE；在微調(diào)（SFT）任務上，USP 表現(xiàn)與 MAE 相當，表 5 總結(jié)了分類結(jié)果： 

進一步，論文在 ADE20 數(shù)據(jù)集上進行了分割性能評估。表 6 顯示了 USP 在單尺度 mIoU 指標上的表現(xiàn)，相比 MAE 提升了 0.5%。

消融實驗

研究者進行了全面的消融實驗，以探討 USP 設計中不同組件的影響。例如 VAE、輸入分辨率、掩蔽率等。更多的消融實驗說明見原論文。

討論

VAE 在圖像理解任務中的作用

在圖像分類任務中應用 VAE 并不是最理想的選擇。因為 VAE 的核心目標是在保證重建能力的同時盡可能壓縮信息，而原始圖像本身是無損的，因此直接在原始圖像上進行分類可能更高效。然而，我們的實驗表明，如果使用高質(zhì)量的 VAE 進行編碼，圖像分類任務的性能至少可以達到與標準方法相當?shù)乃健?/span>

研究者認為，這種現(xiàn)象的主要原因是：

• VAE 的潛空間編碼仍然能保留足夠的判別信息，即使經(jīng)過壓縮，仍能支持良好的分類表現(xiàn)。
• VAE 提供了一種對抗噪聲的方式，通過潛空間中的信息提取，模型可能學習到更魯棒的特征。

工作機制（對比 REPA）

為了更深入地理解 USP，研究者在 DiT-XL/2 訓練過程中，對不同層的線性探測性能進行了分析（見圖 4）。比較了以下幾種情況：

1.DiT-XL/2 預訓練后的線性探測性能（“Pre.”）

2.DiT-XL/2 生成微調(diào)后的線性探測性能（“Ft.”）

3.SiT-XL/2 模型的線性探測性能

4.SiT-XL/2 在應用 REPA 方法后的線性探測性能

主要發(fā)現(xiàn)：

• 與 REPA 不同，USP 不依賴額外的對齊損失，而是通過精心設計的初始化策略，讓模型自動找到最適合線性分類的層。
• 經(jīng)過 40 萬步訓練后，USP 的第 20 層成為最優(yōu)的線性分類層，這表明 USP 能夠自適應地優(yōu)化表征學習，在分類和生成任務之間找到平衡點。
• REPA 通過人為設計的表征對齊方式來增強判別能力，但這種方法可能會限制生成模型的潛力。

這些實驗表明，USP 的初始化策略比基于表征對齊的方法更高效，并且更適用于同時提升分類和生成任務的統(tǒng)一框架。

對基礎(chǔ)范式的另一種視角

論文中使用經(jīng)過 800 輪預訓練且掩碼比例為 0.75 的 ViT-Large 模型研究了圖像修復任務。如圖 5 所示，USP 在圖像修復方面顯著優(yōu)于 MAE，凸顯了強表征能力對有效修復的重要性。

這一結(jié)果與下面基于擴散的生成框架相契合，表明過度使用監(jiān)督標簽微調(diào)編碼器以增強判別能力，并不會顯著提升圖像生成效果。

為了實證這一假設，論文中采用了一個監(jiān)督微調(diào)模型，模型在 ImageNet 驗證集上的準確率達 82.6%，并使用其初始化 DiT-B/2 進行 400K 步訓練。表 8 總結(jié)了實驗結(jié)果。其性能明顯低于預訓練，進一步驗證了對該范式的分析。