文章鏈接：https://arxiv.org/pdf/2410.19310

亮點(diǎn)直擊

• 創(chuàng)新性的一步蒸餾框架：提出了一種名為“流生成匹配 (FGM)”的概率框架，通過一步蒸餾的方式將流匹配模型的采樣過程簡(jiǎn)化為單步生成，大幅提高計(jì)算效率，保持了原始模型的生成質(zhì)量。
• 對(duì)流匹配模型的蒸餾優(yōu)化：FGM 彌補(bǔ)了流匹配模型在蒸餾方面的研究空白，實(shí)現(xiàn)了流匹配模型的一步采樣，克服了該模型在采樣過程中的計(jì)算資源限制。
• 出色的生成表現(xiàn)：在 CIFAR10 和大規(guī)模文本生成圖像等任務(wù)上，F(xiàn)GM 在單步生成中取得了優(yōu)異的生成表現(xiàn)，與多步生成模型的性能相當(dāng)，并刷新了流匹配模型的一步生成記錄。
• 提升流模型的實(shí)用性：FGM 不僅擴(kuò)展了蒸餾技術(shù)的應(yīng)用場(chǎng)景，還顯著提升了流匹配模型的實(shí)際應(yīng)用潛力，特別適合需要快速高效采樣的應(yīng)用場(chǎng)景。

總結(jié)速覽

解決的問題
在AIGC領(lǐng)域中，流匹配模型因其堅(jiān)實(shí)的理論基礎(chǔ)和強(qiáng)大的大規(guī)模生成能力而取得了顯著成功。然而，這些模型的采樣過程對(duì)計(jì)算資源要求極高，因?yàn)樗鼈冃枰嗖綌?shù)值常微分方程 (ODEs) 的計(jì)算。

提出的方案
本文提出了一種名為“流生成匹配 (Flow Generator Matching, FGM)”的創(chuàng)新方法，通過將流匹配模型的采樣加速為一步生成，降低了對(duì)計(jì)算資源的需求，同時(shí)保證了模型性能。

應(yīng)用的技術(shù)
FGM 模型在 CIFAR10 無(wú)條件生成基準(zhǔn)測(cè)試上獲得了優(yōu)異的成績(jī)，創(chuàng)造了3.08的最新 Fréchet Inception Distance (FID) 分?jǐn)?shù)，超越了原有的50步流匹配模型。此外，通過 FGM 對(duì) Stable Diffusion 3 進(jìn)行了蒸餾，生成了一步文本生成圖像的 MM-DiT-FGM 模型。

達(dá)到的效果
在 GenEval 基準(zhǔn)測(cè)試上，MM-DiT-FGM 展現(xiàn)出卓越的生成效果，在單步生成的高效性下，生成質(zhì)量可與多步模型相媲美。

流生成匹配

本節(jié)介紹流生成匹配(Flow Generator Matching, FGM)，這是一種專為流匹配模型設(shè)計(jì)的一步蒸餾設(shè)計(jì)的通用方法。首先定義問題設(shè)置和符號(hào)，然后介紹我們的匹配目標(biāo)函數(shù)以及 FGM 如何最小化該目標(biāo)。最后，將 FGM 與現(xiàn)有的流蒸餾方法進(jìn)行比較，突出我們方法的經(jīng)驗(yàn)和理論優(yōu)勢(shì)。

問題設(shè)置

問題表述。本文的框架基于一個(gè)預(yù)訓(xùn)練的流匹配模型，該模型能夠準(zhǔn)確地逼近邊際向量場(chǎng) 。流  在噪聲分布和數(shù)據(jù)分布之間建立了聯(lián)系。同時(shí)，也知道條件轉(zhuǎn)移 ，這意味著 。假設(shè)預(yù)訓(xùn)練的流匹配模型對(duì)數(shù)據(jù)分布有足夠好的逼近，即  是真實(shí)的數(shù)據(jù)分布。

本文的目標(biāo)是訓(xùn)練一個(gè)一步生成器模型 ，直接將隨機(jī)噪聲  轉(zhuǎn)換為樣本 。令  表示學(xué)生模型生成樣本  的分布， 表示通過  轉(zhuǎn)移的邊際概率路徑。

該學(xué)生邊際概率路徑隱式地引導(dǎo)了生成該路徑的流向量場(chǎng) ，該向量場(chǎng)未知且難以處理。

不可處理的目標(biāo)。一步流生成匹配的目標(biāo)是使學(xué)生分布  與數(shù)據(jù)分布  匹配。為此，考慮將邊際向量場(chǎng)  與預(yù)訓(xùn)練的向量場(chǎng)  進(jìn)行匹配，從而使得分布  和  能夠彼此匹配。

本節(jié)定義流生成匹配的目標(biāo)?；谥暗挠懻摚繕?biāo)是最小化隱式向量場(chǎng)  和預(yù)訓(xùn)練流模型向量場(chǎng)  之間的期望  距離，表示為

注意到樣本  依賴于參數(shù) 。在必要時(shí)，可以使用  來(lái)強(qiáng)調(diào)這種參數(shù)依賴關(guān)系。

顯然，當(dāng)且僅當(dāng)所有誘導(dǎo)向量場(chǎng)一致時(shí)，，即  幾乎處處對(duì)  成立。因此，這意味著  幾乎處處對(duì)  成立，這表明兩個(gè)分布  幾乎處處對(duì)  匹配。然而，盡管最小化目標(biāo)可以得到一個(gè)強(qiáng)大的一步生成器，但它是不可處理的，因?yàn)槲覀儾恢?nbsp; 和生成器分布  之間的關(guān)系。下面介紹本文的主要貢獻(xiàn)：一個(gè)具有理論保證的、與（4.1）等價(jià)且可處理的訓(xùn)練目標(biāo)。

可處理的目標(biāo)

本文的目標(biāo)是優(yōu)化參數(shù)  以最小化目標(biāo)。然而，隱式向量場(chǎng)  是未知且不可處理的，因此無(wú)法直接最小化該目標(biāo)。不過，通過對(duì)  求損失函數(shù)的梯度，可以得到

其中， 和  定義為

公式 (4.3) 中的梯度考慮了關(guān)于參數(shù)  的所有導(dǎo)數(shù)。

請(qǐng)注意，第一個(gè)梯度  可以通過停止對(duì)  的 -梯度來(lái)獲得，即 。這導(dǎo)致了一個(gè)替代損失函數(shù)，其梯度與  一致。

然而，第二個(gè)梯度 (4.5) 包含一個(gè)不可處理的項(xiàng) 。對(duì)于學(xué)生生成器，我們只能從條件概率路徑中有效地采樣，但向量場(chǎng)  及其關(guān)于  的梯度未知。幸運(yùn)的是，在本文中，我們提出了如下定理 4.2，使得學(xué)生向量場(chǎng)的 -梯度更加可處理。在此之前，我們需要首先引入一個(gè)新的流乘積恒等式（Flow Product Identity），這是我們的貢獻(xiàn)之一，即定理 4.1。

定理 4.1 （流乘積恒等式） 設(shè)  為一個(gè)向量值函數(shù)，使用 4.1 節(jié)中的符號(hào)，在溫和的條件下，該恒等式成立：

接下來(lái)，展示可以引入一個(gè)等效的可處理?yè)p失函數(shù)，該損失函數(shù)與不可處理的損失函數(shù) (4.1) 在參數(shù)梯度上相同，如定理 4.2 所示。

定理 4.2. 如果分布  滿足某些廣義正則性條件，則對(duì)于所有不依賴于參數(shù)的向量值函數(shù) ，對(duì)于所有參數(shù) ，該等式成立：

公式 (4.8) 表明，無(wú)法處理的梯度  的期望可以與來(lái)自學(xué)生模型的可微樣本的可處理期望進(jìn)行互換。

恒等式 (4.8) 的直接結(jié)果是，梯度 Grad2() 與以下可處理?yè)p失函數(shù) (4.9) 一致，該損失函數(shù)在生成器向量中對(duì)  施加了停止梯度操作 ：

將 (4.6) 和 (4.9) 結(jié)合到 (4.3) 中，得到一個(gè)等效的損失函數(shù)，用于最小化原始目標(biāo)，即：

定義在 (4.6) 和 (4.9) 中的 L1(θ) 和 L2(θ)。這產(chǎn)生了通過最小化損失函數(shù) (4.10) 提出的流生成匹配 (FGM) 目標(biāo)。下算法 1 總結(jié)了流生成匹配的偽算法，將預(yù)訓(xùn)練的流匹配模型蒸餾成一個(gè)單步生成器。

與擴(kuò)散蒸餾的區(qū)別 FGM 從單步擴(kuò)散蒸餾中獲得靈感，旨在最小化分布發(fā)散，然而，所得到的理論與單步擴(kuò)散蒸餾的理論本質(zhì)上是不同的。FGM 和單步擴(kuò)散蒸餾之間最顯著的區(qū)別在于，流匹配并不意味著對(duì)概率密度進(jìn)行顯式建模，如擴(kuò)散模型所做。因此，分布發(fā)散的定義不能應(yīng)用于流模型及其蒸餾。然而，F(xiàn)GM 通過直接處理流匹配目標(biāo)而不是分布發(fā)散來(lái)克服這個(gè)問題。主要的見解是我們提出的顯式-隱式梯度等效理論繞過了不可處理的流匹配目標(biāo)，從而產(chǎn)生具有理論保證的強(qiáng)大實(shí)用算法。定理 4.2 也可能為未來(lái)對(duì)流匹配模型的研究帶來(lái)新的貢獻(xiàn)。

與其他流蒸餾方法的比較 目前有少數(shù)現(xiàn)有工作嘗試將流模型加速到單步或少步生成模型。一致性流匹配 (CFM)是最近的一項(xiàng)工作，將預(yù)訓(xùn)練的流模型蒸餾為一個(gè)或兩個(gè)步的模型。盡管 CFM 顯示出不錯(cuò)的結(jié)果，但在理論和實(shí)際方面與我們的 FGM 不同。首先，CFM 背后的理論建立在流模型的軌跡一致性上，這直接從一致性模型中推廣。相反，我們的 FGM 是以流匹配目標(biāo)為出發(fā)點(diǎn)，試圖用真值教師流來(lái)訓(xùn)練單步生成器的隱式流，并具有理論保證。在實(shí)際方面，在 CIFAR10 生成上，展示了訓(xùn)練的單步 FGM 模型在流模型中達(dá)到了 3.08 的新的 SoTA FID，超過了 CFM 最好的 2 步生成結(jié)果 5.34 的 FID。如此強(qiáng)大的實(shí)證表現(xiàn)使得 FGM 成為在標(biāo)準(zhǔn)基準(zhǔn)上加速流匹配模型的可靠解決方案。除了簡(jiǎn)單的 CIFAR10 生成，在第 5.2 節(jié)中，我們還使用 FGM 蒸餾領(lǐng)先的大規(guī)模文本到圖像流模型，獲得了幾乎沒有性能下降的非常強(qiáng)大的單步文本到圖像模型。

實(shí)驗(yàn)

這里進(jìn)行了實(shí)驗(yàn)以評(píng)估 FGM 的有效性和靈活性。實(shí)驗(yàn)涵蓋了標(biāo)準(zhǔn)評(píng)估基準(zhǔn)、無(wú)條件 CIFAR10 圖像生成和使用Stable Diffusion 3 (SD3)的大規(guī)模文本到圖像生成。這些實(shí)驗(yàn)展示了 FGM 在構(gòu)建高效單步生成器的能力，同時(shí)保持高質(zhì)量樣本。

單步 CIFAR10 生成

實(shí)驗(yàn)設(shè)置 首先在 CIFAR10 數(shù)據(jù)集上評(píng)估 FGM 的有效性，這是生成模型性能的標(biāo)準(zhǔn)測(cè)試平臺(tái)。在 CIFAR10 上使用 ReFlow 目標(biāo) (3.6) 預(yù)訓(xùn)練流匹配模型進(jìn)行條件和無(wú)條件生成。參考了用于 EDM 模型的神經(jīng)網(wǎng)絡(luò)架構(gòu)。以 512 的批量大小訓(xùn)練條件和無(wú)條件模型 20,000 張圖像，得到的內(nèi)部訓(xùn)練流模型顯示出 CIFAR10 無(wú)條件 FID 為 2.52，生成步為 300，這略微低于原始 ReFlow 模型，其 FID 為 2.58，使用 127 個(gè)生成步。然而，在下表 1 中，發(fā)現(xiàn)這樣的略微較差模型并不會(huì)影響強(qiáng)單步生成器的蒸餾。

這些流模型作為流生成匹配 (FGM) 的教師模型。然后應(yīng)用 FGM 從流模型中蒸餾單步生成器。通過FID評(píng)估生成圖像的質(zhì)量。較低的 FID 分?jǐn)?shù)表示樣本質(zhì)量和多樣性更高。

注意損失 (4.11) 和損失 (4.12) 共同構(gòu)成 FGM 損失的完整參數(shù)梯度。我們發(fā)現(xiàn)這兩個(gè)損失在使用多層感知機(jī)進(jìn)行 2D 數(shù)據(jù)集生成時(shí)效果很好。在實(shí)踐中，發(fā)現(xiàn)使用損失 (4.11) 在 CIFAR10 模型上會(huì)導(dǎo)致不穩(wěn)定，這與 Poole et al. (2022) 的觀察相似，即其雅可比項(xiàng)的條件數(shù)可能不適定。因此，在訓(xùn)練時(shí)我們不使用損失 (4.11)，而是觀察到良好的性能。

用預(yù)訓(xùn)練流模型初始化生成器 受擴(kuò)散蒸餾技術(shù)的啟發(fā)，用預(yù)訓(xùn)練流模型初始化單步生成器?；貞浟髂Ｐ偷挠?xùn)練目標(biāo) (3.6)，預(yù)訓(xùn)練流模型 θ 大致預(yù)測(cè)從隨機(jī)噪聲到數(shù)據(jù)的方向。因此，我們使用預(yù)訓(xùn)練流來(lái)構(gòu)建我們的單步生成器。具體來(lái)說(shuō)，構(gòu)建單步生成器：

θ 是生成器的可學(xué)習(xí)參數(shù)，而 是預(yù)先確定的最佳時(shí)間步。

定量評(píng)估 使用 Fretchet Inception Distance（FID）評(píng)估每個(gè)模型，這是評(píng)估 CIFAR10 數(shù)據(jù)集上圖像生成結(jié)果的黃金標(biāo)準(zhǔn)。上表1和表2總結(jié)了生成模型在 CIFAR10 數(shù)據(jù)集上的 FID。在無(wú)條件生成任務(wù)中，我們的教師流模型在50個(gè)生成步和100個(gè)生成步下分別獲得了 3.67 和 2.93 的 FID。FGM 模型在僅使用一個(gè)生成步的情況下取得了 3.08 的 FID，顯著超越了教師模型在50個(gè)生成步下的結(jié)果，幅度達(dá)到了16%。在 CIFAR10 條件生成任務(wù)中，我們的一步 FGM 模型獲得了 2.58 的 FID，超越了教師流模型在100個(gè)生成步下的 2.87 FID。在 CIFAR10 生成基準(zhǔn)上的結(jié)果展示了 FGM 的優(yōu)越性能，它能夠在顯著幅度上超越多步教師流模型。

CIFAR-10 生成任務(wù)相對(duì)簡(jiǎn)單。進(jìn)行了實(shí)驗(yàn)，訓(xùn)練大型的一步文本到圖像生成器，采用來(lái)自頂級(jí)性能的基于Transformer的流模型進(jìn)行蒸餾。

文本到圖像生成實(shí)驗(yàn)：

下面展示了通過 FGM 蒸餾的一步 T2I 生成器在其他行業(yè)級(jí)模型中展現(xiàn)出最先進(jìn)的結(jié)果。在進(jìn)一步的實(shí)驗(yàn)中，F(xiàn)GM 被應(yīng)用于從頂尖的流匹配模型（如Stable Diffusion 3 Medium）中訓(xùn)練單步文本到圖像生成器。用于訓(xùn)練的大型數(shù)據(jù)集包含大約 200 萬(wàn)條記錄，增強(qiáng)了模型生成高質(zhì)量圖像的能力。盡管在使用某些損失函數(shù)時(shí)訓(xùn)練動(dòng)態(tài)存在初始不穩(wěn)定性，F(xiàn)GM 仍然保持了與多步模型的競(jìng)爭(zhēng)力，展示了其在處理復(fù)雜文本輸入生成圖像方面的魯棒性。

總的來(lái)說(shuō)，提出的 FGM 框架為生成模型的效率和有效性帶來(lái)了有希望的進(jìn)展，為流匹配和文本到圖像生成技術(shù)的未來(lái)發(fā)展鋪平了道路。

定性評(píng)估 對(duì)所提出的蒸餾方法的定性評(píng)估，以分析其性能。下圖2展示了多個(gè)樣本輸出，比較了教師模型 Hyper-SD3和 Flash-SD3的方法。結(jié)果顯示，即使在僅使用單步生成的情況下，視覺質(zhì)量依然很高，尤其是在細(xì)節(jié)和顏色再現(xiàn)方面。特別是，一步的 MM-DiT-FGM 在每個(gè)生成圖像上展現(xiàn)了美學(xué)上的光影效果。與現(xiàn)有的蒸餾方法相比，我們的模型在生成質(zhì)量上達(dá)到了可比水平，但成本顯著降低。這種優(yōu)勢(shì)使得 FGM 在需要實(shí)時(shí)交互的應(yīng)用中變得切實(shí)可行。

GAN 損失的整合 顯然，純 FGM 算法在訓(xùn)練時(shí)并不依賴任何圖像數(shù)據(jù)。近年來(lái)，許多研究表明，在蒸餾中引入 GAN 損失對(duì)于提高生成圖像的高頻細(xì)節(jié)是有益的。因此，也將 GAN 損失與 FGM 結(jié)合用于訓(xùn)練單步文本到圖像模型，并發(fā)現(xiàn)了其益處。

在訓(xùn)練過程中，觀察到在 FGM 效率較低的某些噪聲調(diào)度間隔中，GAN 損失能夠提供有效的梯度，從而改善模型輸出的質(zhì)量。因此，GAN 損失的一個(gè)重要優(yōu)勢(shì)是能夠補(bǔ)償 FGM 在某些噪聲調(diào)度中的效率不足，從而對(duì)我們的損失函數(shù)形成補(bǔ)充。

結(jié)論

本文介紹了流生成匹配（FGM），一種強(qiáng)有力的概率性單步蒸餾方法，專為流匹配模型而設(shè)計(jì)。本文建立了 FGM 的理論基礎(chǔ)，并驗(yàn)證了其在單步 CIFAR10 生成和大規(guī)模單步文本到圖像生成方面的強(qiáng)大經(jīng)驗(yàn)表現(xiàn)。

盡管 FGM 具有堅(jiān)實(shí)的理論基礎(chǔ)和強(qiáng)大的經(jīng)驗(yàn)表現(xiàn)，但它仍然存在一些局限性。第一個(gè)局限性是目前 FGM 仍需額外的流模型來(lái)近似生成器引起的流向量。這一要求增加了蒸餾的內(nèi)存開銷，并在預(yù)訓(xùn)練的流模型和生成器的模型大小較大時(shí)可能帶來(lái)挑戰(zhàn)。其次，F(xiàn)GM 是一種完全不依賴圖像數(shù)據(jù)的方法，這意味著在蒸餾時(shí)不需要真實(shí)的圖像數(shù)據(jù)。然而，眾所周知的一個(gè)論點(diǎn)是，持續(xù)引入高質(zhì)量圖像數(shù)據(jù)對(duì)于提高文本到圖像生成模型的性能至關(guān)重要。希望未來(lái)的工作能夠探索如何將數(shù)據(jù)整合到蒸餾過程中。

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

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