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

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

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

謝賽寧進(jìn)一步做出了解釋?zhuān)珼iT 雖然取得了長(zhǎng)足的進(jìn)步，但大多數(shù)模型仍然依賴(lài)于 2021 年的舊版 SD-VAE 作為其潛空間基礎(chǔ)。這就帶來(lái)了以下幾個(gè)主要問(wèn)題：

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

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

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

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

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

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

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

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

從網(wǎng)友的反饋來(lái)看，大家非?？春?RAE 的前景，預(yù)計(jì)可以為生成模型帶來(lái)新的可能性。

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

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

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

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

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

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

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

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

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

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

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

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

重建、擴(kuò)展性與表征能力

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

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

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

為 RAE 駕馭擴(kuò)散 Transformer

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

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

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

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

當(dāng)直接在 RAE 的潛變量上進(jìn)行訓(xùn)練時(shí)：

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

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

擴(kuò)展 DiT 寬度以匹配 Token 維度

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

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

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

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

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

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

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

噪聲增強(qiáng)解碼

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

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

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

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

實(shí)驗(yàn)結(jié)果

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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