摘要
尽管在图像生成方面取得了显著成功，我们观察到扩散概率模型（Diffusion Probabilistic Models, DPMs）往往缺乏对图像中物体各部分之间上下文关系的建模能力，导致学习过程缓慢。为解决这一问题，我们提出了一种掩码扩散Transformer（Masked Diffusion Transformer, MDT），通过引入一种掩码潜在表示建模机制，显式增强DPMs在图像语义部件之间进行上下文关系学习的能力。在训练过程中，MDT在潜在空间中对部分标记（tokens）进行掩码处理，并设计了一种非对称扩散Transformer，从未掩码的标记中预测被掩码的标记，同时保持扩散生成过程的完整性。该方法能够仅凭不完整的上下文输入重建图像的完整信息，从而有效学习图像标记之间的关联关系。为进一步提升效率，我们提出了改进版本MDTv2，采用更高效的宏观网络结构与训练策略。实验结果表明，MDTv2在图像生成性能上达到新SOTA水平，例如在ImageNet数据集上取得了1.58的FID分数，且训练速度比此前的SOTA模型DiT快逾10倍。相关源代码已开源，地址为：https://github.com/sail-sg/MDT。