最近，無(wú)分類(lèi)器的指導(dǎo)擴(kuò)散模型(classifier-free guided diffusion models)在高分辨率圖像生成方面非常有效，并且已經(jīng)被廣泛用于大規(guī)模擴(kuò)散框架，包括DALL-E 2、GLIDE和Imagen。

然而，無(wú)分類(lèi)器指導(dǎo)擴(kuò)散模型的一個(gè)缺點(diǎn)是它們?cè)谕评頃r(shí)的計(jì)算成本很高。因?yàn)樗鼈冃枰u(píng)估兩個(gè)擴(kuò)散模型——一個(gè)類(lèi)別條件模型(class-conditional model) 和一個(gè)無(wú)條件模型(unconditional model)，而且需要評(píng)估數(shù)百次。

為了解決這個(gè)問(wèn)題，斯坦福大學(xué)和谷歌大腦的學(xué)者提出使用兩步蒸餾(two-step distillation)的方法來(lái)提升無(wú)分類(lèi)器指導(dǎo)擴(kuò)散模型的采樣效率。

論文地址：https://arxiv.org/abs/2210.03142

如何將無(wú)分類(lèi)器指導(dǎo)擴(kuò)散模型提煉成快速采樣的模型？

首先，對(duì)于一個(gè)預(yù)先訓(xùn)練好的無(wú)分類(lèi)器指導(dǎo)模型，研究者首先學(xué)習(xí)了一個(gè)單一的模型，來(lái)匹配條件模型和無(wú)條件模型的組合輸出。

隨后，研究者逐步將這個(gè)模型蒸餾成一個(gè)采樣步驟更少的擴(kuò)散模型。

可以看到，在ImageNet 64x64和CIFAR-10上，這種方法能夠在視覺(jué)上生成與原始模型相當(dāng)?shù)膱D像。

只需4個(gè)采樣步驟，就能獲得與原始模型相當(dāng)?shù)腇ID/IS分?jǐn)?shù)，而采樣速度卻高達(dá)256倍。

可以看到，通過(guò)改變指導(dǎo)權(quán)重w，研究者蒸餾的模型能夠在樣本多樣性和質(zhì)量之間進(jìn)行權(quán)衡。而且只用一個(gè)取樣步驟，就能獲得視覺(jué)上愉悅的結(jié)果。

擴(kuò)散模型的背景

通過(guò)來(lái)自數(shù)據(jù)分布的樣本x，噪聲調(diào)度函數(shù)研究者通過(guò)最小化加權(quán)均方差來(lái)訓(xùn)練了具有參數(shù)θ的擴(kuò)散模型。

其中是信噪比，和是預(yù)先指定的加權(quán)函數(shù)。

一旦訓(xùn)練了擴(kuò)散模型，就可以使用離散時(shí)間DDIM采樣器從模型中采樣。

具體來(lái)說(shuō)，DDIM采樣器從 z1 ～ N (0,I)開(kāi)始，更新如下

其中，N是采樣步驟的總數(shù)。使用，會(huì)生成最終樣本。

無(wú)分類(lèi)器指導(dǎo)是一種有效的方法，可以顯著提高條件擴(kuò)散模型的樣本質(zhì)量，已經(jīng)廣泛應(yīng)用于包括GLIDE，DALL·E 2和Imagen。

它引入了一個(gè)指導(dǎo)權(quán)重參數(shù)來(lái)衡量樣本的質(zhì)量和多樣性。為了生成樣本，無(wú)分類(lèi)器指導(dǎo)在每個(gè)更新步驟都會(huì)使用作為預(yù)測(cè)模型，來(lái)評(píng)估條件擴(kuò)散模型和聯(lián)合訓(xùn)練的。

由于每次采樣更新都需要評(píng)估兩個(gè)擴(kuò)散模型，因此使用無(wú)分類(lèi)器指導(dǎo)進(jìn)行采樣通常很昂貴。

為了解決這個(gè)問(wèn)題，研究者使用了漸進(jìn)式蒸餾(progressive distillation)  ，這是一種通過(guò)重復(fù)蒸餾提高擴(kuò)散模型采樣速度的方法。

在以前，這種方法不能直接被直接用在引導(dǎo)模型的蒸餾上，也不能在確定性DDIM采樣器以外的采樣器上使用。而在這篇論文中，研究者解決了這些問(wèn)題。

蒸餾無(wú)分類(lèi)器的指導(dǎo)擴(kuò)散模型

他們的辦法是，將無(wú)分類(lèi)器的指導(dǎo)擴(kuò)散模型進(jìn)行蒸餾。

對(duì)于一個(gè)訓(xùn)練有素的教師引導(dǎo)模型，他們采取了兩個(gè)步驟。

第一步，研究者引入了一個(gè)連續(xù)時(shí)間的學(xué)生模型，它具有可學(xué)習(xí)的參數(shù)η1，來(lái)匹配教師模型在任意時(shí)間步長(zhǎng)t ∈ [0, 1] 的輸出。指定一系列他們有興趣的指導(dǎo)強(qiáng)度后，他們使用以下目標(biāo)來(lái)優(yōu)化學(xué)生模型。

其中。

為了結(jié)合指導(dǎo)權(quán)重w，研究者引入了w條件模型，其中w作為學(xué)生模型的輸入。為了更好地捕捉特征，他們將傅里葉嵌入應(yīng)用w，然后用Kingma等人使用的時(shí)間步長(zhǎng)的方式，把它合并到擴(kuò)散模型的主干中。

由于初始化在性能中起著關(guān)鍵作用，研究者初始化學(xué)生模型時(shí)，使用的是與教師條件模型相同的參數(shù)（除了新引入的與w-conditioning相關(guān)的參數(shù)）。

第二步，研究者設(shè)想了一個(gè)離散的時(shí)間步長(zhǎng)場(chǎng)景，并且通過(guò)每次將采樣步數(shù)減半，逐步將學(xué)習(xí)模型從第?步蒸餾成具有可學(xué)習(xí)參數(shù)η2、步?更少的學(xué)?模型。

其中，N表?采樣步驟的數(shù)量，對(duì)于和，研究者開(kāi)始訓(xùn)練學(xué)生模型，讓它用一步來(lái)匹配教師模型的兩步DDIM采樣的輸出（例如：從t/N到t - 0.5/N，從t - 0.5/N到t - 1/N）。

將教師模型中的2N個(gè)步驟蒸餾成學(xué)生模型中的N個(gè)步驟以后，我們可以將新的N-step學(xué)生模型作為新的教師模型，然后重復(fù)同樣的過(guò)程，將教師模型蒸餾成N/2-step的學(xué)生模型。在每?步，研究者都會(huì)?教師模型的參數(shù)來(lái)初始化學(xué)?模型。

N-step的確定性和隨機(jī)采樣

?旦模型被訓(xùn)練出來(lái)，對(duì)于，研究者就可以通過(guò)DDIM更新規(guī)則來(lái)執(zhí)行采樣。研究者注意到，對(duì)于蒸餾模型，這個(gè)采樣過(guò)程在給定初始化的情況下是確定的。

另外，研究者也可以進(jìn)行N步的隨機(jī)采樣。使用兩倍于原始步長(zhǎng)的確定性采樣步驟（ 即與N/2-step確定性采樣器相同），然后使用原始步長(zhǎng)進(jìn)行一次隨機(jī)步回（即用噪聲擾動(dòng)）。

，當(dāng)t > 1/N時(shí)，可用以下的更新規(guī)則——

其中，。

當(dāng)t=1/N時(shí)，研究者使用確定性更新公式，從得出。

值得注意的是，我們注意到，與確定性的采樣器相比，執(zhí)行隨機(jī)采樣需要在稍微不同的時(shí)間步長(zhǎng)內(nèi)評(píng)估模型，并且需要對(duì)邊緣情況的訓(xùn)練算法進(jìn)行小的修改。

其他蒸餾?法

還有一個(gè)直接將漸進(jìn)式蒸餾應(yīng)?于引導(dǎo)模型的方法，即遵循教師模型的結(jié)構(gòu)，直接將學(xué)?模型蒸餾成?個(gè)聯(lián)合訓(xùn)練的條件和?條件模型。研究者嘗試了之后，發(fā)現(xiàn)此?法效果不佳。

實(shí)驗(yàn)和結(jié)論

模型實(shí)驗(yàn)在兩個(gè)標(biāo)準(zhǔn)數(shù)據(jù)集上進(jìn)行：ImageNet（64*64）和 CIFAR 10。

實(shí)驗(yàn)中探索了指導(dǎo)權(quán)重w的不同范圍，并觀察到所有的范圍都有可比性，因此使用[wmin, wmax] = [0, 4]進(jìn)行實(shí)驗(yàn)。使用信噪比損失訓(xùn)練第一步和第二步模型。

基線標(biāo)準(zhǔn)包括DDPM ancestral采樣和DDIM采樣。

為了更好地理解如何納入指導(dǎo)權(quán)重w，使用一個(gè)固定的w值訓(xùn)練的模型作為參照。

為了進(jìn)行公平比較，實(shí)驗(yàn)對(duì)所有的方法使用相同的預(yù)訓(xùn)練教師模型。使用U-Net（Ronneberger等人，2015）架構(gòu)作為基線，并使用相同的U-Net主干，引入嵌入了w的結(jié)構(gòu)作為兩步學(xué)生模型。

上圖為所有方法在ImageNet 64x64上的表現(xiàn)。其中D和S分別代表確定性和隨機(jī)性采樣器。

在實(shí)驗(yàn)中，以指導(dǎo)區(qū)間w∈[0, 4]為條件的模型訓(xùn)練，與w為固定值的模型訓(xùn)練表現(xiàn)相當(dāng)。在步驟較少時(shí)，我們的方法明顯優(yōu)于DDIM基線性能，在8到16個(gè)步驟下基本達(dá)到教師模型的性能水平。

由FID和IS分?jǐn)?shù)評(píng)估的ImageNet 64x64采樣質(zhì)量

由FID和IS評(píng)分評(píng)估的CIFAR-10采樣質(zhì)量

我們還對(duì)教師模型的編碼過(guò)程進(jìn)行蒸餾，并進(jìn)行了風(fēng)格轉(zhuǎn)移的實(shí)驗(yàn)。具體來(lái)說(shuō)，為了在兩個(gè)領(lǐng)域A和B之間進(jìn)行風(fēng)格轉(zhuǎn)換，用在領(lǐng)域A上訓(xùn)練的擴(kuò)散模型對(duì)領(lǐng)域A的圖像進(jìn)行編碼，然后用在領(lǐng)域B上訓(xùn)練的擴(kuò)散模型進(jìn)行解碼。

由于編碼過(guò)程可以理解為顛倒了的DDIM的采樣過(guò)程，我們對(duì)具有無(wú)分類(lèi)器指導(dǎo)的編碼器和解碼器都進(jìn)行了蒸餾，并與DDIM編碼器和解碼器進(jìn)行比較，如上圖所示。我們還探討了對(duì)引導(dǎo)強(qiáng)度w的改動(dòng)對(duì)性能的影響。

總之，我們提出的引導(dǎo)擴(kuò)散模型的蒸餾方法，以及一種隨機(jī)采樣器，從蒸餾后的模型中采樣。從經(jīng)驗(yàn)上看，我們的方法只用了一個(gè)步驟就能實(shí)現(xiàn)視覺(jué)上的高體驗(yàn)采樣，只用8到16個(gè)步驟就能獲得與教師相當(dāng)?shù)腇ID/IS分?jǐn)?shù)。