機器之心報道

機器之心編輯部

存在 10 多年后，VAE（變分自編碼器）時代終于要淘汰了嗎？

就在今天，紐約大學(xué)助理教授謝賽寧團隊放出了新作 ——VAE 的替代解決方案 ——RAE（Representation Autoencoders，表征自編碼器）

他表示，三年前，DiT（Diffusion Transformer） 用基于 Transformer 的去噪骨干網(wǎng)絡(luò)取代了傳統(tǒng)的 U-Net。那時候就知道，笨重的 VAE 遲早也會被淘汰。如今，時機終于到了。

謝賽寧進一步做出了解釋，DiT 雖然取得了長足的進步，但大多數(shù)模型仍然依賴于 2021 年的舊版 SD-VAE 作為其潛空間基礎(chǔ)。這就帶來了以下幾個主要問題：

• 過時的骨干網(wǎng)絡(luò)使架構(gòu)比實際需要的更復(fù)雜：SD-VAE 的計算量約為 450 GFLOPs，而一個簡單的 ViT-B 編碼器只需要大約 22 GFLOPs。
• 過度壓縮的潛空間（只有 4 個通道）限制了可存儲的信息量：人們常說壓縮帶來智能，但這里并非如此：VAE 式壓縮實際上作用有限，幾乎和原始的三通道像素一樣受限。
• 表征能力弱：由于僅使用重建任務(wù)進行訓(xùn)練，VAE 學(xué)到的特征很弱（線性探針精度約 8%），這會導(dǎo)致模型收斂更慢、生成質(zhì)量下降。我們現(xiàn)在已經(jīng)很清楚 —— 表征質(zhì)量直接影響生成質(zhì)量，而 SD-VAE 并不是為此而設(shè)計的。

因此，謝賽寧團隊將預(yù)訓(xùn)練的表征編碼器（如 DINO、SigLIP、MAE）與訓(xùn)練好的解碼器相結(jié)合，以取代傳統(tǒng)的 VAE，形成了一種新的結(jié)構(gòu) —— 表征自編碼器（RAE）。這種模型既能實現(xiàn)高質(zhì)量的重建，又能提供語義豐富的潛空間，同時具備可擴展的 Transformer 架構(gòu)特性

由于這些潛空間通常是高維的，一個關(guān)鍵的挑戰(zhàn)在于如何讓 DiT 能夠在其中高效地運行。從原理上來說，將 DiT 適配到這些高維語義潛空間是可行的，但需要經(jīng)過精心的設(shè)計。最初的 DiT 是為緊湊的 SD-VAE 潛空間而設(shè)計的，當面對高維潛空間時會遇到多方面的困難，包括 Transformer 結(jié)構(gòu)問題、噪聲調(diào)度問題、解碼器魯棒性問題。

為此，研究者提出了一種新的 DiT 變體 ——DiT^DH，它受到了 DDT 的啟發(fā)，但出發(fā)點不同。該變體在標準 DiT 架構(gòu)的基礎(chǔ)上，引入一個輕量、淺層但寬度較大的頭部（head）結(jié)構(gòu)，使擴散模型在不顯著增加二次計算成本的前提下擴展網(wǎng)絡(luò)寬度。

這一設(shè)計在高維 RAE 潛空間中進一步提升了 DiT 的訓(xùn)練效果，在 ImageNet 數(shù)據(jù)集上取得了優(yōu)異的圖像生成效果：在 256×256 分辨率下，無引導(dǎo)條件下的 FID 為 1.51；在 256×256 和 512×512 分辨率下，有引導(dǎo)條件下的 FID 均為 1.13。

因此，RAE 展現(xiàn)出了明顯的優(yōu)勢，應(yīng)當成為 DiT 訓(xùn)練的全新默認方案

當然，RAE 的模型和 PyTorch 代碼全部開源。這項工作的一作為一年級博士生 Boyang Zheng，其本科畢業(yè)于上海交通大學(xué) ACM 班。

• 論文標題：Diffusion Transformers with Representation Autoencoders
• 論文地址：https://arxiv.org/abs/2510.11690
• 項目主頁：https://rae-dit.github.io/
• 代碼：https://github.com/bytetriper/RAE
• HuggingFace：https://huggingface.co/collections/nyu-visionx/rae-68ecb57b8bfbf816c83cce15

從網(wǎng)友的反饋來看，大家非常看好 RAE 的前景，預(yù)計可以為生成模型帶來新的可能性。

基于凍結(jié)編碼器的高保真重建

研究者挑戰(zhàn)了一個普遍的假設(shè)，即像 DINOv2 和 SigLIP2 這類預(yù)訓(xùn)練表征編碼器不適合重建任務(wù)，因為它們 “強調(diào)高層語義，而忽略了底層細節(jié)” 。

該研究證明，只要解碼器訓(xùn)練得當，凍結(jié)的表征編碼器實際上可以作為擴散潛在空間的強大編碼器。RAE 將凍結(jié)的預(yù)訓(xùn)練表征編碼器與一個基于 ViT 的解碼器配對，其重建效果與 SD-VAE 相當甚至更優(yōu)。

更重要的是，RAE 緩解了 VAE 的根本局限性，后者的潛在空間被高度壓縮（例如，SD-VAE 將的圖像映射到的潛在表征，這限制了重建的保真度，更關(guān)鍵的是，也限制了表征的質(zhì)量。

用于 RAE 解碼器的訓(xùn)練方案如下：

首先，給定一個尺寸為 3×H×W 的輸入圖像 x，并使用一個預(yù)先訓(xùn)練好且凍結(jié)的表征編碼器 E。該編碼器的 patch 大小為 p_e，隱藏層大小為 d。經(jīng)過編碼器處理后，輸入圖像被轉(zhuǎn)換為個 token，每個 token 都有 d 個通道。

接著，一個 patch 大小為 p_d 的 ViT 解碼器 D 會接收這些 token，并將它們映射回像素空間，重建出圖像。重建圖像的輸出形狀為。在默認情況下，設(shè)置 p_d = p_e，從而使重建結(jié)果與輸入的分辨率相匹配。

在所有針對 256×256 圖像的實驗中，編碼器均產(chǎn)生 256 個 token。這個數(shù)量與多數(shù)先前基于 DiT 且使用 SD-VAE 潛在表征進行訓(xùn)練的模型的 token 數(shù)量相符。

最后，在訓(xùn)練解碼器 D 時，遵循了 VAE 的常見做法，采用了 L1 損失、LPIPS 損失和對抗性損失相結(jié)合的優(yōu)化目標：

研究者從不同的預(yù)訓(xùn)練范式中選擇了三個代表性的編碼器：

• DINOv2-B (p_e=14,d=768)，一個自監(jiān)督自蒸餾模型；
• SigLIP2-B (p_e=16,d=768)，一個語言監(jiān)督模型；
• MAE-B (p_e=16,d=768)，一個掩碼自編碼器。

對于 DINOv2，還研究了不同模型尺寸 S、B、L (d=384,768,1024)。除非另有說明，研究者在所有 RAE 中都使用 ViT-XL 解碼器。研究者使用在重建的 ImageNet 驗證集上計算的 FID 分數(shù)作為衡量重建質(zhì)量的主要指標，記為 rFID。

重建、擴展性與表征能力

如表 1a 所示，使用凍結(jié)編碼器的 RAE 在重建質(zhì)量 (rFID) 上一致優(yōu)于 SD-VAE。例如，使用 MAE-B/16 的 RAE 達到了 0.16 的 rFID，明顯勝過 SD-VAE，并挑戰(zhàn)了表征編碼器無法恢復(fù)像素級細節(jié)的假設(shè)。

接下來，研究了編碼器和解碼器的擴展性行為。如表 1c 所示，在 DINOv2-S、B 和 L 三種尺寸下，重建質(zhì)量保持穩(wěn)定，這表明即使是小型的表征編碼器模型也保留了足夠的底層細節(jié)以供解碼。在解碼器方面（表 1b），增加其容量能夠持續(xù)提升 rFID：從 ViT-B 的 0.58 提升到 ViT-XL 的 0.49。重要的是，ViT-B 的性能已經(jīng)超過 SD-VAE，而其 GFLOPs 效率要高出 14 倍；ViT-XL 則以僅為 SD-VAE 三分之一的計算成本進一步提升了質(zhì)量。

研究者還在表 1d 中通過在 ImageNet-1K 上的線性探測來評估表征質(zhì)量。因為 RAE 使用凍結(jié)的預(yù)訓(xùn)練編碼器，它們直接繼承了底層表征編碼器的表征能力。相比之下，SD-VAE 僅實現(xiàn)了約 8% 的準確率。

為 RAE 駕馭擴散 Transformer

在 RAE 已展示出良好重建質(zhì)量的基礎(chǔ)上，研究者進一步探討了其在潛空間的可擴散性。

在正式進入生成實驗之前，研究者首先固定編碼器，以研究不同編碼器下的生成能力。表 1a 顯示，MAE、SigLIP2 和 DINOv2 的重建誤差（rFID）均低于 SD-VAE，其中 MAE 的重建表現(xiàn)最好。

然而，研究者指出：僅有重建質(zhì)量好并不意味著生成質(zhì)量高。在實際實驗中，DINOv2 在圖像生成任務(wù)中的表現(xiàn)最強。因此，除非特別說明，后續(xù)實驗都將默認使用 DINOv2 作為編碼器。在模型架構(gòu)上，研究者使用了 LightningDiT 作為基礎(chǔ)網(wǎng)絡(luò)，它是 DiT 的一種改進版本。

然而，出乎意料的是，標準的擴散模型訓(xùn)練方法在 RAE 潛空間中完全失效（見表 2）。

當直接在 RAE 的潛變量上進行訓(xùn)練時：

• 小規(guī)模的模型（如 DiT-S）會徹底訓(xùn)練失敗，無法生成有效結(jié)果；
• 較大的模型（如 DiT-XL）雖然能夠訓(xùn)練，但其表現(xiàn)仍然遠遜于在 SD-VAE 潛空間上訓(xùn)練的同等規(guī)模模型。

為了研究這一觀察結(jié)果，研究者提出了下面幾個假設(shè)：

擴展 DiT 寬度以匹配 Token 維度

為分析擴散 Transformer (DiT) 在 RAE 潛變量上的訓(xùn)練動態(tài)，研究人員進行了一項簡化實驗，旨在通過 DiT 重建由 RAE 編碼的單個圖像。實驗通過固定模型深度并改變其寬度（隱藏維度 d）發(fā)現(xiàn)，當模型寬度小于 Token 維度 n (d < n=768) 時，樣本質(zhì)量和訓(xùn)練損失表現(xiàn)均很差。然而，一旦寬度匹配或超過 Token 維度 (d ≥ n)，樣本質(zhì)量便會急劇提升至近乎完美，同時訓(xùn)練損失也迅速收斂。

為排除這種性能提升僅是模型總?cè)萘吭黾拥慕Y(jié)果，對照實驗將寬度固定為較小值 (d=384) 并將深度加倍。結(jié)果顯示，模型性能并未改善，圖像依然充滿瑕疵，且損失無法收斂。這表明，要使 DiT 在 RAE 的潛空間中成功生成，其模型寬度必須匹配或超過 RAE 的 Token 維度。

這一要求似乎與數(shù)據(jù)流形具有較低內(nèi)在維度的普遍認知相悖。研究者推斷，這源于擴散模型的內(nèi)在機制：在訓(xùn)練過程中持續(xù)向數(shù)據(jù)注入高斯噪聲，實際上將數(shù)據(jù)流形的支撐集擴展至整個空間，使其成為一個「滿秩流形」。因此，模型容量必須與完整的數(shù)據(jù)維度成比例，而非其較低的內(nèi)在維度。

該猜想得到了理論下界 L≥(n?d)/n 的支持，該公式與實驗結(jié)果高度吻合。研究人員通過將不同寬度的 DiT 模型 (S/B/L) 與具有相應(yīng) Token 維度的 DINOv2 編碼器 (S/B/L) 配對，在更真實的場景中進一步驗證了此結(jié)論：模型僅在自身寬度不小于編碼器 Token 維度時才能有效收斂。

維度相關(guān)的噪聲調(diào)度偏移

先前研究已證實，擴散模型訓(xùn)練中的最優(yōu)噪聲調(diào)度與輸入數(shù)據(jù)的空間分辨率相關(guān)。本文將此概念從空間分辨率推廣至有效數(shù)據(jù)維度，即 Token 數(shù)量與 Token 維度的乘積。其核心在于，高斯噪聲會同等地作用于所有維度，因此 RAE 潛變量的高維度（與傳統(tǒng) VAE 或像素的低通道數(shù)不同）在相同的噪聲水平下能保留更多信息，從而需要調(diào)整噪聲注入的策略。

為此，研究者采用了 Esser et al. (2024) 的調(diào)度偏移方法，通過一個維度相關(guān)的縮放因子 α=m/n 來調(diào)整噪聲時間步長（其中 m 為 RAE 的有效數(shù)據(jù)維度，n 為基準維度）。實驗結(jié)果表明，應(yīng)用此維度自適應(yīng)的噪聲調(diào)度帶來了顯著的性能提升，證明了在高維潛空間中訓(xùn)練擴散模型時進行此項調(diào)整的必要性。

噪聲增強解碼

RAE 解碼器通?；谝唤M離散、干凈的潛變量進行訓(xùn)練。然而，擴散模型在推理時生成的潛變量往往帶有噪聲或與訓(xùn)練分布存在偏差，這會給解碼器帶來分布外 (OOD) 挑戰(zhàn)，從而降低最終的樣本質(zhì)量。

為緩解這一問題，研究者提出了噪聲增強解碼方案。該方法在訓(xùn)練解碼器時，向原始的干凈潛變量 z 中注入了加性高斯噪聲 n～N (0,σ2I)。此過程通過平滑潛在分布，增強了解碼器對擴散模型產(chǎn)生的更密集、更連續(xù)的輸出空間的泛化能力。為進一步正則化訓(xùn)練并提升魯棒性，噪聲的標準差 σ 也被隨機化。

這一技術(shù)帶來了預(yù)期的權(quán)衡：通過提升對 OOD 潛變量的魯棒性，模型的生成指標 (gFID) 得以改善，但由于注入的噪聲會去除部分精細細節(jié)，重建指標 (rFID) 會略微下降。

最終，將上述所有技術(shù)（模型寬度匹配、噪聲調(diào)度偏移及噪聲增強解碼）相結(jié)合，一個在 RAE 潛變量上訓(xùn)練的 DiT-XL 模型在 720 個 epoch 后實現(xiàn)了 2.39 的 gFID。這一成果在收斂速度上大幅超越了先前基于 VAE 潛變量的擴散模型（相比 SiT-XL 實現(xiàn) 47 倍訓(xùn)練加速）以及近期的表示對齊方法（相比 REPA-XL 實現(xiàn) 16 倍訓(xùn)練加速），為高效生成模型的訓(xùn)練樹立了新的標桿。

實驗結(jié)果

在標準的 DiT 架構(gòu)中，處理高維的 RAE 潛變量通常需要擴大整個主干網(wǎng)絡(luò)的寬度，而這會導(dǎo)致計算開銷激增。

為了解決這一問題，研究者借鑒了 DDT 的設(shè)計思想，引入了 DDT head，一個淺層但寬度較大的 Transformer 模塊，專門用于去噪任務(wù)。通過將該模塊附加到標準的 DiT 上，模型能夠在不顯著增加計算量的情況下有效提升網(wǎng)絡(luò)寬度。

研究者將這種增強后的架構(gòu)稱為 DiT^DH。

其中，DiT^DH 的收斂速度比 DiT 快，并且，DiT^DH 在計算效率（FLOPs）方面顯著優(yōu)于 DiT，如圖 6a 所示。

此外，DiT^DH 在不同規(guī)模的 RAE 上依然保持性能優(yōu)勢。

如表 6 所示，DiT^DH 在所有情況下都穩(wěn)定優(yōu)于 DiT，并且隨著編碼器規(guī)模的增大，其優(yōu)勢也隨之擴大。例如，在使用 DINOv2-L 時，DiT^DH 將 FID 從 6.09 降低至 2.73。

研究者將這種魯棒性歸功于 DDT head 的設(shè)計。較大的編碼器會生成更高維度的潛變量，這會放大 DiT 的寬度瓶頸問題。而 DiT^DH 通過滿足寬度需求，同時保持特征表示緊湊，有效地解決了這一問題。

此外，DDT head 還能過濾掉高維 RAE 潛變量中更容易出現(xiàn)的噪聲信息，從而進一步提升模型性能與穩(wěn)定性。

收斂性。如圖 6b 所示，研究者繪制了 DiT^DH-XL 的訓(xùn)練收斂曲線，實驗結(jié)果顯示：

• 當訓(xùn)練計算量達到約 5 × 101? GFLOPs 時，DiT^DH-XL 的表現(xiàn)已經(jīng)超越 REPA-XL、MDTv2-XL 和 SiT-XL 等模型。
• 在 5 × 1011 GFLOPs 時，DiT^DH-XL 實現(xiàn)了全場最佳 FID，而所需計算量僅為這些基線模型的 1/40。

換句話說，DiT^DH-XL 不僅收斂速度更快，而且在相同或更低的計算預(yù)算下能達到更優(yōu)性能，展現(xiàn)出極高的計算效率與訓(xùn)練穩(wěn)定性。

擴展性（Scaling）。研究者將 DiT^DH 與近年來不同規(guī)模的擴散模型進行了比較。結(jié)果如圖 6c 所示：

• 隨著 DiT^DH 模型規(guī)模的增加，其 FID 分數(shù)持續(xù)提升，表現(xiàn)出良好的可擴展性；
• 最小的模型 DiT^DH-S 已能取得 6.07 的 FID 分數(shù)，性能甚至超過了體量更大的 REPA-XL；
• 當模型從 DiT^DH-S 擴展到 DiT^DH-B 時，F(xiàn)ID 由 6.07 變?yōu)?3.38，超越了所有相似規(guī)模甚至更大規(guī)模的以往模型；
• 進一步擴展到 DiT^DH-XL 后，性能繼續(xù)提升，在僅 80 個訓(xùn)練周期（epochs）下取得了 2.16 的 FID，創(chuàng)下了新的 SOTA 紀錄。

最后，研究者對 DiT^DH-XL（該系列中性能最強的模型）與近期多款最先進的擴散模型進行了定量性能對比。結(jié)果顯示：本文方法大大優(yōu)于所有先前的擴散模型，在 256×256 下創(chuàng)下了新的最先進的 FID 分數(shù)：無指導(dǎo)時為 1.51，有指導(dǎo)時為 1.13。在 512×512 上，經(jīng)過 400 次 epoch 訓(xùn)練，DiT^DH-XL 在有指導(dǎo)的情況下進一步實現(xiàn)了 1.13 的 FID，超過了 EDM-2 之前的最佳性能（1.25）。

圖 7 為可視化結(jié)果，模型能夠生成多種類別和場景下的圖像，反映出其強大的內(nèi)容理解與泛化能力；圖像細節(jié)逼真、紋理自然，與 ImageNet 的真實樣本相當。

了解更多內(nèi)容，請參考原論文。