Stable Diffusion 3論文終於釋出,架構細節大揭秘,對復現Sora有幫助?

机器之心發表於2024-03-06

在眾多前沿成果都不再透露技術細節之際,Stable Diffusion 3 論文的釋出顯得相當珍貴。

Stable Diffusion 3 的論文終於來了!

這個模型於兩週前釋出,採用了與 Sora 相同的 DiT(Diffusion Transformer)架構,一經發布就引起了不小的轟動。

與之前的版本相比,Stable Diffusion 3 生成的圖在質量上實現了很大改進,支援多主題提示,文字書寫效果也更好了(明顯不再亂碼)。Stable Diffusion 3論文終於釋出,架構細節大揭秘,對復現Sora有幫助?

Stability AI 表示,Stable Diffusion 3 是一個模型系列,引數量從 800M 到 8B 不等。這個引數量意味著,它可以在很多行動式裝置上直接跑,大大降低了 AI 大模型的使用門檻。

在最新發布的論文中,Stability AI 表示,在基於人類偏好的評估中,Stable Diffusion 3 優於當前最先進的文字到影像生成系統,如 DALL・E 3、Midjourney v6 和 Ideogram v1。不久之後,他們將公開該研究的實驗資料、程式碼和模型權重

圖片

在論文中,Stability AI 透露了關於 Stable Diffusion 3 的更多細節。

圖片

  • 論文標題:Scaling Rectified Flow Transformers for High-Resolution Image Synthesis

  • 論文連結:https://stabilityai-public-packages.s3.us-west-2.amazonaws.com/Stable+Diffusion+3+Paper.pdf

架構細節

對於文字到影像的生成,Stable Diffusion 3 模型必須同時考慮文字和影像兩種模式。因此,論文作者稱這種新架構為 MMDiT,意指其處理多種模態的能力。與之前版本的 Stable Diffusion 一樣,作者使用預訓練模型來推導合適的文字和影像表徵。具體來說,他們使用了三種不同的文字嵌入模型 —— 兩種 CLIP 模型和 T5—— 來編碼文字表徵,並使用改進的自編碼模型來編碼影像 token。

圖片

Stable Diffusion 3 模型架構。

圖片

改進的多模態擴散 transformer:MMDiT 塊。

SD3 架構基於 Sora 核心研發成員 William Peebles 和紐約大學電腦科學助理教授謝賽寧合作提出的 DiT。由於文字嵌入和影像嵌入在概念上有很大不同,因此 SD3 的作者對兩種模態使用兩套不同的權重。如上圖所示,這相當於為每種模態設定了兩個獨立的 transformer,但將兩種模態的序列結合起來進行注意力運算,從而使兩種表徵都能在各自的空間內工作,同時也將另一種表徵考慮在內。

圖片

在訓練過程中測量視覺保真度和文字對齊度時,作者提出的 MMDiT 架構優於 UViT 和 DiT 等成熟的文字到影像骨幹。

透過這種方法,資訊可以在影像和文字 token 之間流動,從而提高模型的整體理解能力,並改善所生成輸出的文字排版。正如論文中所討論的那樣,這種架構也很容易擴充套件到影片等多種模式。

圖片

得益於 Stable Diffusion 3 改進的提示遵循能力,新模型有能力製作出聚焦於各種不同主題和質量的影像,同時還能高度靈活地處理影像本身的風格。

圖片

透過 re-weighting 改進 Rectified Flow

Stable Diffusion 3 採用 Rectified Flow(RF)公式,在訓練過程中,資料和噪聲以線性軌跡相連。這使得推理路徑更加平直,從而減少了取樣步驟。此外,作者還在訓練過程中引入了一種新的軌跡取樣計劃。他們假設,軌跡的中間部分會帶來更具挑戰性的預測任務,因此該計劃給予軌跡中間部分更多權重。他們使用多種資料集、指標和取樣器設定進行比較,並將自己提出的方法與 LDM、EDM 和 ADM 等 60 種其他擴散軌跡進行了測試。結果表明,雖然以前的 RF 公式在少步取樣情況下效能有所提高,但隨著步數的增加,其相對效能會下降。相比之下,作者提出的重新加權 RF 變體能持續提高效能。

圖片

擴充套件 Rectified Flow Transformer 模型

作者利用重新加權的 Rectified Flow 公式和 MMDiT 骨幹對文字到影像的合成進行了擴充套件(scaling)研究。他們訓練的模型從帶有 450M 個引數的 15 個塊到帶有 8B 個引數的 38 個塊不等,並觀察到驗證損失隨著模型大小和訓練步驟的增加而平穩降低(上圖的第一行)。為了檢驗這是否轉化為對模型輸出的有意義改進,作者還評估了自動影像對齊指標(GenEval)和人類偏好分數(ELO)(上圖第二行)。結果表明,這些指標與驗證損失之間存在很強的相關性,這表明後者可以很好地預測模型的整體效能。此外,scaling 趨勢沒有顯示出飽和的跡象,這讓作者對未來繼續提高模型效能持樂觀態度。

靈活的文字編碼器

透過移除用於推理的記憶體密集型 4.7B 引數 T5 文字編碼器,SD3 的記憶體需求可顯著降低,而效能損失卻很小。如圖所示,移除該文字編碼器不會影響視覺美感(不使用 T5 時的勝率為 50%),只會略微降低文字一致性(勝率為 46%)。不過,作者建議在生成書面文字時加入 T5,以充分發揮 SD3 的效能,因為他們觀察到,如果不加入 T5,生成排版的效能下降幅度更大(勝率為 38%),如下圖所示:

圖片

只有在呈現涉及許多細節或大量書面文字的非常複雜的提示時,移除 T5 進行推理才會導致效能顯著下降。上圖顯示了每個示例的三個隨機樣本。

模型效能

作者將 Stable Diffusion 3 的輸出影像與其他各種開源模型(包括 SDXL、SDXL Turbo、Stable Cascade、Playground v2.5 和 Pixart-α)以及閉源模型(如 DALL-E 3、Midjourney v6 和 Ideogram v1)進行了比較,以便根據人類反饋來評估效能。在這些測試中,人類評估員從每個模型中獲得輸出示例,並根據模型輸出在多大程度上遵循所給提示的上下文(prompt following)、在多大程度上根據提示渲染文字(typography)以及哪幅影像具有更高的美學質量(visual aesthetics)來選擇最佳結果。

圖片

以 SD3 為基準,這個圖表概述了它在基於人類對視覺美學、提示遵循和文字排版的評估中的勝率。

從測試結果來看,作者發現 Stable Diffusion 3 在上述所有方面都與當前最先進的文字到影像生成系統相當,甚至更勝一籌。

在消費級硬體上進行的早期未最佳化推理測試中,最大的 8B 引數 SD3 模型適合 RTX 4090 的 24GB VRAM,使用 50 個取樣步驟生成解析度為 1024x1024 的影像需要 34 秒。

圖片

此外,在最初發布時,Stable Diffusion 3 將有多種變體,從 800m 到 8B 引數模型不等,以進一步消除硬體障礙。

圖片

圖片

更多細節請參考原論文。

參考連結:https://stability.ai/news/stable-diffusion-3-research-paper

相關文章