史丹佛/Google大腦：兩次蒸餾，引導擴散模型採樣提速256倍！

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

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

為了解決這個問題，史丹佛大學和Google大腦的學者提出使用兩步驟蒸餾(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分數。

以上是史丹佛/Google大腦：兩次蒸餾，引導擴散模型採樣提速256倍！的詳細內容。更多資訊請關注PHP中文網其他相關文章！