時令 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

昔日風光無限的VAE，終于被宣判“退役”？

謝賽寧團隊最新研究給出了答案——VAE的時代結束，RAE將接力前行。

其中表征自編碼器RAE（Representation Autoencoders）是一種用于擴散Transformer（DiT）訓練的新型自動編碼器，其核心設計是用預訓練的表征編碼器（如DINO、SigLIP、MAE 等）與訓練后的輕量級解碼器配對，從而替代傳統擴散模型中依賴的VAE（變分自動編碼器）。

這種新結構不僅能提供高質量重建結果，還具備語義豐富的潛空間，同時支持可擴展的基于變換器的架構。

該方法在無需額外表示對齊損失的情況下，實現了更快的收斂速度。通過采用配備輕量級寬型DDT頭部的DiT變體，他們在ImageNet上取得強勁的圖像生成效果：

• 256×256分辨率下，無引導（no guidance）FID= 1.51；
• 256×256和512×512分辨率下，有引導（with guidance）FID=1.13。

下面具體來看。

VAE退役，RAE當立

如今，Diffusion Transformer雖已取得長足發(fā)展，但多數模型仍依賴2021年的舊版SD-VAE構建潛空間。

這引發(fā)了幾大核心問題：

1、過時的骨干網絡，讓架構過于復雜。SD-VAE約需450 GFLOPs運算量，而簡易的ViT-B編碼器僅需22 GFLOPs。

2、過度壓縮的潛空間（只有4個通道），嚴重限制信息容量。常言道壓縮催生智能，但此處不然：VAE式壓縮收效甚微，其信息承載能力與原始3通道像素幾乎無異。

3、薄弱的表征能力。僅依賴重建訓練的模式使VAE學得的特征質量低下（線性探測精度約8%），最終拖慢收斂速度并損害生成質量?，F有研究已表明：表征質量直接決定生成效果。而SD-VAE的設計初衷并未涵蓋此目標。

謝賽寧曾以為語義編碼器主要捕獲高層次抽象表征而會舍棄細粒度視覺細節(jié)，但他現在意識到這個想法是錯誤的。

針對上述問題，研究團隊采用預訓練表征編碼器（如基于標準化ViT架構的DINO、SigLIP和MAE）與訓練好的解碼器相結合，得到了RAE——

無需額外訓練或對齊階段，沒有輔助損失函數，也不引入重新壓縮的適配層。

只需獲取預訓練語義編碼器，使用L1+LPIPS+GAN損失訓練解碼器即可。

盡管看起來架構如此簡潔，但RAE在重建質量上卻能超越SD-VAE。

有意思的是，謝賽寧還以為擴散模型在高維空間中很難高效去噪，但他承認自己又錯了。

由于RAE的潛空間本質上是高維的，擴散Transformer確實需要一些適配，但只需三個非常簡單的調整，它們的表現就能出乎意料地好。

1、寬DiT設計：要使擴散正常運作，變換器寬度d必須至少等于潛表征維度n。若不滿足此條件，模型甚至無法過擬合單個樣本。

2、噪聲調度：依賴分辨率的噪聲調度調整早已用于高分辨率圖像生成。同理，調整噪聲調度可使擴散模型平滑適應增加的輸入通道維度。

3、噪聲解碼器：為提升解碼器對潛空間微小擴散誤差的魯棒性，他們在解碼器訓練中注入微量噪聲。這使解碼器能優(yōu)雅處理重建表征中的細微瑕疵。

憑借這些簡單調整，團隊訓練的DiT-XL模型已超越REPA，且無需引入任何輔助損失或額外訓練階段。

采用RAE時，收斂速度比基于SD-VAE的REPA快達16倍。

事實表明，模型確實需要足夠的寬度，但單純依靠暴力擴展DiT寬度很快就會變得低效且不切實際。

為此，他們引入了一個簡單而有效的技巧，以在RAE框架內提升DiT的可擴展性。這個思路雖與解耦擴散訓練（DDT）存在松散關聯，但他們的出發(fā)點截然不同。

在新架構中，原始DiT作為條件化骨干網絡，驅動一個極寬但極淺的擴散頭部。該頭部以含噪潛變量x_t為輸入，直接預測速度向量。

借助RAE潛變量，DiTDH在訓練計算量和模型大小方面的擴展效率，均優(yōu)于基于RAE的標準DiT以及基于VAE的傳統方法。

論文鏈接：https://t.co/FGOAP3Eg5m
參考鏈接：https://x.com/sainingxie/status/1977936742763094289