Boyang Zheng Nanye Ma Shengbang Tong Saining Xie

摘要

潜在生成建模——即利用预训练的自编码器将像素映射到潜在空间以支持扩散过程——已成为扩散Transformer（DiT）的标准策略。然而，自编码器组件几乎未取得实质性进展。目前大多数DiT仍依赖于原始的VAE编码器，这带来了诸多局限：过时的骨干网络损害了架构的简洁性，低维潜在空间限制了信息容量，而仅基于重建任务的训练方式导致表征能力薄弱，最终制约了生成质量。在本工作中，我们探索用预训练的表征编码器（如DINO、SigLIP、MAE）结合训练好的解码器替代VAE，构建我们称之为表征自编码器（Representation Autoencoders, RAEs）的新架构。这类模型不仅能够实现高质量的重建，还具备语义丰富的潜在空间，同时支持可扩展的基于Transformer的架构设计。由于这些潜在空间通常具有高维度，一个关键挑战是如何使扩散Transformer在其中高效运行。我们系统分析了该难题的根源，提出了具有理论依据的解决方案，并通过实证加以验证。所提出的方法在无需额外表征对齐损失的情况下实现了更快的收敛速度。采用配备轻量级、宽通道DDT头部的DiT变体，我们在ImageNet数据集上取得了优异的图像生成效果：在256×256分辨率下无引导时FID为1.51，有引导时在256×256和512×512分辨率下FID均达到1.13。结果表明，RAE具有显著优势，应成为扩散Transformer训练的新默认范式。