【AIGC】Stable Diffusion的建模思想、訓練預測方式快速

在這篇博客中，將會用機器學習入門級描述，來介紹Stable Diffusion的關鍵原理。目前，網(wǎng)絡上的使用教程非常多，本篇中不會介紹如何部署、使用或者微調SD模型。也會盡量精簡語言，無公式推導，旨在理解思想。讓有機器學習基礎的朋友，可以快速了解SD模型的重要部分。如有理解錯誤，請不吝指正。

大綱

1. 關鍵概念
2. 模型結構及關鍵組件
3. 訓練和預測方式

關鍵概念

名詞解釋

Stable Diffusion

之所以叫Stable，是因為金主公司叫StabilityAI。

其基礎模型是Latent Diffusion Model（LDM），也是本文主要介紹的部分。

模型任務

1. text-2-img：輸入文本描述、輸出圖像
2. img-2-img：輸入圖片及其他文本描述，輸出圖像

總的來說，不論是輸入是文字還是圖片，都可以稱為是“condition”，用于指引圖像生成的“方向”。因此，SD模型的任務，可以統(tǒng)稱為是cond-2-img任務。

模型結構與關鍵組件

模型結構

LDM論文結構圖，初看時會有點懵，但稍微理解后還是非常清晰準確的。先初步介紹幾個大的模塊。建議把這張圖截圖固定在屏幕上，再繼續(xù)瀏覽下面的內容。

整體輸入輸出

上圖中最左側的$x$和$x$是模型的輸入與輸出，形如$[W,H,C]$的三維張量，代表一張圖像的寬、高和通道數(shù)。

需要注意，這里的輸入$x$，并不是模型img-2-img中的輸入圖像，而是模型訓練時的原始圖像輸入。img-2-img的輸入圖像，是上圖中最右側的Conditioning模塊中的images。

像素空間與隱空間

所謂空間，可以理解為數(shù)據(jù)的表示形式，通常有著不同的坐標軸。

• 像素空間（Pixel Space），上圖左側，紅框部分。通常是人眼可以識別的圖像內容。
• 隱空間（Latent Space），上圖中央，綠框部分。通常是人眼無法識別的內容，但包含的信息量與像素空間相近。

像素空間到隱空間

輸入的圖像$x$，經(jīng)過Encoder（圖中藍色的$\mathcal{E}$），轉換為另一種shape的張量$z$，即稱為隱空間。

從壓縮角度理解：圖像經(jīng)過轉換后，產(chǎn)生的新張量是人眼無法識別的。但其包含的信息量相差不大，數(shù)據(jù)尺寸卻大幅縮小，因此可以看做是一種圖像數(shù)據(jù)壓縮方式。

隱空間到像素空間

經(jīng)過模型處理后的隱向量輸出$z$（特指綠框左下角的$z$），經(jīng)過Decoder（圖中藍色的$\mathcal{D}$），轉換回像素空間。

隱空間Diffusion操作

對應圖中綠色Latent Space框的上半部分，包括以下三步：

1. 圖像經(jīng)過Encoder壓縮后，得到隱向量表示$z=\mathcal{E}(x)$隱向量
2. 從1~1000的均勻分布中，隨機采樣一個整數(shù)$T$，稱為擴散步數(shù)
3. 對向量$z$加$T$次高斯噪聲，滿足分布$N(0,{\beta }_t)$，得到$z_T$向量

在這個操作中，有一些有趣的特性：

噪聲收斂

加噪聲次數(shù)足夠多時，理論上會得到一組符合高斯分布的噪聲。利用這個特性，在預測階段我們就不需要執(zhí)行Diffusion操作，只需要采樣一組高斯分布的噪聲，即代表了$z_T$。

高斯噪聲可加性

當我們需要得到任意時刻的$z_T$時，可以直接從$z_0$以及一系列${\beta }_t$計算得到，只需要采樣一次噪聲。這部分的具體公式推導，可以參考由淺入深了解Diffusion Model - 知乎 (zhihu.com)。

隱空間Denoising操作

對應圖中綠色框的下半部分，包括以下步驟：

1. 輸入$z_t,t,cond$給U-Net結構，預測出一個噪聲${?}_{\theta }(z_t,t,cond)$，shape與$z_t$一致
2. 使$z_{t?1}=z_t?{?}_{\theta }(z_t,t,cond)$，重復上一步驟，直至獲得$z_0$隱向量
3. 使用Decoder得到輸出圖像，$x=\mathcal{D}(z_0)$

條件Conditioning

對應圖中最右邊灰白色框，輸入類型包括text、images等。在Conditioning模塊中，會執(zhí)行以下步驟：

1. 這些“附加信息”會通過對應的編碼器${\tau }_{\theta }$，轉換成向量表示
2. 轉換后的向量，會輸入給U-Net，作為其中Attention模塊的K、V輸入，輔助噪聲的預測

在這個模塊中，有幾個有趣的問題：

為什么需要Conditioning

由于“噪聲收斂”特性，當噪聲加得比較多時，$z_T$已經(jīng)趨近于一個“純噪聲”了，但訓練過程需要比對輸入圖像$x$和輸出圖像$x$的相似度。如何從一個“純噪聲”，還原回與輸入圖像相似的圖像，就必須要給模型提供額外的信息指引，這就是Conditioning的作用。

關鍵組件

VAE（Variational Auto Encoders）

在LDM中，如何將原始圖片“壓縮”轉換至隱空間，經(jīng)過處理再轉換回來，即使用VAE的Encoder和Decoder。這個模塊是預訓練好的，在LDM訓練時固定住參數(shù)。

原理

1. 原始張量輸入，經(jīng)過非常簡單的網(wǎng)絡結構，轉換成較小的張量
2. 在Latent張量上，加一點點噪聲擾動
3. 用對稱的簡單網(wǎng)絡結構，還原回原始大小
4. 對比輸入前后的張量是否相似

特點

1. 網(wǎng)絡計算復雜度比較低
2. Encoder和Decoder可以分開使用
3. 無監(jiān)督訓練，不需要標注輸入的label
4. 有了噪聲擾動之后，Latent Space的距離具有實際物理含義，可以實現(xiàn)例如“（滿杯水+空杯子）/ 2 = 半杯水”的操作

CLIP

文本信息如何轉換成張量，靠的是CLIP模塊。這個模塊是預訓練好的，在LDM訓練時固定住參數(shù)。

訓練方式

圖像以及它的描述文本，經(jīng)過各自的Encoder轉換為向量表示，希望轉換后的向量距離相近。經(jīng)過訓練后，文本描述可以映射到向量空間的一個點，其代表的物理含義與原始圖像相近。

假設無預訓練

開個腦洞，假如沒有這個模塊，直接將文本token化后，去Embedding Table中查表作為文本張量，理論上也是可以訓練的，只不過收斂速度會慢很多。

因此，這里使用一個預訓練text-2-embedding模塊，主要目的是加速訓練。CLIP的訓練數(shù)據(jù)集，也選擇了和LDM的數(shù)據(jù)集的同一個（LAION-5B的子集），語義更一致。

模型標識解釋

我們經(jīng)常會看到類似“ViT-L/14”的模型名，表示一種CLIP的結構。具體的，ViT表示Vision Transformer，L表示Large（此外還有Base、Huge），14表示訓練時把圖像劃分成14*14個子圖序列輸入給Transformer。

U-Net

作為LDM的核心組件，U-Net是模型訓練過程中，唯一需要參數(shù)更新的部分。在這個結構中，輸入是帶有噪聲的隱向量$z_t$、當前的時間戳$t$，文本等Conditioning的張量表示$E$，輸出是$z_t$中的噪聲預測。

模型任務

U-Net的任務，就是從$z_t$中預測出噪聲部分${?}_t$，從而得到降低噪聲后的$z_{t?1}=z_t?{?}_t$，直到獲得$z_0$。下圖是一個可視化示意圖，實際上，我們去噪的$z_t$是隱向量空間的數(shù)據(jù)，人眼無法識別。

模型結構

U-Net大致上可以分為三塊：降采樣層、中間層、上采樣層。之所以叫U-Net，是因為它的模型結構類似字母U。

降采樣層

1. 時間戳$t$轉換為向量形式。用的是“Attention is All you Need”論文的Transformer方法，通過sin和cos函數(shù)再經(jīng)過兩個Linear進行變換
2. 初始化輸入$X=conv(concat(z_t,E))$，其中$conv$是卷積，$E$是Conditioning
3. 重復以下步驟（a~c）多次，將輸入尺寸降至目標尺寸（如上圖的$4\times 4$）
   1. 重復以下兩步多次，訓練多個ResBlock和SpatialTransformer層，輸入值$X$的尺寸不變
      1. 輸入上一層的輸出$X$和時間戳向量，給ResBlock
      2. ResBlock的輸出，與$E$一起輸入給SpatialTransformer，在這里考慮到text等信息
   2. 重復多次3~4步，
   3. 通過卷積或Avg-Pooling進行降采樣，縮小$X$的尺寸

中間層

很簡單，ResBlock + SpatialTransformer + ResBlock，輸入$X$尺寸不變。

上采樣層

大部分步驟與降采樣層一致，只有以下兩點不同

1. 輸入$X$需要拼上對應降采樣層的輸出，稱為skip connection，對應U-Net結構圖中橫向的箭頭
2. 把降采樣步驟，換成使用卷積或插值（interpolate）方式來上采樣，使得$X$的尺寸增大

輸出

上采樣層的輸出，會經(jīng)過normalization + SiLU + conv，得到U-Net的最終輸出，即噪聲的預測值，尺寸保持與輸入$z_t$一致。

訓練方式

模型更新方式

LDM模型需要訓練的部分，只有U-Net的參數(shù)。訓練的方式，可以簡單總結為：

1. 輸入一張圖片$x$，以及它的文本描述等Conditioning，一個隨機的整數(shù)$T$步
2. 經(jīng)過Encoder壓縮、Diffusion加噪聲，得到$z_T$隱向量
3. 結合Conditioning，使用U-Net，進行$T$次去噪，得到預測值$z_0$向量
4. 使用Decoder還原回$x$，計算$x$與$x$之間的差距（KL散度），得到模型更新的loss

模型預測方式

1. 隨機一個高斯噪聲，作為$z_T$向量
2. 輸入text等Conditioning，使用U-Net進行指定次數(shù)$T$的去噪操作
3. 使用Decoder還原回$x$，得到圖像輸出

訓練、預測過程，在論文中的偽代碼為下圖所示。

展望

下一篇文章，將會討論以下幾個更深入的內容：

1. ControlNet、LoRA等插件的實現(xiàn)
2. 各種Conditioning Context是如何轉換為張量的
3. 訓練的數(shù)據(jù)集情況

參考

The Illustrated Stable Diffusion – Jay Alammar – Visualizing machine learning one concept at a time. (jalammar.github.io)

【原創(chuàng)】萬字長文講解Stable Diffusion的AI繪畫基本技術原理 - 知乎 (zhihu.com)

Diffusion Models：生成擴散模型 (yinglinzheng.netlify.app)

由淺入深了解Diffusion Model - 知乎 (zhihu.com)

How does Stable Diffusion work? - Stable Diffusion Art (stable-diffusion-art.com)

[2006.11239] Denoising Diffusion Probabilistic Models (arxiv.org)

CompVis/latent-diffusion: High-Resolution Image Synthesis with Latent Diffusion Models (github.com)文章來源地址http://www.zghlxwxcb.cn/news/detail-481675.html

到了這里，關于【AIGC】Stable Diffusion原理快速上手，模型結構、關鍵組件、訓練預測方式的文章就介紹完了。如果您還想了解更多內容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關文章，希望大家以后多多支持TOY模板網(wǎng)！