DiT作為效果驚艷的Sora的核心技術(shù)之一，利用Difffusion Transfomer 將生成模型擴(kuò)展到更大的模型規(guī)模，從而實(shí)現(xiàn)高質(zhì)量的圖像生成。

然而，更大的模型規(guī)模導(dǎo)致訓(xùn)練成本飆升。

為此，來(lái)自Sea AI Lab、南開(kāi)大學(xué)、昆侖萬(wàn)維2050研究院的顏水成和程明明研究團(tuán)隊(duì)在ICCV 2023提出的Masked Diffusion Transformer利用mask modeling表征學(xué)習(xí)策略通過(guò)學(xué)習(xí)語(yǔ)義表征信息來(lái)大幅加速Diffusion Transfomer的訓(xùn)練速度，并實(shí)現(xiàn)SoTA的圖像生成效果。

圖片

論文地址：https://arxiv.org/abs/2303.14389

GitHub地址：https://github.com/sail-sg/MDT

近日，Masked Diffusion Transformer V2再次刷新SoTA, 相比DiT的訓(xùn)練速度提升10倍以上，并實(shí)現(xiàn)了ImageNet benchmark 上 1.58的FID score。

最新版本的論文和代碼均已開(kāi)源。

背景

盡管以DiT 為代表的擴(kuò)散模型在圖像生成領(lǐng)域取得了顯著的成功，但研究者發(fā)現(xiàn)擴(kuò)散模型往往難以高效地學(xué)習(xí)圖像中物體各部分之間的語(yǔ)義關(guān)系，這一局限性導(dǎo)致了訓(xùn)練過(guò)程的低收斂效率。

圖片

例如上圖所示，DiT在第50k次訓(xùn)練步驟時(shí)已經(jīng)學(xué)會(huì)生成狗的毛發(fā)紋理，然后在第200k次訓(xùn)練步驟時(shí)才學(xué)會(huì)生成狗的一只眼睛和嘴巴，但是卻漏生成了另一只眼睛。

即使在第300k次訓(xùn)練步驟時(shí)，DiT生成的狗的兩只耳朵的相對(duì)位置也不是非常準(zhǔn)確。

這一訓(xùn)練學(xué)習(xí)過(guò)程揭示了擴(kuò)散模型未能高效地學(xué)習(xí)到圖像中物體各部分之間的語(yǔ)義關(guān)系，而只是獨(dú)立地學(xué)習(xí)每個(gè)物體的語(yǔ)義信息。

研究者推測(cè)這一現(xiàn)象的原因是擴(kuò)散模型通過(guò)最小化每個(gè)像素的預(yù)測(cè)損失來(lái)學(xué)習(xí)真實(shí)圖像數(shù)據(jù)的分布，這個(gè)過(guò)程忽略了圖像中物體各部分之間的語(yǔ)義相對(duì)關(guān)系，因此導(dǎo)致模型的收斂速度緩慢。

方法：Masked Diffusion Transformer

受到上述觀察的啟發(fā)，研究者提出了Masked Diffusion Transformer (MDT) 提高擴(kuò)散模型的訓(xùn)練效率和生成質(zhì)量。

MDT提出了一種針對(duì)Diffusion Transformer 設(shè)計(jì)的mask modeling表征學(xué)習(xí)策略，以顯式地增強(qiáng)Diffusion Transformer對(duì)上下文語(yǔ)義信息的學(xué)習(xí)能力，并增強(qiáng)圖像中物體之間語(yǔ)義信息的關(guān)聯(lián)學(xué)習(xí)。

圖片

如上圖所示，MDT在保持?jǐn)U散訓(xùn)練過(guò)程的同時(shí)引入mask modeling學(xué)習(xí)策略。通過(guò)mask部分加噪聲的圖像token，MDT利用一個(gè)非對(duì)稱Diffusion Transformer (Asymmetric Diffusion Transformer) 架構(gòu)從未被mask的加噪聲的圖像token預(yù)測(cè)被mask部分的圖像token，從而同時(shí)實(shí)現(xiàn)mask modeling 和擴(kuò)散訓(xùn)練過(guò)程。

在推理過(guò)程中，MDT仍保持標(biāo)準(zhǔn)的擴(kuò)散生成過(guò)程。MDT的設(shè)計(jì)有助于Diffusion Transformer同時(shí)具有mask modeling表征學(xué)習(xí)帶來(lái)的語(yǔ)義信息表達(dá)能力和擴(kuò)散模型對(duì)圖像細(xì)節(jié)的生成能力。

具體而言，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)

圖片

如上圖所示，Asymmetric Diffusion Transformer架構(gòu)包含encoder、side-interpolater（輔助插值器）和decoder。

圖片

在訓(xùn)練過(guò)程中，Encoder只處理未被mask的token；而在推理過(guò)程中，由于沒(méi)有mask步驟，它會(huì)處理所有token。

因此，為了保證在訓(xùn)練或推理階段，decoder始終能處理所有的token，研究者們提出了一個(gè)方案：在訓(xùn)練過(guò)程中，通過(guò)一個(gè)由DiT block組成的輔助插值器（如上圖所示），從encoder的輸出中插值預(yù)測(cè)出被mask的token，并在推理階段將其移除因而不增加任何推理開(kāi)銷。

MDT的encoder和decoder在標(biāo)準(zhǔn)的DiT block中插入全局和局部位置編碼信息以幫助預(yù)測(cè)mask部分的token。

Asymmetric Diffusion Transformer V2

圖片

如上圖所示，MDTv2通過(guò)引入了一個(gè)針對(duì)Masked Diffusion過(guò)程設(shè)計(jì)的更為高效的宏觀網(wǎng)絡(luò)結(jié)構(gòu)，進(jìn)一步優(yōu)化了diffusion和mask modeling的學(xué)習(xí)過(guò)程。

這包括在encoder中融合了U-Net式的long-shortcut，在decoder中集成了dense input-shortcut。

其中，dense input-shortcut將添加噪后的被mask的token送入decoder，保留了被mask的token對(duì)應(yīng)的噪聲信息，從而有助于diffusion過(guò)程的訓(xùn)練。

此外，MDT還引入了包括采用更快的Adan優(yōu)化器、time-step相關(guān)的損失權(quán)重，以及擴(kuò)大掩碼比率等更優(yōu)的訓(xùn)練策略來(lái)進(jìn)一步加速M(fèi)asked Diffusion模型的訓(xùn)練過(guò)程。

實(shí)驗(yàn)結(jié)果

ImageNet 256基準(zhǔn)生成質(zhì)量比較

圖片

上表比較了不同模型尺寸下MDT與DiT在ImageNet 256基準(zhǔn)下的性能對(duì)比。

顯而易見(jiàn)，MDT在所有模型規(guī)模上都以較少的訓(xùn)練成本實(shí)現(xiàn)了更高的FID分?jǐn)?shù)。

MDT的參數(shù)和推理成本與DiT基本一致，因?yàn)檎缜拔乃榻B的，MDT推理過(guò)程中仍保持與DiT一致的標(biāo)準(zhǔn)的diffusion過(guò)程。

對(duì)于最大的XL模型，經(jīng)過(guò)400k步驟訓(xùn)練的MDTv2-XL/2，顯著超過(guò)了經(jīng)過(guò)7000k步驟訓(xùn)練的DiT-XL/2，F(xiàn)ID分?jǐn)?shù)提高了1.92。在這一setting下，結(jié)果表明了MDT相對(duì)DiT有約18倍的訓(xùn)練加速。

對(duì)于小型模型，MDTv2-S/2 仍然以顯著更少的訓(xùn)練步驟實(shí)現(xiàn)了相比DiT-S/2顯著更好的性能。例如同樣訓(xùn)練400k步驟，MDTv2以39.50的FID指標(biāo)大幅領(lǐng)先DiT 68.40的FID指標(biāo)。

更重要的是，這一結(jié)果也超過(guò)更大模型DiT-B/2在400k訓(xùn)練步驟下的性能（39.50 vs 43.47）。

ImageNet 256基準(zhǔn)CFG生成質(zhì)量比較

圖片

我們還在上表中比較了MDT與現(xiàn)有方法在classifier-free guidance下的圖像生成性能。

MDT以1.79的FID分?jǐn)?shù)超越了以前的SOTA DiT和其他方法。MDTv2進(jìn)一步提升了性能，以更少的訓(xùn)練步驟將圖像生成的SOTA FID得分推至新低，達(dá)到1.58。

與DiT類似，我們?cè)谟?xùn)練過(guò)程中沒(méi)有觀察到模型的FID分?jǐn)?shù)在繼續(xù)訓(xùn)練時(shí)出現(xiàn)飽和現(xiàn)象。

MDT在PaperWithCode的leaderboard上刷新SoTA

收斂速度比較

圖片

上圖比較了ImageNet 256基準(zhǔn)下，8×A100 GPU上DiT-S/2基線、MDT-S/2和MDTv2-S/2在不同訓(xùn)練步驟/訓(xùn)練時(shí)間下的FID性能。

得益于更優(yōu)秀的上下文學(xué)習(xí)能力，MDT在性能和生成速度上均超越了DiT。MDTv2的訓(xùn)練收斂速度相比DiT提升10倍以上。

MDT在訓(xùn)練步驟和訓(xùn)練時(shí)間方面大相比DiT約3倍的速度提升。MDTv2進(jìn)一步將訓(xùn)練速度相比于MDT提高了大約5倍。

例如，MDTv2-S/2僅需13小時(shí)（15k步驟）就展示出比需要大約100小時(shí)（1500k步驟）訓(xùn)練的DiT-S/2更好的性能，這揭示了上下文表征學(xué)習(xí)對(duì)于擴(kuò)散模型更快的生成學(xué)習(xí)至關(guān)重要。

總結(jié)&討論

MDT通過(guò)在擴(kuò)散訓(xùn)練過(guò)程中引入類似于MAE的mask modeling表征學(xué)習(xí)方案，能夠利用圖像物體的上下文信息重建不完整輸入圖像的完整信息，從而學(xué)習(xí)圖像中語(yǔ)義部分之間的關(guān)聯(lián)關(guān)系，進(jìn)而提升圖像生成的質(zhì)量和學(xué)習(xí)速度。

研究者認(rèn)為，通過(guò)視覺(jué)表征學(xué)習(xí)增強(qiáng)對(duì)物理世界的語(yǔ)義理解，能夠提升生成模型對(duì)物理世界的模擬效果。這正與Sora期待的通過(guò)生成模型構(gòu)建物理世界模擬器的理念不謀而合。希望該工作能夠激發(fā)更多關(guān)于統(tǒng)一表征學(xué)習(xí)和生成學(xué)習(xí)的工作。

參考資料：

https://arxiv.org/abs/2303.14389