35張圖，直觀理解Stable Diffusion

語言: CN / TW / HK

時間 2023-01-20 15:01:12 OneFlow深度學習框架

作者｜Jay Alammar

翻譯｜楊婷、徐佳渝

最近，AI影象生成引人注目，它能夠根據文字描述生成精美影象，這極大地改變了人們的影象創作方式。 Stable Diffusion 作為一款高效能模型，它生成的影象質量更高、執行速度更快、消耗的資源以及記憶體佔用更小，是AI影象生成領域的里程碑。

在接觸了AI影象生成以後，你可能會好奇這些模型背後的工作原理。

下面是對Stable Diffusion工作原理的概述。

Stable Diffusion用途多樣，是一款多功能模型。首先它可以根據文字生成影象(text2img)。上圖是從文字輸入到影象生成的示例。除此之外，我們還可以使用Stable Diffusion來替換、更改影象（這時我們需要同時輸入文字和影象）。

下面是Stable Diffusion的內部結構，瞭解內部結構可以讓我們更好地理解Stable Diffusion的組成、各組成部分的互動方式、以及各種影象生成選項/引數的含義。

1
Stable Diffusion的組成

Stable Diffusion並不是一個單一模型，而是由多個部分和模型一起構成的系統。

從內部來看，首先我們可以看到一個文字理解元件，這個元件將文字資訊轉化為數字表示（numeric representation）以捕捉文字意圖。

這部分主要對ML進行大概介紹，文章後續還會講解更多細節。可以說這個文字理解元件（文字編碼器）是一個特殊的Transformer語言模型（嚴格來說它是一個CLIP模型的文字編碼器）。將文字輸入到 Clip 文字編碼器得到特徵列表，對於文字中的每一個word/token 都有會得到一個向量特徵。

然後將文字特徵作為影象生成器的輸入，影象生成器又由幾部分組成。

影象生成器兩步驟：

1-影象資訊建立器（Image information creator）

影象資訊建立器是Stable Diffusion特有的關鍵部分，也是其效能遠超其他模型的原因。

影象資訊建立器執行多個step生成影象資訊。Stable Diffusion介面（interfaces）和庫（libraries）的step引數一般預設為50或100。

影象資訊建立器完全在影象資訊空間（亦稱潛在空間）上執行，這讓Stable Diffusion比以前在畫素空間（pixel space）上執行的擴散模型速度更快。從技術上講，影象資訊建立器由UNet神經網路和排程演算法組成。

“擴散”一詞描述了影象資訊建立器中發生的事情。因為影象資訊建立器對資訊作了逐步處理，所以影象解碼器（image decoder）才能隨後產出高質量影象。

2-影象解碼器（Image Decoder）

影象解碼器根據影象資訊建立器的資訊繪製圖像，它只用在過程結束時執行一次，以生成最終的畫素影象。

這樣就構成了Stable Diffusion的三個主要組成部分，每個部分都有自己的神經網路：

ClipText：用於文字編碼。輸入: 文字。輸出: 77個token embeddings向量，每個向量有768維。
UNet+排程程式：在資訊（潛在）空間中逐步處理資訊。輸入: 文字embeddings和一個初始化的多維陣列（結構化的數字列表，也稱為張量）組成的噪聲。輸出：經過處理的資訊陣列。
自動編碼解碼器（Autoencoder Decoder）：使用經過處理的資訊陣列繪製最終影象。輸入：經過處理的資訊陣列（維數：(4,64,64)）輸出: 生成的影象（維數：（3,512,512），即（紅/綠/藍；寬，高）)。

到底何為擴散（Diffusion）？

擴散是發生在粉色區域影象資訊建立器元件中的過程。這一部分有一個表示輸入文字的token embeddings和一個隨機初始化的影象資訊陣列，這些陣列也被稱為latents，在這個過程中會產生一個資訊陣列，影象解碼器（Image Decoder）使用這個資訊陣列生成最終影象。

擴散是循序漸進逐步發生的，每一步都會增加更多的相關資訊。為了更加直觀地瞭解這個過程，我們可以檢查隨機latents陣列，看它是否轉化為了視覺噪音（visual noise）。在這種情況下，視覺檢查（Visual inspection）是通過影象解碼器進行的。

擴散分為多個step，每個step都在輸入的latents陣列上執行，並且會產生另一個latents陣列，這個陣列更類似於輸入文字以及模型在模型訓練時的所有影象中獲取的所有視覺資訊。

我們可以對一組這樣的latents陣列執行視覺化，看看每一步都添加了什麼資訊。這一過程令人歎為觀止。

影片連結：

在這種情況下，步驟2和4之間發生了一些特別有意思的事情，就好像輪廓是從噪音中浮現出來的。

3
Diffusion的工作原理

擴散模型影象生成的核心是強大的計算機視覺模型。在足夠大的資料集的基礎上，這些模型可以學會很多複雜運算。擴散模型通過如下方式建構問題來實現影象生成：

假設我們有一個影象，我們首先生成一些噪音（noise），然後將這些噪音新增到影象上。

我們可以將這看成是一個訓練示例。之後我們使用同樣的公式去建立更多的訓練示例，然後用這些示例去訓練影象生成模型的中心元件。

雖然這個例子展示了從影象（總量0，沒有噪音）到總噪音（總量4，總噪音）的一些噪音值，但是我們可以輕鬆控制向影象中新增的噪音，因此我們可以將其分為數十個step，為資料集中的每個影象建立數十個訓練示例。

有了這個資料集，我們可以訓練噪音預測器（noise predictor），並最終得到一個在特定配置下執行時可以建立影象的預測器。接觸過ML的人會覺得訓練步驟非常熟悉：

接下來我們來看看Stable Diffusion是如何生成影象的。

4
通過降噪繪圖

經過訓練的噪音預測器可以對噪音影象進行降噪處理，並且可以預測噪音。

因為樣本噪音（sampled noise）被預測，所以如果我們從影象中去掉這個樣本，我們得到的影象就會更接近模型訓練的影象。（這個影象不是確切的影象本身，而是影象分佈，也就是影象的畫素排列，在畫素排列中天空通常是藍色的，高於地面，人有兩隻眼睛，貓有尖耳朵並且總是懶洋洋的）。

如果訓練資料集中的影象比較美觀，比如說Stable Diffusion訓練的LAION Aesthetics ，那麼訓練出來的影象的可觀賞性也會更高。如果我們在logo影象上對其進行訓練，那麼我們最終會得到一個logo生成模型。

這裡總結了擴散模型處理影象生成的過程，主要如論文Denoising Diffusion Probabilistic Models所述。相信你對擴散的含義有了一定的瞭解，知道了Stable Diffusion、Dall-E 2和谷歌Imagen的主要元件。

值得注意的是，到目前為止我們所描述的擴散過程，沒有使用任何文字資料，只需執行模型就能生成精美影象。不過我們無法控制影象的內容，它可能是一座金字塔，也可能是一隻貓。接下來，我們將討論如何將文字資訊融入擴散過程以控制圖片型別。

5
速度提升：在壓縮（Latent）資料中擴散

為了加快影象生成過程，Stable Diffusion論文沒有在畫素影象上進行執行，而是在影象的壓縮版本上執行。論文將這稱為前往潛在空間(Departure to Latent Space)。

壓縮（隨後是解壓縮/繪圖）通過編碼器完成。自動編碼器使用Image Encoder將影象壓縮排潛空間，然後使用Image Decoder再對壓縮資訊進行重構。

正向擴散在潛空間上完成。噪聲資訊應用於潛空間，而不是應用於畫素圖象。因此，訓練噪聲預測器（noise predictor）實際上是為了預測壓縮表示（compressed representation）上的噪音，這個壓縮表示也被稱為潛空間（latent space）。

正向擴散是使用Image Encoder生成影象資料，來訓練噪聲預測器。訓練一旦完成，就可以執行反向擴散，使用Image Decoder生成影象。

LDM/Stable Diffusion論文的圖3中提及了這兩個過程：

上圖還顯示了“conditioning”元件，這個元件在本例中是描述模型生成影象的文字提示詞（text prompts）。接下來，我們繼續探討文字元件。

6
文字編碼器：一種Transformer語言模型

Transformer語言模型作為語言理解元件，能夠接受文字提示詞，生成token embeddings。Stable Diffusion模型使用的是ClipText（基於GPT的模型），而論文中採用的是BERT。

Imagen論文表明，語言模型的選擇相當重要。相較於較大的影象生成元件，較大的語言模型元件對生成影象的質量影響更大。

較大的/更好的語言模型對影象生成模型的質量有巨大的影響。資料來源：Saharia等人所著論文Google Imagen中的圖A.5。

早期的Stable Diffusion模型僅使用了OpenAI釋出的預訓練模型ClipText。未來模型可能轉向新發布的更大的CLIP變體OpenCLIP。（更新於2022 年 11 月，詳情見 Stable Diffusion V2 uses OpenClip 。與僅含有630萬文字模型引數的ClipText相比，OpenCLIP文字模型引數多達3.54億。）

7
如何訓練CLIP

CLIP模型是在影象和影象說明資料集上訓練的。我們可以設想這樣一個數據集，它裡面有4億張影象以及這些影象說明的材料。

影象及影象說明資料集

實際上，CLIP是在網路上抓取的帶有“alt”標籤的影象上訓練的。CLIP是影象編碼器和文字編碼器的結合。簡單來說，訓練CLIP就是分別對影象和影象文字說明進行編碼。

然後，使用餘弦相似度來比較生成的embeddings。剛開始訓練時，即使文字正確描述了影象，相似度也會很低。

我們更新了這兩個模型，這樣下次嵌入它們時就可以得到相似的embeddings。

通過在資料集上重複此操作並使用大的batch size，最終使編碼器能夠生成影象和文字說明相似的embeddings。如 word2vec ，訓練過程也需要包含不匹配的影象和文字說明作為負樣本，以得到較低的相似度分數。

8
將文字資訊融入影象生成過程

為了使文字融入影象生成，我們須調整噪聲預測器來輸入文字。

現在，在資料集中加入文字。因為我們是在潛空間中執行，所以輸入的影象和預測的噪聲都處於潛空間中。

為了更好地理解UNet中文字tokens的使用方式，下面我們將進一步探究UNet模型。

Unet噪聲預測器的Layers(未使用文字)

首先來看沒有使用文字的UNet，其輸入和輸出如下：

可以看到：

UNet是一系列用於轉換latents陣列的layers
每一layer都對前一個layer的輸出進行操作
Some of the outputs are fed (via residual connections) into the processing later in the network
通過殘差連線（residual connections），將網路前面的layer輸出送入到後面的layer進行處理
時間步長被轉化為embedding向量，在網路層中使用

Unet噪聲預測器中的Layers (帶文字)

現在讓我們看看如何改變該系統以增加對文字的關注度。

為了給文字輸入提供支援，也就是專業上所說的文字條件（text conditioning），我們需要在系統的 ResNet blocks之間新增一個注意力層（attention layer）。

文字資訊不直接由ResNet處理，而是通過注意力層將這些文字表示融入到latents中。這樣，下一個ResNet就能在處理過程中利用融入的文字資訊。

9
總結

希望本文能幫助你深入瞭解Stable Diffusion的運作機制。雖然還涉及到許多其他概念，但是隻要熟悉了以上板塊，這些概念就會變得很容易理解。下面是一些我認為很有用的資源。

資源

http://www.youtube.com/shorts/qL6mKRyjK-0
http://huggingface.co/blog/stable_diffusion
http://huggingface.co/blog/annotated-diffusion
http://www.youtube.com/watch?v=J87hffSMB60
http://www.youtube.com/watch?v=ltLNYA3lWAQ
http://ommer-lab.com/research/latent-diffusion-models/
http://lilianweng.github.io/posts/2021-07-11-diffusion-models/
http://www.youtube.com/watch?v=_7rMfsA24Ls

（本文在遵循CC BY-NC-SA 4.0協議的基礎上由OneFlow編譯釋出，譯文轉載請聯絡獲得授權。原文：Alammar, J (2018). The Illustrated Transformer [Blog post]. http://jalammar.github.io/illustrated-stable-diffusion/）

其他人都在看

歡迎Star、試用OneFlow最新版本：http://github.com/Oneflow-Inc/oneflow/

本文分享自微信公眾號 - OneFlow（OneFlowTechnology）。
如有侵權，請聯絡 [email protected] 刪除。
本文參與“OSC源創計劃”，歡迎正在閱讀的你也加入，一起分享。

「其他文章」