如何改進(jìn)和加速擴(kuò)散模型采樣的方法2

這是 NVIDIA 研究人員如何改進(jìn)和加速擴(kuò)散模型采樣的系列文章的一部分，擴(kuò)散模型是一種新的、強(qiáng)大的生成模型。 Part 1 介紹了擴(kuò)散模型作為深層生成模型的一個(gè)強(qiáng)大類，并研究了它們?cè)诮鉀Q生成性學(xué)習(xí)三重困境中的權(quán)衡。

雖然擴(kuò)散模型同時(shí)滿足 生成性學(xué)習(xí)三位一體 的第一和第二個(gè)要求，即高樣本質(zhì)量和多樣性，但它們?nèi)狈鹘y(tǒng) GAN 的采樣速度。在這篇文章中，我們回顧了 NVIDIA 最近開發(fā)的三種技術(shù)，它們克服了擴(kuò)散模型中緩慢采樣的挑戰(zhàn)。

潛空間擴(kuò)散模型

擴(kuò)散模型的采樣速度較慢的主要原因之一是，從簡單的高斯噪聲分布到具有挑戰(zhàn)性的多模態(tài)數(shù)據(jù)分布的映射非常復(fù)雜。最近， NVIDIA 推出了 基于潛在分?jǐn)?shù)的生成模型 （ LSGM ），這是一種新的框架，可以在潛在空間而不是直接在數(shù)據(jù)空間中訓(xùn)練擴(kuò)散模型。

在 LSGM 中，我們利用變分自動(dòng)編碼器（ VAE ）框架將輸入數(shù)據(jù)映射到一個(gè)潛在空間，并在那里應(yīng)用擴(kuò)散模型。然后，擴(kuò)散模型的任務(wù)是對(duì)數(shù)據(jù)集潛在嵌入的分布進(jìn)行建模，這在本質(zhì)上比數(shù)據(jù)分布更簡單。

新的數(shù)據(jù)合成是通過從簡單的基分布中提取嵌入，然后迭代去噪，然后使用解碼器將該嵌入轉(zhuǎn)換為數(shù)據(jù)空間來實(shí)現(xiàn)的（圖 1 ）。

圖 1 基于潛在分?jǐn)?shù)的生成模型

圖 1 顯示，在基于潛在分?jǐn)?shù)的生成模型（ LSGM ）中：

合成速度

通過先用高斯先驗(yàn)對(duì) VAE 進(jìn)行預(yù)訓(xùn)練，可以使數(shù)據(jù)分布的潛在編碼接近高斯先驗(yàn)分布，這也是擴(kuò)散模型的基本分布。擴(kuò)散模型只需對(duì)剩余的不匹配進(jìn)行建模，從而形成一個(gè)簡單得多的模型，從中采樣變得更容易、更快。

可以相應(yīng)地調(diào)整潛在空間。例如，我們可以使用分層潛變量，并僅在其中的一個(gè)子集上或以較小的分辨率應(yīng)用擴(kuò)散模型，從而進(jìn)一步提高合成速度。

表現(xiàn)力

訓(xùn)練一個(gè)規(guī)則的擴(kuò)散模型可以看作是直接在數(shù)據(jù)上訓(xùn)練一個(gè)神經(jīng)網(wǎng)絡(luò)。然而，之前的研究發(fā)現(xiàn)，用潛在變量增強(qiáng)神經(jīng)微分方程以及其他類型的生成模型通常可以提高它們的表達(dá)能力。

我們期望通過將擴(kuò)散模型與潛在變量框架相結(jié)合，獲得類似的表現(xiàn)力收益。

定制編碼器和解碼器

在潛在空間中使用擴(kuò)散模型時(shí)，可以使用精心設(shè)計(jì)的編碼器和解碼器在潛在空間和數(shù)據(jù)空間之間映射，進(jìn)一步提高合成質(zhì)量。因此， LSGM 方法可以自然地應(yīng)用于非連續(xù)數(shù)據(jù)。

原則上， LSGM 可以通過使用編碼器和解碼器網(wǎng)絡(luò)，輕松地對(duì)文本、圖形和類似的離散或分類數(shù)據(jù)類型等數(shù)據(jù)進(jìn)行建模，這些網(wǎng)絡(luò)將這些數(shù)據(jù)轉(zhuǎn)換為連續(xù)的潛在表示并返回。

直接對(duì)數(shù)據(jù)進(jìn)行操作的常規(guī)擴(kuò)散模型無法輕松對(duì)此類數(shù)據(jù)類型進(jìn)行建模。標(biāo)準(zhǔn)擴(kuò)散框架僅適用于連續(xù)數(shù)據(jù)，這些數(shù)據(jù)可以逐漸擾動(dòng)并以有意義的方式生成。

后果

在實(shí)驗(yàn)上， LSGM 在 CIFAR-10 和 CelebA-HQ-256 這兩個(gè)廣泛使用的圖像生成基準(zhǔn)數(shù)據(jù)集上實(shí)現(xiàn)了最先進(jìn)的 Fr é chet Inception 距離（ FID ），這是量化視覺圖像質(zhì)量的標(biāo)準(zhǔn)度量。在這些數(shù)據(jù)集上，它優(yōu)于先前的生成模型，包括 GANs 。

在 CelebA-HQ-256 上， LSGM 的合成速度比以前的擴(kuò)散模型快兩個(gè)數(shù)量級(jí)。在對(duì) CelebA-HQ-256 數(shù)據(jù)建模時(shí)， LSGM 只需要 23 次神經(jīng)網(wǎng)絡(luò)調(diào)用，而之前在數(shù)據(jù)空間上訓(xùn)練的擴(kuò)散模型通常依賴數(shù)百次或數(shù)千次網(wǎng)絡(luò)調(diào)用。

臨界阻尼朗之萬擴(kuò)散

擴(kuò)散模型中的一個(gè)關(guān)鍵因素是固定前向擴(kuò)散過程，以逐漸擾動(dòng)數(shù)據(jù)。與數(shù)據(jù)本身一起，它唯一地決定了去噪模型學(xué)習(xí)的難度。因此，我們能否設(shè)計(jì)一種特別容易去噪的前向擴(kuò)散，從而實(shí)現(xiàn)更快、更高質(zhì)量的合成？

擴(kuò)散模型中使用的擴(kuò)散過程在統(tǒng)計(jì)學(xué)和物理學(xué)等領(lǐng)域得到了很好的研究，它們?cè)诟鞣N抽樣應(yīng)用中都很重要。受這些領(lǐng)域的啟發(fā)，我們最近提出了 臨界阻尼朗之萬擴(kuò)散 （ CLD ）。

在 CLD 中，必須擾動(dòng)的數(shù)據(jù)與可被視為 velocities 的輔助變量耦合，這與物理學(xué)中的速度相似，因?yàn)樗鼈兓旧厦枋隽藬?shù)據(jù)向擴(kuò)散模型的基本分布移動(dòng)的速度。

就像一個(gè)落在山頂上的球，在相對(duì)直接的路徑上迅速滾動(dòng)到山谷中，積累一定的速度，這種受物理啟發(fā)的技術(shù)有助于數(shù)據(jù)快速平穩(wěn)地?cái)U(kuò)散。描述 CLD 的正向擴(kuò)散 SDE 如下所示：

這里，xt表示數(shù)據(jù)，vt表示速度。m、t和b是決定擴(kuò)散以及速度和數(shù)據(jù)之間耦合的參數(shù)。dwt是一個(gè)高斯白噪聲過程，負(fù)責(zé)噪聲注入，如公式所示。

CLD 可以解釋為兩個(gè)不同術(shù)語的組合。首先是一個(gè) Ornstein-Uhlenback 過程，這是一種特殊的噪聲注入過程，作用于速度變量vt。

其次，在哈密頓動(dòng)力學(xué)中，數(shù)據(jù)和速度相互耦合，因此注入速度的噪聲也會(huì)影響數(shù)據(jù)xt。哈密頓動(dòng)力學(xué)提供了物理系統(tǒng)力學(xué)的基本描述，比如前面提到的例子中滾下山的球。

圖 2 顯示了一個(gè)簡單的一維玩具問題的數(shù)據(jù)和速度如何在 CLD 中擴(kuò)散：

圖 2 在嚴(yán)重阻尼的朗之萬擴(kuò)散中，數(shù)據(jù) xT增加了速度 vT.擴(kuò)散耦合 xT和 vT在聯(lián)合數(shù)據(jù)速度空間中運(yùn)行（紅色表示概率）。噪聲只注入 vT.這將導(dǎo)致數(shù)據(jù) x 的平滑擴(kuò)散軌跡（綠色）T.

在擴(kuò)散開始時(shí)，我們從簡單的高斯分布中提取一個(gè)隨機(jī)速度，然后在聯(lián)合數(shù)據(jù)速度空間中進(jìn)行完全擴(kuò)散。當(dāng)觀察數(shù)據(jù)的演變（圖中右下角）時(shí)，模型的擴(kuò)散方式比之前的擴(kuò)散方式要平滑得多。

直觀地說，這也應(yīng)該使去噪和反轉(zhuǎn)生成過程更容易。我們只在擴(kuò)散參數(shù)m和t的特定選擇下獲得這種行為，特別是T*T=4M。這種結(jié)構(gòu)在物理學(xué)上被稱為臨界阻尼，對(duì)應(yīng)于更廣泛的隨機(jī)動(dòng)力系統(tǒng) Langevin dynamics 的一個(gè)特例，因此被稱為臨界阻尼 Langevin 擴(kuò)散。

我們還可以可視化圖像在正向擴(kuò)散和生成期間如何在高維聯(lián)合數(shù)據(jù)速度空間中演化：

圖 3 CLD 的正向擴(kuò)散和反向時(shí)間合成過程

在圖 3 的頂部，我們可視化了一維數(shù)據(jù)分布和速度如何在聯(lián)合數(shù)據(jù)速度空間中擴(kuò)散，以及生成如何以相反的方向進(jìn)行。我們對(duì)三種不同的擴(kuò)散軌跡進(jìn)行了采樣，并在右側(cè)顯示了到數(shù)據(jù)和速度空間的投影。在底部，我們將相應(yīng)的擴(kuò)散和合成過程可視化，以生成圖像。我們看到速度在中間時(shí)間對(duì)數(shù)據(jù)進(jìn)行“編碼”t。

在培訓(xùn)生成性擴(kuò)散模型時(shí)使用 CLD 有兩個(gè)關(guān)鍵優(yōu)勢(shì)：

更簡單的評(píng)分函數(shù)和培訓(xùn)目標(biāo)

使用定制 SDE 解算器加速采樣

更簡單的評(píng)分函數(shù)和培訓(xùn)目標(biāo)

在常規(guī)擴(kuò)散模型中，神經(jīng)網(wǎng)絡(luò)的任務(wù)是學(xué)習(xí)擴(kuò)散數(shù)據(jù)分布的得分函數(shù)

相比之下，在基于 CLD 的模型中，我們的任務(wù)是學(xué)習(xí)

即給定數(shù)據(jù)的速度的條件分?jǐn)?shù)函數(shù)。這是只向速度變量中注入噪聲的結(jié)果。

然而，由于速度總是遵循比數(shù)據(jù)本身更平滑的分布，這是一個(gè)更容易學(xué)習(xí)的問題?；?CLD 的擴(kuò)散模型中使用的神經(jīng)網(wǎng)絡(luò)可以更簡單，同時(shí)仍能實(shí)現(xiàn)較高的生成性能。與此相關(guān)，我們還可以針對(duì)基于 CLD 的擴(kuò)散模型制定改進(jìn)的、更穩(wěn)定的訓(xùn)練目標(biāo)。

使用定制 SDE 解算器加速采樣

要集成 CLD 的反向時(shí)間合成 SDE ，可以導(dǎo)出定制的 SDE 解算器，以便對(duì) CLD 中產(chǎn)生的更平滑的正向擴(kuò)散進(jìn)行更有效的去噪。這會(huì)加速合成。

在實(shí)驗(yàn)上，對(duì)于廣泛使用的 CIFAR-10 圖像建?；鶞?zhǔn)，對(duì)于類似的神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)和采樣計(jì)算預(yù)算， CLD 在合成質(zhì)量上優(yōu)于以前的擴(kuò)散模型。此外， CLD 為生成性 SDE 量身定制的 SDE 解算器在生成速度方面明顯優(yōu)于 Euler – Maruyama 等解算器，后者是一種解決擴(kuò)散模型中 SDE 的常用方法。有關(guān)更多信息，請(qǐng)參閱 基于分?jǐn)?shù)的臨界阻尼朗之萬擴(kuò)散生成模型 。

我們已經(jīng)證明，只要仔細(xì)設(shè)計(jì)固定正向擴(kuò)散過程，就可以改進(jìn)擴(kuò)散模型。

擴(kuò)散算子去噪

到目前為止，我們已經(jīng)討論了如何通過將訓(xùn)練數(shù)據(jù)移動(dòng)到平滑的潛在空間（如 LSGM ）來加速擴(kuò)散模型的采樣，或者通過使用輔助速度變量來增加數(shù)據(jù)，以及設(shè)計(jì)改進(jìn)的前向擴(kuò)散過程（如基于 CLD 的擴(kuò)散模型）。

然而，加速擴(kuò)散模型采樣的最直觀的方法之一是直接減少反向過程中的去噪步驟。在這一部分中，我們回到離散時(shí)間擴(kuò)散模型，在數(shù)據(jù)空間中進(jìn)行訓(xùn)練，并分析在減少去噪步驟的數(shù)量和執(zhí)行大步驟時(shí)，去噪過程的行為。

在最近的 study 中，我們觀察到擴(kuò)散模型通常假設(shè)反向合成過程中學(xué)習(xí)到的去噪分布P0（xt-1！xt）可以近似為高斯分布。然而，眾所周知，高斯假設(shè)只適用于許多小的去噪步驟的無窮小極限，這最終導(dǎo)致擴(kuò)散模型的合成緩慢。

當(dāng)反向生成過程使用較大的步長（去噪步驟較少）時(shí)，我們需要一個(gè)非高斯、多峰分布來建模去噪分布P0（xt-1！xt)。

直觀地說，在圖像合成中，多峰分布產(chǎn)生于多個(gè)看似合理且干凈的圖像可能對(duì)應(yīng)于同一個(gè)噪聲圖像的事實(shí)。由于這種多模性，簡單地減少去噪步驟的數(shù)量，同時(shí)在去噪分布中保持高斯假設(shè)，會(huì)損害發(fā)電質(zhì)量。

Figure 5. （ 頂部） 一維數(shù)據(jù)分布 q （ x ）的演化0） 根據(jù)正向擴(kuò)散過程。（下）在固定 x 軸條件下真實(shí)去噪分布的可視化5以不同的顏色顯示不同的步長。

在圖 5 中，小步距（以黃色顯示）的真實(shí)去噪分布接近高斯分布。然而，隨著步長的增加，它變得更加復(fù)雜和多模態(tài)。

受上述觀察結(jié)果的啟發(fā)，我們建議使用表達(dá)性多峰分布參數(shù)化去噪分布，以實(shí)現(xiàn)大步長去噪。特別是，我們引入了一種新的生成模型 去噪擴(kuò)散 GAN ，在該模型中，使用條件 GAN 對(duì)去噪分布進(jìn)行建模（圖 6 ）。

圖 6 去噪擴(kuò)散過程

生成性去噪擴(kuò)散模型通常假設(shè)去噪分布可以用高斯分布建模。這一假設(shè)僅適用于小的去噪步驟，實(shí)際上，這意味著合成過程中有數(shù)千個(gè)去噪步驟。

在我們的去噪擴(kuò)散算法中，我們使用多模態(tài)和復(fù)雜條件算法來表示去噪模型，使我們能夠在兩個(gè)步驟中高效地生成數(shù)據(jù)。

使用對(duì)抗性訓(xùn)練設(shè)置對(duì)去噪擴(kuò)散裝置進(jìn)行訓(xùn)練（圖 7 ）。給定一個(gè)訓(xùn)練圖像x0，我們使用前向高斯擴(kuò)散過程從xt-1和xt兩個(gè)連續(xù)步驟的擴(kuò)散樣本中取樣。

給定xt，我們的條件去噪 GAN 首先隨機(jī)生x‘0，然后使用可處理的后驗(yàn)分布q（xt-1！xt，x’0）通過加回噪聲生成x‘t-1。訓(xùn)練鑒別器來區(qū)分實(shí)際（xt-1，xt）和生成的（x’t-1，xt）對(duì)，并提供反饋以學(xué)習(xí)條件去噪算法。

在訓(xùn)練之后，我們通過從噪聲中采樣并使用我們的去噪擴(kuò)散生成器在幾個(gè)步驟中迭代去噪來生成新實(shí)例。

圖 7 去噪擴(kuò)散算子的訓(xùn)練過程

我們訓(xùn)練了一個(gè)條件 GAN 發(fā)生器，利用擴(kuò)散過程中不同步驟的對(duì)抗性損失對(duì)輸入xt進(jìn)行去噪。

與傳統(tǒng)干草相比的優(yōu)勢(shì)

與我們通過去噪迭代生成樣本的模型相比，為什么不訓(xùn)練一個(gè)可以使用傳統(tǒng)設(shè)置一次性生成樣本的 GAN 呢？與傳統(tǒng)的 GaN 相比，我們的模型有幾個(gè)優(yōu)點(diǎn)。

眾所周知， GAN 會(huì)遭受訓(xùn)練不穩(wěn)定和模式崩潰的影響。一些可能的原因包括難以從復(fù)雜分布中一次性直接生成樣本，以及當(dāng)鑒別器僅查看干凈樣本時(shí)存在過度擬合問題。

相比之下，由于xt上的強(qiáng)條件作用，我們的模型將生成過程分解為幾個(gè)條件去噪擴(kuò)散步驟，其中每個(gè)步驟對(duì)建模相對(duì)簡單。擴(kuò)散過程平滑了數(shù)據(jù)分布，使鑒別器不太可能過度擬合。

我們觀察到，我們的模型具有更好的訓(xùn)練穩(wěn)定性和模式覆蓋率。在圖像生成中，我們觀察到我們的模型實(shí)現(xiàn)了與擴(kuò)散模型競(jìng)爭(zhēng)的樣本質(zhì)量和模式覆蓋率，同時(shí)只需要兩個(gè)去噪步驟。 與常規(guī)擴(kuò)散模型相比，它的采樣速度提高了 2000 倍。我們還發(fā)現(xiàn)，我們的模型在樣本多樣性方面顯著優(yōu)于最先進(jìn)的傳統(tǒng) GAN ，同時(shí)在樣本保真度方面具有競(jìng)爭(zhēng)力。

圖 8 基于擴(kuò)散的不同生成模型的樣本質(zhì)量與采樣時(shí)間

圖 8 顯示了與 CIFAR-10 圖像建?；鶞?zhǔn)的不同基于擴(kuò)散的生成模型的采樣時(shí)間相比，樣本質(zhì)量（通過 Fr é chet Inception 距離測(cè)量；越低越好）。與其他擴(kuò)散模型相比，去噪擴(kuò)散 GaN 在保持相似合成質(zhì)量的同時(shí)實(shí)現(xiàn)了幾個(gè)數(shù)量級(jí)的加速。

結(jié)論

擴(kuò)散模型是一類很有前途的深層生成模型，因?yàn)樗鼈兘Y(jié)合了高質(zhì)量的合成、強(qiáng)大的多樣性和模式覆蓋。這與常規(guī) GAN 等方法形成對(duì)比，后者很受歡迎，但樣本多樣性有限。擴(kuò)散模型的主要缺點(diǎn)是合成速度慢。

在本文中，我們介紹了 NVIDIA 最近開發(fā)的三種技術(shù)，它們成功地解決了這一挑戰(zhàn)。有趣的是，他們每個(gè)人都從不同的角度處理問題，分析擴(kuò)散模型的不同組成部分：

潛空間擴(kuò)散模型 基本上簡化了數(shù)據(jù)本身，首先將其嵌入平滑的潛在空間，在那里可以訓(xùn)練更有效的擴(kuò)散模型。

臨界阻尼朗之萬擴(kuò)散 是一種改進(jìn)的前向擴(kuò)散過程，特別適合于更簡單、更快的去噪和生成。

擴(kuò)散算子去噪 通過表達(dá)性多峰去噪分布，直接學(xué)習(xí)顯著加速的反向去噪過程。

我們相信，擴(kuò)散模型非常適合克服生成性學(xué)習(xí)的三重困境，尤其是在使用本文中強(qiáng)調(diào)的技術(shù)時(shí)。原則上，這些技術(shù)也可以結(jié)合使用。

事實(shí)上，擴(kuò)散模型已經(jīng)在深層生成性學(xué)習(xí)方面取得了重大進(jìn)展。我們預(yù)計(jì)，它們可能會(huì)在圖像和視頻處理、 3D 內(nèi)容生成和數(shù)字藝術(shù)以及語音和語言建模等領(lǐng)域得到實(shí)際應(yīng)用。它們還將用于藥物發(fā)現(xiàn)和材料設(shè)計(jì)等領(lǐng)域，以及其他各種重要應(yīng)用。我們認(rèn)為，基于擴(kuò)散的方法有可能推動(dòng)下一代領(lǐng)先的生成模型。

最后但并非最不重要的一點(diǎn)是，我們是 2022 年 6 月 19 日在美國路易斯安那州新奧爾良舉行的 計(jì)算機(jī)視覺與模式識(shí)別 （ CVPR ）會(huì)議期間舉辦的擴(kuò)散模型、其基礎(chǔ)和應(yīng)用教程組織委員會(huì)的成員。

關(guān)于作者

Arash Vahdat 是 NVIDIA research 的首席研究科學(xué)家，專攻計(jì)算機(jī)視覺和機(jī)器學(xué)習(xí)。在加入 NVIDIA 之前，他是 D-Wave 系統(tǒng)公司的研究科學(xué)家，從事深度生成學(xué)習(xí)和弱監(jiān)督學(xué)習(xí)。在 D-Wave 之前，阿拉什是西蒙·弗雷澤大學(xué)（ Simon Fraser University ， SFU ）的一名研究人員，他領(lǐng)導(dǎo)了深度視頻分析的研究，并教授大數(shù)據(jù)機(jī)器學(xué)習(xí)的研究生課程。阿拉什在格雷格·莫里（ Greg Mori ）的指導(dǎo)下獲得了 SFU 的博士和理學(xué)碩士學(xué)位，致力于視覺分析的潛變量框架。他目前的研究領(lǐng)域包括深層生成學(xué)習(xí)、表征學(xué)習(xí)、高效神經(jīng)網(wǎng)絡(luò)和概率深層學(xué)習(xí)。

Karsten Kreis 是 NVIDIA 多倫多人工智能實(shí)驗(yàn)室的高級(jí)研究科學(xué)家。在加入 NVIDIA 之前，他在 D-Wave Systems 從事深度生成建模工作，并與他人共同創(chuàng)立了變分人工智能，這是一家利用生成模型進(jìn)行藥物發(fā)現(xiàn)的初創(chuàng)公司?？ㄋ闺隈R克斯·普朗克光科學(xué)研究所獲得量子信息理論理學(xué)碩士學(xué)位，并在馬克斯·普朗克聚合物研究所獲得計(jì)算和統(tǒng)計(jì)物理博士學(xué)位。目前，卡斯滕的研究重點(diǎn)是開發(fā)新的生成性學(xué)習(xí)方法，以及將深層生成模型應(yīng)用于計(jì)算機(jī)視覺、圖形和數(shù)字藝術(shù)等領(lǐng)域的問題。

審核編輯：郭婷