一、結(jié)論寫在前面

最近的一些方法已經(jīng)顯示出將昂貴的擴(kuò)散模型蒸餾到高效的單步生成器中的前景。其中，分布匹配蒸餾(DMD)能夠生成與教師模型在分布上匹配的單步生成器，即蒸餾過程不強(qiáng)制與教師模型的采樣軌跡一一對(duì)應(yīng)。然而，為了在實(shí)踐中確保穩(wěn)定訓(xùn)練，DMD需要使用教師模型通過多步確定性采樣器生成的大量噪聲-圖像對(duì)計(jì)算一個(gè)額外的回歸損失。這不僅在大規(guī)模文本到圖像合成中計(jì)算代價(jià)高昂，而且還限制了學(xué)生模型的質(zhì)量，使其過于緊密地綁定到教師模型的原始采樣路徑。

論文提出了DMD2，一組技術(shù)來(lái)解決這一限制并改進(jìn)DMD訓(xùn)練。首先，論文消除了回歸損失和昂貴數(shù)據(jù)集構(gòu)建的需求。論文表明，由此導(dǎo)致的不穩(wěn)定性是由于"假"評(píng)價(jià)器無(wú)法充分準(zhǔn)確地估計(jì)生成樣本的分布。因此，論文提出了一種兩時(shí)間尺度更新規(guī)則作為補(bǔ)救措施。其次，論文將GAN損失整合到蒸餾過程中，區(qū)分生成樣本和真實(shí)圖像。這使論文能夠在真實(shí)數(shù)據(jù)上訓(xùn)練學(xué)生模型，從而緩解教師模型中不完美的"真實(shí)"分?jǐn)?shù)估計(jì)，并因此提高質(zhì)量。第三，論文介紹了一種新的訓(xùn)練程序，允許學(xué)生模型進(jìn)行多步采樣，并通過在訓(xùn)練時(shí)模擬推理時(shí)的生成器樣本來(lái)解決之前工作中訓(xùn)練-推理輸入不匹配的問題。

總的來(lái)說，論文的改進(jìn)在單步圖像生成上設(shè)置了新的基準(zhǔn)，ImageNet-64×64上的FID分?jǐn)?shù)為1.28，零樣本COCO 2014上為8.35，盡管推理成本降低了500倍，但仍優(yōu)于原始教師模型。此外，論文展示了論文的方法可以生成百萬(wàn)像素級(jí)的圖像，通過對(duì)SDXL進(jìn)行蒸餾，在少步方法中展現(xiàn)出卓越的視覺質(zhì)量，并超過了教師模型。論文發(fā)布了論文的代碼和預(yù)訓(xùn)練模型。

二、論文的簡(jiǎn)單介紹

2.1 論文的背景

擴(kuò)散模型在視覺生成任務(wù)中達(dá)到了前所未有的質(zhì)量水平。但其采樣過程通常需要數(shù)十次迭代去噪步驟，每次迭代都需通過神經(jīng)網(wǎng)絡(luò)進(jìn)行前向傳播。這使得高分辨率文本到圖像的合成既緩慢又昂貴。為解決這一問題，研究者們開發(fā)了多種蒸餾方法，旨在將一個(gè)教師擴(kuò)散模型轉(zhuǎn)化為一個(gè)高效、僅需少數(shù)步驟的學(xué)生生成器。然而，這些方法往往導(dǎo)致生成質(zhì)量下降，因?yàn)閷W(xué)生模型通常是通過損失函數(shù)來(lái)學(xué)習(xí)教師模型的成對(duì)噪聲到圖像映射，但在完美模仿教師行為方面存在困難。             

              圖1：由論文的4步生成器從SDXL蒸餾得到的1024x1024樣本。請(qǐng)放大查看細(xì)節(jié)    

盡管如此，應(yīng)該注意的是，旨在匹配分布的損失函數(shù)，如GAN [21] 或 DMD [22] 損失，并不負(fù)擔(dān)精確學(xué)習(xí)從噪聲到圖像的具體路徑的復(fù)雜性，因?yàn)樗鼈兊哪繕?biāo)是在分布上與教師模型對(duì)齊——通過最小化學(xué)生和教師輸出分布之間的Jensen-Shannon（JS）或近似的Kullback-Leibler（KL）散度。

特別是DMD [22]在蒸餾Stable Diffusion 1.5時(shí)展現(xiàn)出了最先進(jìn)的結(jié)果，但相比基于GAN的方法[23-29]來(lái)說，它受到的研究還不夠深入。一個(gè)可能的原因是DMD仍然需要一個(gè)額外的回歸損失來(lái)確保穩(wěn)定訓(xùn)練。反過來(lái)，這就需要通過運(yùn)行教師模型的全部采樣步驟來(lái)創(chuàng)建數(shù)百萬(wàn)個(gè)噪聲-圖像對(duì)，對(duì)于文本到圖像合成來(lái)說代價(jià)是特別高昂的?；貧w損失也抵消了DMD的無(wú)配對(duì)分布匹配目標(biāo)的關(guān)鍵優(yōu)勢(shì)，因?yàn)樗鼘?dǎo)致學(xué)生的質(zhì)量被教師所限制。    

              圖2：由論文從SDXL蒸餾出的4步生成器產(chǎn)生的1024x1024樣本。請(qǐng)放大查看細(xì)節(jié)    

2.2 改進(jìn)的分布匹配蒸餾

論文重新審視了DMD算法中的多個(gè)設(shè)計(jì)選擇，并確定了顯著的改進(jìn)。

圖3：論文的方法將一個(gè)成本高昂的擴(kuò)散模型（灰色，右側(cè)）提煉成一個(gè)一步或多步生成器（紅色，左側(cè)）。論文的訓(xùn)練交替進(jìn)行兩個(gè)步驟：1. 使用隱式分布匹配目標(biāo)的梯度（紅色箭頭）和GAN損失（綠色）優(yōu)化生成器；2. 訓(xùn)練一個(gè)得分函數(shù)（藍(lán)色）來(lái)模擬生成器產(chǎn)生的“假”樣本的分布，以及一個(gè)GAN判別器（綠色)來(lái)區(qū)分假樣本和真實(shí)圖像。學(xué)生生成器可以是一步或多步模型，如圖所示，具有中間步驟輸入

2.2.1 移除回歸損失：真正的分布匹配與更易于大規(guī)模訓(xùn)練

DMD[22]中使用的回歸損失[16]確保了模式覆蓋和訓(xùn)練穩(wěn)定性，但正如論文在第3節(jié)中討論的，它使得大規(guī)模蒸餾變得繁瑣，并且與分布匹配的理念相沖突，因此本質(zhì)上限制了蒸餾生成器的性能，使其無(wú)法超越教師模型。論文的第一個(gè)改進(jìn)是移除這個(gè)損失。

2.2.2 通過雙時(shí)間尺度更新規(guī)則穩(wěn)定純分布匹配

從DMD中簡(jiǎn)單地省略回歸目標(biāo)，如式(3)所示，會(huì)導(dǎo)致訓(xùn)練不穩(wěn)定并顯著降低質(zhì)量（見表3）。例如，論文觀察到生成的樣本的平均亮度以及其他統(tǒng)計(jì)量波動(dòng)很大，沒有收斂到一個(gè)穩(wěn)定點(diǎn)（見附錄C）。論文將這種不穩(wěn)定性歸因于假擴(kuò)散模型wfake中的近似誤差，它沒有準(zhǔn)確跟蹤假得分，因?yàn)樗窃谏善鞯姆瞧椒€(wěn)輸出分布上動(dòng)態(tài)優(yōu)化的。這導(dǎo)致了近似誤差和偏置生成器梯度（如[30]中也討論的）。

論文采用Heusel等人[59]啟發(fā)下的雙時(shí)間尺度更新規(guī)則來(lái)解決這一問題。具體而言，論文以不同頻率訓(xùn)練ufake和生成器G，確保ufake精確跟蹤生成器的輸出分布。論文發(fā)現(xiàn)，每進(jìn)行一次生成器更新，進(jìn)行5次虛假評(píng)分更新，不使用回歸損失，能夠提供良好的穩(wěn)定性，并且在ImageNet上與原始DMD的質(zhì)量相匹配（見表3），同時(shí)實(shí)現(xiàn)更快的收斂。進(jìn)一步的分析包含在附錄C中。

2.2.3 利用GAN損失和真實(shí)數(shù)據(jù)超越教師模型

到目前為止，論文的模型在不需昂貴的數(shù)據(jù)集構(gòu)建的情況下，實(shí)現(xiàn)了與DMD相當(dāng)?shù)挠?xùn)練穩(wěn)定性和性能（見表3）。然而，蒸餾生成器與教師擴(kuò)散模型之間仍存在性能差距。論文推測(cè)這一差距可能歸因于DMD中使用的真實(shí)評(píng)分函數(shù)的近似誤差，這會(huì)傳遞到生成器并導(dǎo)致次優(yōu)結(jié)果。由于DMD的蒸餾模型從未接受過真實(shí)數(shù)據(jù)訓(xùn)練，因此無(wú)法從這些誤差中恢復(fù)。

論文通過在論文的流程中加入額外的GAN目標(biāo)來(lái)解決這一問題，其中判別器被訓(xùn)練以區(qū)分真實(shí)圖像和論文生成器生成的圖像。通過使用真實(shí)數(shù)據(jù)訓(xùn)練，GAN分類器不受教師模型的限制，可能使論文的學(xué)生生成器在樣本質(zhì)量上超越它。論文將GAN分類器整合到DMD中遵循極簡(jiǎn)主義設(shè)計(jì)：論文在虛假擴(kuò)散去噪器的瓶頸層之上添加了一個(gè)分類分支（見圖3）。分類分支和UNet中的上游編碼器特征通過最大化標(biāo)準(zhǔn)非飽和GAN目標(biāo)進(jìn)行訓(xùn)練：

其中D是判別器，是前向擴(kuò)散過程（即噪聲注入），其噪聲水平對(duì)應(yīng)于時(shí)間步。

2.2.4 多步生成器

通過提出的改進(jìn)措施，論文能夠在ImageNet和COCO上匹配教師擴(kuò)散模型的性能（見表1和表5）。然而，論文發(fā)現(xiàn)，對(duì)于像SDXL這樣的大規(guī)模模型，由于模型容量有限以及學(xué)習(xí)從噪聲到高度多樣化和詳細(xì)圖像的直接映射的復(fù)雜優(yōu)化景觀，將其提煉成一步生成器仍然具有挑戰(zhàn)性。這促使論文將DMD擴(kuò)展以支持多步采樣。

論文固定一個(gè)預(yù)定的調(diào)度，使用個(gè)時(shí)間步，在訓(xùn)練和推理期間保持不變。在推理過程中，每個(gè)步驟交替進(jìn)行去噪和噪聲注入步驟，遵循一致性模型[9]，以提高樣本質(zhì)量。

2.2.5 避免訓(xùn)練/推理不匹配的多步生成器模擬

論文通過在訓(xùn)練期間用當(dāng)前學(xué)生生成器運(yùn)行幾個(gè)步驟產(chǎn)生的噪聲合成圖像替換噪聲真實(shí)圖像來(lái)解決這個(gè)問題，類似于論文的推理流程。這是可行的，因?yàn)榕c教師擴(kuò)散模型不同，論文的生成器只運(yùn)行幾步。然后，論文的生成器對(duì)這些模擬圖像進(jìn)行去噪，輸出由提出的損失函數(shù)進(jìn)行監(jiān)督。使用噪聲合成圖像避免了不匹配并提高了整體性能（見第5.3節(jié)）。

2.2.6匯總一切

總之，論文的蒸餾方法解決了DMD 對(duì)預(yù)計(jì)算噪聲-圖像對(duì)的嚴(yán)格要求。它進(jìn)一步整合了GAN的優(yōu)勢(shì)，并支持多步生成器。如圖3所示，從預(yù)訓(xùn)練的擴(kuò)散模型開始，論文交替優(yōu)化生成器Gθ以最小化原始分布匹配目標(biāo)以及GAN目標(biāo)，并優(yōu)化使用假數(shù)據(jù)的去噪分?jǐn)?shù)匹配目標(biāo)和GAN分類損失來(lái)優(yōu)化假分?jǐn)?shù)估計(jì)器μfake。為確保假分?jǐn)?shù)估計(jì)準(zhǔn)確且穩(wěn)定，盡管是在線優(yōu)化，論文以比生成器更高的頻率(5步比1步)更新它。

2.3論文的效果

論文使用幾個(gè)基準(zhǔn)評(píng)估論文的方法DMD2，包括在ImageNet-64x64 上的條件類圖像生成，以及在COCO 2014 上使用各種教師模型進(jìn)行文本到圖像合成。

2.3.1 類別條件圖像生成

表1比較了論文的模型與最近在ImageNet-64x64上的基準(zhǔn)模型。通過單次前向傳播，論文的方法顯著超越了現(xiàn)有的蒸餾技術(shù)，甚至超越了使用ODE采樣器的教師模型[52]。論文將這一顯著性能歸功于移除了DMD的回歸損失，這消除了由ODE采樣器施加的性能上限，以及論文增加的GAN項(xiàng)，這減輕了教師擴(kuò)散模型分?jǐn)?shù)近似誤差的不利影響。

2.3.2文本到圖像合成

論文在零樣本COCO 2014上評(píng)估DMD2的文本到圖像生成性能。論文的生成器通過蒸餾SDXL和SD v1.5進(jìn)行訓(xùn)練，分別使用來(lái)自L.AION-Aesthetics[58]的300萬(wàn)提示的子集。此外，論文還從LAION-Aesthetic收集了50萬(wàn)張圖像作為GAN判別器的訓(xùn)練數(shù)據(jù)。表2總結(jié)了SDXL模型的蒸餾結(jié)果。論文的4步生成器產(chǎn)生了高質(zhì)量且多樣的樣本，實(shí)現(xiàn)了19.32的FID分?jǐn)?shù)和0.332的CLIP分?jǐn)?shù)，與教師擴(kuò)散模型在圖像質(zhì)量和提示一致性上相媲美。

為了進(jìn)一步驗(yàn)證論文方法的有效性，論文進(jìn)行了一項(xiàng)廣泛的用戶研究，比較論文的模型輸出與教師模型和現(xiàn)有蒸餾方法的輸出。論文使用了LADD之后的PartiPrompts的一個(gè)子集128個(gè)提示。對(duì)于每次比較，論文要求一組隨機(jī)的五名評(píng)估者選擇視覺上更吸引人的圖像，以及更好地代表文本提示的圖像。關(guān)于人類評(píng)估的詳細(xì)信息包含在附錄H中。如圖5所示，論文的模型在用戶偏好上遠(yuǎn)超基線方法。值得注意的是，論文的模型在圖像質(zhì)量上超越其教師模型的樣本占比達(dá)到249%，并且在提示對(duì)齊上達(dá)到可比性，同時(shí)需要的前向傳播次數(shù)減少了25倍（4 vs 100）。               

圖5：用戶研究比較了論文蒸餾的模型與其教師模型以及競(jìng)爭(zhēng)性的蒸餾基線[23， 27， 31]。所有蒸餾模型使用4個(gè)采樣步驟，教師模型使用50個(gè)。論文的模型在圖像質(zhì)量和提示對(duì)齊方面均達(dá)到了最佳性能

2.3.3 消融研究

表3在ImageNet上消融了論文提出的方法的不同組件。簡(jiǎn)單地從原始DMD中移除ODE回歸損失導(dǎo)致由于訓(xùn)練不穩(wěn)定而降級(jí)的FID為3.48（見附錄C中的進(jìn)一步分析）。然而，結(jié)合論文的雙時(shí)間尺度更新規(guī)則（TTUR）緩解了這種性能下降，無(wú)需額外的數(shù)據(jù)集構(gòu)建即可匹配DMD基線性能。添加論文的GAN損失進(jìn)一步將FID提高了1.1個(gè)點(diǎn)。論文集成的方案超越了僅使用GAN（沒有分布匹配目標(biāo)）的性能，并且將雙時(shí)間尺度更新規(guī)則添加到單獨(dú)的GAN中并沒有改善它，突出了在統(tǒng)一框架中結(jié)合分布匹配與GAN的有效性。

在表4中，論文分析了GAN項(xiàng)、分布匹配目標(biāo)和反向模擬（對(duì)將SDXL模型蒸餾為四步生成器的影響。定性結(jié)果如圖7所示。在沒有GAN損失的情況下，論文的基線模型生成的圖像過度飽和且過度平滑（圖7第三列）。同樣，移除分布匹配目標(biāo)將論文的方法簡(jiǎn)化為純粹的基于GAN的方法，這面臨著訓(xùn)練穩(wěn)定性的挑戰(zhàn)。

此外，純粹的基于GAN的方法也缺乏一種自然的方式來(lái)整合無(wú)分類器指導(dǎo)，這對(duì)于高質(zhì)量的文本到圖像合成至關(guān)重要。因此，在背景中，戴著太陽(yáng)鏡的駱駝坐在太空船甲板上的照片，盡管基于GAN的方法通過緊密匹配真實(shí)分布實(shí)現(xiàn)了最低的FID，但它們?cè)谖谋緦?duì)齊和美學(xué)質(zhì)量方面顯著表現(xiàn)不佳（圖7第二列）。同樣，省略反向模擬會(huì)導(dǎo)致圖像質(zhì)量下降，這一點(diǎn)通過退化的塊FID分?jǐn)?shù)得以體現(xiàn)。。

圖6：論文的模型、SDXL教師模型以及選定競(jìng)爭(zhēng)方法[23， 27， 31]之間的視覺對(duì)比。所有蒸餾模型使用4個(gè)采樣步驟，而教師模型使用50個(gè)采樣步驟并結(jié)合無(wú)分類器指導(dǎo)。所有圖像均使用相同的噪聲和文本提示生成。論文的模型生成的圖像在真實(shí)感和文本對(duì)齊方面表現(xiàn)更優(yōu)。（放大查看細(xì)節(jié)。）更多對(duì)比見附錄圖10

論文標(biāo)題：Improved Distribution Matching Distillation for Fast Image Synthesis

論文鏈接：??https://arxiv.org/pdf/2405.14867???

本文轉(zhuǎn)載自 ??AI帝國(guó)??，作者： 無(wú)影寺