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

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

這篇工作提出了基于“score”的生成式模型，和DDPM擴(kuò)散模型有著千絲萬(wàn)縷的聯(lián)系，后面宋飏博士發(fā)表中 ICLR2021 上的另一篇工作將 NCSN 和 DDPM 在 SDE 視角下進(jìn)行了很好的統(tǒng)一。

本文經(jīng)自動(dòng)駕駛之心公眾號(hào)授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

簡(jiǎn)介

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

宋飏博士在博客中提到，score-based generative model 的主要優(yōu)點(diǎn)是：

有媲美 GAN 的生成質(zhì)量，但是不需要對(duì)抗訓(xùn)練
靈活的模型架構(gòu)、精確的對(duì)數(shù)似然計(jì)算
可以用于解決逆問(wèn)題，并且不需要重新訓(xùn)練網(wǎng)絡(luò)

從生成模型到 score-based model

生成模型

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

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

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

為了實(shí)現(xiàn)對(duì)原數(shù)據(jù)分布的擬合，可以通過(guò)最大似然估計(jì)來(lái)訓(xùn)練這個(gè)公式要求

這個(gè)公式要求是一個(gè)歸一化的概率密度函數(shù)——這一點(diǎn)其實(shí)會(huì)給我們帶來(lái)很多麻煩。因?yàn)槿绻笫菤w一化的，就必須知道歸一化常數(shù)，但是對(duì)于任意的來(lái)說(shuō)往往是一個(gè)難以求解的量。因此，為了讓上述基于最大似然估計(jì)的訓(xùn)練能夠進(jìn)行，似然模型往往只能采用特定的網(wǎng)絡(luò)結(jié)構(gòu)（比如在自回歸模型中采用常規(guī)卷積結(jié)構(gòu)，在 normalizing flow models 中采用 invertible networks 結(jié)構(gòu)）讓可解，或者通過(guò)其他方式去近似（比如 VAE 中的變分推理，contrastive divergence 中的 MCMC），但是往往計(jì)算復(fù)雜度很高。

為了解決上述問(wèn)題，宋飏博士等人提出了 score-based model，基本思想是通過(guò)對(duì)概率密度函數(shù)的梯度進(jìn)行建模，而不是直接對(duì)概率密度本身進(jìn)行建模，來(lái)間接表征數(shù)據(jù)的分布，并規(guī)避難以求解的問(wèn)題。

score-based model

“score function”，或者簡(jiǎn)單的稱為“score”，其實(shí)就是論文題目中提到的“Gradients of the Data Distribution”，更具體來(lái)說(shuō)其實(shí)是概率密度函數(shù)的對(duì)數(shù)的梯度，即：用來(lái)對(duì)其進(jìn)行建模/擬合的模型就叫做——score-based model，記作。和上面直接建模概率密度函數(shù)p(x)時(shí)不同，score-based model 的優(yōu)勢(shì)是不會(huì)受到歸一化常數(shù)的制約，可以直接用上面提到的可學(xué)習(xí)的實(shí)值函數(shù)（energy-based model）來(lái)建模，不用考慮歸一化常數(shù)的問(wèn)題（因?yàn)榍髮?dǎo)時(shí)相關(guān)項(xiàng)變?yōu)榱?0）：這樣我們就可以靈活的選擇不同類型的網(wǎng)絡(luò)架構(gòu)，不受歸一化常數(shù)是否易求解的制約。

我們也可以給 score 的物理意義一個(gè)直觀的解釋——它是數(shù)據(jù)的對(duì)數(shù)概率密度增長(zhǎng)最快的方向。

一個(gè)數(shù)據(jù)分布的梯度場(chǎng)

觀察上面一個(gè)數(shù)據(jù)分布的對(duì)數(shù)概率密度梯度場(chǎng)可以發(fā)現(xiàn)，梯度方向匯聚到了兩個(gè)中心區(qū)域，這代表了這兩個(gè)區(qū)域是對(duì)數(shù)概率密度最大的地方——也就是此數(shù)據(jù)域的分布的核心、最能代表這個(gè)數(shù)據(jù)域先驗(yàn)的區(qū)域。

回想一下，對(duì)于生成模型來(lái)說(shuō)，我們期望生成的圖像采樣于哪些概率分布區(qū)域呢？當(dāng)然就是概率密度最大的區(qū)域！因?yàn)檫@些區(qū)域是最符合數(shù)據(jù)域先驗(yàn)的，有助于生成“最真”的圖像。所以如果我們能夠通過(guò) score-based model 估計(jì)出來(lái)數(shù)據(jù)分布的 score，并實(shí)現(xiàn)在其指示的概率密度大的地方進(jìn)行采樣生成，就可以得到理想的生成結(jié)果。

如何用 score-based model 生成樣本——朗之萬(wàn)采樣

假設(shè)我們已經(jīng)獲得了 score-based model，我們要如何實(shí)現(xiàn)采樣生成的步驟呢？

1 朗之萬(wàn)采樣示意圖

朗之萬(wàn)動(dòng)力學(xué)方程是描述物理學(xué)中布朗運(yùn)動(dòng)的微分方程，相關(guān)推導(dǎo)涉及到伊藤積分等一系列隨機(jī)過(guò)程相關(guān)的數(shù)學(xué)，還是比較復(fù)雜的。

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

上面我們已經(jīng)解決了從訓(xùn)練好的 score-based model 進(jìn)行采樣生成樣本的問(wèn)題，但是如何設(shè)計(jì)和訓(xùn)練 score-based model 則是這一切的基礎(chǔ)和關(guān)鍵。

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

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

這個(gè)公式里面的是比較麻煩的地方，因?yàn)槲覀冸m然有訓(xùn)練集的所有數(shù)據(jù)，但是很難直接算出來(lái)它的解析的概率分布，所以需要想辦法避開(kāi)它，經(jīng)過(guò)一系列推導(dǎo)，我們可以把 score matching 目標(biāo)函數(shù)簡(jiǎn)化為：

上式中，我們需要計(jì)算向量函數(shù)的偏導(dǎo)，即計(jì)算雅克比（Jocabian）矩陣，對(duì)的每一維輸出都要計(jì)算一次導(dǎo)數(shù)（如下圖，一般是個(gè)神經(jīng)網(wǎng)絡(luò)，可以用自動(dòng)微分功能計(jì)算），運(yùn)算量比較大，限制了 score matching 對(duì)深度網(wǎng)絡(luò)和高維數(shù)據(jù)的處理能力。

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

還有一種不用計(jì)算導(dǎo)數(shù)的 score matching 方法——Denoising Score Matching (DSM)，這種方法是作者對(duì) NCSN 采取的默認(rèn)方法。詳細(xì)推導(dǎo)見(jiàn) DSM 原論文的 Sec4 和 Appendix，這里簡(jiǎn)述下基本思路：

在這個(gè)分布中，原來(lái)的x相當(dāng)于均值，相當(dāng)于方差。我們能夠直接寫出這個(gè)分布的解析形式（高斯分布的概率密度函數(shù)），并求出：

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

到目前為止，我們從概率分布的表示、估計(jì)到采樣的過(guò)程基本已經(jīng)完成了，整個(gè)“Score-based generative modeling”的過(guò)程可以總結(jié)如下圖:

Score-based generative modeling: score matching + Langevin dynamics

解決實(shí)現(xiàn)中的存在的問(wèn)題——Annealed Langevin dynamics

雖然經(jīng)過(guò)上面的設(shè)計(jì)與推導(dǎo)，score-based model 的基本方案已經(jīng)有了，但是在具體實(shí)現(xiàn)中還有一些問(wèn)題需要解決。這些問(wèn)題的根源主要在于“流形假設(shè)”（manifold hypothesis）。

流形假設(shè)指出，現(xiàn)實(shí)世界的數(shù)據(jù)分布傾向于聚集在內(nèi)嵌于一個(gè)高維空間（也叫 ambinet space）的低維流形上。這個(gè)假設(shè)對(duì)于大部分?jǐn)?shù)據(jù)集（自然也包括這里用來(lái)訓(xùn)練 score-based model 的數(shù)據(jù)集）都是成立的，它也是流形學(xué)習(xí)的理論基礎(chǔ)。

這些問(wèn)題可能導(dǎo)致 score-based model 在訓(xùn)練時(shí) loss 不收斂，從而無(wú)法得到理想的生成效果。下面對(duì)這兩個(gè)問(wèn)題產(chǎn)生的具體影響做進(jìn)一步介紹，并給出作者的解決方案。

低概率密度區(qū)域 score 估計(jì)不準(zhǔn)確

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

作者發(fā)現(xiàn)上述的基礎(chǔ)方案在實(shí)際測(cè)試時(shí)，在數(shù)據(jù)密度低的區(qū)域，score 估計(jì)會(huì)不準(zhǔn)確。這是因?yàn)閿?shù)據(jù)分布的低概率密度區(qū)域意味著對(duì)應(yīng)的數(shù)據(jù)樣本比較少、權(quán)重較低，所以在有限的訓(xùn)練樣本下（往往無(wú)法充分覆蓋到概率密度很低的區(qū)域的數(shù)據(jù)樣本），訓(xùn)練過(guò)程往往也很難充分的學(xué)習(xí)/擬合到相應(yīng)區(qū)域的準(zhǔn)確的 score。

此外，從目標(biāo)函數(shù)出發(fā)，可以看到上面的目標(biāo)函數(shù)中p(x)其實(shí)是作為距離的權(quán)重因子的，所以對(duì)于數(shù)據(jù)分布的概率密度較小的區(qū)域,p(x)也比較小，也就意味著對(duì)應(yīng)的樣本在目標(biāo)函數(shù)中的貢獻(xiàn)比較低。

另外，由于朗之萬(wàn)動(dòng)力學(xué)采樣步驟的初始值是從任意分布中隨機(jī)產(chǎn)生的，它大概率會(huì)落在真實(shí)數(shù)據(jù)分布 p(x)（以及訓(xùn)練集數(shù)據(jù)分布的低密度區(qū)域甚至是分布之外，這種從一開(kāi)始就出現(xiàn)的不匹配也會(huì)導(dǎo)致朗之萬(wàn)采樣過(guò)程難以生成理想的樣本。

“忽視”數(shù)據(jù)分布的比例

使用朗之萬(wàn)采樣會(huì)忽視掉其實(shí)右上方的區(qū)域應(yīng)該權(quán)重更大

作者還發(fā)現(xiàn)，直接用朗之萬(wàn)采樣會(huì)導(dǎo)致“忽視”數(shù)據(jù)分布比例的問(wèn)題，從而讓學(xué)到的分布和實(shí)際的分布可能不符。如上圖所示，仍以雙高斯混合分布為例，在實(shí)際的數(shù)據(jù)分布 (a) 中，右上部分的分布對(duì)應(yīng)的樣本比較多，左下部分的比較少；而生成的結(jié)果 (b) 確是兩種分布對(duì)應(yīng)的樣本一樣多，并沒(méi)有反映出原始各分布之間的關(guān)系。這是因?yàn)槔手f(wàn)采樣過(guò)程中，在根據(jù)數(shù)據(jù)梯度“前進(jìn)”時(shí)，選擇向左下的高密度區(qū)域走還是向右上的高密度區(qū)域走，往往可能是隨機(jī)的，這也是朗之萬(wàn)采樣使用貪心思想的問(wèn)題?。

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

設(shè)計(jì)思路

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

對(duì)于第一個(gè)問(wèn)題——“低數(shù)據(jù)區(qū)域 score 難以計(jì)算”：當(dāng)用高斯噪聲去擾動(dòng)訓(xùn)練數(shù)據(jù)時(shí)，一方面，因?yàn)榧拥臄_動(dòng)噪聲服從高斯分布，能夠覆蓋整個(gè)概率空間的，所以在一定程度上可以“彌補(bǔ)”原始訓(xùn)練集分布范圍有限的問(wèn)題，讓其不至于被局限在一個(gè)低維流形上；另一方面，如果所加噪聲足夠大，也會(huì)讓新產(chǎn)生的加噪數(shù)據(jù)集的分布中，原來(lái)低概率密度的區(qū)域被填充，從而能夠提高該區(qū)域 score 估計(jì)的準(zhǔn)確性。通過(guò)使用不同的噪聲強(qiáng)度進(jìn)行加噪，最終可以獲得一系列噪聲擾動(dòng)的數(shù)據(jù)分布，它們將收斂到真實(shí)的數(shù)據(jù)分布。

被擾動(dòng)后的數(shù)據(jù)

對(duì)于第二個(gè)問(wèn)題——“忽視數(shù)據(jù)分布的比例”，在加噪擾動(dòng)后，數(shù)據(jù)分布比例高的區(qū)域，被采樣的概率也會(huì)增大，這也就意味著采樣的時(shí)候更可能向比例高的分布接近，從而能夠讓生成模型的分布和原數(shù)據(jù)分布有相同的比例特性?。

從左往右依次為原始數(shù)據(jù)分布，沒(méi)加噪聲擾動(dòng)的擬合結(jié)果和加了噪聲擾動(dòng)之后的結(jié)果

具體實(shí)現(xiàn)

具體實(shí)現(xiàn)中應(yīng)該如何控制所加噪聲的程度呢？噪聲加的大，能有效“填充”概率密度低的區(qū)域，緩解 score 估計(jì)不準(zhǔn)確的問(wèn)題，但是這樣同時(shí)也會(huì)導(dǎo)致加噪之后的數(shù)據(jù)分布嚴(yán)重偏離原始數(shù)據(jù)分布；反之，噪聲加的小，雖然能夠保持加噪前后數(shù)據(jù)分布的一致性，但是不能很好的“填充”低密度區(qū)域，解決前述存在的問(wèn)題。

引入不同程度的噪聲擾動(dòng)原始數(shù)據(jù)（第一行），然后聯(lián)合估計(jì)它們的score（第二行）

得到訓(xùn)練完成的 noise-conditional score-based model之后，就可以按照 i=L,L?1,…,1的遞減次序，串聯(lián)進(jìn)行多輪不同噪聲程度下的朗之萬(wàn)動(dòng)力學(xué)采樣，并最終生成新樣本了。這種采樣方法因?yàn)樵酵螅肼曉叫。ㄍ嘶穑砸步型嘶鹄手f(wàn)動(dòng)力學(xué)采樣（annealed Langevin dynamics）。

2 利用噪聲尺度逐漸減小的退火朗之萬(wàn)采樣，生成的樣本逐漸落到目標(biāo)分布區(qū)域（從右往左）

直觀來(lái)看，按照噪聲遞減的順序來(lái)采樣是因?yàn)?，一開(kāi)始噪聲先大一些，能夠讓數(shù)據(jù)先移動(dòng)到高密度區(qū)域，之后噪聲再小一些，可以讓數(shù)據(jù)在高密度區(qū)域能更精準(zhǔn)地移動(dòng)到更加符合數(shù)據(jù)分布的位置。

總結(jié)一下退火朗之萬(wàn)動(dòng)力學(xué)采樣（annealed Langevin dynamics）的算法流程：

退火朗之萬(wàn)動(dòng)力學(xué)采樣（annealed Langevin dynamics）的算法流程

Trick

一些具體實(shí)現(xiàn)中的技巧

效果展示

可以看到，原始的沒(méi)有加噪聲擾動(dòng)的 sliced score matching 實(shí)際訓(xùn)練時(shí)容易出現(xiàn) loss 不收斂的問(wèn)題（下圖左），而加噪聲擾動(dòng)之后，就能很好的收斂了（下圖右）

對(duì)數(shù)據(jù)加噪（右）loss就很穩(wěn)定

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

3CelebA數(shù)據(jù)集

4CIFAR-10數(shù)據(jù)集

討論

DDPM 和 NCSN 有什么區(qū)別？

也就是說(shuō) score 的方向就是噪聲?的反方向，在采樣生成時(shí)沿著 score 的方向走，就是往噪聲的反方向走，這樣就會(huì)走回原來(lái)數(shù)據(jù)樣本的位置。所以 NCSN 也可以看成是一個(gè)去噪的過(guò)程。

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