斯坦福大學(xué)聯(lián)合谷歌大腦使用「兩步蒸餾方法」提升無(wú)分類器指導(dǎo)的采樣效率，在生成樣本質(zhì)量和采樣速度上都有非常亮眼的表現(xiàn)。

去噪擴(kuò)散概率模型（DDPM）在圖像生成、音頻合成、分子生成和似然估計(jì)領(lǐng)域都已經(jīng)實(shí)現(xiàn)了 SOTA 性能。同時(shí)無(wú)分類器（classifier-free）指導(dǎo)進(jìn)一步提升了擴(kuò)散模型的樣本質(zhì)量，并已被廣泛應(yīng)用在包括 GLIDE、DALL·E 2 和 Imagen 在內(nèi)的大規(guī)模擴(kuò)散模型框架中。

然而，無(wú)分類器指導(dǎo)的一大關(guān)鍵局限是它的采樣效率低下，需要對(duì)兩個(gè)擴(kuò)散模型評(píng)估數(shù)百次才能生成一個(gè)樣本。這一局限阻礙了無(wú)分類指導(dǎo)模型在真實(shí)世界設(shè)置中的應(yīng)用。盡管已經(jīng)針對(duì)擴(kuò)散模型提出了蒸餾方法，但目前這些方法不適用無(wú)分類器指導(dǎo)擴(kuò)散模型。

為了解決這一問(wèn)題，近日斯坦福大學(xué)和谷歌大腦的研究者在論文《On Distillation of Guided Diffusion Models》中提出使用兩步蒸餾（two-step distillation）方法來(lái)提升無(wú)分類器指導(dǎo)的采樣效率。

在第一步中，他們引入單一學(xué)生模型來(lái)匹配兩個(gè)教師擴(kuò)散模型的組合輸出；在第二步中，他們利用提出的方法逐漸地將從第一步學(xué)得的模型蒸餾為更少步驟的模型。

利用提出的方法，單個(gè)蒸餾模型能夠處理各種不同的指導(dǎo)強(qiáng)度，從而高效地對(duì)樣本質(zhì)量和多樣性進(jìn)行權(quán)衡。此外為了從他們的模型中采樣，研究者考慮了文獻(xiàn)中已有的確定性采樣器，并進(jìn)一步提出了隨機(jī)采樣過(guò)程。

研究者在 ImageNet 64x64 和 CIFAR-10 上進(jìn)行了實(shí)驗(yàn)，結(jié)果表明提出的蒸餾模型只需 4 步就能生成在視覺上與教師模型媲美的樣本，并且在更廣泛的指導(dǎo)強(qiáng)度上只需 8 到 16 步就能實(shí)現(xiàn)與教師模型媲美的 FID/IS 分?jǐn)?shù)，具體如下圖 1 所示。

此外，在 ImageNet 64x64 上的其他實(shí)驗(yàn)結(jié)果也表明了，研究者提出的框架在風(fēng)格遷移應(yīng)用中也表現(xiàn)良好。

方法介紹

接下來(lái)本文討論了蒸餾無(wú)分類器指導(dǎo)擴(kuò)散模型的方法（ distilling a classifier-free guided diffusion model）。給定一個(gè)訓(xùn)練好的指導(dǎo)模型，即教師模型之后本文分兩步完成。

第一步引入一個(gè)連續(xù)時(shí)間學(xué)生模型，該模型具有可學(xué)習(xí)參數(shù)η_1，以匹配教師模型在任意時(shí)間步 t∈[0,1] 處的輸出。給定一個(gè)優(yōu)化范圍 [w_min, w_max]，對(duì)學(xué)生模型進(jìn)行優(yōu)化：

其中，。為了合并指導(dǎo)權(quán)重 w，本文引入了一個(gè) w - 條件模型，其中 w 作為學(xué)生模型的輸入。為了更好地捕捉特征，本文還對(duì) w 應(yīng)用傅里葉嵌入。此外，由于初始化在模型性能中起著關(guān)鍵作用，因此本文初始化學(xué)生模型的參數(shù)與教師模型相同。

在第二步中，本文將離散時(shí)間步（discrete time-step）考慮在內(nèi)，并逐步將第一步中的蒸餾模型轉(zhuǎn)化為步數(shù)較短的學(xué)生模型，其可學(xué)習(xí)參數(shù)為η_2，每次采樣步數(shù)減半。設(shè) N 為采樣步數(shù)，給定 w ~ U[w_min, w_max] 和 t∈{1，…， N}，然后根據(jù) Salimans & Ho 等人提出的方法訓(xùn)練學(xué)生模型。在將教師模型中的 2N 步蒸餾為學(xué)生模型中的 N 步之后，之后使用 N 步學(xué)生模型作為新的教師模型，這個(gè)過(guò)程不斷重復(fù)，直到將教師模型蒸餾為 N/2 步學(xué)生模型。

N 步可確定性和隨機(jī)采樣：一旦模型訓(xùn)練完成，給定一個(gè)指定的 w ∈ [w_min, w_max]，然后使用 DDIM 更新規(guī)則執(zhí)行采樣。

實(shí)際上，本文也可以執(zhí)行 N 步隨機(jī)采樣，使用兩倍于原始步長(zhǎng)的確定性采樣步驟，然后使用原始步長(zhǎng)向后執(zhí)行一個(gè)隨機(jī)步驟 。對(duì)于，當(dāng) t > 1/N 時(shí)，本文使用以下更新規(guī)則

實(shí)驗(yàn)

實(shí)驗(yàn)評(píng)估了蒸餾方法的性能，本文主要關(guān)注模型在 ImageNet 64x64 和 CIFAR-10 上的結(jié)果。他們探索了指導(dǎo)權(quán)重的不同范圍，并觀察到所有范圍都具有可比性，因此實(shí)驗(yàn)采用 [w_min, w_max] = [0, 4]。圖 2 和表 1 報(bào)告了在 ImageNet 64x64 上所有方法的性能。

本文還進(jìn)行了如下實(shí)驗(yàn)。具體來(lái)說(shuō)，為了在兩個(gè)域 A 和 B 之間執(zhí)行風(fēng)格遷移，本文使用在域 A 上訓(xùn)練的擴(kuò)散模型對(duì)來(lái)自域 A 的圖像進(jìn)行編碼，然后使用在域 B 上訓(xùn)練的擴(kuò)散模型進(jìn)行解碼。由于編碼過(guò)程可以理解為反向 DDIM 采樣過(guò)程，本文在無(wú)分類器指導(dǎo)下對(duì)編碼器和解碼器進(jìn)行蒸餾，并與下圖 3 中的 DDIM 編碼器和解碼器進(jìn)行比較。