文章地址：https://arxiv.org/pdf/2405.05224

?

擴散模型是一個強大的生成框架，但推理過程非常昂貴。現(xiàn)有的加速方法通常會在極低步驟情況下犧牲圖像質(zhì)量，或者在復(fù)雜條件下失敗。這項工作提出了一個新穎的蒸餾框架Imagine Flash，旨在通過僅使用一到三個步驟實現(xiàn)高保真度和多樣化樣本生成。方法包括三個關(guān)鍵組成部分：

反向蒸餾，通過校準學生模型自身的反向軌跡來緩解訓練推理差異；
遷移重構(gòu)損失，根據(jù)當前時間步動態(tài)調(diào)整知識傳遞；
噪聲校正，一種增強樣本質(zhì)量的推理時技術(shù)，通過解決噪聲預(yù)測中的奇異性來改善樣本質(zhì)量。

通過大量實驗證明，我們的方法在定量指標和人類評估方面優(yōu)于現(xiàn)有競爭對手。值得注意的是，我們的方法僅使用三個去噪步驟就實現(xiàn)了與教師模型相當?shù)男阅?，實現(xiàn)了高效的高質(zhì)量生成。

生成的效果圖

介紹

隨著去噪擴散模型（DMs）的出現(xiàn)，生成建模已經(jīng)發(fā)生了一場范式轉(zhuǎn)變。這些模型在各個領(lǐng)域樹立了新的基準，提供了前所未有的逼真性和多樣性，同時確保了穩(wěn)定的訓練。然而，去噪過程的順序性質(zhì)帶來了顯著挑戰(zhàn)。從DMs中抽樣是一個耗時和昂貴的過程，所需時間在很大程度上取決于兩個因素：（i）每步神經(jīng)網(wǎng)絡(luò)評估的延遲，以及 （ii）去噪步驟的總數(shù)。

大量的研究工作致力于加速抽樣過程。對于文本到圖像的合成，提出的方法涵蓋了廣泛的技術(shù)，包括高階求解器、用于曲率降低的修改擴散公式，以及引導、步驟和一致性蒸餾。這些方法取得了令人印象深刻的改進，在接近10步的區(qū)間內(nèi)達到了非常高的質(zhì)量。最近，利用蒸餾和對抗性損失兩者的混合方法[17,31,39]將界限推動到了不到五步。盡管這些方法在簡單提示和簡單樣式（如動畫）上取得了令人印象深刻的質(zhì)量，但它們在真實照片般的圖像上，特別是在復(fù)雜的文本條件下，樣本質(zhì)量下降。

上述方法中的一個共同主題是試圖將少步驟的學生模型與復(fù)雜的教師路徑對齊，盡管學生模型具有明顯較低的容量（即步驟）。我們認識到這是一個限制，因此提出了一種新穎的蒸餾框架，旨在使教師模型沿著自己的擴散路徑改進學生模型?？傊?，本文貢獻有三個方面：

• 首先，方法引入了反向蒸餾，這是一種設(shè)計用于在其自身上游反向軌跡上校準學生模型的蒸餾過程，從而減少了訓練和推理分布之間的差距，并確保在所有時間步長上訓練期間沒有數(shù)據(jù)泄漏。
• 其次，提出了一種遷移重構(gòu)損失，它動態(tài)地調(diào)整來自教師模型的知識遷移。具體而言，該損失被設(shè)計為在高時間步長時從教師模型中提煉全局結(jié)構(gòu)信息，而在低時間步長時專注于渲染細粒度細節(jié)和高頻成分。這種自適應(yīng)方法使學生能夠有效地模擬教師的生成過程在擴散軌跡的不同階段。
• 最后，提出了噪聲校正，這是一種在推理時修改，通過解決在初始采樣步驟中存在的噪聲預(yù)測模型中的奇異性來增強樣本質(zhì)量的技術(shù)。這種無需訓練的技術(shù)減輕了對比度和顏色強度的降級，這種降級通常會在極低數(shù)量的去噪步驟中操作時出現(xiàn)。

通過將這三個新穎組件協(xié)同結(jié)合起來，將我們的蒸餾框架應(yīng)用到基線擴散模型Emu中，從而實現(xiàn)了Imagine Flash，在極低步長區(qū)間內(nèi)實現(xiàn)了高質(zhì)量的生成，而不會影響樣本質(zhì)量或條件保真度（下圖2）。通過大量實驗和人類評估，展示了我們的方法在各種任務(wù)和模態(tài)之間實現(xiàn)了抽樣效率和生成質(zhì)量之間的有利折衷。

方法

Imagine Flash，這是一種新穎的蒸餾技術(shù)，旨在快速進行文本到圖像的生成，它建立在Emu 的基礎(chǔ)上，但不局限于此。與原始的Emu模型不同，后者需要至少50個神經(jīng)函數(shù)評估（NFEs）才能生成高質(zhì)量的樣本，Imagine Flash僅需少量NFEs即可實現(xiàn)可比較的結(jié)果。提出的蒸餾方法包括三個新穎的關(guān)鍵組成部分：

• 反向蒸餾，這是一種蒸餾過程，確保了在所有時間點t上的訓練期間零數(shù)據(jù)泄漏。
• 遷移重構(gòu)損失（SRL），這是一種自適應(yīng)損失函數(shù)，旨在最大化從教師模型中的知識遷移。
• 噪聲校正，這是一種無需訓練的推理修改，用于改善在噪聲預(yù)測模式下訓練的少步驟方法的樣本質(zhì)量。

接下來，假設(shè)可以訪問一個預(yù)訓練的擴散模型，它預(yù)測噪聲估計ε。這個教師模型可以在圖像空間或潛在空間中操作。我們的目標是將的知識蒸餾到一個學生模型中，同時減少抽樣步驟的總數(shù)，并提供額外允許中每一步的高質(zhì)量增加。如果模型使用無分類器的引導(cfg)，那么也將這些知識蒸餾到學生模型中，并且消除cfg的需要。

反向蒸餾

總的來說，反向蒸餾消除了所有時間步驟t上的信息泄漏，防止模型依賴于真值信號。這是通過在訓練期間模擬推理過程來實現(xiàn)的，這也可以解釋為在其自身的上游反向路徑上校準學生模型。

SRL：遷移重構(gòu)損失

通過反向擴散進行圖像生成的過程中，早期階段（其中t接近T）對于構(gòu)建圖像的整體結(jié)構(gòu)和組成至關(guān)重要。相反，后期階段（其中t接近0）對于添加高級細節(jié)至關(guān)重要。借鑒這一觀察結(jié)果，對默認的知識蒸餾損失進行了改進，鼓勵學生模型學習教師模型的結(jié)構(gòu)組成和細節(jié)渲染能力。這涉及將教師去噪的起始點從學生的起始點t移開，因此我們將這種方法稱為遷移重構(gòu)損失（SRL）。下圖4概述了提出的Loss。

這種方法鼓勵學生在早期后向步驟中優(yōu)先提取結(jié)構(gòu)知識，并在最后后向步驟中集中精力生成更精細的細節(jié)。

噪聲校正

觀察到，這個小的修改可以顯著改善估計的顏色，使色調(diào)更加生動飽滿。當推理步驟數(shù)量較少時，這種效應(yīng)尤為明顯。

實驗

為了確保公平性，將Emu模型作為所有實驗的基礎(chǔ)。Emu是一個擁有27億參數(shù)和分辨率為768×768的最先進模型。將我們的結(jié)果與之前的蒸餾方法進行比較，例如Step Distillation、LCM和ADD，并將它們直接應(yīng)用在Emu上。所有模型都在一個委托的圖像數(shù)據(jù)集上進行了替換訓練。由于ADD訓練沒有公開可用的代碼，根據(jù)論文中提供的細節(jié)自行實現(xiàn)了它。

與最新技術(shù)的定量比較

使用FID 、CLIP分數(shù) 和CompBench將Imagine Flash與之前的方法進行比較。FID和CLIP分別衡量圖像質(zhì)量和提示對齊度，并根據(jù)來自COCO2017 的5k樣本分割進行評估，遵循[31]中的評估協(xié)議。CompBench是一個基準，單獨衡量屬性綁定（顏色、形狀和紋理）和對象關(guān)系（空間、非空間和復(fù)雜）。

在CompBench驗證集中為每個提示生成2張圖像（總共300個提示）。對于LCM和Imagine Flash，計算1、2和3步的指標。對于ADD，計算4步的指標，因為該方法專門針對4步推理進行了調(diào)優(yōu)和配置，以確保公平比較。還評估了4步的Step Distillation，以提供更直接的比較。表1顯示了結(jié)果。

我們的3步Imagine Flash在FID方面優(yōu)于Step Distillation和ADD，即使使用了少一步。它在1、2和3步的FID上也比LCM低。我們的3步模型的CLIP分數(shù)高于所有變體的ADD和LCM，并與4步Step Distillation模型的得分（30.2）相匹配。與Step Distillation和ADD相比，它們相應(yīng)地降低了FID分別為10.1和3.4，而我們的3步和2步Imagine Flash則保持了FID并略有改善。

對于CompBench，我們的1、2或3步Imagine Flash在所有類別中都優(yōu)于以前的方法，除了顏色外，其中4步Step Distillation和ADD的得分與我們的類似。這突顯了Imagine Flash的優(yōu)越提示對齊性。

與最新技術(shù)的定性比較

在圖5中，展示了Imagine Flash與當前最先進技術(shù)（SOTA）的定性比較：Step Distillation、LCM和ADD，它們都對相同的基線Emu模型進行了蒸餾，以進行公平比較。

觀察到ADD生成的圖像比Step Distillation和LCM更清晰，這是因為它使用了對抗損失。雖然Imagine Flash和ADD都使用了鑒別器，但Imagine Flash生成的圖像比ADD更銳利和更詳細。Imagine Flash的增強銳度和細節(jié)是由我們提出的SRL導致的，它有效地精煉了學生預(yù)測的高頻細節(jié)，如下圖6的最后一行所示。

另一方面，對于ADD，目標圖像可能展示出明顯不同的色譜，出現(xiàn)顏色偽影（見圖6），并且顏色在訓練迭代過程中可能出現(xiàn)不可預(yù)測的波動。我們假設(shè)，為了在期望中最小化L2重構(gòu)損失，ADD模型最好通過預(yù)測接近零的顏色值來達到目標，從而導致圖像蒼白和輪廓模糊。除了改善局部細節(jié)外，SRL還可以糾正學生的文本對齊錯誤，如圖6右側(cè)所示（1步），其中小熊貓被轉(zhuǎn)換回了狗。

與公開模型的比較

還將Imagine Flash的性能與ADD-LDMXL和Lightning-LDMXL發(fā)布的公開模型進行了比較。為此，按照前面詳細說明的方法計算CLIP和FID分數(shù)，并與基準模型進行相對增益/下降的比較。請在附錄A中找到表格。我們的方法在文本對齊能力上與ADD和Lightning保持類似，但在FID增加方面顯示出更為有利的結(jié)果，尤其是對于兩步和三步。

此外，進行了廣泛的人類評估。為此，使用了三個推理步驟為OUI數(shù)據(jù)集中隨機抽樣的1,000個提示生成了所有方法的圖像。成對的圖像呈現(xiàn)給42名受過訓練的人類標注者中的五名，他們的任務(wù)是為更具視覺吸引力的圖像投票。通過多數(shù)投票匯總的結(jié)果顯示在表2中，明顯表明了對Imagine Flash的偏好。

消融研究

對Imagine Flash進行了定量和定性消融，以評估所提出的反向蒸餾、SRL和噪聲校正的效果。定量評估如表3所示，而補充視覺消融如圖7所示

結(jié)論

Imagine Flash，這是一種新穎的蒸餾框架，能夠利用擴散模型實現(xiàn)高保真度的少步驟圖像生成。方法包括三個關(guān)鍵組成部分：反向蒸餾以減少訓練推理差異，遷移重構(gòu)損失（SRL）動態(tài)地調(diào)整每個時間步的知識傳遞，以及噪聲校正以增強初始樣本質(zhì)量。

通過大量實驗，Imagine Flash取得了顯著的成果，僅使用三個去噪步驟就與預(yù)訓練的教師模型的性能相匹配，并始終超越現(xiàn)有方法。這種前所未有的抽樣效率結(jié)合了高樣本質(zhì)量和多樣性，使我們的模型非常適用于實時生成應(yīng)用。

我們的工作為超高效的生成建模鋪平了道路。未來的方向包括擴展到其他模態(tài)，如視頻和3D，進一步降低抽樣預(yù)算，并將我們的方法與互補的加速技術(shù)相結(jié)合。通過實現(xiàn)即時的高保真度生成，Imagine Flash為實時創(chuàng)意工作流程和交互式媒體體驗開啟了新的可能性。

更多生成的圖像

本文轉(zhuǎn)自 AI生成未來 ，作者：Meta     

原文鏈接:??https://mp.weixin.qq.com/s/osiR6U4srAidT0E_YwUVDA??