“U-Net 已死，Transformer 成為擴(kuò)散模型新 SOTA 了！”
就在 ChatGPT 占盡 AI 圈風(fēng)頭時(shí)，紐約大學(xué)謝賽寧的圖像生成模型新論文橫空出世，收獲一眾同行驚訝的聲音。
論文創(chuàng)意性地將 Transformer 與擴(kuò)散模型融合，在計(jì)算效率和生成效果上均超越了基于 U-Net 的經(jīng)典模型 ADM 和 LDM，打破了 U-Net 統(tǒng)治擴(kuò)散模型的“普遍認(rèn)知”。
網(wǎng)友給這對(duì)新組合命名也是腦洞大開：
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)。
如今新研究表明，U-Net 并非不可用 Transformer 替代。
“U-Net 并非不可替代”
論文提出的新架構(gòu)名叫 Diffusion Transformers（DiTs）。
架構(gòu)保留了很多 ViT 的特性，其中整體架構(gòu)如圖左（包含多個(gè) DiT 模塊），具體的 DiT 模塊組成如圖右：
更右邊的兩個(gè)灰色框的模塊，則是 DiT 架構(gòu)的“變體”。主要是探討在條件輸入下，不同的架構(gòu)是否能對(duì)信息進(jìn)行更好的處理，包括交叉注意力等。
最終結(jié)果表明，還是層歸一化（Layer Normalization）更好用，這里最終選用了 Adaptive Layer Normalization（自適應(yīng)層歸一化）的方法。
對(duì)于這篇論文研究的目的，作者表示希望探討擴(kuò)散模型中不同架構(gòu)選擇的重要性，以及也是給將來生成模型的評(píng)估做一個(gè)評(píng)判標(biāo)準(zhǔn)。
先說結(jié)果 —— 作者認(rèn)為，U-Net 的歸納偏置（inductive bias），對(duì)于擴(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。
不過，DiT 還不是一個(gè)文本生成圖像模型 —— 目前只能基于訓(xùn)練標(biāo)簽生成對(duì)應(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 測(cè)量結(jié)果可以看出，就像其他領(lǐng)域一樣，增加模型大小和減少輸入標(biāo)記數(shù)量可以大大提高 DiT 的性能。
FID 是計(jì)算真實(shí)圖像和生成圖像的特征向量之間距離的一種度量，越小越好。
換句話說，較大的 DiTs 模型相對(duì)于較小的模型是計(jì)算效率高的，而且較大的模型比較小的模型需要更少的訓(xùn)練計(jì)算來達(dá)到給定的 FID。
其中，Gflop 最高的模型是 DiT-XL / 2，它使用最大的 XL 配置，patch 大小為 2，當(dāng)訓(xùn)練時(shí)間足夠長(zhǎng)時(shí)，DiT-XL / 2 就是里面的最佳模型。
于是在接下來，研究人員就專注于 DiT-XL / 2，他們?cè)?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ì)算效率為 119 Gflops，相比而言 LDM-4 是 103 Gflops，ADM-U 則是 742 Gflops。
同樣，在 512x512 分辨率下，DiT-XL / 2 也將 ADM-U 之前獲得的最佳 FID 3.85 降至了 3.04。
不過此時(shí) ADM-U 的計(jì)算效率是 2813 Gflops，而 XL / 2 只有 525 Gflops。
研究作者
本篇論文作者為 UC 伯克利的 William Peebles 和紐約大學(xué)的謝賽寧。
William Peebles，目前是 UC 伯克利的四年級(jí)博士生，本科畢業(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)然，對(duì)于這次 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
參考鏈接：
[1]https://twitter.com/ethanCaballero/status/1605621603135471616
[2]https://www.wpeebles.com/DiT
[3]https://paperswithcode.com/paper/scalable-diffusion-models-with-transformers#code
本文來自微信公眾號(hào)：量子位 （ID：QbitAI），作者：羿閣 蕭簫
廣告聲明：文內(nèi)含有的對(duì)外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時(shí)間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。