MDTv2開源，Sora 核心組件 DiT 訓練提速 10 倍

　　由顏水成和程明明領(lǐng)銜的研究團隊在Sora核心技術(shù)上進行了重要的升級，推出了Masked Diffusion Transformer V2。該模型在ImageNet benchmark上取得了驚人的成績，F(xiàn)ID score達到1.58，刷新了State-of-the-Art(SoTA)。

　　作為Sora的核心技術(shù)之一，Diffusion Transformer(DiT)在圖像生成領(lǐng)域取得了顯著的成功，但其訓練成本隨著模型規(guī)模的增大而飆升。研究者發(fā)現(xiàn)，擴散模型難以高效地學習圖像中物體各部分之間的語義關(guān)系，導致訓練過程低收斂效率。為了解決這一問題，顏水成和程明明研究團隊提出了Masked Diffusion Transformer(MDT)，通過引入mask modeling表征學習策略，顯著提升了DiT的訓練效率。

　　MDT采用了針對Diffusion Transformer設(shè)計的mask modeling表征學習策略，增強了模型對上下文語義信息的學習能力，并加強了圖像中物體之間語義信息的關(guān)聯(lián)學習。通過在擴散訓練過程中引入類似于MAE的mask modeling表征學習方案，MDT能夠重建不完整輸入圖像的完整信息，提升圖像生成的質(zhì)量和學習速度。

　　具體而言，MDT通過VAE encoder將圖片映射到latent空間，并在latent空間中進行處理以節(jié)省計算成本。在訓練過程中，MDT首先mask掉部分加噪聲后的圖像token，并將剩余的token送入Asymmetric Diffusion Transformer來預測去噪聲后的全部圖像token。Asymmetric Diffusion Transformer架構(gòu)包含encoder、side-interpolater和decoder，在訓練和推理階段分別處理未被mask的token和所有token，確保了模型的一致性。

　　最新版本的MDT，即MDTv2，通過更為高效的宏觀網(wǎng)絡(luò)結(jié)構(gòu)進一步優(yōu)化了diffusion和mask modeling的學習過程，實現(xiàn)了訓練速度的顯著提升。在ImageNet256基準下，MDTv2相較于DiT，不僅在訓練速度上提升了10倍以上，而且在生成質(zhì)量上取得了更高的FID分數(shù)。MDTv2-S/2在400k步驟訓練下，F(xiàn)ID指標為39.50，明顯領(lǐng)先于DiT-S/2的68.40。

　　總體而言，MDT的創(chuàng)新設(shè)計在擴散模型訓練中引入了有效的語義信息學習，提高了圖像生成的質(zhì)量和訓練速度。研究者認為，通過視覺表征學習增強對物理世界的語義理解，有望提升生成模型對物理世界的模擬效果。這一工作符合Sora的期望，通過生成模型構(gòu)建物理世界模擬器的理念，為未來的表征學習和生成學習研究提供了有力的啟示。