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

然而，無分類器指導(dǎo)擴(kuò)散模型的一個缺點是它們在推理時的計算成本很高。因為它們需要評估兩個擴(kuò)散模型——一個類別條件模型(class-conditional model) 和一個無條件模型(unconditional model)，而且需要評估數(shù)百次。

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

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

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

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

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

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

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

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

擴(kuò)散模型的背景

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

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

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

具體來說，DDIM采樣器從 z1 ～ N (0,I)開始，更新如下

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

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

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

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

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

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

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

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

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

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

其中。

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

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

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

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

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

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

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

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

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

其中，。

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

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

其他蒸餾?法

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

實驗和結(jié)論

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

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

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

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

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

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

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

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

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

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

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

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