機(jī)器學(xué)習(xí)中如何使用擴(kuò)散模型產(chǎn)生高保真度圖像

發(fā)布人：Google Research Brain 團(tuán)隊(duì)研究員 Jonathan Ho 和軟件工程師 Chitwan Saharia

自然圖像合成是范圍很廣的一類機(jī)器學(xué)習(xí) （ML） 任務(wù)，會(huì)面臨各種類型的應(yīng)用所帶來(lái)的許多設(shè)計(jì)挑戰(zhàn)。其中之一就是圖像超分辨率，即訓(xùn)練一個(gè)模型來(lái)將低分辨率的圖像，轉(zhuǎn)換為細(xì)節(jié)豐富的高分辨率圖像（例如 RAISR）。超分辨率的應(yīng)用有很多，從恢復(fù)古舊的家族畫像，到改善醫(yī)療成像系統(tǒng) （Super-resolution imaging），均有涉獵。另一個(gè)類似的圖像合成任務(wù)是類條件下的圖像生成，即訓(xùn)練一個(gè)模型，并根據(jù)輸入的類別標(biāo)簽生成樣本圖像。最終生成的樣本圖像可用于提高圖像分類、分割等下游模型的性能。例如，GAN 經(jīng)常受到不穩(wěn)定的訓(xùn)練和模式崩潰的影響，而自回歸模型通常受到合成速度慢的影響。

RAISR

https://ai.googleblog.com/2016/11/enhance-raisr-sharp-images-with-machine.html

一般來(lái)說(shuō)，這些圖像合成任務(wù)由深度生成模型完成，如 GAN、VAE，以及自回歸模型。然而，通過(guò)訓(xùn)練讓這些模型在困難的高分辨率數(shù)據(jù)集上，合成高質(zhì)量的樣本時(shí)，它們的表現(xiàn)均有不足。例如，GAN 經(jīng)常受到不穩(wěn)定的訓(xùn)練和模式崩潰的影響，而自回歸模型通常受到合成速度慢的影響。

GAN

https://arxiv.org/abs/1406.2661

VAE

https://arxiv.org/abs/1312.6114

自回歸模型

https://arxiv.org/abs/1601.06759

但是，2015 年首次提出的擴(kuò)散模型 （Diffusion Model） 由于其訓(xùn)練的穩(wěn)定性及在圖像和音頻生成方面的樣本質(zhì)量結(jié)果很有潛力，最近又重新引起了關(guān)注。因此，與其他類型的深度生成模型相比，此模型在得失方面有了不錯(cuò)的權(quán)衡。擴(kuò)散模型的工作原理是通過(guò)逐步添加高斯噪聲來(lái)破壞訓(xùn)練數(shù)據(jù)，慢慢抹去數(shù)據(jù)中的細(xì)節(jié)，直到變成純?cè)肼?。然后再?xùn)練一個(gè)神經(jīng)網(wǎng)絡(luò)來(lái)逆轉(zhuǎn)這個(gè)破壞過(guò)程。運(yùn)行這個(gè)反向破壞過(guò)程，可以通過(guò)逐漸去噪，來(lái)從純?cè)肼曋泻铣蓴?shù)據(jù)，直到產(chǎn)生一個(gè)干凈的樣本。這個(gè)合成過(guò)程可以被解釋為一種優(yōu)化算法，會(huì)按照數(shù)據(jù)密度的梯度來(lái)產(chǎn)生可能的樣本。

擴(kuò)散模型

https://arxiv.org/abs/1503.03585

可以被解釋為

https://arxiv.org/abs/2006.11239

按照數(shù)據(jù)密度的梯度

https://arxiv.org/abs/1907.05600

今天，我們會(huì)介紹兩種相互聯(lián)系的方法。它們將推動(dòng)擴(kuò)散模型的圖像合成質(zhì)量界限，即通過(guò)重復(fù)優(yōu)化獲得的超級(jí)分辨率 （SR3） 和用于類條件合成的模型，此模型又名為級(jí)聯(lián)擴(kuò)散模型 （CDM）。經(jīng)證明，通過(guò)擴(kuò)大擴(kuò)散模型的規(guī)模以及精心選擇的數(shù)據(jù)增強(qiáng)技術(shù)，我們可以取得比現(xiàn)有方法更出色的結(jié)果。具體來(lái)說(shuō)，SR3 可以實(shí)現(xiàn)強(qiáng)大的圖像超分辨率結(jié)果，在人工評(píng)估方面的表現(xiàn)優(yōu)于 GAN。CDM 生成的高保真 ImageNet 樣本在 FID 得分 （Fréchet inception distance） 和分類準(zhǔn)確率得分上都大大超過(guò)了 BigGAN-deep 和 VQ-VAE2。

分類準(zhǔn)確率得分

https://arxiv.org/abs/1905.10887

BigGAN-deep

https://arxiv.org/abs/1809.11096

VQ-VAE2

https://arxiv.org/abs/1906.00446

SR3：圖像超分辨率

SR3 是一個(gè)超分辨率擴(kuò)散模型，它通過(guò)輸入低分辨率圖像，從純?cè)肼曋袠?gòu)建相應(yīng)的高分辨率圖像。該模型利用圖像破壞過(guò)程為訓(xùn)練對(duì)象，在這個(gè)過(guò)程中，噪聲被逐步添加到高分辨率圖像中，直到只剩下純?cè)肼暈橹埂Ｈ缓螅鼘⒛孓D(zhuǎn)這一過(guò)程，從純?cè)肼曢_始，逐步去除噪聲，再通過(guò)指導(dǎo)輸入的低分辨率圖像，達(dá)到目標(biāo)分布。

SR3

http://iterative-refinement.github.io

通過(guò)大規(guī)模的訓(xùn)練，在擴(kuò)縮到所輸入低分辨率圖像的 4-8 倍時(shí)，SR3 在人臉和自然圖像的超分辨率任務(wù)中取得了強(qiáng)大的基準(zhǔn)結(jié)果。這些超分辨率模型可以進(jìn)一步級(jí)聯(lián)起來(lái)，增加有效的超分辨率縮放系數(shù)，例如，通過(guò)將 64x64→256x256 和 256x256→1024x1024 的人臉超分辨率模型堆疊在一起，來(lái)執(zhí)行 64x64→1024x1024 的超分辨率任務(wù)。

以人工評(píng)估研究為例，我們比較了 SR3 與現(xiàn)有方法。通過(guò)開展雙選項(xiàng)必選實(shí)驗(yàn) （Two-alternative forced choice），我們要求受訪者在引用的高分辨率圖像和模型輸出之間做出選擇，并回答問(wèn)題（猜一猜哪張圖像是攝像機(jī)拍攝的？）。我們通過(guò)混淆率（評(píng)分者選擇模型輸出而不是參考圖像的次數(shù)百分比，一個(gè)完美的算法能夠達(dá)到 50% 的混淆率）來(lái)衡量模型的性能。這項(xiàng)研究的結(jié)果如下圖所示。

上：我們?cè)?16x16→128x128 的人臉任務(wù)中實(shí)現(xiàn)了接近 50% 的混淆率，優(yōu)于一流的人臉超分辨率方法 PULSE 和 FSRGAN。下：在 64x64 → 256x256 的自然圖像這個(gè)更困難的任務(wù)中，我們也達(dá)到了 40% 的混淆率，表現(xiàn)水準(zhǔn)遠(yuǎn)超回歸基線

PULSE

https://arxiv.org/abs/2003.03808

FSRGAN

https://arxiv.org/abs/1711.10703

CDM：類條件 ImageNet 生成

在展示了 SR3 處理自然圖像超分辨率方面的有效性之后，我們便更進(jìn)一步——將這些 SR3 模型用于類條件圖像生成。CDM 是利用 ImageNet 數(shù)據(jù)訓(xùn)練的類條件擴(kuò)散模型，可以生成高分辨率的自然圖像。由于 ImageNet 是一個(gè)困難的高熵?cái)?shù)據(jù)集，我們將 CDM 構(gòu)建為多個(gè)擴(kuò)散模型的級(jí)聯(lián)。這種級(jí)聯(lián)方法會(huì)將多個(gè)空間分辨率的生成模型串聯(lián)起來(lái)：一個(gè)擴(kuò)散模型以低分辨率生成數(shù)據(jù)，搭配一連串的 SR3 超分辨率擴(kuò)散模型，生成圖像的分辨率逐漸提高到最高。眾所周知，級(jí)聯(lián)可以改善高分辨率數(shù)據(jù)的質(zhì)量和訓(xùn)練速度，這一點(diǎn)在以前的研究（例如自回歸模型和 VQ-VAE-2）以及同時(shí)進(jìn)行的擴(kuò)散模型研究中都有所體現(xiàn)。正如以下定量結(jié)果所證明的那樣，CDM 進(jìn)一步突出了級(jí)聯(lián)在擴(kuò)散模型中對(duì)樣本質(zhì)量的有效性和對(duì)下游任務(wù)（如圖像分類）的有用性。

級(jí)聯(lián)流水線的示例，其中包括一系列擴(kuò)散模型：第一個(gè)模型用于生成低分辨率的圖像，然后其余的模型負(fù)責(zé)執(zhí)行上采樣，以最終生成高分辨率的圖像。此為用于生成類條件 ImageNet 的流水線，以 32x32 分辨率的類條件擴(kuò)散模型開始，然后是使用 SR3 的 2 倍和 4 倍的類條件超分辨率

從我們的 256x256 級(jí)聯(lián)類條件 ImageNet 模型中選擇的生成圖像

同時(shí)進(jìn)行

https://arxiv.org/abs/2102.09672

研究

https://arxiv.org/abs/2105.05233

將 SR3 模型納入級(jí)聯(lián)流水線的同時(shí)，我們還引入了一種新的數(shù)據(jù)增強(qiáng)技術(shù)，稱之為 “條件增強(qiáng)”，它可以進(jìn)一步優(yōu)化 CDM 的樣本質(zhì)量結(jié)果。雖然 CDM 中的超分辨率模型利用數(shù)據(jù)集中的原始圖像訓(xùn)練而成，但在生成過(guò)程中，它們需要以超分辨率對(duì)低分辨率基礎(chǔ)模型生成的圖像進(jìn)行處理，而這些圖像與原始圖像相比，質(zhì)量可能不夠高。這就導(dǎo)致了超分辨率模型的訓(xùn)練和測(cè)試水平不對(duì)等。條件增強(qiáng)是指對(duì)級(jí)聯(lián)流水線中每個(gè)超分辨率模型的低分辨率輸入圖像進(jìn)行數(shù)據(jù)增強(qiáng)。在我們的研究中，這些增強(qiáng)包括高斯噪聲和高斯模糊，防止各個(gè)超分辨率模型對(duì)其低分辨率的條件輸入過(guò)度擬合，最終讓 CDM 得到更好的高分辨率樣本質(zhì)量。

總的來(lái)說(shuō)，CDM 生成的高保真樣本在類條件 ImageNet 生成的 FID 得分和分類準(zhǔn)確率得分方面都優(yōu)于 BigGAN-deep 和 VQ-VAE-2。不同于 ADM 和 VQ-VAE-2 等其他模型，CDM 是一個(gè)純粹的生成模型，不會(huì)使用分類器來(lái)提高樣本質(zhì)量。有關(guān)樣本質(zhì)量的定量結(jié)果如下所示。

就不使用額外分類器來(lái)提高樣本質(zhì)量的方法而言，類條件 ImageNet FID 在 256x256 分辨率下的得分。BigGAN-deep 的分?jǐn)?shù)記錄自其在最佳截?cái)嘀禃r(shí)的表現(xiàn)（數(shù)值越低越好）

ImageNet 分類準(zhǔn)確率在 256x256 分辨率下的得分，用以衡量利用生成數(shù)據(jù)訓(xùn)練的分類器的驗(yàn)證集準(zhǔn)確率。與現(xiàn)有的方法相比，使用 CDM 生成的數(shù)據(jù)取得了明顯的進(jìn)步，縮小了真實(shí)數(shù)據(jù)和生成數(shù)據(jù)之間的分類準(zhǔn)確率的差距（數(shù)值越高越好）

結(jié)論

在 SR3 和 CDM 的助力下，我們已經(jīng)將擴(kuò)散模型的性能提升到了超分辨率和類條件 ImageNet 生成基準(zhǔn)的一流水平。我們很高興能進(jìn)一步探索擴(kuò)散模型在各種生成性建模問(wèn)題上的極限。有關(guān)我們研究的更多信息，請(qǐng)參閱通過(guò)迭代優(yōu)化實(shí)現(xiàn)圖像超分辨和用于生成高保真圖像的級(jí)聯(lián)擴(kuò)散模型。

責(zé)任編輯：haq