還記得火爆全網的圖解Transformer嗎？

最近這位大佬博主Jay Alammar在博客上對大火的Stable Diffusion模型也撰寫了一篇圖解，讓你從零開始徹底搞懂圖像生成模型的原理，還配有超詳細的視頻講解！

文章鏈接：https://jalammar.github.io/illustrated-stable-diffusion/

視頻鏈接：https://www.youtube.com/watch?v=MXmacOUJUaw

圖解Stable Diffusion

AI模型最新展現出的圖像生成能力遠遠超出人們的預期，直接根據文字描述就能創造出具有驚人視覺效果的圖像，其背后的運行機制顯得十分神秘與神奇，但確實影響了人類創造藝術的方式。

Stable Diffusion的發布是AI圖像生成發展過程中的一個里程碑，相當于給大眾提供了一個可用的高性能模型，不僅生成的圖像質量非常高，運行速度快，并且有資源和內存的要求也較低。

相信只要試過AI圖像生成的人都會想了解它到底是如何工作的，這篇文章就將為你揭開Stable Diffusion工作原理的神秘面紗。

Stable Diffusion從功能上來說主要包括兩方面：1）其核心功能為僅根據文本提示作為輸入來生成的圖像（text2img）；2）你也可以用它對圖像根據文字描述進行修改（即輸入為文本+圖像）。

下面將使用圖示來輔助解釋Stable Diffusion的組件，它們之間如何交互，以及圖像生成選項及參數的含義。

Stable Diffusion組件

Stable Diffusion是一個由多個組件和模型組成的系統，而非單一的模型。

當我們從模型整體的角度向模型內部觀察時，可以發現，其包含一個文本理解組件用于將文本信息翻譯成數字表示（numeric representation），以捕捉文本中的語義信息。

雖然目前還是從宏觀角度分析模型，后面才有更多的模型細節，但我們也可以大致推測這個文本編碼器是一個特殊的Transformer語言模型（具體來說是CLIP模型的文本編碼器）。

模型的輸入為一個文本字符串，輸出為一個數字列表，用來表征文本中的每個單詞/token，即將每個token轉換為一個向量。

然后這些信息會被提交到圖像生成器（image generator）中，它的內部也包含多個組件。

圖像生成器主要包括兩個階段：

1. Image information creator

這個組件是Stable Diffusion的獨家秘方，相比之前的模型，它的很多性能增益都是在這里實現的。

該組件運行多個steps來生成圖像信息，其中steps也是Stable Diffusion接口和庫中的參數，通常默認為50或100。

圖像信息創建器完全在圖像信息空間（或潛空間）中運行，這一特性使得它比其他在像素空間工作的Diffusion模型運行得更快；從技術上來看，該組件由一個UNet神經網絡和一個調度（scheduling）算法組成。

擴散（diffusion）這個詞描述了在該組件內部運行期間發生的事情，即對信息進行一步步地處理，并最終由下一個組件（圖像解碼器）生成高質量的圖像。

2. 圖像解碼器

圖像解碼器根據從圖像信息創建器中獲取的信息畫出一幅畫，整個過程只運行一次即可生成最終的像素圖像。

可以看到，Stable Diffusion總共包含三個主要的組件，其中每個組件都擁有一個獨立的神經網絡：

1）Clip Text用于文本編碼。

輸入：文本

輸出：77個token嵌入向量，其中每個向量包含768個維度

2）UNet + Scheduler在信息（潛）空間中逐步處理/擴散信息。

輸入：文本嵌入和一個由噪聲組成的初始多維數組（結構化的數字列表，也叫張量tensor）。

輸出：一個經過處理的信息陣列

3）自編碼解碼器（Autoencoder Decoder），使用處理過的信息矩陣繪制最終圖像的解碼器。

輸入：處理過的信息矩陣，維度為（4, 64, 64）

輸出：結果圖像，各維度為（3，512，512），即（紅/綠/藍，寬，高）

什么是Diffusion？

擴散是在下圖中粉紅色的圖像信息創建器組件中發生的過程，過程中包含表征輸入文本的token嵌入，和隨機的初始圖像信息矩陣（也稱之為latents），該過程會還需要用到圖像解碼器來繪制最終圖像的信息矩陣。

整個運行過程是step by step的，每一步都會增加更多的相關信息。

為了更直觀地感受整個過程，可以中途查看隨機latents矩陣，并觀察它是如何轉化為視覺噪聲的，其中視覺檢查（visual inspection）是通過圖像解碼器進行的。

整個diffusion過程包含多個steps，其中每個step都是基于輸入的latents矩陣進行操作，并生成另一個latents矩陣以更好地貼合「輸入的文本」和從模型圖像集中獲取的「視覺信息」。

將這些latents可視化可以看到這些信息是如何在每個step中相加的。

整個過程就是從無到有，看起來相當激動人心。

步驟2和4之間的過程轉變看起來特別有趣，就好像圖片的輪廓是從噪聲中出現的。

Diffusion的工作原理

使用擴散模型生成圖像的核心思路還是基于已存在的強大的計算機視覺模型，只要輸入足夠大的數據集，這些模型可以學習任意復雜的操作。

假設我們已經有了一張圖像，生成產生一些噪聲加入到圖像中，然后就可以將該圖像視作一個訓練樣例。

使用相同的操作可以生成大量訓練樣本來訓練圖像生成模型中的核心組件。

上述例子展示了一些可選的噪聲量值，從原始圖像(級別0，不含噪聲)到噪聲全部添加(級別4) ，從而可以很容易地控制有多少噪聲添加到圖像中。

所以我們可以將這個過程分散在幾十個steps中，對數據集中的每張圖像都可以生成數十個訓練樣本。

基于上述數據集，我們就可以訓練出一個性能極佳的噪聲預測器，每個訓練step和其他模型的訓練相似。當以某一種確定的配置運行時，噪聲預測器就可以生成圖像。

移除噪聲，繪制圖像

經過訓練的噪聲預測器可以對一幅添加噪聲的圖像進行去噪，也可以預測添加的噪聲量。

由于采樣的噪聲是可預測的，所以如果從圖像中減去噪聲，最后得到的圖像就會更接近模型訓練得到的圖像。

得到的圖像并非是一張精確的原始圖像，而是分布（distribution），即世界的像素排列，比如天空通常是藍色的，人有兩只眼睛，貓有尖耳朵等等，生成的具體圖像風格完全取決于訓練數據集。

不止Stable Diffusion通過去噪進行圖像生成，DALL-E 2和谷歌的Imagen模型都是如此。

需要注意的是，到目前為止描述的擴散過程還沒有使用任何文本數據生成圖像。因此，如果我們部署這個模型的話，它能夠生成很好看的圖像，但用戶沒有辦法控制生成的內容。

在接下來的部分中，將會對如何將條件文本合并到流程中進行描述，以便控制模型生成的圖像類型。

加速：在壓縮數據上擴散

為了加速圖像生成的過程，Stable Diffusion并沒有選擇在像素圖像本身上運行擴散過程，而是選擇在圖像的壓縮版本上運行，論文中也稱之為「Departure to Latent Space」。

整個壓縮過程，包括后續的解壓、繪制圖像都是通過自編碼器完成的，將圖像壓縮到潛空間中，然后僅使用解碼器使用壓縮后的信息來重構。

前向擴散（forward diffusion）過程是在壓縮latents完成的，噪聲的切片（slices）是應用于latents上的噪聲，而非像素圖像，所以噪聲預測器實際上是被訓練用來預測壓縮表示（潛空間）中的噪聲。

前向過程，即使用使用自編碼器中的編碼器來訓練噪聲預測器。一旦訓練完成后，就可以通過運行反向過程（自編碼器中的解碼器）來生成圖像。

前向和后向過程如下所示，圖中還包括了一個conditioning組件，用來描述模型應該生成圖像的文本提示。

文本編碼器：一個Transformer語言模型

模型中的語言理解組件使用的是Transformer語言模型，可以將輸入的文本提示轉換為token嵌入向量。發布的Stable Diffusion模型使用 ClipText (基于 GPT 的模型) ，這篇文章中為了方便講解選擇使用 BERT模型。

Imagen論文中的實驗表明，相比選擇更大的圖像生成組件，更大的語言模型可以帶來更多的圖像質量提升。

早期的Stable Diffusion模型使用的是OpenAI發布的經過預訓練的 ClipText 模型，而在Stable Diffusion V2中已經轉向了最新發布的、更大的CLIP模型變體OpenClip.

CLIP是怎么訓練的？

CLIP需要的數據為圖像及其標題，數據集中大約包含4億張圖像及描述。

數據集通過從網上抓取的圖片以及相應的「alt」標簽文本來收集的。

CLIP 是圖像編碼器和文本編碼器的組合，其訓練過程可以簡化為拍攝圖像和文字說明，使用兩個編碼器對數據分別進行編碼。

然后使用余弦距離比較結果嵌入，剛開始訓練時，即使文本描述與圖像是相匹配的，它們之間的相似性肯定也是很低的。

隨著模型的不斷更新，在后續階段，編碼器對圖像和文本編碼得到的嵌入會逐漸相似。

通過在整個數據集中重復該過程，并使用大batch size的編碼器，最終能夠生成一個嵌入向量，其中狗的圖像和句子「一條狗的圖片」之間是相似的。

就像在 word2vec 中一樣，訓練過程也需要包括不匹配的圖片和說明的負樣本，模型需要給它們分配較低的相似度分數。

文本信息喂入圖像生成過程

為了將文本條件融入成為圖像生成過程的一部分，必須調整噪聲預測器的輸入為文本。

所有的操作都是在潛空間上，包括編碼后的文本、輸入圖像和預測噪聲。

為了更好地了解文本token在 Unet 中的使用方式，還需要先了解一下 Unet模型。

Unet 噪聲預測器中的層（無文本）

一個不使用文本的diffusion Unet，其輸入輸出如下所示：

在模型內部，可以看到：

1. Unet模型中的層主要用于轉換latents；

2. 每層都是在之前層的輸出上進行操作；

3. 某些輸出（通過殘差連接）將其饋送到網絡后面的處理中

4. 將時間步轉換為時間步長嵌入向量，可以在層中使用。

Unet 噪聲預測器中的層（帶文本）

現在就需要將之前的系統改裝成帶文本版本的。

主要的修改部分就是增加對文本輸入（術語：text conditioning）的支持，即在ResNet塊之間添加一個注意力層。

需要注意的是，ResNet塊沒有直接看到文本內容，而是通過注意力層將文本在latents中的表征合并起來，然后下一個ResNet就可以在這一過程中利用上文本信息。