Stable Diffusion背後的技術：高效、高解析又易控制的Latent Diffusion Model

Stable Diffusion背後的技術：高效、高解析又易控制的Latent Diffusion Model

近年，生成式模型 (generative model) 用於圖像生成展現了驚人的成果，**最知名的莫過於基於 diffusion model 的 Stable Diffusion 與 **Midjourney 。雖然目前還無法完整地指出 Midjourney 背後的方法，但透過Stable Diffusion 團隊的技術論文 與 diffusion model 的發展，還是可以大致推論這些 state-of-the-art 方法的概念。

早在 2020年 DDPM 中已經證明 diffusion model 是生成表述能力最好的模型之一，但其效率與解析度仍然是其劣勢。而Stable Diffusion 透過 perceptual image encoder 與 decoder，除了訓練 diffusion model 的過程可以在較低的解析度完成，在生成高解析度圖像上的能力也更好。另一方面，透過 domain specific expert 的結合，** Stable Diffusion 具備了良好的可控制性，可以執行結合其他語意的任務 (如文字生成圖像)。**

Cover made with Canva ! (image source )

Tip

文章難度：★★★★☆

閱讀建議： 本篇文章主要就 Stable Diffusion 開發團隊 Runway 於 2022 年 CVPR 發表的 “High-Resolution Image Synthesis with Latent Diffusion Models ” 進行現行 diffusion model 的技術介紹，其中會參雜到許多相關研究，如 DDPM 、 DDIM 、 VQGAN 等方法都會在文章中展開介紹。整體上，文章自 generative model 與 diffusion model 基礎開始，後續才會介紹 Stable Diffusion 的核心 latent diffusion model 。雖然 diffusion model 是相當數學的模型，這篇文章會盡量減少數學描述，但保留相關重要觀念。

推薦背景知識：Machine Learning, Deep Learning, Generative Model, VAE (Variational Auto-Encoder), GAN (Generative Adversarial Network), Normalizing Flow, Maximize Likelihood Estimation (MLE), Density Estimation, Markov Chain, Score Matching, Variational Inference, Vector Quantization, VQGAN.

Stable Diffusion [1] 方法是基於 diffusion model 重要里程碑DDPM (Denoising Diffusion Probabilistic Models) [2] 改良的 generative model，因此以下先簡單介紹一下 generative model 與 DDPM 的概念。

Diffusion Model

生成模型 (Generative Model) 網路一直都是 machine learing中熱門的題目。除了過往的幾年中**常被提起的 VAE (Variational Auto-Encoder) [3] 與 GAN (Generative Adversarial Network) [4]，**近年崛起、或者說復辟成功的其中一個方法就是擴散模型 (Diffusion Model) [5]。

以統計的角度來看，

生成模型 (generative model)的目標是 modeling目標資料的分布。

比如說，我們手邊有一萬張人臉的照片，但我們希望可以獲得更多人臉的照片，或者說想要依照某些情境生成更多人臉的照片。而使用 generative model 來 modeling 目標資料的分布，可以確保生成出來的是人臉的照片，而不是貓狗的照片，或是四不像的照片。

不考慮條件 (unconditional) 的 generative model 的概念大致如下圖，首先會從一個已知的分布中採樣，將樣本放到 generative model 生成目標資料，最後要讓這些生成資料的分布與真實資料的分布越接近越好。

Generative model的概念。(資料來源)

Generative model 有許多不同的實現方法 (或者說流派)。其中包含 VAE (Variational Auto Encoder) 、 GAN (Generative Adversarial Network) 、 ARM (Auto-Regressive Model) 、flow-based model [6] 與這次要介紹的 diffusion model 等。一些分類的分支可以參考下圖。

Generative Model的一些分支。

Diffusion model的核心精神是學習一個逐步denoise的過程。

在訓練模型時，逐步地將原圖加入雜訊。具體上來說，使用一個高斯分布 (Gaussian distribution) 一次又一次地在原圖上打上很小的雜訊，然後讓網路來學習如何reverse這個雜訊。

Diffusion model的精神是學習逐步denoise。(資料來源)

在數學上 diffusion model 是被定義為一個參數化的 Markov chain，在DDPM的原論文中是這樣描述：`“Diffusion probabilistic model is parameterized Markov chain trained using variational inference to produce samples matching the data after finite time.”

聽起來好像有點複雜，但實際上 diffusion model 的實作多半反而是很單純的。網路輸入通常為要被 denoise 的圖片以及現在是 denoise 的第幾個步數。而輸出通常是這步 denoise 完的圖片，或是說預測這步被加入的 noise。

可以用這個角度簡單理解大部分的diffusion model。(資料來源)

目前的 diffusion model 大多都是延續 DDPM (Denoising Diffusion Probabilistic Models) 精巧設計的機率模型。 DDPM 簡化後的訓練目標非常簡單，**拿過程打的 noise ，就可以直接計算 L2 loss **。

概念性解釋diffusion model的訓練目標。(資料來源)

實際上 DDPM 在訓練時也不需要一步一步加 noise，因為多個 Gaussian distribution 的疊加還是 Gaussian distribution 。因此，計算多個加 Gaussian noise 的動作得以一步完成。

Stable Diffusion

雖然 Diffusion Model 發展在 DDPM 後，具備與其他 generative model 同等甚至更好的圖像生成能力，但還是有兩個主要問題：

  • 生成圖像本身還是需要通過多個 iteration，速度偏慢。
  • 訓練時需要在原解析度上計算 loss ，對於硬體與訓練時間要求高。另外，也讓 diffusion model 較難生成高解析度圖像。

在生成圖像加速上，後續雖然有一些透過 sampling 機制 (如DDIM [7]) 或者是 hierarchical 生成 (如Google Imagen [8, 9] 的前身 [10]) 來加速圖像生成，但在訓練時需要在原解析度上計算 loss 這件事，一直沒有一個較好的解法。也讓學習很大一部分重點放在 pixel-space，相較難集中於 sematic 資訊。

而這就是 Stable Diffusion 一個重要的特色，

透過 perceptual image encoder & decoder，讓 diffusion model 訓練可以在較低的解析度完成，大幅降低訓練所需硬體與時間。

Stable Diffusion 中使用的 latent diffusion model。(資料來源)

這邊說的 Stable Diffusion 指的都是官方 GitHub 中引述，由Stability AI andRunway 合作的“High-Resolution Image Synthesis with Latent Diffusion Models” 論文 [11] 中描述的方法。實際上，這篇論文八成已無法代表現在 Stable Diffusion 的完整方法了，但是核心精神應該是差不多的。

從大方向來看， Stable Diffusion 可以分成三個重要元素：

  • Perceptual image compression: 先將圖像透過 perceptual image encoder & decoder (VQ-VAE ** [12] and VQ-GAN [13] style) 降低解析度後，直接在降解析度的圖像上，或者說特徵圖上，進行 DDPM 的訓練。**
  • Latent diffusion models: 基本上就是 DDPM 的描述方法，不過這個 DM 運行在latent space，因此論文稱為 latent diffusion model。
  • Conditioning mechanisms: SD論文內設計了利用 domain specific models 抽取 semantic 資訊後，再使用 attention 機制 [14] 與原本抽取圖像的 latent 結合。可以達到很泛用又有效的 conditioning。
Stable Diffusion 架構與運作。(資料來源 )

結合這三個就是構成 Stable Diffusion 的核心想法，** perceptual image compression 後的 latent diffusion model 讓 Stable Diffusion 可以有效率地生成高解析度圖片。Conditioning mechanisms 則讓 Stable Diffusion 可以具備良好的可控制性，完成各式各樣結合其他描述的生成任務。**

一般簡單的 autoencoder 雖然很容易收斂，也可以有效地壓縮 latent 大小，但是 decoder 重建的圖形實常是有一些模糊的，甚至有一些不真實的 artifact 產生。

常見 autoencoder 的概念。(資料來源)

因為 Stable Diffusion 的 diffusion model 是在 latent space 運作，如果decoder 品質不好會直接影響到生成圖像的品質。為了讓讓 decode 的圖像保持真持、並且減少模糊的情況，**Stable Diffusion 使用基於 VQGAN 的 autoencoder 設計。**簡單來說：

向量量化 (vector quantization) 其實在訊號處理法上其實已經存在幾十年了，簡單說就是找附近既定的點，來當作一個區間的代表。

在 VQGAN 中，使用稱為 codebook 的矩陣描述這些既定的點的向量數值，通常大小設計為 code 數量N乘上 code 長度n_z。然後先使用 CNN 針對影像 (H×W×3) 抽取特徵z後 (h×w×n_z，n_z為 code 的長度)，與 codebook 中所有 code 比較差異，最後用差異最小的那組 code 取代，得到z_q。

在 generative model 領域 VQ 原本被用於 VQVAE，利用 VQ 來避免訓練 VAE 時產生的 posterior collapse。

接下來就是 decode 成圖片，然後丟去計算 perceptual loss 與 patch-based adversarial objective。

  • Perceptual loss: 大概念就是透過先訓練好網路，**對原始圖片與重建圖片抽取特徵後，利用特徵計算 loss，要求兩邊抽出來的特徵越接近越好。**這個做法比起單純對兩張圖片計算 L2 ，更能考慮到影像整體。
  • Patch-based adversarial objective: 訓練discriminator 判斷每個 patch 屬於真實還是生成的圖片 patch，而非單純對整張圖片進行分類。
VQGAN 整體架構。(資料來源)

在 diffusion model 本體上， Stable Diffusion 基本上可以說是 DDPM 設計如出一轍。差別是，原本 DDPM 的 loss 計算是基於影像x_t。

而 Stable Diffution 則是基於 perceptual image encoder 抽完的 latentz_t。

Conditioning mechanisms 是 Stable Diffusion 可控性的重要元素，在辦法結合 condition 的情況下，模型才能夠做出各種符合人類想像的應用，比如說以文字為條件的 text to image 與以物體範圍為條件的 object removal。

Stable Diffusion 雖然主要還是使用 U-Net 的設計，但在 conditioning 的部分修改並加入 attention 機制。具體上控制的 condition (如 semantic map 或 text等)透過 domain specific encoder 抽出固定維度特徵後，在使用 attention 與 convolution 抽的 feature 做融合。

現今通用的 attention 描述。 (資料來源)

Stable Diffusion 設計 Q 項的輸入 convolutional feature ， K 跟 V 項的輸入則是 domain specific encoder 抽的 feature。

Domain specific encoder 通常都是額外先訓練好的，比如說在 Stable Diffusion 使用的文字 encoderCLIP [17]。

Synthesis Results

先撇除技術結果，現在 Stable Diffusion 或 Midjourney 生成的圖片真的都極為強大，並且可以生成各種風格：真實相片、電影海報、漫畫、藝術都是可以做出來的。

Stable Diffusion 生成的圖片。 (資料來源 )
Midjourney 生成的圖片。(資料來源 )

實際上目前那些假新聞用的，或是 18 禁的圖片通常都是 Stable Diffusion 做出來的，因為 Midjourney 目前是比較封閉式的應用情境，針對可能產生爭議的 prompt 也都會直接封鎖。而 Stable Diffusion 比較自由，甚至也都有開源的程式與 UI 可以拉下來自己部屬。

在生成圖片時，除了從 Gaussian distribution 中 sample 初始點外， condition 就是這些 generative model 最有意思的地方。前面說過， Stable Diffusion 是透過 domain specific encoder 抽出固定維度特徵後，再使用 attention 與 convolution 抽的 feature 做融合。

因此，只要透過各種 domain 的 plug-in ，比如說描述語言與圖像關聯的CLIP (Contrastive Language-Image Pre-Training) [17] 或者是泛用的ControlNet [18]，就可以做出各式各樣的應用。

Stable Diffusion 用於 text to image。(資料來源 )
Stable Diffusion 用於 layout to image。(資料來源 )
Stable Diffusion 用於 semantic synthesis。(資料來源 )
Stable Diffusion 用於 object removal。(資料來源 )

Experiments Results

Stable Diffusion 論文中的實驗數據，其實有點複雜，這邊只就最核心的部分，也就是 latent diffusion 這件是來看效果。因為大部分的實驗數據都是使用DDIM (Denoising Diffusion Implicit Models) ** [7] 的 sampling 技巧**，並在**評估上使用 **FID score [19]，因此先解釋一下這兩個名詞。

FID (Fréchet inception distance) [19] 是在圖像生成領域最常使用的評估 metric ，首見於 NeurIPS 2017 的論文 “GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium” 中。

FID 的概念其實跟 perceptual loss 很像

透過訓練好的網路，從圖像中抽 latent 進行比較。

真實與生成的圖象抽出的 latent 會分別形成兩個 distribution，然後計算這兩個 distribution 的距離就是 FID score。因此， FID score 數值越小代表生成品質越好。

FID計算概念。(資料來源)

至於兩個未知的 distribution 如何計算 distance 確實是個難題。在 FID 原本設計的方法，是直接把這兩個 distribution 視為 Gaussian。這個手法雖然很粗糙，但是在實務上確實是有效的。而 FID score 也因此至今都是圖像生成最熱門的 evaluation metric 之一。

雖然在 DDPM 論文中疊加多個 Gaussian noise 的動作可以在一步內完成，但是在生成圖像時，還是需要一步一步地進行 denoise。因為 DDPM 是一個parameterized Markov chain，當前的任何狀態轉變的機率，都與過去狀態完全無關。

DDPM 定義 diffusion model 為 parameterized Markov chain 。(資料來源)

DDIM 則將 DDPM 生成樣本時隨機的行為，變成了確定的行為。

DDIM 將 DDPM 生成樣本時隨機的行為，變成了確定的行為。 (資料來源 )

**DDIM 可以將 DDPM 多個 iteration 的 sampling 一步完成，大幅加速生成速度。理論上 DDIM 可以視為 DDPM 的 relaxation，或者說 DDPM 是 DDIM 的一種特例。**DDIM 並不一定是 Markov chain ，當下狀態轉變的機率可以基於原始的狀態。

DDIM 定義的 diffusion model 並不一定是 Markov chain 。(資料來源 )

聽起來有點複雜，總之，

用 DDPM 訓練的模型，可以直接套用 DDIM 的 sampling 技巧來加速生成速度。

Stable Diffusion 論文中關於 latent diffusion model (LDM) 的數據以 LDM-X 來表示不同的下採樣率 (downsample rate)。當X=1即沒有 downsample ，在 pixel-space 下進行 diffusion precess。

在訓練上可以看到有趣的現象是，平均表現最好 (又快圖像品質又好) 的是 LDM-4 到 LDM-16。 Pixel-space 下的 LDM-1 反而不僅慢，連生成的品質也比較差。在 Inference (生成圖片) 上，表現最好的則是 LDM-4 與 LDM-8 。總之，整體上來看，在 latent space 下做 diffusion 的結果真的不只是快而已，連品質都可以有顯著地提升。

Stable Diffusion 用於 semantic synthesis。(資料來源 )
Stable Diffusion 用於 semantic synthesis。(資料來源 )

Diffusion model 近期真的是一個超級火紅的題目，不過有鑒於他的生成品質與可控制性，這其實是相當合理，也相當期待未來可以有什麼更驚人的突破。在寫這篇文章的過程，也發現 OpenAI 持續在這邊放出砲火，比如說這個月 (2023年4月) 的Consistency Models [20] 。之後如果有發現什麼重要突破，有空也會盡量補充上來。

另外，如果覺得這篇文章還是有點晦澀難懂，建議可以看一下李弘毅老師近期的影片，有兩支影片分別用十來分鐘概快速地介紹 diffusion model 與 stable diffusion 相關的概念。這兩支影片 (DM影片 ** [21] / SD影片 [22] ) 將大部分地數學難題藏起來，聚焦於簡化的實現，非常地淺顯易懂。理解概念後，再回來看這篇文章或是其他介紹 diffusion model 的技術文章 (之前寫的DDPM介紹 / Lil’Log 的 DM Blog [23] / Ari Seff 的 DM 影片 [24] / Tidiane的 DDIM 影片 [25])**，會更融會貫通。

好了～這篇文章就先到這邊。老話一句，Deep Learning領域每年都會有大量高質量的論文產出，說真的要跟緊不是一件容易的事。所以我的觀點可能也會存在瑕疵，若有發現什麼錯誤或值得討論的地方，歡迎回覆文章或來信一起討論 :)

Reference

  1. Stable Diffusion Public Release[Official Web]
  2. Denoising Diffusion Probabilistic Model[NeurIPS 2020] [Project]
  3. Auto-Encoding Variational Bayes[ICLR 2014]
  4. Generative Adversarial Networks[NIPS 2014]
  5. Deep Unsupervised Learning using Nonequilibrium Thermodynamics[ICML 2015]
  6. Variational Inference with Normalizing Flows[ICML 2015]
  7. Denoising Diffusion Implicit Models[ICLR 2021]
  8. Google Imagen[Official Web]
  9. Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding[NeurIPS 2022]
  10. Cascaded Diffusion Models for High Fidelity Image Generation[arXiv 2021]
  11. High-Resolution Image Synthesis with Latent Diffusion Models[CVPR 2022]
  12. Neural Discrete Representation Learning[NeurIPS 2017]
  13. Taming Transformers for High-Resolution Image Synthesis[CVPR 2021]
  14. Attention Is All You Need[NIPS 2017]
  15. Perceptual Losses for Real-Time Style Transfer and Super-Resolution[ECCV 2016]
  16. High-Resolution Image Synthesis and Semantic Manipulation with Conditional GANs[CVPR 2018]
  17. Learning Transferable Visual Models From Natural Language Supervision[ICML 2021]
  18. Adding Conditional Control to Text-to-Image Diffusion Models[arXiv 2023]
  19. GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium[NeurIPS 2017]
  20. Consistency Models[arXiv 2023]
  21. 【生成式AI】淺談圖像生成模型 Diffusion Model 原理[Youtube]
  22. 【生成式AI】Stable Diffusion、DALL-E、Imagen 背後共同的套路[Youtube]
  23. What are Diffusion Models?[Blog Lil’Log]
  24. What are Diffusion Models?[Youtube]
  25. Faster Diffusion — presentation of the Denoising Diffusion Implicit Models paper[Youtube]
comments powered by Disqus