MDTv2开源，Sora 核心组件 DiT 训练提速 10 倍

　　由颜水成和程明明领衔的研究团队在Sora核心技术上进行了重要的升级，推出了Masked Diffusion Transformer V2。该模型在ImageNet benchmark上取得了惊人的成绩，FID score达到1.58，刷新了State-of-the-Art(SoTA)。

　　作为Sora的核心技术之一，Diffusion Transformer(DiT)在图像生成领域取得了显著的成功，但其训练成本随着模型规模的增大而飙升。研究者发现，扩散模型难以高效地学习图像中物体各部分之间的语义关系，导致训练过程低收敛效率。为了解决这一问题，颜水成和程明明研究团队提出了Masked Diffusion Transformer(MDT)，通过引入mask modeling表征学习策略，显著提升了DiT的训练效率。

　　MDT采用了针对Diffusion Transformer设计的mask modeling表征学习策略，增强了模型对上下文语义信息的学习能力，并加强了图像中物体之间语义信息的关联学习。通过在扩散训练过程中引入类似于MAE的mask modeling表征学习方案，MDT能够重建不完整输入图像的完整信息，提升图像生成的质量和学习速度。

　　具体而言，MDT通过VAE encoder将图片映射到latent空间，并在latent空间中进行处理以节省计算成本。在训练过程中，MDT首先mask掉部分加噪声后的图像token，并将剩余的token送入Asymmetric Diffusion Transformer来预测去噪声后的全部图像token。Asymmetric Diffusion Transformer架构包含encoder、side-interpolater和decoder，在训练和推理阶段分别处理未被mask的token和所有token，确保了模型的一致性。

　　最新版本的MDT，即MDTv2，通过更为高效的宏观网络结构进一步优化了diffusion和mask modeling的学习过程，实现了训练速度的显著提升。在ImageNet256基准下，MDTv2相较于DiT，不仅在训练速度上提升了10倍以上，而且在生成质量上取得了更高的FID分数。MDTv2-S/2在400k步骤训练下，FID指标为39.50，明显领先于DiT-S/2的68.40。

　　总体而言，MDT的创新设计在扩散模型训练中引入了有效的语义信息学习，提高了图像生成的质量和训练速度。研究者认为，通过视觉表征学习增强对物理世界的语义理解，有望提升生成模型对物理世界的模拟效果。这一工作符合Sora的期望，通过生成模型构建物理世界模拟器的理念，为未来的表征学习和生成学习研究提供了有力的启示。