機器之心專欄
機器之心編輯部
該論文提出了一個為多模態設計的概率建模框架 UniDiffuser，除了單向的文生圖，還能實現圖生文、圖文聯合生成、無條件圖文生成、圖文改寫等多種功能。
據悉 GPT-4 將於本周發布，多模態將成為其一大亮點。當前的大語言模型正在成為理解各種模態的通用接口，能夠根據不同模態信息來給出回復文本，但大語言模型生成的內容也僅僅局限於文本。另一方面，當前的擴散模型 DALL・E 2、Imagen、Stable Diffusion 等在視覺創作上掀起一場革命，但這些模型僅僅支持文到圖的單一跨模態功能，離通用式生成模型還有一定距離。而多模態大模型將能夠打通各種模態能力，實現任意模態之間轉化，被認為是通用式生成模型的未來發展方向。
清華大學計算機系朱軍教授帶領的 TSAIL 團隊近期公開的一篇論文《One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale》，率先發布了對多模態生成式模型的一些探索工作，實現了任意模態之間的相互轉化。
論文連結：https://ml.cs.tsinghua.edu.cn/diffusion/unidiffuser.pdf
開原始碼：https://github.com/thu-ml/unidiffuser
該論文提出了一個為多模態設計的概率建模框架 UniDiffuser，並採用該團隊提出的基於 transformer 的網絡架構 U-ViT，在開源的大規模圖文數據集 LAION-5B 上訓練了一個十億參數量的模型，使得一個底層模型能夠高質量地完成多種生成任務（圖 1）。簡單來講，除了單向的文生圖，還能實現圖生文、圖文聯合生成、無條件圖文生成、圖文改寫等多種功能，大幅提升文圖內容的生產效率，也進一步提升了生成式模型的應用想像力。
該論文一作鮑凡目前博士在讀，是此前 Analytic-DPM 的提出者，憑藉在擴散模型方面的優秀工作榮獲 ICLR 2022 的 outstanding paper award（目前唯一一篇大陸單位獨立完成的獲獎論文）。
此外，機器之心之前還報導過 TSAIL 團隊提出的DPM-Solver 快速算法，目前仍是擴散模型最快的生成算法。多模態大模型正是該團隊在深度概率模型的算法和原理方面上長期深入積累的一個集中展示。該工作的合作者包括人民大學高瓴人工智慧學院的李崇軒、北京智源研究院的曹越等。
值得注意的是，該項目的論文和代碼均已開源。
效果展示
如下的圖 8 展示了 UniDiffuser 在圖文聯合生成的效果：
如下的圖 9 展示了 UniDiffuser 在文到圖上的效果：
如下的圖 10 展示了 UniDiffuser 在圖到文上的效果：
如下的圖 11 展示了 UniDiffuser 在無條件圖像生成上的效果：
如下的圖 12 展示了 UniDiffuser 在圖像改寫上的效果：
如下的圖 15 展示了 UniDiffuser 能夠實現在圖文兩個模態之間的來回跳躍 ：
如下圖 16 展示了 UniDiffuser 能對真實的兩張圖像進行插值：
方法概覽
研究團隊將針對通用生成式模型的設計劃分成了兩個子問題：
- 概率建模框架：是否能尋找到一個概率建模框架，能同時建模出模態之間所有的分布，例如圖文之間的邊緣分布、條件分布、聯合分布等？
- 網絡架構：是否能設計出一個統一的網絡架構，來支持各種不同模態的輸入？
概率建模框架
針對概率建模框架，研究團隊提出 UniDiffuser，一個基於擴散模型的概率建模框架。UniDiffuser 能夠顯示地建模多模態數據中包括邊緣分布、條件分布、聯合分布在內的所有分布。研究團隊發現，關於不同分布的擴散模型學習都可以統一成一個視角：首先向兩個模態的數據分別加入某種大小的噪聲，然後再預測兩個模態數據上的噪聲。其中兩個模態數據上的噪聲大小決定了具體的分布。例如，將文本的噪聲大小設置為 0，則對應了文生圖的條件分布；將文本噪聲大小設置為最大值，則對應了無條件圖像生成的分布；將圖文噪聲大小設置為相同，則對應了圖文的聯合分布。根據該統一的視角，UniDiffuser 只需要將原始擴散模型的訓練算法做少許的修改，便能同時學習上述的所有分布 — 如下圖所示，UniDiffuser 同時向所有模態加噪而非單個模態，輸入所有模態對應的噪聲大小，以及預測所有模態上的噪聲。
以雙模態為例子，最終的訓練目標函數如下所示：
在訓練後，通過向噪聲預測網絡設置兩個模態合適的時間，UniDiffuser 能夠實現無條件、條件以及聯合生成。例如將文本的時間設置為 0，可以實現文到圖生成；將文本的時間設置為最大值，可以實現無條件圖像生成；將圖文時間設置為相同值，可以實現圖文聯合生成。
下面羅列了 UniDiffuser 的訓練和採樣算法，可見這些算法相對原始的擴散模型均只做了微小的改動，易於實現。
此外，由於 UniDiffuser 同時建模了條件分布和無條件分布，因此 UniDiffuser 天然地支持 classifier-free guidance。下面的圖 3 展示了 UniDiffuser 的條件生成和聯合生成在不同的 guidance scale 下的效果：
網絡架構
針對網絡架構，研究團隊提出使用基於 transformer 的架構來參數化噪聲預測網絡。具體地，研究團隊採用了最近提出的 U-ViT 架構。U-ViT 將所有的輸入都視作 token，並在 transformer 塊之間加入了 U 型連接。研究團隊也採用了 Stable Diffusion 的策略，將不同模態的數據都轉換到了隱空間再進行擴散模型的建模。值得注意的是，U-ViT 架構同樣來自該研究團隊，並且已被開源在 https://github.com/baofff/U-ViT。
實驗結果
UniDiffuser 首先和 Versatile Diffusion 進行了比較。Versatile Diffusion 是過去的一個基於多任務框架的多模態擴散模型。首先 UniDiffuser 和 Versatile Diffusion 進行了文到圖上的效果比較。如下面的圖 5 所示，在不同的 classifier-free guidance scale 下，UniDiffuser 在 CLIP Score 和 FID 指標上均要好於 Versatile Diffusion。
然後 UniDiffuser 和 Versatile Diffusion 進行了圖到文上的效果比較。如下面的圖 6 所示，UniDiffuser 在圖到文上有更好的 CLIP Score。
UniDiffuser 也和專用的文到圖模型在 MS-COCO 上進行了 zero-shot FID 的比較。如下面的表 1 所示，UniDiffuser 可以和專用的文到圖模型取得可比的效果。