羿閣 蕭簫 發自 凹非寺
量子位 | 公眾號 QbitAI
“U-Net已死，Transformer成為擴散模型新SOTA了！”
就在ChatGPT佔盡AI圈風頭時，紐約大學謝賽寧的影象生成模型新論文橫空出世，收穫一眾同行驚訝的聲音。
△MILA在讀ML博士生Ethan Caballero
論文創意性地將Transformer與擴散模型融合，在計算效率和生成效果上均
超越了基於U-Net的經典模型
ADM和LDM，打破了U-Net統治擴散模型的“普遍認知”。
網友給這對新組合命名也是腦洞大開：
All we need is U-Transformer
希望他們沒有錯過Transffusion這個名字。
要知道，這幾年雖然Transformer佔盡風頭，但U-Net在擴散模型領域仍然一枝獨秀——
無論是“前任王者”DALL·E2還是“新晉生成AI”Stable Diffusion，都沒有使用Transformer作為影象生成架構。
△英偉達AI科學家Jim Fan
如今新研究表明，U-Net並非不可用Transformer替代。
“U-Net並非不可替代”
論文提出的新架構名叫Diffusion Transformers
（DiTs）
。
架構保留了很多ViT的特性，其中整體架構如圖左
（包含多個DiT模組）
，具體的DiT模組組成如圖右：
更右邊的兩個灰色框的模組，則是DiT架構的“變體”。主要是探討在條件輸入下，不同的架構是否能對資訊進行更好的處理，包括交叉注意力等。
最終結果表明，還是層歸一化
（Layer Normalization）
更好用，這裡最終選用了Adaptive Layer Normalization
（自適應層歸一化）
的方法。
對於這篇論文研究的目的，作者表示希望探討擴散模型中不同架構選擇的重要性，以及也是給將來生成模型的評估做一個評判標準。
先說結果——作者認為，U-Net的歸納偏置
（inductive bias）
，對於擴散模型效能提升不是必須的。
與之相反，他們能“輕鬆地”
（readily）
被Transformer的標準架構取代。
有網友發現，DALL·E和DALL·E2似乎都有用到Transformer。
這篇論文和它們的差異究竟在哪裡？
事實上，DALL·E雖然是Transformer，但
並非擴散模型
，本質是基於VQVAE架構實現的；
至於DALL·E2和Stable Diffusion，雖然都分別將Transformer用在了CLIP和文字編碼器上，但關鍵的影象生成用的還是U-Net。
△經典U-Net架構
不過，DiT還不是一個文字生成影象模型——目前只能基於訓練標籤生成對應的新影象。
雖然生成的圖片還帶著股“ImageNet風”，不過英偉達AI科學家
Jim Fan
認為，將它改造成想要的風格和加上文字生成功能，都不是難點。
如果將標籤輸入調整成其他向量、乃至於文字嵌入，就能很快地將DiT改造成一個文生圖模型：
Stable-DiT馬上就要來了！
所以DiTs在生成效果和運算速率上，相比其他影象生成模型究竟如何？
在ImageNet基準上取得SOTA
為了驗證DiTs的最終效果，研究者將DiTs沿“模型大小”和“輸入標記數量”兩個軸進行了縮放。
具體來說，他們嘗試了四種不同模型深度和寬度的配置：DiT-S、DiT-B、DiT-L和DiT-XL，在此基礎上又分別訓練了3個潛塊大小為8、4和2的模型，總共是12個模型。
從FID測量結果可以看出，就像其他領域一樣，增加模型大小和減少輸入標記數量可以大大提高DiT的效能。
FID是計算真實影象和生成影象的特徵向量之間距離的一種度量，越小越好。
換句話說，較大的DiTs模型相對於較小的模型是計算效率高的，而且較大的模型比較小的模型需要更少的訓練計算來達到給定的FID。
其中，Gflop最高的模型是DiT-XL/2，它使用最大的XL配置，patch大小為2，當訓練時間足夠長時，DiT-XL/2就是裡面的最佳模型。
於是在接下來，研究人員就專注於DiT-XL/2，他們在ImageNet上訓練了兩個版本的DiT-XL/2，解析度分別為256x256和512x512，步驟分別為7M和3M。
當使用無分類器指導時，DiT-XL/2比之前的擴散模型資料都要更好，取得SOTA效果：
在256x256解析度下，DiT-XL/2將之前由LDM實現的最佳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和紐約大學的謝賽寧。
William Peebles，目前是UC伯克利的四年級博士生，本科畢業於麻省理工學院。研究方向是深度學習和人工智慧，重點是深度生成模型。
之前曾在META、Adobe、英偉達實習過，這篇論文就是在Meta實習期間完成。
謝賽寧，紐約大學計算機科學系助理教授，之前曾是Meta FAIR研究員，本科就讀於上海交通大學ACM班，博士畢業於UC聖迭戈分校。
謝賽寧讀博士時曾在FAIR實習，期間與何愷明合作完成ResNeXt，是該論文的一作，之前何愷明一作論文MAE他也有參與。
當然，對於這次Transformer的表現，也有研究者們表示“U-Net不服”。
例如三星AI Lab科學家Alexia Jolicoeur-Martineau就表示：
U-Net仍然充滿生機，我相信只需要經過細小調整，有人能將它做得比Transformer更好。
看來，影象生成領域很快又要掀起新的“較量風暴”了。
論文地址：
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
炒股開戶享福利，入金抽188元紅包，100%中獎！
開啟App看更多精彩內容
