蘋果“套娃”式擴散模型,訓練步數(shù)減少七成!
蘋果的一項最新研究,大幅提高了擴散模型在高分辨率圖像上性能。
利用這種方法,同樣分辨率的圖像,訓練步數(shù)減少了超過七成。
在1024×1024的分辨率下,圖片畫質(zhì)直接拉滿,細節(jié)都清晰可見。
蘋果把這項成果命名為MDM,DM就是擴散模型(Diffusion Model)的縮寫,而第一個M則代表了套娃(Matryoshka)。
就像真的套娃一樣,MDM在高分辨率過程中嵌套了低分辨率過程,而且是多層嵌套。
高低分辨率擴散過程同時進行,極大降低了傳統(tǒng)擴散模型在高分辨率過程中的資源消耗。
對于256×256分辨率的圖像,在批大?。╞atch size)為1024的環(huán)境下,傳統(tǒng)擴散模型需要訓練150萬步,而MDM僅需39萬,減少了超七成。
另外,MDM采用了端到端訓練,不依賴特定數(shù)據(jù)集和預訓練模型,在提速的同時依然保證了生成質(zhì)量,而且使用靈活。
不僅可以畫出高分辨率的圖像,還能合成16×2562的視頻。
有網(wǎng)友評論到,蘋果終于把文本連接到圖像中了。
那么,MDM的“套娃”技術(shù),具體是怎么做的呢?
整體與漸進相結(jié)合
在開始訓練之前,需要將數(shù)據(jù)進行預處理,高分辨率的圖像會用一定算法重新采樣,得到不同分辨率的版本。
然后就是利用這些不同分辨率的數(shù)據(jù)進行聯(lián)合UNet建模,小UNet處理低分辨率,并嵌套進處理高分辨率的大UNet。
通過跨分辨率的連接,不同大小的UNet之間可以共用特征和參數(shù)。
MDM的訓練則是一個循序漸進的過程。
雖然建模是聯(lián)合進行的,但訓練過程并不會一開始就針對高分辨率進行,而是從低分辨率開始逐步擴大。
這樣做可以避免龐大的運算量,還可以讓低分辨率UNet的預訓練可以加速高分辨率訓練過程。
訓練過程中會逐步將更高分辨率的訓練數(shù)據(jù)加入總體過程中,讓模型適應(yīng)漸進增長的分辨率,平滑過渡到最終的高分辨率過程。
不過從整體上看,在高分辨率過程逐步加入之后,MDM的訓練依舊是端到端的聯(lián)合過程。
在不同分辨率的聯(lián)合訓練當中,多個分辨率上的損失函數(shù)一起參與參數(shù)更新,避免了多階段訓練帶來的誤差累積。
每個分辨率都有對應(yīng)的數(shù)據(jù)項的重建損失,不同分辨率的損失被加權(quán)合并,其中為保證生成質(zhì)量,低分辨率損失權(quán)重較大。
在推理階段,MDM采用的同樣是并行與漸進相結(jié)合的策略。
此外,MDM利還采用了預訓練的圖像分類模型(CFG)來引導生成樣本向更合理的方向優(yōu)化,并為低分辨率的樣本添加噪聲,使其更貼近高分辨率樣本的分布。
那么,MDM的效果究竟如何呢?
更少參數(shù)匹敵SOTA
圖像方面,在ImageNet和CC12M數(shù)據(jù)集上,MDM的FID(數(shù)值越低效果越好)和CLIP表現(xiàn)都顯著優(yōu)于普通擴散模型。
其中FID用于評價圖像本身的質(zhì)量,CLIP則說明了圖像和文本指令之間的匹配程度。
和DALL E、IMAGEN等SOTA模型相比,MDM的表現(xiàn)也很接近,但MDM的訓練參數(shù)遠少于這些模型。
不僅是優(yōu)于普通擴散模型,MDM的表現(xiàn)也超過了其他級聯(lián)擴散模型。
消融實驗結(jié)果表明,低分辨率訓練的步數(shù)越多,MDM效果增強就越明顯;另一方面,嵌套層級越多,取得相同的CLIP得分需要的訓練步數(shù)就越少。
而關(guān)于CFG參數(shù)的選擇,則是一個多次測試后再FID和CLIP之間權(quán)衡的結(jié)果(CLIP得分高相對于CFG強度增大)。