Part1: Overview of Diffusion Process

本文將會概括性地介紹\(\textit{Diffusion Process}\)算法與實踐，主要參考論文《Denoising Diffusion Probabilistic Models》。它的一些改進與優化，將“擴散方法”帶入主流視野。

而具體的數學推導部分，請參考其它系列文章。整個系列有相對完整的公式推導，若正文中有涉及到的省略部分，皆額外整理在Part4，並會在正文中會指明具體位置。

• Part1: Overview of Diffusion Process
• Part2: DDPM as Example of Variational Inference
• Part3: Dive into DDPM
• Part4: Appendix

Image Synthesis

在正式引入\(\textit{diffusion}\)前，希望先簡單介紹圖像生成的相關背景。一張圖像由很多個像素點組成，對於彩色圖像，每個像素點由三個0至255的整數表達，比如[255, 255, 255]代表像素點對應的顏色爲白色。而一張512*512的圖像就意味着共有26w左右的像素點。

圖像生成的目標是學習像素及像素間的概率分佈。結合具體的例子來理解概率分佈，上圖爲拍攝於戶外的照片。
假定它是由比較好的模型生成得到的圖像。看到“圖像上方是一片天空”，因爲它是天空，所以上方的像素點不會是大面積的綠色覆蓋。在生成天空區域時進行數值採樣，採樣到淡白色或者淡藍色對應數值的概率，遠大於採樣到綠色對應數值的概率。

通過常規的機器學習方法，直接學習上述提到的概率分佈很難，同時計算也很複雜。\(\textit{diffusion}\)的提出爲學習圖像的概率分佈提供新的思路和解決辦法。

Diffusion Process

圖像生成領域的\(\textit{Diffusion Process}\)最早在2015年發表的論文《Deep Unsupervised Learning using Nonequilibrium Thermodynamics》提到，作者認爲它是一種適配於深度學習背景下的圖片生成框架。其靈感來源於統計物理學中的非平衡態熱力學\((\textit{non-equilibrium})\)。在該領域，\(\textit{diffusion}\)並非新名詞，用於描述分子運動的現象。

Diffusion is the movement of a substance from a region of high concentration to a region of low concentration without bulk motion.

上圖中，起初紅色分子都聚集在液體的左上角，而當趨於穩定後，紅色分子遍佈在整個液體中。

那在圖像生成領域，論文作者獲取到的靈感是什麼？
圖像生成的其中一個難點在於不知曉圖像服從何種分佈，如此統計機器學習無法派上用場。故\(\textit{diffusion}\)的思路是，在\(\textit{forward process}\)中，通過緩慢、持續地往原數據上增加其它分佈的數據，將圖片原本分佈破壞掉，轉變到其它分佈；同時，存在一個\(\textit{reverse process}\)，即通過深度模型學習「如何將破壞掉的分佈復原回破壞前的分佈」。
前後兩個過程都是緩慢且逐步進行的，天然地可以用馬爾可夫鏈\((\textit{Markov chain})\)建模。並且，選擇的其它分佈是常見分佈，比如高斯分佈或二項分佈等。這類分佈的特性爲人們熟知，可以方便地進行計算。

上圖第一行描述二維數據分佈從\(t = 0\)至\(t = T\)的變化情況，數據分佈初始時爲Swiss Roll，在逐步經過Gaussian的diffusion後，最終完全轉變到服從Gaussian。

下面將圍繞以下幾部分展開介紹，符號定義與論文《Denoising Diffusion Probabilistic Models》保持一致：

1. Forward Process
2. Reverse Process
   • Loss Function
   • Model Architecture
3. Train and Sample

Forward Process

給定一張圖片，將其初始狀態記爲\(\mathbf{x}_0\)，逐漸地對其增加服從其它分佈的數據，如服從高斯分佈，重複若干輪後，原圖片的特徵逐漸被銷燬，完全難以分辨。爲了方便，下文記“其它分佈的數據”爲噪聲。

若對上圖經過緩慢地加噪，其變換情況如下圖所示。

整個緩慢的加噪過程建模爲馬爾可夫過程，狀態間的轉變\(\mathbf{x_{t-1}} \rightarrow \mathbf{x_{t}}\)服從高斯分佈：

\[q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)=\mathcal{N}\left(\mathbf{x}_t ; \sqrt{1-\beta_t} \mathbf{x}_{t-1}, \beta_t \mathbf{I}\right) \quad \tag{1} \]

故時刻\(t\)的圖片\(\mathbf{x_t}\)，由\(\mathbf{x_{t-1}}\)經過\((2)\)式得到，\((2)\)式的推導見Part4中高斯分佈的性質1：

\[\mathbf{x}_t =\sqrt{1-\beta_t} \mathbf{x}_{t-1}+\sqrt{\beta_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, \mathbf{I}) \tag{2}\]

其中，\(\beta_t\)是個參數，爲0到1的小數，表示時刻\(t\)噪聲增加的幅度。有過相關實踐的同學，應該會在模型文件或者代碼中看到scheduler的字樣，它便是用來產生\(\beta_t\)的模塊。

對\((2)\)式中的項\(\beta_t\)改寫，定義：

\[\begin{aligned} \alpha_t & =1-\beta_t \\ \bar{\alpha}_t & =\prod_{i=1}^t \alpha_i \end{aligned} \]

則通過推導（見Part4中推導一，可以得到\((3)\)式:

\[\begin{aligned} \mathbf{x}_t =\sqrt{\bar{\alpha}_t} \mathbf{x}_0+\sqrt{1-\bar{\alpha}_t} \epsilon \end{aligned}, \quad \epsilon \sim \mathcal{N}(0, \mathbf{I})\tag{3} \]

這是一個很好的性質，意味着任意時刻的狀態\(\mathbf{x}_t\)皆可由\(\mathbf{x}_0\)基於一次運算得到。在後續實際進行訓練時，自然地與隨機梯度下降\((\textit{Stochastic Gradient Descent})\)結合使用。

並且，當\(t\)足夠長，發現\((3)\)式的第一項趨近於0，而\(\mathbf{x}_t \approx \epsilon\)。這也印證，經過逐步緩慢的“擴散過程”後，原數據的分佈會被破壞，轉而服從所加入噪聲的分佈。

Reverse Process

當\(\textit{forward process}\)進行完畢後，認爲\(\mathbf{x}_T\)基本看作是一張噪聲圖片，它的像素點服從標準高斯分佈。
$$p(\mathbf{x}_T) \sim \mathcal{N}(0, \mathbf{I})$$
前文提到\(diffusion\)方法需要在reverse階段對被破壞的分佈進行復原，如視頻所示：

在這一過程中，深度學習模型介入，基於\((4)\)式表達這一過程：

\[\begin{gathered}p_\theta\left(x_0\right):=\int p_\theta\left(x_{0: T}\right) d x_{1: T} \\ p_\theta\left(\mathbf{x}_{0: T}\right):=p\left(\mathbf{x}_T\right) \prod_{t=1}^T p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right), \quad p(\mathbf{x}_T) \sim \mathcal{N}(\mathbf{x}_T;0, \mathbf{I}) \\ p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right):=\mathcal{N}\left(\mathbf{x}_{t-1} ; \boldsymbol{\mu}_\theta\left(\mathbf{x}_t, t\right), \boldsymbol{\Sigma}_\theta\left(\mathbf{x}_t, t\right)\right) \end{gathered}\tag{4} \]

其中，\(\theta\)對應模型的參數。

Loss Function

擴散模型的損失函數基於變分推斷\((\textit{variational inference})\)得到。簡單來說，最初的目標是最大化對數似然估計，但直接對\(\ln p_\theta\left(x_0\right)\)無從做起。基於變分推斷的思想，引入該目標的下界ELBO，通過最大化該下界實現計算。具體推導ELBO的方式請查看Part2。

經過一系列推導，最後簡化得到的損失函數如下所示：
$$L = \left|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_\theta\left(\sqrt{\bar{\alpha}_t} \mathbf{x}_0+\sqrt{1-\bar{\alpha}_t} \boldsymbol{\epsilon}, t\right)\right|^2
\tag{5}$$

目標是給定任意時刻\(t\)的狀態\(\mathbf{x}_t\)，預測此狀態中所包含的噪聲\(\epsilon_{\theta}(\mathbf{x}_t, t)\)，並與真實增加的噪聲\(\epsilon\)比較。\((5)\)式完全就是平方差損失函數\((\textit{mean squared error})\)。

Model Architecture

基於\((5)\)式，不難發現模型的輸入與輸出皆爲圖片格式的數據。在DDPM中，使用U-Net作爲該部分模型的結構。下圖爲U-Net原始的結構框架：


DDPM中使用的UNet在原基礎上，對於每個Block加入Attention以及Residual等模塊。

拿上圖所示的UNet舉例，其包含Down、Mid以及Up三部分，體現在圖中時，左邊包含有紅色向下箭頭的部分皆屬於Down，最底下包含數字1024的爲Mid，而右邊屬於Up部分。不難看出，UNet的形狀是鏡像對稱的，但是Down與Up輸入的張量維度卻是不同，主要原因在於：Up中所有的子模塊輸入，除了接收上一層傳遞過來的張量，還會接收對應層級Down子模塊傳遞過來的輸出。拿Up部分最低一層56x56x1024來說，其中56x56x512來自較低一層傳遞而來，另一部分的56x56x512從左邊的64x64x512經過copy+resize而來，兩者concate之後構成了當前層的輸入。

更加具體的代碼實現可以參考。

Train and Sample

這樣一來，模型的訓練算法如下所示：


而採樣算法如下所示：

其中，\(\sigma_t\)在具體實現中一般設置爲\(\sqrt{\beta_t}\)，是在diffusion process時該時間步\(t\)對應的標準差。

Summary

本文相對粗曠地介紹了“擴散過程”。前半階段，主要基於論文《Deep Unsupervised Learning using Nonequilibrium Thermodynamics》引入\(\textit{diffusion}\)的具體思想；後半階段，以《Denoising Diffusion Probabilistic Models》爲例，引入其訓練與採樣的算法步驟。

對於\(\textit{diffusion}\)，其最終在代碼層面的實現不算複雜；實際上，比較令人着迷的是其完備的數學推導，以及與其它隱變量模型\((\textit{latent models})\)的關聯。

Reference

• Deep Unsupervised Learning using Nonequilibrium Thermodynamics
• Denoising Diffusion Probabilistic Models
• U-net Implementation
• Diffusion Wikipedia
• denoise視頻來源