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

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

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

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

如何將無分類器指導擴散模型提煉成快速采樣的模型？

首先，對于一個預先訓練好的無分類器指導模型，研究者首先學習了一個單一的模型，來匹配條件模型和無條件模型的組合輸出。

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

可以看到，在ImageNet 64x64和CIFAR-10上，這種方法能夠在視覺上生成與原始模型相當的圖像。

只需4個采樣步驟，就能獲得與原始模型相當的FID/IS分數，而采樣速度卻高達256倍。

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

擴散模型的背景

通過來自數據分布的樣本x，噪聲調度函數研究者通過最小化加權均方差來訓練了具有參數θ的擴散模型。

其中是信噪比，和是預先指定的加權函數。

一旦訓練了擴散模型，就可以使用離散時間DDIM采樣器從模型中采樣。

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

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

無分類器指導是一種有效的方法，可以顯著提高條件擴散模型的樣本質量，已經廣泛應用于包括GLIDE，DALL·E 2和Imagen。

它引入了一個指導權重參數來衡量樣本的質量和多樣性。為了生成樣本，無分類器指導在每個更新步驟都會使用作為預測模型，來評估條件擴散模型和聯合訓練的。

由于每次采樣更新都需要評估兩個擴散模型，因此使用無分類器指導進行采樣通常很昂貴。

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

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

蒸餾無分類器的指導擴散模型

他們的辦法是，將無分類器的指導擴散模型進行蒸餾。

對于一個訓練有素的教師引導模型，他們采取了兩個步驟。

第一步，研究者引入了一個連續時間的學生模型，它具有可學習的參數η1，來匹配教師模型在任意時間步長t ∈ [0, 1] 的輸出。指定一系列他們有興趣的指導強度后，他們使用以下目標來優化學生模型。

其中。

為了結合指導權重w，研究者引入了w條件模型，其中w作為學生模型的輸入。為了更好地捕捉特征，他們將傅里葉嵌入應用w，然后用Kingma等人使用的時間步長的方式，把它合并到擴散模型的主干中。

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

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

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

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

N-step的確定性和隨機采樣

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

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

，當t > 1/N時，可用以下的更新規則——

其中，。

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

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

其他蒸餾?法

還有一個直接將漸進式蒸餾應?于引導模型的方法，即遵循教師模型的結構，直接將學?模型蒸餾成?個聯合訓練的條件和?條件模型。研究者嘗試了之后，發現此?法效果不佳。

實驗和結論

模型實驗在兩個標準數據集上進行：ImageNet（64*64）和 CIFAR 10。

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

基線標準包括DDPM ancestral采樣和DDIM采樣。

為了更好地理解如何納入指導權重w，使用一個固定的w值訓練的模型作為參照。

為了進行公平比較，實驗對所有的方法使用相同的預訓練教師模型。使用U-Net（Ronneberger等人，2015）架構作為基線，并使用相同的U-Net主干，引入嵌入了w的結構作為兩步學生模型。

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

在實驗中，以指導區間w∈[0, 4]為條件的模型訓練，與w為固定值的模型訓練表現相當。在步驟較少時，我們的方法明顯優于DDIM基線性能，在8到16個步驟下基本達到教師模型的性能水平。

由FID和IS分數評估的ImageNet 64x64采樣質量

由FID和IS評分評估的CIFAR-10采樣質量

我們還對教師模型的編碼過程進行蒸餾，并進行了風格轉移的實驗。具體來說，為了在兩個領域A和B之間進行風格轉換，用在領域A上訓練的擴散模型對領域A的圖像進行編碼，然后用在領域B上訓練的擴散模型進行解碼。

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

總之，我們提出的引導擴散模型的蒸餾方法，以及一種隨機采樣器，從蒸餾后的模型中采樣。從經驗上看，我們的方法只用了一個步驟就能實現視覺上的高體驗采樣，只用8到16個步驟就能獲得與教師相當的FID/IS分數。