本文經(jīng)AI新媒體量子位（公眾號ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。

“U-Net已死，Transformer成為擴散模型新SOTA了！”

就在ChatGPT占盡AI圈風(fēng)頭時，紐約大學(xué)謝賽寧的圖像生成模型新論文橫空出世，收獲一眾同行驚訝的聲音。

△MILA在讀ML博士生Ethan Caballero

論文創(chuàng)意性地將Transformer與擴散模型融合，在計算效率和生成效果上均超越了基于U-Net的經(jīng)典模型ADM和LDM，打破了U-Net統(tǒng)治擴散模型的“普遍認(rèn)知”。

網(wǎng)友給這對新組合命名也是腦洞大開：

All we need is U-Transformer

希望他們沒有錯過Transffusion這個名字。

要知道，這幾年雖然Transformer占盡風(fēng)頭，但U-Net在擴散模型領(lǐng)域仍然一枝獨秀——

無論是“前任王者”DALL·E2還是“新晉生成AI”Stable Diffusion，都沒有使用Transformer作為圖像生成架構(gòu)。

△英偉達(dá)AI科學(xué)家Jim Fan

如今新研究表明，U-Net并非不可用Transformer替代。

“U-Net并非不可替代”

論文提出的新架構(gòu)名叫Diffusion Transformers（DiTs）。

架構(gòu)保留了很多ViT的特性，其中整體架構(gòu)如圖左（包含多個DiT模塊），具體的DiT模塊組成如圖右：

更右邊的兩個灰色框的模塊，則是DiT架構(gòu)的“變體”。主要是探討在條件輸入下，不同的架構(gòu)是否能對信息進(jìn)行更好的處理，包括交叉注意力等。

最終結(jié)果表明，還是層歸一化（Layer Normalization）更好用，這里最終選用了Adaptive Layer Normalization（自適應(yīng)層歸一化）的方法。

對于這篇論文研究的目的，作者表示希望探討擴散模型中不同架構(gòu)選擇的重要性，以及也是給將來生成模型的評估做一個評判標(biāo)準(zhǔn)。

先說結(jié)果——作者認(rèn)為，U-Net的歸納偏置（inductive bias），對于擴散模型性能提升不是必須的。

與之相反，他們能“輕松地”（readily）被Transformer的標(biāo)準(zhǔn)架構(gòu)取代。

有網(wǎng)友發(fā)現(xiàn)，DALL·E和DALL·E2似乎都有用到Transformer。

這篇論文和它們的差異究竟在哪里？

事實上，DALL·E雖然是Transformer，但并非擴散模型，本質(zhì)是基于VQVAE架構(gòu)實現(xiàn)的；

至于DALL·E2和Stable Diffusion，雖然都分別將Transformer用在了CLIP和文本編碼器上，但關(guān)鍵的圖像生成用的還是U-Net。

△經(jīng)典U-Net架構(gòu)

不過，DiT還不是一個文本生成圖像模型——目前只能基于訓(xùn)練標(biāo)簽生成對應(yīng)的新圖像。

雖然生成的圖片還帶著股“ImageNet風(fēng)”，不過英偉達(dá)AI科學(xué)家Jim Fan認(rèn)為，將它改造成想要的風(fēng)格和加上文本生成功能，都不是難點。

如果將標(biāo)簽輸入調(diào)整成其他向量、乃至于文本嵌入，就能很快地將DiT改造成一個文生圖模型：

Stable-DiT馬上就要來了！

所以DiTs在生成效果和運算速率上，相比其他圖像生成模型究竟如何？

在ImageNet基準(zhǔn)上取得SOTA

為了驗證DiTs的最終效果，研究者將DiTs沿“模型大小”和“輸入標(biāo)記數(shù)量”兩個軸進(jìn)行了縮放。

具體來說，他們嘗試了四種不同模型深度和寬度的配置：DiT-S、DiT-B、DiT-L和DiT-XL，在此基礎(chǔ)上又分別訓(xùn)練了3個潛塊大小為8、4和2的模型，總共是12個模型。

從FID測量結(jié)果可以看出，就像其他領(lǐng)域一樣，增加模型大小和減少輸入標(biāo)記數(shù)量可以大大提高DiT的性能。

FID是計算真實圖像和生成圖像的特征向量之間距離的一種度量，越小越好。

換句話說，較大的DiTs模型相對于較小的模型是計算效率高的，而且較大的模型比較小的模型需要更少的訓(xùn)練計算來達(dá)到給定的FID。

其中，Gflop最高的模型是DiT-XL/2，它使用最大的XL配置，patch大小為2，當(dāng)訓(xùn)練時間足夠長時，DiT-XL/2就是里面的最佳模型。

于是在接下來，研究人員就專注于DiT-XL/2，他們在ImageNet上訓(xùn)練了兩個版本的DiT-XL/2，分辨率分別為256x256和512x512，步驟分別為7M和3M。

當(dāng)使用無分類器指導(dǎo)時，DiT-XL/2比之前的擴散模型數(shù)據(jù)都要更好，取得SOTA效果：

在256x256分辨率下，DiT-XL/2將之前由LDM實現(xiàn)的最佳FID-50K從3.60降至了2.27。

并且與基線相比，DiTs模型本身的計算效率也很高：

DiT-XL/2的計算效率為119 Gflops，相比而言LDM-4是103 Gflops，ADM-U則是742 Gflops。

同樣，在512x512分辨率下，DiT-XL/2也將ADM-U之前獲得的最佳FID 3.85降至了3.04。

不過此時ADM-U的計算效率是2813 Gflops，而XL/2只有525 Gflops。

研究作者

本篇論文作者為UC伯克利的William Peebles和紐約大學(xué)的謝賽寧。

William Peebles，目前是UC伯克利的四年級博士生，本科畢業(yè)于麻省理工學(xué)院。研究方向是深度學(xué)習(xí)和人工智能，重點是深度生成模型。

之前曾在Meta、Adobe、英偉達(dá)實習(xí)過，這篇論文就是在Meta實習(xí)期間完成。

謝賽寧，紐約大學(xué)計算機科學(xué)系助理教授，之前曾是Meta FAIR研究員，本科就讀于上海交通大學(xué)ACM班，博士畢業(yè)于UC圣迭戈分校。

謝賽寧讀博士時曾在FAIR實習(xí)，期間與何愷明合作完成ResNeXt，是該論文的一作，之前何愷明一作論文MAE他也有參與。

當(dāng)然，對于這次Transformer的表現(xiàn)，也有研究者們表示“U-Net不服”。

例如三星AI Lab科學(xué)家Alexia Jolicoeur-Martineau就表示：

U-Net仍然充滿生機，我相信只需要經(jīng)過細(xì)小調(diào)整，有人能將它做得比Transformer更好。

看來，圖像生成領(lǐng)域很快又要掀起新的“較量風(fēng)暴”了。

?論文地址：???https://arxiv.org/abs/2212.09748v1?