統一自監督預訓練！視覺模型權重無縫遷移下游任務，SiT收斂提速近47倍

作者：機器之心 2025-03-17 09:12:00

來自高德地圖的研究者提出了統一自監督預訓練（USP, Unified Self-Supervised Pretraining）。

最近的研究強調了擴散模型與表征學習之間的相互作用。擴散模型的中間表征可用于下游視覺任務，同時視覺模型表征能夠提升擴散模型的收斂速度和生成質量。然而，由于輸入不匹配和 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 預訓練權重能夠無縫適配到下游分類和生成任務，而不引入額外計算開銷或存儲需求。