RAE：Diffusion Transformers with Representation Autoencoders
abstract
- 作者认为原来的假设/实验结论是错误的：（1）语义建模和重建目标是互相冲突的；（2）diffusion很难预测高维空间特征
method
- 之前的方法假设diffusion 很难学习高维空间的内容
- VAE + diffusion
- 以重构损失为目标训练VAE，将图片压缩到低维度空间，但是token 数目多，比如2562564（token=256*256, dim=4)
- diffusion 预测vae-emb，然后通过VAE-Decoder 重建
- RAE+Diffusion
- 预训练的语义encoder，高维度编码，少量token，比如768*256（token=256,dim=768）
- diffusion预测语义embedding，然后通过一个单独训练的VIT decoder 重建
- 为了能够让diffusion预测高维特征，做出的改进：
- 匹配模型宽度与潜在维度→ 要求DiT的隐藏层宽度 ≥ 潜在token的通道数（否则无法有效拟合）。
- 维度感知的噪声调度调整（Dimension-dependent shift）→ 原有的噪声调度策略针对像素或低维VAE设计，在高维空间下需重新校准时间步。
- 噪声增强解码器训练（Noise-augmented decoding）→ 在训练解码器时加入噪声，使其能更好地处理扩散模型输出的“非完美”潜在变量。
此外，为了更高效地扩展模型宽度而不增加过多计算成本，作者引入了一个新架构 DiTDH（受DDT启发），它包含一个浅但极宽的“去噪头”（wide diffusion head），专门负责处理高维输入。
实验结果
- 在ImageNet上实现了新的SOTA性能：
- 256×256分辨率，无引导生成 FID = 1.51，有AutoGuidance时 FID = 1.13（256×256 和 512×512）
- 训练收敛速度比现有方法快数十倍（如比SiT快47倍）
- 所有指标均超越此前基于VAE的方法