“U-Net已死,Transformer成為擴(kuò)散模型新SOTA了!”
(資料圖片)
就在ChatGPT占盡AI圈風(fēng)頭時(shí),紐約大學(xué)謝賽寧的圖像生成模型新論文橫空出世,收獲一眾同行驚訝的聲音。
△MILA在讀ML博士生Ethan Caballero
論文創(chuàng)意性地將Transformer與擴(kuò)散模型融合,在計(jì)算效率和生成效果上均超越了基于U-Net的經(jīng)典模型ADM和LDM,打破了U-Net統(tǒng)治擴(kuò)散模型的“普遍認(rèn)知”。
網(wǎng)友給這對新組合命名也是腦洞大開:
All we need is U-Transformer
希望他們沒有錯(cuò)過Transffusion這個(gè)名字。
要知道,這幾年雖然Transformer占盡風(fēng)頭,但U-Net在擴(kuò)散模型領(lǐng)域仍然一枝獨(dú)秀——
無論是“前任王者”DALL E2還是“新晉生成AI”Stable Diffusion,都沒有使用Transformer作為圖像生成架構(gòu)。
△英偉達(dá)AI科學(xué)家Jim Fan
如今新研究表明,U-Net并非不可用Transformer替代?!癠-Net并非不可替代”
論文提出的新架構(gòu)名叫Diffusion Transformers(DiTs)。
架構(gòu)保留了很多ViT的特性,其中整體架構(gòu)如圖左(包含多個(gè)DiT模塊),具體的DiT模塊組成如圖右:
更右邊的兩個(gè)灰色框的模塊,則是DiT架構(gòu)的“變體”。主要是探討在條件輸入下,不同的架構(gòu)是否能對信息進(jìn)行更好的處理,包括交叉注意力等。
最終結(jié)果表明,還是層歸一化(Layer Normalization)更好用,這里最終選用了Adaptive Layer Normalization(自適應(yīng)層歸一化)的方法。
對于這篇論文研究的目的,作者表示希望探討擴(kuò)散模型中不同架構(gòu)選擇的重要性,以及也是給將來生成模型的評估做一個(gè)評判標(biāo)準(zhǔn)。
先說結(jié)果——作者認(rèn)為,U-Net的歸納偏置(inductive bias),對于擴(kuò)散模型性能提升不是必須的。
與之相反,他們能“輕松地”(readily)被Transformer的標(biāo)準(zhǔn)架構(gòu)取代。
有網(wǎng)友發(fā)現(xiàn),DALL E和DALL E2似乎都有用到Transformer。
這篇論文和它們的差異究竟在哪里?
事實(shí)上,DALL E雖然是Transformer,但并非擴(kuò)散模型,本質(zhì)是基于VQVAE架構(gòu)實(shí)現(xiàn)的;
至于DALL E2和Stable Diffusion,雖然都分別將Transformer用在了CLIP和文本編碼器上,但關(guān)鍵的圖像生成用的還是U-Net。
△經(jīng)典U-Net架構(gòu)
不過,DiT還不是一個(gè)文本生成圖像模型——目前只能基于訓(xùn)練標(biāo)簽生成對應(yīng)的新圖像。
雖然生成的圖片還帶著股“ImageNet風(fēng)”,不過英偉達(dá)AI科學(xué)家Jim Fan認(rèn)為,將它改造成想要的風(fēng)格和加上文本生成功能,都不是難點(diǎn)。
如果將標(biāo)簽輸入調(diào)整成其他向量、乃至于文本嵌入,就能很快地將DiT改造成一個(gè)文生圖模型:
Stable-DiT馬上就要來了!
所以DiTs在生成效果和運(yùn)算速率上,相比其他圖像生成模型究竟如何?在ImageNet基準(zhǔn)上取得SOTA
為了驗(yàn)證DiTs的最終效果,研究者將DiTs沿“模型大小”和“輸入標(biāo)記數(shù)量”兩個(gè)軸進(jìn)行了縮放。
具體來說,他們嘗試了四種不同模型深度和寬度的配置:DiT-S、DiT-B、DiT-L和DiT-XL,在此基礎(chǔ)上又分別訓(xùn)練了3個(gè)潛塊大小為8、4和2的模型,總共是12個(gè)模型。
從 FID測量結(jié)果可以看出,就像其他領(lǐng)域一樣,增加模型大小和減少輸入標(biāo)記數(shù)量可以大大提高DiT的性能。
FID是計(jì)算真實(shí)圖像和生成圖像的特征向量之間距離的一種度量,越小越好。
換句話說,較大的DiTs模型相對于較小的模型是計(jì)算效率高的,而且較大的模型比較小的模型需要更少的訓(xùn)練計(jì)算來達(dá)到給定的FID。
其中,Gflop最高的模型是DiT-XL/2,它使用最大的XL配置,patch大小為2,當(dāng)訓(xùn)練時(shí)間足夠長時(shí),DiT-XL/2就是里面的最佳模型。
于是在接下來,研究人員就專注于DiT-XL/2,他們在ImageNet上訓(xùn)練了兩個(gè)版本的DiT-XL/2,分辨率分別為256x256和512x512,步驟分別為7M和3M。
當(dāng)使用無分類器指導(dǎo)時(shí),DiT-XL/2比之前的擴(kuò)散模型數(shù)據(jù)都要更好,取得SOTA效果:
在 256x256分辨率下,DiT-XL/2將之前由LDM實(shí)現(xiàn)的最佳FID-50K從3.60降至了2.27。
并且與基線相比,DiTs模型本身的計(jì)算效率也很高:
DiT-XL/2的計(jì)算效率為119Gflops,相比而言LDM-4是103Gflops,ADM-U則是742Gflops。
同樣,在512x512分辨率下,DiT-XL/2也將ADM-U之前獲得的最佳FID3.85降至了3.04。
不過此時(shí)ADM-U的計(jì)算效率是2813Gflops,而XL/2只有525Gflops。
研究作者
本篇論文作者為UC伯克利的William Peebles和紐約大學(xué)的謝賽寧。
William Peebles,目前是UC伯克利的四年級博士生,本科畢業(yè)于麻省理工學(xué)院。研究方向是深度學(xué)習(xí)和人工智能,重點(diǎn)是深度生成模型。
之前曾在Meta、Adobe、英偉達(dá)實(shí)習(xí)過,這篇論文就是在Meta實(shí)習(xí)期間完成。
謝賽寧,紐約大學(xué)計(jì)算機(jī)科學(xué)系助理教授,之前曾是Meta FAIR研究員,本科就讀于上海交通大學(xué)ACM班,博士畢業(yè)于UC圣迭戈分校。
謝賽寧讀博士時(shí)曾在FAIR實(shí)習(xí),期間與何愷明合作完成ResNeXt,是該論文的一作,之前何愷明一作論文MAE他也有參與。
當(dāng)然,對于這次Transformer的表現(xiàn),也有研究者們表示“U-Net不服”。
例如三星AI Lab科學(xué)家Alexia Jolicoeur-Martineau就表示:
U-Net仍然充滿生機(jī),我相信只需要經(jīng)過細(xì)小調(diào)整,有人能將它做得比Transformer更好。
看來,圖像生成領(lǐng)域很快又要掀起新的“較量風(fēng)暴”了。
論文地址:
https://arxiv.org/abs/2212.09748v1
標(biāo)簽: 計(jì)算效率 紐約大學(xué) 標(biāo)記數(shù)量
熱門