Diffusion擴散模型

關于擴散模型（Diffusion Models）有很多種理解，本文的介紹是基于denoising diffusion probabilistic model （DDPM），DDPM已經在（無）條件圖像/音頻/視頻生成領域取得了較多顯著的成果，現有的比較受歡迎的的例子包括由OpenAI主導的GLIDE和DALL-E 2、由海德堡大學主導的潛在擴散和由Google Brain主導的圖像生成。

實際上生成模型的擴散概念已經在（Sohl-Dickstein et al., 2015）中介紹過。然而，直到（Song et al., 2019）（斯坦福大學）和（Ho et al., 2020）（在Google Brain）才各自獨立地改進了這種方法。

本文是在Phil Wang基于PyTorch框架的復現的基礎上（而它本身又是基于TensorFlow實現），遷移到MindSpore AI框架上實現的。

模型簡介

什么是Diffusion Model？

如果將Diffusion與其他生成模型（如Normalizing Flows、GAN或VAE）進行比較，它并沒有那么復雜，它們都將噪聲從一些簡單分布轉換為數據樣本，Diffusion也是從純噪聲開始通過一個神經網絡學習逐步去噪，最終得到一個實際圖像。 Diffusion對于圖像的處理包括以下兩個過程：

• 我們選擇的固定（或預定義）正向擴散過程?𝑞𝑞?：它逐漸將高斯噪聲添加到圖像中，直到最終得到純噪聲

• 一個學習的反向去噪的擴散過程?𝑝𝜃𝑝𝜃?：通過訓練神經網絡從純噪聲開始逐漸對圖像去噪，直到最終得到一個實際的圖像

由?𝑡𝑡?索引的正向和反向過程都發生在某些有限時間步長?𝑇𝑇（DDPM作者使用?𝑇=1000𝑇=1000）內。從𝑡=0𝑡=0開始，在數據分布中采樣真實圖像?𝐱0𝑥0（本文使用一張來自ImageNet的貓圖像形象的展示了diffusion正向添加噪聲的過程），正向過程在每個時間步長?𝑡𝑡?都從高斯分布中采樣一些噪聲，再添加到上一個時刻的圖像中。假定給定一個足夠大的?𝑇𝑇?和一個在每個時間步長添加噪聲的良好時間表，您最終會在?𝑡=𝑇𝑡=𝑇?通過漸進的過程得到所謂的各向同性的高斯分布。

擴散模型實現原理

Diffusion 前向過程

所謂前向過程，即向圖片上加噪聲的過程。雖然這個步驟無法做到圖片生成，但這是理解diffusion model以及構建訓練樣本至關重要的一步。 首先我們需要一個可控的損失函數，并運用神經網絡對其進行優化。

設?𝑞(𝑥0)𝑞(𝑥0)?是真實數據分布，由于?𝑥0～𝑞(𝑥0)𝑥0～𝑞(𝑥0)?，所以我們可以從這個分布中采樣以獲得圖像?𝑥0𝑥0?。接下來我們定義前向擴散過程?𝑞(𝑥𝑡|𝑥𝑡?1)𝑞(𝑥𝑡|𝑥𝑡?1)?，在前向過程中我們會根據已知的方差?0<𝛽1<𝛽2<...<𝛽𝑇<10<𝛽1<𝛽2<...<𝛽𝑇<1?在每個時間步長 t 添加高斯噪聲，由于前向過程的每個時刻 t 只與時刻 t-1 有關，所以也可以看做馬爾科夫過程：

回想一下，正態分布（也稱為高斯分布）由兩個參數定義：平均值?𝜇𝜇?和方差?𝜎2≥0𝜎2≥0?。基本上，在每個時間步長?𝑡𝑡?處的產生的每個新的（輕微噪聲）圖像都是從條件高斯分布中繪制的，其中

我們可以通過采樣然后設置

請注意，?𝛽𝑡𝛽𝑡?在每個時間步長?𝑡𝑡?（因此是下標）不是恒定的：事實上，我們定義了一個所謂的“動態方差”的方法，使得每個時間步長的?𝛽𝑡𝛽𝑡?可以是線性的、二次的、余弦的等（有點像動態學習率方法）。

因此，如果我們適當設置時間表，從?𝐱0𝑥0?開始，我們最終得到?𝐱1,...,𝐱𝑡,...,𝐱𝑇𝑥1,...,𝑥𝑡,...,𝑥𝑇，即隨著?𝑡𝑡?的增大?𝐱𝑡𝑥𝑡?會越來越接近純噪聲，而?𝐱𝑇𝑥𝑇?就是純高斯噪聲。

那么，如果我們知道條件概率分布?𝑝(𝐱𝑡?1|𝐱𝑡)𝑝(𝑥𝑡?1|𝑥𝑡)?，我們就可以反向運行這個過程：通過采樣一些隨機高斯噪聲?𝐱𝑇𝑥𝑇，然后逐漸去噪它，最終得到真實分布?𝐱0𝑥0?中的樣本。但是，我們不知道條件概率分布?𝑝(𝐱𝑡?1|𝐱𝑡)𝑝(𝑥𝑡?1|𝑥𝑡)?。這很棘手，因為需要知道所有可能圖像的分布，才能計算這個條件概率。

Diffusion 逆向過程

為了解決上述問題，我們將利用神經網絡來近似（學習）這個條件概率分布?𝑝𝜃(𝐱𝑡?1|𝐱𝑡)𝑝𝜃(𝑥𝑡?1|𝑥𝑡)?, 其中?𝜃𝜃?是神經網絡的參數。如果說前向過程(forward)是加噪的過程，那么逆向過程(reverse)就是diffusion的去噪推斷過程，而通過神經網絡學習并表示?𝑝𝜃(𝐱𝑡?1|𝐱𝑡)𝑝𝜃(𝑥𝑡?1|𝑥𝑡)?的過程就是Diffusion 逆向去噪的核心。

現在，我們知道了需要一個神經網絡來學習逆向過程的（條件）概率分布。我們假設這個反向過程也是高斯的，任何高斯分布都由2個參數定義：

• 由?𝜇𝜃𝜇𝜃?參數化的平均值

• 由?𝜇𝜃𝜇𝜃?參數化的方差

綜上，我們可以將逆向過程公式化為

其中平均值和方差也取決于噪聲水平?𝑡𝑡?，神經網絡需要通過學習來表示這些均值和方差。

• 注意，DDPM的作者決定保持方差固定，讓神經網絡只學習（表示）這個條件概率分布的平均值?𝜇𝜃𝜇𝜃?。

• 本文我們同樣假設神經網絡只需要學習（表示）這個條件概率分布的平均值?𝜇𝜃𝜇𝜃?。

為了導出一個目標函數來學習反向過程的平均值，作者觀察到?𝑞𝑞?和?𝑝𝜃𝑝𝜃?的組合可以被視為變分自動編碼器(VAE)。因此，變分下界（也稱為ELBO）可用于最小化真值數據樣本?𝐱0𝑥0?的似然負對數（有關ELBO的詳細信息，請參閱VAE論文（Kingma等人，2013年）），該過程的ELBO是每個時間步長的損失之和?𝐿=𝐿0+𝐿1+...+𝐿𝑇𝐿=𝐿0+𝐿1+...+𝐿𝑇?，其中，每項的損失?𝐿𝑡𝐿𝑡?（除了?𝐿0𝐿0?）實際上是2個高斯分布之間的KL發散，可以明確地寫為相對于均值的L2-loss!

如Sohl-Dickstein等人所示，構建Diffusion正向過程的直接結果是我們可以在條件是?𝐱0𝑥0?（因為高斯和也是高斯）的情況下，在任意噪聲水平上采樣?𝐱𝑡𝑥𝑡?，而不需要重復應用?𝑞𝑞?去采樣?𝐱𝑡𝑥𝑡?，這非常方便。使用

我們就有

這意味著我們可以采樣高斯噪聲并適當地縮放它，然后將其添加到?𝐱0𝑥0?中，直接獲得?𝐱𝑡𝑥𝑡?。

請注意，𝛼ˉ𝑡𝛼ˉ𝑡?是已知?𝛽𝑡𝛽𝑡?方差計劃的函數，因此也是已知的，可以預先計算。這允許我們在訓練期間優化損失函數?𝐿𝐿?的隨機項。或者換句話說，在訓練期間隨機采樣?𝑡𝑡?并優化?𝐿𝑡𝐿𝑡?。

正如Ho等人所展示的那樣，這種性質的另一個優點是可以重新參數化平均值，使神經網絡學習（預測）構成損失的KL項中噪聲的附加噪聲。這意味著我們的神經網絡變成了噪聲預測器，而不是（直接）均值預測器。其中，平均值可以按如下方式計算：

最終的目標函數?𝐿𝑡𝐿𝑡?如下 （隨機步長 t 由?(𝜖～𝑁(0,𝐈))(𝜖～𝑁(0,𝐼))?給定）：

在這里，?𝐱0𝑥0?是初始（真實，未損壞）圖像，?𝜖𝜖?是在時間步長?𝑡𝑡?采樣的純噪聲，𝜖𝜃(𝐱𝑡,𝑡)𝜖𝜃(𝑥𝑡,𝑡)是我們的神經網絡。神經網絡是基于真實噪聲和預測高斯噪聲之間的簡單均方誤差（MSE）進行優化的。

訓練算法現在如下所示：

換句話說：

• 我們從真實未知和可能復雜的數據分布中隨機抽取一個樣本?𝑞(𝐱0)𝑞(𝑥0)

• 我們均勻地采樣11和𝑇𝑇之間的噪聲水平𝑡𝑡（即，隨機時間步長）

• 我們從高斯分布中采樣一些噪聲，并使用上面定義的屬性在?𝑡𝑡?時間步上破壞輸入

• 神經網絡被訓練以基于損壞的圖像?𝐱𝑡𝑥𝑡?來預測這種噪聲，即基于已知的時間表?𝐱𝑡𝑥𝑡?上施加的噪聲

實際上，所有這些都是在批數據上使用隨機梯度下降來優化神經網絡完成的。

U-Net神經網絡預測噪聲

神經網絡需要在特定時間步長接收帶噪聲的圖像，并返回預測的噪聲。請注意，預測噪聲是與輸入圖像具有相同大小/分辨率的張量。因此，從技術上講，網絡接受并輸出相同形狀的張量。那么我們可以用什么類型的神經網絡來實現呢？

這里通常使用的是非常相似的自動編碼器，您可能還記得典型的"深度學習入門"教程。自動編碼器在編碼器和解碼器之間有一個所謂的"bottleneck"層。編碼器首先將圖像編碼為一個稱為"bottleneck"的較小的隱藏表示，然后解碼器將該隱藏表示解碼回實際圖像。這迫使網絡只保留bottleneck層中最重要的信息。

在模型結構方面，DDPM的作者選擇了U-Net，出自（Ronneberger et al.，2015）（當時，它在醫學圖像分割方面取得了最先進的結果）。這個網絡就像任何自動編碼器一樣，在中間由一個bottleneck組成，確保網絡只學習最重要的信息。重要的是，它在編碼器和解碼器之間引入了殘差連接，極大地改善了梯度流。

可以看出，U-Net模型首先對輸入進行下采樣（即，在空間分辨率方面使輸入更小），之后執行上采樣。

構建Diffusion模型

DDPM論文表明擴散模型是（非）條件圖像有希望生成的方向。自那以后，diffusion得到了（極大的）改進，最明顯的是文本條件圖像生成。下面，我們列出了一些重要的（但遠非詳盡無遺的）后續工作：

• 改進的去噪擴散概率模型(Nichol et al., 2021)：發現學習條件分布的方差（除平均值外）有助于提高性能

• 用于高保真圖像生成的級聯擴散模型([Ho et al., 2021)：引入級聯擴散，它包括多個擴散模型的流水線，這些模型生成分辨率提高的圖像，用于高保真圖像合成

• 擴散模型在圖像合成上擊敗了GANs(Dhariwal et al., 2021)：表明擴散模型通過改進U-Net體系結構以及引入分類器指導，可以獲得優于當前最先進的生成模型的圖像樣本質量

• 無分類器擴散指南([Ho et al., 2021)：表明通過使用單個神經網絡聯合訓練條件和無條件擴散模型，不需要分類器來指導擴散模型

• 具有CLIP Latents (DALL-E 2) 的分層文本條件圖像生成 (Ramesh et al., 2022)：在將文本標題轉換為CLIP圖像嵌入之前使用，然后擴散模型將其解碼為圖像

• 具有深度語言理解的真實文本到圖像擴散模型（ImageGen）(Saharia et al., 2022)：表明將大型預訓練語言模型（例如T5）與級聯擴散結合起來，對于文本到圖像的合成很有效

請注意，此列表僅包括在撰寫本文，即2022年6月7日之前的重要作品。

目前，擴散模型的主要（也許唯一）缺點是它們需要多次正向傳遞來生成圖像（對于像GAN這樣的生成模型來說，情況并非如此）。然而，有正在進行中的研究表明只需要10個去噪步驟就能實現高保真生成。