作者:京東零售 劉巖
擴散模型講解
前沿
人工智能生成內容(AI Generated Content,AIGC)近年來成為了非常前沿的一個研究方向,生成模型目前有四個流派,分別是生成對抗網路(Generative Adversarial Models,GAN),變分自編碼器(Variance Auto-Encoder,VAE),標準化流模型(Normalization Flow, NF)以及這里要介紹的擴散模型(Diffusion Models,DM),擴散模型是受到熱力學中的一個分支,它的思想來源是非平衡熱力學(Non-equilibrium thermodynamics),擴散模型的演算法理論基礎是通過變分推斷(Variational Inference)訓練引數化的馬爾可夫鏈(Markov Chain),它在許多任務上展現了超過GAN等其它生成模型的效果,例如最近非常火熱的OpenAI的DALL-E 2,Stability.ai的Stable Diffusion等,這些效果驚艷的模型擴散模型的理論基礎便是我們這里要介紹的提出擴散模型的文章[1]和非常重要的DDPM[2],擴散模型的實作并不復雜,但其背后的數學原理卻非常豐富,在這里我會介紹這些重要的數學原理,但省去了這些公式的推導計算,如果你對這些推導感興趣,可以學習參考文獻[4,5,11]的相關內容,我在這里主要以一個相對簡單的角度來講解擴散模型,幫助你快速入門這個非常重要的生成演算法,
1. 背景知識: 生成模型
生成模型的本質是通過一個已知的概率模型來擬合所給的資料樣本,也就是說,我們往往需要通過模型得到一個帶引數的分布,即如果訓練資料的分布是\(p_\text{data}(x)\),生成樣本的分布為\(p_\text{model}(x)\),我們希望得到的分布和訓練資料的分布盡可能相似,目前生成模型主要有圖1所示的四類,其中GAN的原理是通過判別器和生成器的互相博弈來讓生成器生成足以以假亂真的影像,VAE的原理是通過一個編碼器將輸入影像編碼成特征向量,它用來學習高斯分布的均值和方差,而解碼器則可以將特征向量轉化為生成影像,它側重于學習生成能力,流模型是從一個簡單的分布開始,通過一系列可逆的轉換函式將分布轉化成目標分布,擴散模型先通過正向程序將噪聲逐漸加入到資料中,然后通過反向程序預測每一步加入的噪聲,通過將噪聲去掉的方式逐漸還原得到無噪聲的影像,擴散模型本質上是一個馬爾可夫架構,只是其中訓練程序用到了深度學習的BP,但它更屬于數學層面的創新,這也就是為什么很多計算機的同學看擴散模型相關的論文會如此費力,
圖1:生成模型的四種型別 [4]
擴散模型中最重要的思想根基是馬爾可夫鏈,它的一個關鍵性質是平穩性,即如果一個概率隨時間變化,那么再馬爾可夫鏈的作用下,它會趨向于某種平穩分布,時間越長,分布越平穩,如圖2所示,當你向一滴水中滴入一滴顏料時,無論你滴在什么位置,只要時間足夠長,最終顏料都會均勻的分布在水溶液中,這也就是擴散模型的前向程序,
圖2:顏料分子在水溶液中的擴散程序
如果我們能夠在擴散的程序顏料分子的位置、移動速度、方向等移動屬性,那么也可以根據正向程序的保存的移動屬性從一杯被溶解了顏料的水中反推顏料的滴入位置,這邊是擴散模型的反向程序,記錄移動屬性的快照便是我們要訓練的模型,
2. 擴散模型
在這一部分我們將集中介紹擴散模型的數學原理以及推導的幾個重要性質,因為推導程序涉及大量的數學知識但是對理解擴散模型本身思想并無太大幫助,所以這里我會省去推導的程序而直接給出結論,但是我也會給出推導程序的出處,對其中的推導程序比較感興趣的請自行查看,
2.1 計算原理
擴散模型簡單的講就是通過神經網路學習從純噪聲資料逐漸對資料進行去噪的程序,它包含兩個步驟,如圖3:
- 固定的前向程序\(q\):在這一步我們逐漸將高斯噪聲添加到影像中,直到得到一個純噪聲的影像;
- 可學習的反向去噪程序\(p_\theta\):在這一步我們從純噪聲影像中逐漸對其進行去噪,直到得到真實的影像,
圖3:DDPM的前向加噪和后向去噪程序
更具體些,對于一個\(T\)步的擴散模型,每一步的索引為\(t\),在前向程序中,我們從一個真實影像\(\boldsymbol x_0\)開始,在每一步我們隨機生成一些高斯噪聲,然后將生成的噪聲逐步加入到輸入影像中,當\(T\)足夠大時,我們得到的加噪后的影像便接近一個高斯噪聲影像,例如DDPM中 \(T=1000\),在后向程序中,我們從噪聲影像\(\boldsymbol x_T\)開始(訓練時是真實影像加噪的結果,采樣時是隨機噪聲),通過一個神經網路學習\(\boldsymbol x_{t-1}\)到\(\boldsymbol x_t\)添加的噪聲,然后通過逐漸去噪的方式得到最后要生成的影像,
2.1.1 前向程序
前向程序即擴散程序指的是向資料中逐漸添加高斯噪聲直到資料完全變成噪聲的程序,假設\(q(\boldsymbol x_0)\)是真實影像的分布,我們可以通過從訓練集的真實影像中隨機采樣一張影像,表示為\(\boldsymbol x \sim q(\boldsymbol x_0)\),那么前向程序\(q(\boldsymbol x_t | \boldsymbol x_{t-1})\)指的是在前向的每一步通過向影像\(\boldsymbol x_{t-1}\)中添加高斯噪聲得到\(\boldsymbol x_t\),我們知道,一個高斯分布由均值\(\mu\)和方差\(\sigma^2\)定義,那么在每一步向影像中添加高斯噪聲的程序表示為式(1),它是一個均值\(\mu_t = \sqrt{1-\beta_t} \boldsymbol x_{t-1}\),方差\(\sigma^2_t = \beta_t\)的高斯分布,
\[q(\boldsymbol x_t | \boldsymbol x_{t-1}) = \mathcal N(\boldsymbol x_t | \sqrt{1 - \beta_t} \boldsymbol x_{t-1}, \beta_t \mathbf I) \tag1 \]具體到每一步的計算時,我們先采樣一個二維標準高斯分布\(\epsilon \sim \mathcal N(\boldsymbol 0, \mathbf I)\),然后通過引數\(\beta_t\)由\(\boldsymbol x_{t-1}\)得到\(\boldsymbol x_t\),表示為\(\boldsymbol x_t = \sqrt{1-\beta_t} \boldsymbol x_{t-1} + \sqrt{\beta_t} \epsilon\),注意這里\(\beta_t\)并不是一個常數值,而是一個隨時間變化的變數,\(\beta_t\)的變化情況在擴散模型中被叫做差異時間表(Variance Schedule),常見的時間表策略有線性時間表(DDPM),平方時間表,cosine時間表[6]等,在DDPM擴散模型的前項程序中,我們需要保證\(T\)足夠大并且\(\beta_t\)的時間表配置合理,才能保證我們最終得到的\(\boldsymbol x_t\)也將是一個高斯噪聲影像,擴散模型的全部前向程序可以表示為從\(t=1\)到\(t=T\)的時刻的馬爾可夫鏈,如式(2),
\[q(\boldsymbol x_{0:T}) = q(\boldsymbol x_0) \prod_{t=1}^T q\left( \boldsymbol x_t | \boldsymbol x_{t-1} \right) \tag2 \]擴散程序一個隱藏的重要特征是我們可以直接基于原始資料\(\boldsymbol x_0\)來對任意\(t\)步的\(\boldsymbol x_t\)進行采樣,這里我們定義\(\bar{a}_t = \prod_{s=1}^t \alpha_t\)以及\(\alpha_t = 1 - \beta_t\),通過重引數(Reparamazation)技巧,我們可以得到服從分布\(q(\boldsymbol x_t| \boldsymbol x_0)\)的任意一個樣本\(\boldsymbol x_t\),我們可以通過反重引數化得到式(3),推導程序見參考文獻[11]的式(61)到式(70),
\[\begin{aligned} \boldsymbol x_t & \sim q(\boldsymbol x_t | \boldsymbol x_0) \\ & = \sqrt{\bar{\alpha_t}} \boldsymbol x_0 + \sqrt{1 - \bar{\alpha_t}}\epsilon \\ & = \mathcal N(\boldsymbol x_t; \sqrt{\bar{\alpha}_t} \boldsymbol x_0, (1 - \bar{\alpha}_t) \mathbf I) \end{aligned} \tag3 \]上面推理反應了一個重要的性質,即\(\boldsymbol x_t\)可以看做原始資料\(\boldsymbol x_0\)和隨機噪聲\(\boldsymbol \epsilon\)的線性組合,其中\(\sqrt{\bar{\alpha}_t}\)和\(\sqrt{1 - \bar{\alpha}_t}\)是組合系數,它們的平方和為\(1\),公式(3)在論文中被叫做“Nice Property”,進一步的,我們可以使用\(\bar \alpha\)來定義差異時間表,例如[6]中的cosine時間表便是這么做的,\(\bar{\alpha}_t\)要比\(\beta\)更直接,例如我們可以通過將\(\bar{\alpha}_T\)設定為一個接近\(0\)的值,來使得最終得到的噪聲更傾向于是一個高斯噪聲,
2.1.2 后向程序
前向程序是將資料噪聲化的程序,那么擴散模型的后向程序\(p(\boldsymbol x_{t-1} | \boldsymbol x_t)\)則是一個去噪程序,即我們先在\(T\)時刻隨機采樣一個二維高斯噪聲,然后逐步進行去噪,最終得到一個和真實影像分布一致的生成影像\(\boldsymbol x_0\),
所以擴散模型的核心是如何進行這個去噪程序,因為我們并不知道\(p(\boldsymbol x_{t-1} | \boldsymbol x_t)\)的具體形式是什么,擴散模型指出,我們可以使用一個神經網路學習這個去噪程序,因為第\(t\)時刻的分布\(\boldsymbol x_t\)是已知的,因此我們這個神經網路的目標是根據\(\boldsymbol x_t\)去學習\(\boldsymbol x_{t-1}\)的概率分布,綜上,擴散模型的后向程序表示為\(p_\theta(\boldsymbol x_{t-1} | \boldsymbol x_t)\),其中\(\theta\)是神經網路的引數,我們可以用SGD等策略對該網路進行優化,
因為前向程序我們添加的噪聲是高斯噪聲,為了簡化模型的訓練難度,我們假設反向的去噪程序去掉的噪聲也是高斯噪聲,因為一個高斯分布是通過均值\(\mu_\theta\)和方差\(\sum_\theta\)決定的,那么\(p_\theta(\boldsymbol x_{t-1} | \boldsymbol x_t)\)可以表示為式(4)的形式,
\[p_\theta(\boldsymbol x_{t-1} | \boldsymbol x_t) = \mathcal N(\boldsymbol x_{t-1}; \mu_\theta(\boldsymbol x_t,t), \sum_\theta(\boldsymbol x_t, t)) \tag4 \]其中均值和方差均是根據模型計算得到的,綜合所有時間步,我們也可以通過馬爾可夫鏈得到擴散模型的后向程序,如式(5),
\[p_\theta(\boldsymbol x_{0:T}) = p(x_T) \prod_{t=1}^T p_\theta (\boldsymbol x_{t-1} | \boldsymbol x_t) \tag5 \]其中\(p(\boldsymbol x_t) = \mathcal N(\boldsymbol x_T; \boldsymbol 0, \mathbf I)\)是隨機采樣的高斯噪聲,\(p_\theta (\boldsymbol x_{t-1} | \boldsymbol x_t)\)是一個均值和方差需要計算的高斯分布,
2.1.3 目標函式
那么問題來了,我們究竟使用什么樣的優化目標才能比較好的預測高斯噪聲的分布呢?一個比較復雜的方式是使用變分自編碼器的最大化證據下界(Evidence Lower Bound, ELBO)的思想來推導,如式(6),推導詳細程序見論文[11]的式(47)到式(58),這里主要用到了貝葉斯定理和琴生不等式,
\[\begin{aligned} \mathcal L & = - \log p(\boldsymbol x) \\ & = - \log \int \frac{p_\theta(\boldsymbol x_{0:T})q(\boldsymbol x_{1:T} | \boldsymbol x_0)}{q(\boldsymbol x_{1:T} | \boldsymbol x_0)} d \boldsymbol x_{1:T} \\ & \leq - \mathbb E_{q(\boldsymbol x_{1:T} | \boldsymbol x_0)} \left[ \frac{p_\theta(\boldsymbol x_{0:T})}{q(\boldsymbol x_{1:T} | \boldsymbol x_0)}\right] \\ & = - \underbrace{\mathbb{E}_{q\left(\boldsymbol{x}_1 \mid \boldsymbol{x}_0\right)}\left[\log p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_0 \mid \boldsymbol{x}_1\right)\right]}_{\text {重構項}} + \underbrace{D_{\mathrm{KL}}\left(q\left(\boldsymbol{x}_T \mid \boldsymbol{x}_0\right) \| p\left(\boldsymbol{x}_T\right)\right)}_{\text {先驗匹配項}} + \sum_{t=2}^T \underbrace{\mathbb{E}_{q\left(\boldsymbol{x}_t \mid \boldsymbol{x}_0\right)}\left[D_{\mathrm{KL}}\left(q\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0\right) \| p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t\right)\right)\right]}_{\text {去噪匹配項}} \end{aligned} \tag6 \]式(6)的推導細節并不重要,我們需要重點關注的是它的最終等式的三個組成部分,下面我們分別介紹它們:
(1) 重構項 \(\mathcal L_0 = \mathbb{E}_{q\left(\boldsymbol{x}_1 \mid \boldsymbol{x}_0\right)}\left[\log p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_0 \mid \boldsymbol{x}_1\right)\right]\),它的作用是對原始資料進行重構,優化的是負log似然,DDPM提供的計算方式是:它首先將離散的影像像素從\([0, 255]\)歸一化到了\([-1,1]\)的范圍,然后用式(3)中估計的$\mathcal N(\boldsymbol x_t; \sqrt{\bar{\alpha}_t} \boldsymbol x_0, (1 - \bar{\alpha}_t) \boldsymbol I) $構建一個離散的解碼器來計算,如式(7),它計算的是高斯分布落在以Ground Truth為中心,且范圍大小為\(2/255\)時的概率積分,即累積分布函式(Cumulative Distribution Function,CDF),
\[\begin{aligned} p_\theta\left(\boldsymbol{x}_0 \mid \boldsymbol{x}_1\right) & =\prod_{i=1}^D \int_{\delta_{-}\left(x_0^i\right)}^{\delta_{+}\left(x_0^i\right)} \mathcal{N}\left(x ; \mu_\theta^i\left(\boldsymbol{x}_1, 1\right), \sigma_1^2\right) d x \\ \delta_{+}(x) & =\left\{\begin{array}{ll} \infty & \text { if } x=1 \\ x+\frac{1}{255} & \text { if } x<1 \end{array} \quad \delta_{-}(x)= \begin{cases}-\infty & \text { if } x=-1 \\ x-\frac{1}{255} & \text { if } x>-1\end{cases} \right. \end{aligned} \tag7 \]其中\(D\)是整張影像,\(i\)是影像上的像素點坐標,
(2) 先驗重構項 \(\mathcal L_T=D_{\mathrm{KL}}\left(q\left(\boldsymbol{x}_T \mid \boldsymbol{x}_0\right) \| p\left(\boldsymbol{x}_T\right)\right)\):它使用KL散度計算了最后的噪聲輸入和標準的高斯先驗的接近程度,因為這一部分沒有可以訓練的引數,我們可以將它視作常數0,
(3) 去噪匹配項 \(\mathcal L_{t-1} = \mathbb{E}_{q\left(\boldsymbol{x}_t \mid \boldsymbol{x}_0\right)}\left[D_{\mathrm{KL}}\left(q\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0\right) \| p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t\right)\right)\right]\):它計算的是真實后驗分布\(q\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t, \boldsymbol{x}_0\right)\)和預測的分布\(p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t\right)\)之間的KL散度,因為我們希望真實的去噪程序和模型預測的去噪程序完全一致,如圖4,
圖4:擴散模型的去噪匹配項在每一步都要擬合噪音的真實后驗分布和估計分布
真實后驗分布可以使用貝葉斯定理進行推導,最終結果如式(8),推導程序見論文[11]的式(71)到式(84),
\[\begin{aligned} q(\boldsymbol x_{t-1} | \boldsymbol x_t, \boldsymbol x_0) & = \frac{q(\boldsymbol x_{t} | \boldsymbol x_{t-1}, \boldsymbol x_0) q(\boldsymbol x_{t-1} | \boldsymbol x_0)}{q(\boldsymbol x_{t} | \boldsymbol x_0)} \\ & \propto \mathcal N \left( \boldsymbol x_{t-1}; \frac{\sqrt{\alpha_t} (1 - \bar{\alpha}_{t-1}) \boldsymbol x_t + \sqrt{\bar{\alpha}_{t-1}}(1 - \alpha_t) \boldsymbol x_0}{ 1- \bar{\alpha}_t}, \frac{(1 - \alpha_t)(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf I \right) \\ & = \mathcal N(\boldsymbol x_{t-1}; \mu_q(\boldsymbol x_t, \boldsymbol x_0), \Sigma_q(t)) \end{aligned} \tag8 \]其中\(\mu_q(\boldsymbol x_t, \boldsymbol x_0) = \frac{\sqrt{\alpha_t} (1 - \bar{\alpha}_{t-1}) \boldsymbol x_t + \sqrt{\bar{\alpha}_{t-1}}(1 - \alpha_t) \boldsymbol x_0}{ 1- \bar{\alpha}_t}\)是均值,\(\Sigma_q(t) = \frac{(1 - \alpha_t)(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf I\)是方差,可以看出均值是和\(\boldsymbol x_t\)和\(\boldsymbol x_0\)相關的,但是方差與資料無關,假設預測的分布也服從正態分布,表示為式(9),
\[p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_t\right) = \mathcal N(\boldsymbol x_{t-1}; \mu_\theta(\boldsymbol x_t, t), \Sigma_q(t)) \tag9 \]為了進一步化簡\(\mathcal L_{t-1}\)我們需要用到式(10)的兩個高斯分布的KL散度計算公式,
\[\begin{aligned} & D_\text{KL}(\mathcal N(\boldsymbol x; \boldsymbol \mu_x, \boldsymbol \Sigma_x), \mathcal N(\boldsymbol y; \boldsymbol \mu_y, \boldsymbol \Sigma_y) \\ = & \frac{1}{2}\left[ \log \frac{|\boldsymbol \Sigma_x|}{|\boldsymbol \Sigma_y|} - d + \text{tr}(\boldsymbol \Sigma_y^{-1} \boldsymbol \Sigma_x) + (\boldsymbol \mu_y - \boldsymbol \mu_x)^\intercal \boldsymbol \sigma_y^{-1}(\boldsymbol \mu_y - \boldsymbol \mu_x)\right]) \end{aligned} \tag{10} \]在兩個分布均是高斯分布的前提下,我們可以使用公式(10)繼續對\(\mathcal L_{t-1}\)進行進行計算,這一部分完整的推導流程參考論文[11]的式(87)到式(92),從這里我們可以看出,當兩個高斯分布方差相同時,求它們之間的KL散度既是求兩個分布的均值的l2距離,
\[\begin{aligned} \mathop{\arg\min}_\theta D_\text{KL}(q(\boldsymbol x_{t-1} | \boldsymbol x_t, \boldsymbol x_0) || p_\theta(\boldsymbol x_{t-1} | \boldsymbol x_t)) = \mathop{\arg\min}_\theta \frac{1}{2\sigma_q^2(t)} \left[\|\boldsymbol \mu_\theta (\boldsymbol x_t, \boldsymbol x_0) - \boldsymbol \mu_q(\boldsymbol x_t, t) \|^2_2\right] \\ \end{aligned} \tag{11} \]通過式(3)我們可以得到\(\boldsymbol x_0 = \frac{\boldsymbol x_t - \sqrt{1 - \bar\alpha_t}\boldsymbol \epsilon_0}{\sqrt {\bar \alpha_t}}\),將它代入到\(\mu_q(\boldsymbol x_t, \boldsymbol x_0) = \frac{\sqrt{\alpha_t} (1 - \bar{\alpha}_{t-1}) \boldsymbol x_t + \sqrt{\bar{\alpha}_{t-1}}(1 - \alpha_t) \boldsymbol x_0}{ 1- \bar{\alpha}_t}\)中,我們可以得到式(12)的最終結果,推導程序見參考文獻(11)的式(116)到式(124),
\[\begin{aligned} \boldsymbol \mu_q(\boldsymbol x_t, \boldsymbol x_0) & = \frac{\sqrt{\alpha_t} (1 - \bar{\alpha}_{t-1}) \boldsymbol x_t + \sqrt{\bar{\alpha}_{t-1}}(1 - \alpha_t) \boldsymbol x_0}{ 1- \bar{\alpha}_t} \\ & = \frac{\sqrt{\alpha_t} (1 - \bar{\alpha}_{t-1}) \boldsymbol x_t + \sqrt{\bar{\alpha}_{t-1}}(1 - \alpha_t) \frac{\boldsymbol x_t - \sqrt{1 - \bar\alpha_t}\boldsymbol \epsilon_0}{\sqrt {\bar \alpha_t}}}{ 1- \bar{\alpha}_t} \\ & = \frac{1}{\sqrt{\alpha_t}} \boldsymbol x_t - \frac{1-\alpha_t}{\sqrt{1 - \bar{\alpha_t}}{\sqrt{\alpha_t}}} \boldsymbol \epsilon_0 \end{aligned} \tag{12} \]同理,我們也可以用這種方式計算$ \mu_\theta (\boldsymbol x_t, \boldsymbol x_0)$,如式(13):
\[\mu_\theta (\boldsymbol x_t, t) = \frac{1}{\sqrt{\alpha_t}} \boldsymbol x_t - \frac{1-\alpha_t}{\sqrt{1 - \bar{\alpha_t}}{\sqrt{\alpha_t}}} \hat {\boldsymbol \epsilon }_\theta(\boldsymbol x_t, t) \tag{13} \]將式(12)和式(13)帶入到式(11)中,我們可以得到:
\[\begin{aligned} & \mathop{\arg\min}_\theta D_\text{KL}(q(\boldsymbol x_{t-1} | \boldsymbol x_t, \boldsymbol x_0) || p_\theta(\boldsymbol x_{t-1} | \boldsymbol x_t)) \\ = \; & \mathop{\arg\min}_\theta \frac{1}{2\sigma^2_q(t)} \frac{(1-\alpha_t)^2}{(1-\bar{\alpha}_t)\alpha_t} \left[ \| \boldsymbol \epsilon_0 - \hat{\boldsymbol \epsilon}_\theta (\boldsymbol x_t, t)\|^2_2 \right] \\ = \; & \mathop{\arg\min}_\theta \frac{1}{2\sigma^2_q(t)} \frac{(1-\alpha_t)^2}{(1-\bar{\alpha}_t)\alpha_t} \left[ \| \boldsymbol \epsilon_0 - \hat{\boldsymbol \epsilon}_\theta (\sqrt{\bar{\alpha}_t} \boldsymbol x_0 + \sqrt{ 1- \bar{\alpha}_t} \boldsymbol \epsilon , t)\|^2_2 \right] \end{aligned} \tag{14} \]最終我們可以將擴散模型的損失函式簡化為式(15)的形式,其中\(\boldsymbol \epsilon_t\)是添加的高斯噪聲,\(\hat{\boldsymbol \epsilon}_\theta(\boldsymbol x_t, t)\)是一個神經網路,用于預測從\(\boldsymbol x_0\)到\(\boldsymbol x_t\)時刻添加的噪聲,
\[\mathcal L_\text{simple} = \mathbb E_{t, \boldsymbol x_0, \epsilon} \| \boldsymbol \epsilon_t - \hat{\boldsymbol \epsilon}_\theta (\sqrt{\bar{\alpha}_t} \boldsymbol x_0 + \sqrt{ 1- \bar{\alpha}_t} \boldsymbol \epsilon , t)\|^2 \tag{15} \]2.1.4 模型訓練
在第2.1.1節我們講到我們可以直接基于原始資料\(\boldsymbol x_0\)來對任意\(t\)步的\(\boldsymbol x_t\)進行采樣,那么在實際訓練程序中,我們不必將所有的時間片都拿來訓練,而采取直接采樣到時刻\(t\),然后得到該時刻的\(\boldsymbol x_t\)并使用神經網路預測添加的噪聲即可,因為擴散模型的\(T\)是一個非常大的值,使用這種方式將大幅提升訓練速度,它的訓練程序為:
- 從分布為\(q(\boldsymbol x_0)\)的資料集隨機采樣一個樣本\(\boldsymbol x_0 \sim q(\boldsymbol x_0)\);
- 從\(1\)到\(T\)中隨機采樣一個值\(t\),用于表示添加噪聲的水平;
- 隨機采樣一個二維高斯噪音\(\epsilon\),然后使用上面介紹的“Nice Property”對\(\boldsymbol x_0\)施加\(t\)級別的噪聲;
- 訓練神經網路根據加噪之后的\(\boldsymbol x_t\)預測作用到\(\boldsymbol x_0\)之上的噪聲,
雖然上面我們介紹了很多內容,并給出了大量公式,但得益于推匯出的幾個重要性質,擴散模型的訓練并不復雜,它的訓練偽代碼見演算法1,
2.1.5 樣本生成
正如我們所介紹的,擴散模型的生成程序是一個反向去噪的程序,它的偽代碼見演算法2,具體的講,我們從\(T\)時刻開始,首先隨機采樣一個高斯噪聲,然后我們使用神經網路預測的噪聲逐漸對其去噪,直到\(0\)時刻停止,通過式(9)我們得到了\(\boldsymbol x_{t-1} \sim p_\theta(\boldsymbol x_{t-1}|\boldsymbol x_t)\),那么我們可以進一步得到從\(\boldsymbol x_t\)到\(\boldsymbol x_{t-1}\)的計算公式,如式(16),
\[\boldsymbol x_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left( \boldsymbol x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}}\epsilon(\boldsymbol x_t, t)\right) + \sigma_t \boldsymbol z \tag{16} \]其中\(\boldsymbol z\sim \mathcal N(\mathbf 0, \mathbf I)\)是一個二維標準高斯分布,
2.2 演算法實作
2.2.1模型結構
DDPM在預測施加的噪聲時,它的輸入是施加噪聲之后的影像,預測內容是和輸入影像相同尺寸的噪聲,所以它可以看做一個Img2Img的任務,DDPM選擇了U-Net[9]作為噪聲預測的模型結構,U-Net是一個U形的網路結構,它由編碼器,解碼器以及編碼器和解碼器之間的跨層連接(殘差連接)組成,其中編碼器將影像降采樣成一個特征,解碼器將這個特征上采樣為目標噪聲,跨層連接用于拼接編碼器和解碼器之間的特征,
圖5:U-Net的網路結構
下面我們介紹DDPM的模型結構的重要組件,首先在U-Net的卷積部分,DDPM使用了寬殘差網路(Wide Residual Network,WRN)[12]作為核心結構,WRN是一個比標準殘差網路層數更少,但是通道數更多的網路結構,也有作者復現發現ConvNeXT作為基礎結構會取得非常顯著的效果提升[13,14],這里我們可以根據訓練資源靈活的調整卷積結構以及具體的層數等超參,因為我們在擴散程序的整個流程中都共享同一套引數,為了區分不同的時間片,作者借鑒了Transformer [15]的位置編碼的思想,采用了正弦位置嵌入對時間\(t\)進行了編碼,這使得模型在預測噪聲時知道它預測的是批次中分別是哪個時間片添加的噪聲,在卷積層之間,DDPM添加了一個注意力層,這里我們可以使用Transformer中提出的自注意力機制或是多頭自注意力機制,[13]則提出了一個線性注意力機制的模塊,它的特點是消耗的時間以及占用的記憶體和序列長度是線性相關的,對比傳統注意力機制的平方相關要高效很多,在進行歸一化時,DDPM選擇了組歸一化(Group Normalization,GN)[16],最后,對于U-Net中的降采樣和上采樣操作,DDPM分別選擇了步長為2的卷積以及反卷積,
確定了這些組件,我們便可以搭建用于DDPM的U-Net的模型了,從第2.1節的介紹我們知道,模型的輸入為形狀為(batch_size, num_channels, height, width)的噪聲影像和形狀為(batch_size,1)的噪聲水平,回傳的是形狀為(batch_size, num_channels, height, width)的預測噪聲,我們搭建的用于噪聲預測的模型結構如下:
- 首先在噪聲影像\(\boldsymbol x_0\)上應用卷積層,并為噪聲水平\(t\)計算時間嵌入;
- 接下來是降采樣階段,采用的模型結構依次是兩個卷積(WRNS或是ConvNeXT)+GN+Attention+降采樣層;
- 在網路的最中間,依次是卷積層+Attention+卷積層;
- 接下來是上采樣階段,它首先會使用Short-cut拼接來自降采樣中同樣尺寸的卷積,再之后是兩個卷積+GN+Attention+上采樣層,
- 最后是使用WRNS或是ConvNeXT作為輸出層的卷積,
U-Net類的forword函式如下面代碼片段所示,完整的實作代碼參照[3],
def forward(self, x, time):
x = self.init_conv(x)
t = self.time_mlp(time) if exists(self.time_mlp) else None
h = []
# downsample
for block1, block2, attn, downsample in self.downs:
x = block1(x, t)
x = block2(x, t)
x = attn(x)
h.append(x)
x = downsample(x)
# bottleneck
x = self.mid_block1(x, t)
x = self.mid_attn(x)
x = self.mid_block2(x, t)
# upsample
for block1, block2, attn, upsample in self.ups:
x = torch.cat((x, h.pop()), dim=1)
x = block1(x, t)
x = block2(x, t)
x = attn(x)
x = upsample(x)
return self.final_conv(x)
2.2.2 前向加噪
擴散模型的前向程序是逐漸向影像中添加噪聲的程序,這個從時刻\(0\)到時刻\(T\)的\(t\)的變化情況叫做差異時間表,DDPM使用的是線性時間表,即我們首先對\(T\)個時間步做均勻拆分,得到\(\beta_t\),接下來我們根據\(\beta_t\)計算我們需要的其它變數,例如\(\alpha\),\(\bar{\alpha}\)等,它們最好被存盤起來以避免重復計算,接下來我們需要準備輸入影像\(\boldsymbol x_0\)以及隨機噪聲\(\epsilon\),其中影像的處理主要包括resize到模型輸入大小以及進行歸一化,DDPM的策略是將它們線性歸一化到\([-1,1]\)之間,隨機噪聲即隨機生成一個和輸入影像相同尺寸的二維高斯噪聲,最后我們根據式(3)將它們合成一張影像即可,圖6是參考文獻[3]中給出的輸入影像依次經過0次,50次,100次,150次以及199次加噪后的效果圖,可以看出隨著逐漸添加噪聲,影像越來越難以區分,直到徹底變成一個二維高斯噪聲,
圖6:一張圖依次經過0次,50次,100次,150次以及199次加噪后的效果圖
根據式(14)我們知道,擴散模型的損失函式計算的是兩張影像的相似性,因此我們可以選擇使用回歸演算法的所有損失函式,以MSE為例,前向程序的核心代碼如下面代碼片段,
def p_losses(denoise_model, x_start, t, noise=None, loss_type="l1"):
# 1. 根據時刻t計算隨機噪聲分布,并對影像x_start進行加噪
x_noisy = q_sample(x_start=x_start, t=t, noise=noise)
# 2. 根據噪聲影像以及時刻t,預測添加的噪聲
predicted_noise = denoise_model(x_noisy, t)
# 3. 對比添加的噪聲和預測的噪聲的相似性
loss = F.mse_loss(noise, predicted_noise)
return loss
2.2.3 樣本生成
根據2.1.5節介紹的樣本生成流程,它的核心代碼片段所示,關于這段代碼的講解我通過注釋添加到了代碼片段中,
@torch.no_grad()
def p_sample(model, x, t, t_index):
betas_t = extract(betas, t, x.shape)
sqrt_one_minus_alphas_cumprod_t = extract(sqrt_one_minus_alphas_cumprod, t, x.shape)
sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)
# 使用式(13)計算模型的均值
model_mean = sqrt_recip_alphas_t * (x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t)
if t_index == 0:
return model_mean
else:
# 獲取保存的方差
posterior_variance_t = extract(posterior_variance, t, x.shape)
noise = torch.randn_like(x)
# 演算法2的第4行
return model_mean + torch.sqrt(posterior_variance_t) * noise
# 演算法2的流程,但是我們保存了所有中間樣本
@torch.no_grad()
def p_sample_loop(model, shape):
device = next(model.parameters()).device
b = shape[0]
# start from pure noise (for each example in the batch)
img = torch.randn(shape, device=device)
imgs = []
for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=timesteps):
img = p_sample(model, img, torch.full((b,), i, device=device, dtype=torch.long), i)
imgs.append(img.cpu().numpy())
return imgs
最后我們看下在人臉影像資料集下訓練的模型,一批隨機噪聲經過逐漸去噪變成人臉影像的示例,
圖7:擴散模型由隨機噪聲通過去噪逐漸生成人臉影像
3. 總結
這里我們以DDPM為例介紹了另一個派系的生成演算法:擴散模型,擴散模型是一個基于馬爾可夫鏈的數學模型,它通過預測每個時間片添加的噪聲來進行模型的訓練,作為近日來引發熱烈討論的ControlNet, Stable Diffusion等模型的底層演算法,我們十分有必要對其有所了解,DDPM的實作并不復雜,這得益于大量數學界大佬通過大量的數學推導將整個擴散程序和反向去噪程序進行了精彩的化簡,這才有了DDPM的大道至簡的實作,DDPM作為一個擴散模型的基石演算法,它有著很多早期演算法的共同問題:
- 采樣速度慢:DDPM的去噪是從時刻\(T\)到時刻\(1\)的一個完整的馬爾可夫鏈的計算,尤其是DDPM還需要一個比較大的\(T\)才能保證比較好的效果,這就導致了DDPM的采樣程序注定是非常慢的;
- 生成效果差:DDPM的效果并不能說是非常好,尤其是對于高解析度影像的生成,這一方面是因為它的計算速度限制了它擴展到更大的模型;另一方面它的設計還有一些問題,例如逐像素的計算損失并使用相同權值而忽略影像中的主體并不是非常好的策略,
- 內容不可控:我們可以看出,DDPM生成的內容完全還是取決于它的訓練集,它并沒有引入一些先驗條件,因此并不能通過控制影像中的細節來生成我們制定的內容,
我們現在已經知道,DDPM的這些問題已大幅得到改善,現在基于擴散模型生成的影像已經達到甚至超過人類多數的畫師的效果,我也會在之后逐漸給出這些優化方案的講解,
Reference
[1] Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International Conference on Machine Learning. PMLR, 2015.
[2] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in Neural Information Processing Systems 33 (2020): 6840-6851.
[3] https://huggingface.co/blog/annotated-diffusion
[4] https://lilianweng.github.io/posts/2021-07-11-diffusion-models/#simplification
[5] https://openai.com/blog/generative-models/
[6] Nichol, Alexander Quinn, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." International Conference on Machine Learning. PMLR, 2021.
[7] Kingma, Diederik P., and Max Welling. "Auto-encoding variational bayes." arXiv preprint arXiv:1312.6114 (2013).
[8] Hinton, Geoffrey E., and Ruslan R. Salakhutdinov. "Reducing the dimensionality of data with neural networks." science 313.5786 (2006): 504-507.
[9] Ronneberger O, Fischer P, Brox T. U-net: Convolutional networks for biomedical image segmentation[C]//International Conference on Medical image computing and computer-assisted intervention. Springer, Cham, 2015: 234-241.
[10] Long, Jonathan, Evan Shelhamer, and Trevor Darrell. "Fully convolutional networks for semantic segmentation." Proceedings of the IEEE conference on computer vision and pattern recognition. 2015.
[11] Luo, Calvin. "Understanding diffusion models: A unified perspective." arXiv preprint arXiv:2208.11970 (2022).
[12] Zagoruyko, Sergey, and Nikos Komodakis. "Wide residual networks." arXiv preprint arXiv:1605.07146 (2016).
[13] https://github.com/lucidrains/denoising-diffusion-pytorch
[14] Liu, Zhuang, et al. "A convnet for the 2020s." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.
[15] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems 30 (2017).
[16] Wu, Yuxin, and Kaiming He. "Group normalization." Proceedings of the European conference on computer vision (ECCV). 2018.
轉載請註明出處,本文鏈接:https://www.uj5u.com/qita/548613.html
標籤:其他
