"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
