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

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

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

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

研究背景

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

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

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

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

1. 高昂的計(jì)算成本：REPA 依賴(lài)于 DINOv2 這樣的大規(guī)模預(yù)訓(xùn)練視覺(jué)模型，其預(yù)訓(xùn)練需要超過(guò) 22,000 GPU 小時(shí)（A100），計(jì)算資源需求極高。

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

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

1. 預(yù)訓(xùn)練是否對(duì)擴(kuò)散模型的訓(xùn)練是可行且必要的？

2. 是否可以找到一種同時(shí)適用于生成和理解任務(wù)的預(yù)訓(xùn)練方法？

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

該論文旨在提出一種簡(jiǎn)單而有效的方法來(lái)解決這些問(wèn)題。

方法設(shè)計(jì)

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

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

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

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

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

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

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

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

實(shí)驗(yàn)設(shè)置

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

在額外實(shí)驗(yàn)中，將預(yù)訓(xùn)練時(shí)長(zhǎng)擴(kuò)展到 1600 輪，以證明 USP 在更長(zhǎng)預(yù)訓(xùn)練時(shí)間上的可擴(kuò)展性。為了與 MAE 進(jìn)行公平比較，本文在 224×224 分辨率上進(jìn)行預(yù)訓(xùn)練，盡管消融實(shí)驗(yàn)表明更高的分辨率可以帶來(lái)更好的性能。

圖像生成實(shí)驗(yàn)

本文在兩種基于 Transformer 的擴(kuò)散模型，DiT 和 SiT 上驗(yàn)證了 USP。評(píng)估均在 ImageNet 256×256 上進(jìn)行，使用 50,000 個(gè)樣本，不使用 CFG。

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

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

圖像理解實(shí)驗(yàn)

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

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

消融實(shí)驗(yàn)

研究者進(jìn)行了全面的消融實(shí)驗(yàn)，以探討 USP 設(shè)計(jì)中不同組件的影響。例如 VAE、輸入分辨率、掩蔽率等。更多的消融實(shí)驗(yàn)說(shuō)明見(jiàn)原論文。

討論

VAE 在圖像理解任務(wù)中的作用

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

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

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

工作機(jī)制（對(duì)比 REPA）

為了更深入地理解 USP，研究者在 DiT-XL/2 訓(xùn)練過(guò)程中，對(duì)不同層的線(xiàn)性探測(cè)性能進(jìn)行了分析（見(jiàn)圖 4）。比較了以下幾種情況：

1.DiT-XL/2 預(yù)訓(xùn)練后的線(xiàn)性探測(cè)性能（“Pre.”）

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

3.SiT-XL/2 模型的線(xiàn)性探測(cè)性能

4.SiT-XL/2 在應(yīng)用 REPA 方法后的線(xiàn)性探測(cè)性能

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

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

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

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

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

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

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