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

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

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

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

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

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

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

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

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

因此，RAE 展現出了明顯的優勢，應當成為 DiT 訓練的全新默認方案。

當然，RAE 的模型和 PyTorch 代碼全部開源。這項工作的一作為一年級博士生 Boyang Zheng，其本科畢業于上海交通大學 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

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

基于凍結編碼器的高保真重建

研究者挑戰了一個普遍的假設，即像 DINOv2 和 SigLIP2 這類預訓練表征編碼器不適合重建任務，因為它們 “強調高層語義，而忽略了底層細節” 。

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

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

用于 RAE 解碼器的訓練方案如下：

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

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

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

最后，在訓練解碼器 D 時，遵循了 VAE 的常見做法，采用了 L1 損失、LPIPS 損失和對抗性損失相結合的優化目標：

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

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

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

重建、擴展性與表征能力

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

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

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

為 RAE 駕馭擴散 Transformer

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

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

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

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

當直接在 RAE 的潛變量上進行訓練時：

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

為了研究這一觀察結果，研究者提出了下面幾個假設：

擴展 DiT 寬度以匹配 Token 維度

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

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

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

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

維度相關的噪聲調度偏移

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

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

噪聲增強解碼

RAE 解碼器通常基于一組離散、干凈的潛變量進行訓練。然而，擴散模型在推理時生成的潛變量往往帶有噪聲或與訓練分布存在偏差，這會給解碼器帶來分布外 (OOD) 挑戰，從而降低最終的樣本質量。

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

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

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

實驗結果

在標準的 DiT 架構中，處理高維的 RAE 潛變量通常需要擴大整個主干網絡的寬度，而這會導致計算開銷激增。

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

研究者將這種增強后的架構稱為 DiT^DH。

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

此外，DiT^DH 在不同規模的 RAE 上依然保持性能優勢。

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

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

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

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

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

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

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

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

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

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

了解更多內容，請參考原論文。