具有表示自编码器的扩散 Transformer

潜在生成模型，其中预训练的自编码器将像素映射到用于扩散过程的潜在空间，已成为扩散Transformer（DiT）的标准策略；然而，自编码器组件几乎没有进化。大多数DiT继续依赖原始VAE编码器，这带来了一些限制：过时的骨干网络，破坏了架构的简洁性；低维度的潜在空间，限制了信息容量；以及纯粹基于重构的训练产生的弱表示，最终限制了生成质量。在这项工作中，我们探索使用预训练的表示编码器（例如，DINO、SigLIP、MAE）与训练好的解码器配对来替换VAE，形成我们称之为表示自编码器（RAE）的模型。这些模型既提供了高质量的重构，又提供了语义丰富的潜在空间，同时允许可扩展的基于Transformer的架构。由于这些潜在空间通常是高维的，一个关键的挑战是如何使扩散Transformer在其中有效地运行。我们分析了这种困难的根源，提出了理论上受启发的解决方案，并在实践中进行了验证。我们的方法在没有辅助表示对齐损失的情况下实现了更快的收敛。使用配备轻量级、宽DDT头的DiT变体，我们在ImageNet上取得了出色的图像生成结果：256x256时FID为1.51（无引导），256x256和512x512时FID为1.13（有引导）。RAE提供了明显的优势，应成为扩散Transformer训练的新默认选项。