什么是diffusion model? 它為什么好用？

作者：凌晗 2024-01-23 11:31:24

這篇工作提出了基于“score”的生成式模型，和DDPM擴散模型有著千絲萬縷的聯系，后面宋飏博士發表中 ICLR2021 上的另一篇工作將 NCSN 和 DDPM 在 SDE 視角下進行了很好的統一。

本文經自動駕駛之心公眾號授權轉載，轉載請聯系出處。

簡介

NCSN （Noise Conditional Score Networks）來自于宋飏博士發表在 NeurIPS2019 上面的文章“Generative Modeling by Estimating Gradients of the Data Distribution”，也是推動擴散模型領域興起的重要工作之一，比 DDPM 發表的還要早。這篇工作提出了基于“score”的生成式模型，和DDPM擴散模型有著千絲萬縷的聯系，后面宋飏博士發表中 ICLR2021 上的另一篇工作將 NCSN 和 DDPM 在 SDE 視角下進行了很好的統一。

宋飏博士在博客中提到，score-based generative model 的主要優點是：

有媲美 GAN 的生成質量，但是不需要對抗訓練
靈活的模型架構、精確的對數似然計算
可以用于解決逆問題，并且不需要重新訓練網絡

從生成模型到 score-based model

生成模型

首先回顧一下，生成模型是什么？

假設我們從一個未知的數據分布p(x)中，獨立的采樣出了一系列的樣本{,,...,}，也就是我們常說的“數據集”。生成模型的目的就是，從這些含有有限樣本的“數據集”出發，去擬合原來的數據分布p(x),從而讓我們能夠繼續合成更多類似的樣本。

為了實現這個目標，我們首先需要想辦法去建模并估計數據的概率分布p(x)。一種方式是，類似于似然模型（likelihood-based models）那樣，直接定義出來一個概率密度函數（probability density function, p.d.f.）或者概率分布函數（probability mass function, p.m.f.）。以概率密度函數為例：

為了實現對原數據分布的擬合，可以通過最大似然估計來訓練這個公式要求

這個公式要求是一個歸一化的概率密度函數——這一點其實會給我們帶來很多麻煩。因為如果要求是歸一化的，就必須知道歸一化常數，但是對于任意的來說往往是一個難以求解的量。因此，為了讓上述基于最大似然估計的訓練能夠進行，似然模型往往只能采用特定的網絡結構（比如在自回歸模型中采用常規卷積結構，在 normalizing flow models 中采用 invertible networks 結構）讓可解，或者通過其他方式去近似（比如 VAE 中的變分推理，contrastive divergence 中的 MCMC），但是往往計算復雜度很高。

為了解決上述問題，宋飏博士等人提出了 score-based model，基本思想是通過對概率密度函數的梯度進行建模，而不是直接對概率密度本身進行建模，來間接表征數據的分布，并規避難以求解的問題。

score-based model

“score function”，或者簡單的稱為“score”，其實就是論文題目中提到的“Gradients of the Data Distribution”，更具體來說其實是概率密度函數的對數的梯度，即：用來對其進行建模/擬合的模型就叫做——score-based model，記作。和上面直接建模概率密度函數p(x)時不同，score-based model 的優勢是不會受到歸一化常數的制約，可以直接用上面提到的可學習的實值函數（energy-based model）來建模，不用考慮歸一化常數的問題（因為求導時相關項變為了 0）：這樣我們就可以靈活的選擇不同類型的網絡架構，不受歸一化常數是否易求解的制約。

我們也可以給 score 的物理意義一個直觀的解釋——它是數據的對數概率密度增長最快的方向。

一個數據分布的梯度場

觀察上面一個數據分布的對數概率密度梯度場可以發現，梯度方向匯聚到了兩個中心區域，這代表了這兩個區域是對數概率密度最大的地方——也就是此數據域的分布的核心、最能代表這個數據域先驗的區域。

回想一下，對于生成模型來說，我們期望生成的圖像采樣于哪些概率分布區域呢？當然就是概率密度最大的區域！因為這些區域是最符合數據域先驗的，有助于生成“最真”的圖像。所以如果我們能夠通過 score-based model 估計出來數據分布的 score，并實現在其指示的概率密度大的地方進行采樣生成，就可以得到理想的生成結果。

如何用 score-based model 生成樣本——朗之萬采樣

假設我們已經獲得了 score-based model，我們要如何實現采樣生成的步驟呢？

1 朗之萬采樣示意圖

朗之萬動力學方程是描述物理學中布朗運動的微分方程，相關推導涉及到伊藤積分等一系列隨機過程相關的數學，還是比較復雜的。

如何獲得 score-based model——score matching

上面我們已經解決了從訓練好的 score-based model 進行采樣生成樣本的問題，但是如何設計和訓練 score-based model 則是這一切的基礎和關鍵。

首先寫出目標函數（損失函數），和似然模型類似，score-based model 也可以把最小化模型和數據分布之間的 Fisher divergence 作為訓練目標：

由于這個目標函數里面存在未知的先驗分布p(x)，所以沒法直接計算。不過有一類叫做“score matching”的方法可以解決這個問題。和訓練似然模型類似（知道歸一化常數的前提下），score matching 方法可以在訓練集上通過隨機梯度下降（stochastic gradient descent, SGD），而不需要對抗訓練，就可以估計出它的目標函數（score matching objective）：

這個公式里面的是比較麻煩的地方，因為我們雖然有訓練集的所有數據，但是很難直接算出來它的解析的概率分布，所以需要想辦法避開它，經過一系列推導，我們可以把 score matching 目標函數簡化為：

上式中，我們需要計算向量函數的偏導，即計算雅克比（Jocabian）矩陣，對的每一維輸出都要計算一次導數（如下圖，一般是個神經網絡，可以用自動微分功能計算），運算量比較大，限制了 score matching 對深度網絡和高維數據的處理能力。

論文中介紹了兩個能夠用于大尺度場景的 score matching 方法，一個是同年（2019）宋飏博士提出并發表的 Sliced Score Matching——使用一個隨機向量v把雅克比矩陣降到一維變成向量：

還有一種不用計算導數的 score matching 方法——Denoising Score Matching (DSM)，這種方法是作者對 NCSN 采取的默認方法。詳細推導見 DSM 原論文的 Sec4 和 Appendix，這里簡述下基本思路：

在這個分布中，原來的x相當于均值，相當于方差。我們能夠直接寫出這個分布的解析形式（高斯分布的概率密度函數），并求出：

從這里我們可以看出，SDE/NCSN 中用網絡去擬合的 score，和 DDPM 中用網絡去擬合的噪聲，其實本質是一樣的。score 所對應的梯度方向，正是的方向，即從加噪圖像還原回清晰圖像的方向，也就是 Denoising Score Matching 和 DDPM 中“Denoise”的含義。

到目前為止，我們從概率分布的表示、估計到采樣的過程基本已經完成了，整個“Score-based generative modeling”的過程可以總結如下圖:

Score-based generative modeling: score matching + Langevin dynamics

解決實現中的存在的問題——Annealed Langevin dynamics

雖然經過上面的設計與推導，score-based model 的基本方案已經有了，但是在具體實現中還有一些問題需要解決。這些問題的根源主要在于“流形假設”（manifold hypothesis）。

流形假設指出，現實世界的數據分布傾向于聚集在內嵌于一個高維空間（也叫 ambinet space）的低維流形上。這個假設對于大部分數據集（自然也包括這里用來訓練 score-based model 的數據集）都是成立的，它也是流形學習的理論基礎。

這些問題可能導致 score-based model 在訓練時 loss 不收斂，從而無法得到理想的生成效果。下面對這兩個問題產生的具體影響做進一步介紹，并給出作者的解決方案。

低概率密度區域 score 估計不準確

score在低密度區域很難估計（圖中以雙高斯混合分布為例）

作者發現上述的基礎方案在實際測試時，在數據密度低的區域，score 估計會不準確。這是因為數據分布的低概率密度區域意味著對應的數據樣本比較少、權重較低，所以在有限的訓練樣本下（往往無法充分覆蓋到概率密度很低的區域的數據樣本），訓練過程往往也很難充分的學習/擬合到相應區域的準確的 score。

此外，從目標函數出發，可以看到上面的目標函數中p(x)其實是作為距離的權重因子的，所以對于數據分布的概率密度較小的區域,p(x)也比較小，也就意味著對應的樣本在目標函數中的貢獻比較低。

另外，由于朗之萬動力學采樣步驟的初始值是從任意分布中隨機產生的，它大概率會落在真實數據分布 p(x)（以及訓練集數據分布的低密度區域甚至是分布之外，這種從一開始就出現的不匹配也會導致朗之萬采樣過程難以生成理想的樣本。

“忽視”數據分布的比例

使用朗之萬采樣會忽視掉其實右上方的區域應該權重更大

作者還發現，直接用朗之萬采樣會導致“忽視”數據分布比例的問題，從而讓學到的分布和實際的分布可能不符。如上圖所示，仍以雙高斯混合分布為例，在實際的數據分布 (a) 中，右上部分的分布對應的樣本比較多，左下部分的比較少；而生成的結果 (b) 確是兩種分布對應的樣本一樣多，并沒有反映出原始各分布之間的關系。這是因為朗之萬采樣過程中，在根據數據梯度“前進”時，選擇向左下的高密度區域走還是向右上的高密度區域走，往往可能是隨機的，這也是朗之萬采樣使用貪心思想的問題?。

解決辦法——Noise Conditional Score Networks （NCSN）

設計思路

為了解決上述問題，作者提出了升級版方法 —— Noise Conditional Score Networks （NCSN）& Annealed Langevin dynamics，其核心是用高斯噪聲去擾動訓練數據，然后估計不同噪聲擾動下的數據分布的 score，并用于朗之萬采樣。這么做有什么好處呢？

對于第一個問題——“低數據區域 score 難以計算”：當用高斯噪聲去擾動訓練數據時，一方面，因為加的擾動噪聲服從高斯分布，能夠覆蓋整個概率空間的，所以在一定程度上可以“彌補”原始訓練集分布范圍有限的問題，讓其不至于被局限在一個低維流形上；另一方面，如果所加噪聲足夠大，也會讓新產生的加噪數據集的分布中，原來低概率密度的區域被填充，從而能夠提高該區域 score 估計的準確性。通過使用不同的噪聲強度進行加噪，最終可以獲得一系列噪聲擾動的數據分布，它們將收斂到真實的數據分布。

被擾動后的數據

對于第二個問題——“忽視數據分布的比例”，在加噪擾動后，數據分布比例高的區域，被采樣的概率也會增大，這也就意味著采樣的時候更可能向比例高的分布接近，從而能夠讓生成模型的分布和原數據分布有相同的比例特性?。

從左往右依次為原始數據分布，沒加噪聲擾動的擬合結果和加了噪聲擾動之后的結果

具體實現

具體實現中應該如何控制所加噪聲的程度呢？噪聲加的大，能有效“填充”概率密度低的區域，緩解 score 估計不準確的問題，但是這樣同時也會導致加噪之后的數據分布嚴重偏離原始數據分布；反之，噪聲加的小，雖然能夠保持加噪前后數據分布的一致性，但是不能很好的“填充”低密度區域，解決前述存在的問題。

引入不同程度的噪聲擾動原始數據（第一行），然后聯合估計它們的score（第二行）

得到訓練完成的 noise-conditional score-based model之后，就可以按照 i=L,L?1,…,1的遞減次序，串聯進行多輪不同噪聲程度下的朗之萬動力學采樣，并最終生成新樣本了。這種采樣方法因為越往后，噪聲越小（退火），所以也叫退火朗之萬動力學采樣（annealed Langevin dynamics）。

2 利用噪聲尺度逐漸減小的退火朗之萬采樣，生成的樣本逐漸落到目標分布區域（從右往左）

直觀來看，按照噪聲遞減的順序來采樣是因為，一開始噪聲先大一些，能夠讓數據先移動到高密度區域，之后噪聲再小一些，可以讓數據在高密度區域能更精準地移動到更加符合數據分布的位置。

總結一下退火朗之萬動力學采樣（annealed Langevin dynamics）的算法流程：

退火朗之萬動力學采樣（annealed Langevin dynamics）的算法流程

Trick

一些具體實現中的技巧

效果展示

可以看到，原始的沒有加噪聲擾動的 sliced score matching 實際訓練時容易出現 loss 不收斂的問題（下圖左），而加噪聲擾動之后，就能很好的收斂了（下圖右）

對數據加噪（右）loss就很穩定

基于退火朗之萬采樣，用 NCSN 生成新樣本的示例：

3CelebA數據集

4CIFAR-10數據集

討論

DDPM 和 NCSN 有什么區別？

也就是說 score 的方向就是噪聲?的反方向，在采樣生成時沿著 score 的方向走，就是往噪聲的反方向走，這樣就會走回原來數據樣本的位置。所以 NCSN 也可以看成是一個去噪的過程。

宋飏在 ICLR 2021 上獲得 Outstanding Paper Award 的論文“Score-based generative modeling through stochastic differential equations”從 SDE 的視角對 NCSN 和 DDPM 給出了更加完善的分析，證明了它們的內在統一性。