谢赛宁新作：VAE退役，RAE当立

时令 发自 凹非寺
量子位 | 公众号 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虽已取得长足发展，但多数模型仍依赖2021年的旧版SD-VAE构建潜空间。

这引发了几大核心问题：

1、过时的骨干网络，让架构过于复杂。SD-VAE约需450 GFLOPs运算量，而简易的ViT-B编码器仅需22 GFLOPs。

2、过度压缩的潜空间（只有4个通道），严重限制信息容量。常言道压缩催生智能，但此处不然：VAE式压缩收效甚微，其信息承载能力与原始3通道像素几乎无异。

3、薄弱的表征能力。仅依赖重建训练的模式使VAE学得的特征质量低下（线性探测精度约8%），最终拖慢收敛速度并损害生成质量。现有研究已表明：表征质量直接决定生成效果。而SD-VAE的设计初衷并未涵盖此目标。

谢赛宁曾以为语义编码器主要捕获高层次抽象表征而会舍弃细粒度视觉细节，但他现在意识到这个想法是错误的。

针对上述问题，研究团队采用预训练表征编码器（如基于标准化ViT架构的DINO、SigLIP和MAE）与训练好的解码器相结合，得到了RAE——

无需额外训练或对齐阶段，没有辅助损失函数，也不引入重新压缩的适配层。

只需获取预训练语义编码器，使用L1+LPIPS+GAN损失训练解码器即可。

尽管看起来架构如此简洁，但RAE在重建质量上却能超越SD-VAE。

有意思的是，谢赛宁还以为扩散模型在高维空间中很难高效去噪，但他承认自己又错了。

由于RAE的潜空间本质上是高维的，扩散Transformer确实需要一些适配，但只需三个非常简单的调整，它们的表现就能出乎意料地好。

1、宽DiT设计：要使扩散正常运作，变换器宽度d必须至少等于潜表征维度n。若不满足此条件，模型甚至无法过拟合单个样本。

2、噪声调度：依赖分辨率的噪声调度调整早已用于高分辨率图像生成。同理，调整噪声调度可使扩散模型平滑适应增加的输入通道维度。

3、噪声解码器：为提升解码器对潜空间微小扩散误差的鲁棒性，他们在解码器训练中注入微量噪声。这使解码器能优雅处理重建表征中的细微瑕疵。

凭借这些简单调整，团队训练的DiT-XL模型已超越REPA，且无需引入任何辅助损失或额外训练阶段。

采用RAE时，收敛速度比基于SD-VAE的REPA快达16倍。

事实表明，模型确实需要足够的宽度，但单纯依靠暴力扩展DiT宽度很快就会变得低效且不切实际。

为此，他们引入了一个简单而有效的技巧，以在RAE框架内提升DiT的可扩展性。这个思路虽与解耦扩散训练（DDT）存在松散关联，但他们的出发点截然不同。

在新架构中，原始DiT作为条件化骨干网络，驱动一个极宽但极浅的扩散头部。该头部以含噪潜变量x_t为输入，直接预测速度向量。

借助RAE潜变量，DiTDH在训练计算量和模型大小方面的扩展效率，均优于基于RAE的标准DiT以及基于VAE的传统方法。

论文链接：https://t.co/FGOAP3Eg5m
参考链接：https://x.com/sainingxie/status/1977936742763094289