昔日风光无限的 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。
下面具体来看。
如今,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
一键三连「点赞」「转发」「小心心」
欢迎在评论区留下你的想法!
— 完 —
年度科技风向标「2025 人工智能年度榜单」评选报名开启啦!我们正在寻找 AI+ 时代领航者 点击了解详情
❤️ 企业、产品、人物 3 大维度,共设立了 5 类奖项,欢迎企业报名参与
一键关注 点亮星标
科技前沿进展每日见
登录后才可以发布评论哦
打开小程序可以发布评论哦