一、結論寫在前面

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

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

總的來說，論文的改進在單步圖像生成上設置了新的基準，ImageNet-64×64上的FID分數為1.28，零樣本COCO 2014上為8.35，盡管推理成本降低了500倍，但仍優于原始教師模型。此外，論文展示了論文的方法可以生成百萬像素級的圖像，通過對SDXL進行蒸餾，在少步方法中展現出卓越的視覺質量，并超過了教師模型。論文發布了論文的代碼和預訓練模型。

二、論文的簡單介紹

2.1 論文的背景

擴散模型在視覺生成任務中達到了前所未有的質量水平。但其采樣過程通常需要數十次迭代去噪步驟，每次迭代都需通過神經網絡進行前向傳播。這使得高分辨率文本到圖像的合成既緩慢又昂貴。為解決這一問題，研究者們開發了多種蒸餾方法，旨在將一個教師擴散模型轉化為一個高效、僅需少數步驟的學生生成器。然而，這些方法往往導致生成質量下降，因為學生模型通常是通過損失函數來學習教師模型的成對噪聲到圖像映射，但在完美模仿教師行為方面存在困難。             

              圖1：由論文的4步生成器從SDXL蒸餾得到的1024x1024樣本。請放大查看細節    

盡管如此，應該注意的是，旨在匹配分布的損失函數，如GAN [21] 或 DMD [22] 損失，并不負擔精確學習從噪聲到圖像的具體路徑的復雜性，因為它們的目標是在分布上與教師模型對齊——通過最小化學生和教師輸出分布之間的Jensen-Shannon（JS）或近似的Kullback-Leibler（KL）散度。

特別是DMD [22]在蒸餾Stable Diffusion 1.5時展現出了最先進的結果，但相比基于GAN的方法[23-29]來說，它受到的研究還不夠深入。一個可能的原因是DMD仍然需要一個額外的回歸損失來確保穩定訓練。反過來，這就需要通過運行教師模型的全部采樣步驟來創建數百萬個噪聲-圖像對，對于文本到圖像合成來說代價是特別高昂的。回歸損失也抵消了DMD的無配對分布匹配目標的關鍵優勢，因為它導致學生的質量被教師所限制。    

              圖2：由論文從SDXL蒸餾出的4步生成器產生的1024x1024樣本。請放大查看細節    

2.2 改進的分布匹配蒸餾

論文重新審視了DMD算法中的多個設計選擇，并確定了顯著的改進。

圖3：論文的方法將一個成本高昂的擴散模型（灰色，右側）提煉成一個一步或多步生成器（紅色，左側）。論文的訓練交替進行兩個步驟：1. 使用隱式分布匹配目標的梯度（紅色箭頭）和GAN損失（綠色）優化生成器；2. 訓練一個得分函數（藍色）來模擬生成器產生的“假”樣本的分布，以及一個GAN判別器（綠色)來區分假樣本和真實圖像。學生生成器可以是一步或多步模型，如圖所示，具有中間步驟輸入

2.2.1 移除回歸損失：真正的分布匹配與更易于大規模訓練

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

2.2.2 通過雙時間尺度更新規則穩定純分布匹配

從DMD中簡單地省略回歸目標，如式(3)所示，會導致訓練不穩定并顯著降低質量（見表3）。例如，論文觀察到生成的樣本的平均亮度以及其他統計量波動很大，沒有收斂到一個穩定點（見附錄C）。論文將這種不穩定性歸因于假擴散模型wfake中的近似誤差，它沒有準確跟蹤假得分，因為它是在生成器的非平穩輸出分布上動態優化的。這導致了近似誤差和偏置生成器梯度（如[30]中也討論的）。

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

2.2.3 利用GAN損失和真實數據超越教師模型

到目前為止，論文的模型在不需昂貴的數據集構建的情況下，實現了與DMD相當的訓練穩定性和性能（見表3）。然而，蒸餾生成器與教師擴散模型之間仍存在性能差距。論文推測這一差距可能歸因于DMD中使用的真實評分函數的近似誤差，這會傳遞到生成器并導致次優結果。由于DMD的蒸餾模型從未接受過真實數據訓練，因此無法從這些誤差中恢復。

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

其中D是判別器，是前向擴散過程（即噪聲注入），其噪聲水平對應于時間步。

2.2.4 多步生成器

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

論文固定一個預定的調度，使用個時間步，在訓練和推理期間保持不變。在推理過程中，每個步驟交替進行去噪和噪聲注入步驟，遵循一致性模型[9]，以提高樣本質量。

2.2.5 避免訓練/推理不匹配的多步生成器模擬

論文通過在訓練期間用當前學生生成器運行幾個步驟產生的噪聲合成圖像替換噪聲真實圖像來解決這個問題，類似于論文的推理流程。這是可行的，因為與教師擴散模型不同，論文的生成器只運行幾步。然后，論文的生成器對這些模擬圖像進行去噪，輸出由提出的損失函數進行監督。使用噪聲合成圖像避免了不匹配并提高了整體性能（見第5.3節）。

2.2.6匯總一切

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

2.3論文的效果

論文使用幾個基準評估論文的方法DMD2，包括在ImageNet-64x64 上的條件類圖像生成，以及在COCO 2014 上使用各種教師模型進行文本到圖像合成。

2.3.1 類別條件圖像生成

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

2.3.2文本到圖像合成

論文在零樣本COCO 2014上評估DMD2的文本到圖像生成性能。論文的生成器通過蒸餾SDXL和SD v1.5進行訓練，分別使用來自L.AION-Aesthetics[58]的300萬提示的子集。此外，論文還從LAION-Aesthetic收集了50萬張圖像作為GAN判別器的訓練數據。表2總結了SDXL模型的蒸餾結果。論文的4步生成器產生了高質量且多樣的樣本，實現了19.32的FID分數和0.332的CLIP分數，與教師擴散模型在圖像質量和提示一致性上相媲美。

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

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

2.3.3 消融研究

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

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

此外，純粹的基于GAN的方法也缺乏一種自然的方式來整合無分類器指導，這對于高質量的文本到圖像合成至關重要。因此，在背景中，戴著太陽鏡的駱駝坐在太空船甲板上的照片，盡管基于GAN的方法通過緊密匹配真實分布實現了最低的FID，但它們在文本對齊和美學質量方面顯著表現不佳（圖7第二列）。同樣，省略反向模擬會導致圖像質量下降，這一點通過退化的塊FID分數得以體現。。

圖6：論文的模型、SDXL教師模型以及選定競爭方法[23， 27， 31]之間的視覺對比。所有蒸餾模型使用4個采樣步驟，而教師模型使用50個采樣步驟并結合無分類器指導。所有圖像均使用相同的噪聲和文本提示生成。論文的模型生成的圖像在真實感和文本對齊方面表現更優。（放大查看細節。）更多對比見附錄圖10

論文標題：Improved Distribution Matching Distillation for Fast Image Synthesis

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

本文轉載自 ??AI帝國??，作者： 無影寺