CVPR'23 最佳論文候選｜提速256倍！蒸餾擴(kuò)散模型生成圖像質(zhì)量媲美教師模型，只需4步！

斯坦福大學(xué)和谷歌大腦的研究者提出使用兩步蒸餾方法來提升無分類器指導(dǎo)的采樣效率。該方法能夠在 ImageNet 64x64 和 CIFAR-10 上使用少至 4 個(gè)采樣步驟生成視覺上與原始模型相當(dāng)?shù)膱D像，實(shí)現(xiàn)與原始模型的采樣速度提高了 256 倍。 >>加入極市CV技術(shù)交流群，走在計(jì)算機(jī)視覺的最前沿

轉(zhuǎn)自《極市平臺》

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

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

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

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

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

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

研究者在 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)良好。

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

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

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

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

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

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

實(shí)驗(yàn)評估了蒸餾方法的性能，本文主要關(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)。具體來說，為了在兩個(gè)域 A 和 B 之間執(zhí)行風(fēng)格遷移，本文使用在域 A 上訓(xùn)練的擴(kuò)散模型對來自域 A 的圖像進(jìn)行編碼，然后使用在域 B 上訓(xùn)練的擴(kuò)散模型進(jìn)行解碼。由于編碼過程可以理解為反向 DDIM 采樣過程，本文在無分類器指導(dǎo)下對編碼器和****進(jìn)行蒸餾，并與下圖 3 中的 DDIM 編碼器和****進(jìn)行比較。

本文還探討了如何修改指導(dǎo)強(qiáng)度 w 以影響性能，如下圖 4 所示。