登陆注册
19869

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

站长网2024-03-13 14:37:151

**划重点:**

1. 🌟 Masked Diffusion Transformer V2MDTv2()在ImageNet benchmark上刷新SoTA,FID score达到1.58。

2. 🚀 颜水成/程明明研究团队通过引入Masked Diffusion Transformer有效提升了DiT的训练速度,并解决了扩散模型在学习语义关系方面的难题。

3. 📄 论文和代码已开源,详细信息可查阅[GitHub地址](https://github.com/sail-sg/MDT)。

由颜水成和程明明领衔的研究团队在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的期望,通过生成模型构建物理世界模拟器的理念,为未来的表征学习和生成学习研究提供了有力的启示。

项目入口:

0001
评论列表
共(0)条