MDTv2開(kāi)源，Sora 核心組件 DiT 訓(xùn)練提速 10 倍

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

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

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

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

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

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