Meta最新成果：前所未有的加速Emu推理 | Imagine Flash：新型蒸餾框架發布 精華
文章地址：https://arxiv.org/pdf/2405.05224
?
擴散模型是一個強大的生成框架，但推理過程非常昂貴。現有的加速方法通常會在極低步驟情況下犧牲圖像質量，或者在復雜條件下失敗。這項工作提出了一個新穎的蒸餾框架Imagine Flash，旨在通過僅使用一到三個步驟實現高保真度和多樣化樣本生成。方法包括三個關鍵組成部分：
反向蒸餾，通過校準學生模型自身的反向軌跡來緩解訓練推理差異；
遷移重構損失，根據當前時間步動態調整知識傳遞；
噪聲校正，一種增強樣本質量的推理時技術，通過解決噪聲預測中的奇異性來改善樣本質量。
通過大量實驗證明，我們的方法在定量指標和人類評估方面優于現有競爭對手。值得注意的是，我們的方法僅使用三個去噪步驟就實現了與教師模型相當的性能，實現了高效的高質量生成。
生成的效果圖
介紹
隨著去噪擴散模型（DMs）的出現，生成建模已經發生了一場范式轉變。這些模型在各個領域樹立了新的基準，提供了前所未有的逼真性和多樣性，同時確保了穩定的訓練。然而，去噪過程的順序性質帶來了顯著挑戰。從DMs中抽樣是一個耗時和昂貴的過程，所需時間在很大程度上取決于兩個因素：（i）每步神經網絡評估的延遲，以及 （ii）去噪步驟的總數。
大量的研究工作致力于加速抽樣過程。對于文本到圖像的合成，提出的方法涵蓋了廣泛的技術，包括高階求解器、用于曲率降低的修改擴散公式，以及引導、步驟和一致性蒸餾。這些方法取得了令人印象深刻的改進，在接近10步的區間內達到了非常高的質量。最近，利用蒸餾和對抗性損失兩者的混合方法[17,31,39]將界限推動到了不到五步。盡管這些方法在簡單提示和簡單樣式（如動畫）上取得了令人印象深刻的質量，但它們在真實照片般的圖像上，特別是在復雜的文本條件下，樣本質量下降。
上述方法中的一個共同主題是試圖將少步驟的學生模型與復雜的教師路徑對齊，盡管學生模型具有明顯較低的容量（即步驟）。我們認識到這是一個限制，因此提出了一種新穎的蒸餾框架，旨在使教師模型沿著自己的擴散路徑改進學生模型。總之，本文貢獻有三個方面：
- 首先，方法引入了反向蒸餾，這是一種設計用于在其自身上游反向軌跡上校準學生模型的蒸餾過程，從而減少了訓練和推理分布之間的差距，并確保在所有時間步長上訓練期間沒有數據泄漏。
- 其次，提出了一種遷移重構損失，它動態地調整來自教師模型的知識遷移。具體而言，該損失被設計為在高時間步長時從教師模型中提煉全局結構信息，而在低時間步長時專注于渲染細粒度細節和高頻成分。這種自適應方法使學生能夠有效地模擬教師的生成過程在擴散軌跡的不同階段。
- 最后，提出了噪聲校正，這是一種在推理時修改，通過解決在初始采樣步驟中存在的噪聲預測模型中的奇異性來增強樣本質量的技術。這種無需訓練的技術減輕了對比度和顏色強度的降級，這種降級通常會在極低數量的去噪步驟中操作時出現。
通過將這三個新穎組件協同結合起來，將我們的蒸餾框架應用到基線擴散模型Emu中，從而實現了Imagine Flash，在極低步長區間內實現了高質量的生成，而不會影響樣本質量或條件保真度（下圖2）。通過大量實驗和人類評估，展示了我們的方法在各種任務和模態之間實現了抽樣效率和生成質量之間的有利折衷。
方法
Imagine Flash，這是一種新穎的蒸餾技術，旨在快速進行文本到圖像的生成，它建立在Emu 的基礎上，但不局限于此。與原始的Emu模型不同，后者需要至少50個神經函數評估（NFEs）才能生成高質量的樣本，Imagine Flash僅需少量NFEs即可實現可比較的結果。提出的蒸餾方法包括三個新穎的關鍵組成部分：
- 反向蒸餾，這是一種蒸餾過程，確保了在所有時間點t上的訓練期間零數據泄漏。
- 遷移重構損失（SRL），這是一種自適應損失函數，旨在最大化從教師模型中的知識遷移。
- 噪聲校正，這是一種無需訓練的推理修改，用于改善在噪聲預測模式下訓練的少步驟方法的樣本質量。
接下來，假設可以訪問一個預訓練的擴散模型，它預測噪聲估計ε。這個教師模型可以在圖像空間或潛在空間中操作。我們的目標是將的知識蒸餾到一個學生模型中，同時減少抽樣步驟的總數，并提供額外允許中每一步的高質量增加。如果模型使用無分類器的引導(cfg)，那么也將這些知識蒸餾到學生模型中，并且消除cfg的需要。
反向蒸餾
總的來說，反向蒸餾消除了所有時間步驟t上的信息泄漏，防止模型依賴于真值信號。這是通過在訓練期間模擬推理過程來實現的，這也可以解釋為在其自身的上游反向路徑上校準學生模型。
SRL：遷移重構損失
通過反向擴散進行圖像生成的過程中，早期階段（其中t接近T）對于構建圖像的整體結構和組成至關重要。相反，后期階段（其中t接近0）對于添加高級細節至關重要。借鑒這一觀察結果，對默認的知識蒸餾損失進行了改進，鼓勵學生模型學習教師模型的結構組成和細節渲染能力。這涉及將教師去噪的起始點從學生的起始點t移開，因此我們將這種方法稱為遷移重構損失（SRL）。下圖4概述了提出的Loss。
這種方法鼓勵學生在早期后向步驟中優先提取結構知識，并在最后后向步驟中集中精力生成更精細的細節。
噪聲校正
觀察到，這個小的修改可以顯著改善估計的顏色，使色調更加生動飽滿。當推理步驟數量較少時，這種效應尤為明顯。
實驗
為了確保公平性，將Emu模型作為所有實驗的基礎。Emu是一個擁有27億參數和分辨率為768×768的最先進模型。將我們的結果與之前的蒸餾方法進行比較，例如Step Distillation、LCM和ADD，并將它們直接應用在Emu上。所有模型都在一個委托的圖像數據集上進行了替換訓練。由于ADD訓練沒有公開可用的代碼，根據論文中提供的細節自行實現了它。
與最新技術的定量比較
使用FID 、CLIP分數 和CompBench將Imagine Flash與之前的方法進行比較。FID和CLIP分別衡量圖像質量和提示對齊度，并根據來自COCO2017 的5k樣本分割進行評估，遵循[31]中的評估協議。CompBench是一個基準，單獨衡量屬性綁定（顏色、形狀和紋理）和對象關系（空間、非空間和復雜）。
在CompBench驗證集中為每個提示生成2張圖像（總共300個提示）。對于LCM和Imagine Flash，計算1、2和3步的指標。對于ADD，計算4步的指標，因為該方法專門針對4步推理進行了調優和配置，以確保公平比較。還評估了4步的Step Distillation，以提供更直接的比較。表1顯示了結果。
我們的3步Imagine Flash在FID方面優于Step Distillation和ADD，即使使用了少一步。它在1、2和3步的FID上也比LCM低。我們的3步模型的CLIP分數高于所有變體的ADD和LCM，并與4步Step Distillation模型的得分（30.2）相匹配。與Step Distillation和ADD相比，它們相應地降低了FID分別為10.1和3.4，而我們的3步和2步Imagine Flash則保持了FID并略有改善。
對于CompBench，我們的1、2或3步Imagine Flash在所有類別中都優于以前的方法，除了顏色外，其中4步Step Distillation和ADD的得分與我們的類似。這突顯了Imagine Flash的優越提示對齊性。
與最新技術的定性比較
在圖5中，展示了Imagine Flash與當前最先進技術（SOTA）的定性比較：Step Distillation、LCM和ADD，它們都對相同的基線Emu模型進行了蒸餾，以進行公平比較。
觀察到ADD生成的圖像比Step Distillation和LCM更清晰，這是因為它使用了對抗損失。雖然Imagine Flash和ADD都使用了鑒別器，但Imagine Flash生成的圖像比ADD更銳利和更詳細。Imagine Flash的增強銳度和細節是由我們提出的SRL導致的，它有效地精煉了學生預測的高頻細節，如下圖6的最后一行所示。
另一方面，對于ADD，目標圖像可能展示出明顯不同的色譜，出現顏色偽影（見圖6），并且顏色在訓練迭代過程中可能出現不可預測的波動。我們假設，為了在期望中最小化L2重構損失，ADD模型最好通過預測接近零的顏色值來達到目標，從而導致圖像蒼白和輪廓模糊。除了改善局部細節外，SRL還可以糾正學生的文本對齊錯誤，如圖6右側所示（1步），其中小熊貓被轉換回了狗。
與公開模型的比較
還將Imagine Flash的性能與ADD-LDMXL和Lightning-LDMXL發布的公開模型進行了比較。為此，按照前面詳細說明的方法計算CLIP和FID分數，并與基準模型進行相對增益/下降的比較。請在附錄A中找到表格。我們的方法在文本對齊能力上與ADD和Lightning保持類似，但在FID增加方面顯示出更為有利的結果，尤其是對于兩步和三步。
此外，進行了廣泛的人類評估。為此，使用了三個推理步驟為OUI數據集中隨機抽樣的1,000個提示生成了所有方法的圖像。成對的圖像呈現給42名受過訓練的人類標注者中的五名，他們的任務是為更具視覺吸引力的圖像投票。通過多數投票匯總的結果顯示在表2中，明顯表明了對Imagine Flash的偏好。
消融研究
對Imagine Flash進行了定量和定性消融，以評估所提出的反向蒸餾、SRL和噪聲校正的效果。定量評估如表3所示，而補充視覺消融如圖7所示
結論
Imagine Flash，這是一種新穎的蒸餾框架，能夠利用擴散模型實現高保真度的少步驟圖像生成。方法包括三個關鍵組成部分：反向蒸餾以減少訓練推理差異，遷移重構損失（SRL）動態地調整每個時間步的知識傳遞，以及噪聲校正以增強初始樣本質量。
通過大量實驗，Imagine Flash取得了顯著的成果，僅使用三個去噪步驟就與預訓練的教師模型的性能相匹配，并始終超越現有方法。這種前所未有的抽樣效率結合了高樣本質量和多樣性，使我們的模型非常適用于實時生成應用。
我們的工作為超高效的生成建模鋪平了道路。未來的方向包括擴展到其他模態，如視頻和3D，進一步降低抽樣預算，并將我們的方法與互補的加速技術相結合。通過實現即時的高保真度生成，Imagine Flash為實時創意工作流程和交互式媒體體驗開啟了新的可能性。
更多生成的圖像
本文轉自 AI生成未來 ，作者：Meta
