新智元報道
編輯:KingHZ
【新智元導讀】何愷明團隊又一力作!這次他們帶來的是「生成模型界的降維打擊」——MeanFlow:無需預訓練、無需蒸餾、不搞課程學習,僅一步函數評估(1-NFE),就能碾壓以往的擴散與流模型!
何愷明有新論文了!
全新的生成模型MeanFlow,最大亮點在于它徹底跳脫了傳統訓練范式——無須預訓練、蒸餾或課程學習,僅通過一次函數評估(1-NFE)即可完成生成。
MeanFlow在ImageNet 256×256上創下3.43 FID分數,實現從零開始訓練下的SOTA性能。
圖1(上):在ImageNet 256×256上從零開始的一步生成結果
在ImageNet 256×256數據集上,MeanFlow在一次函數評估(1-NFE)下達到了3.43的FID分數,性能相比此前同類最佳方法有50%到70%的相對提升(見圖1左)。
此外,MeanFlow訓練過程從零開始,無需預訓練、蒸餾或課程學習。
圖1(左):算力和一次函數評估FID分數
其中iCT、Shortcut和MF都是一次函數評估(1-NFE),而IMM則使用了兩次函數評估(2-NFE)的引導策略。
具體數值見表2。
表2:ImageNet-256×256上的類別條件生成實驗,不同模型的參數、FID得分等統計數據
值得一提的是,作者共有5位,其中4位是華人,均來自CMU和MIT兩所頂校。
其中一作耿正陽,是CMU的博士生,在MIT訪問時完成了這次研究的部分工作。
論文鏈接:https://arxiv.org/abs/2505.13447
在新論文中,研究者提出了系統且高效的一步生成建模框架MeanFlow。
傳統Flow Matching依賴建模瞬時速度場,而MeanFlow首創性地引入平均速度場(Mean Velocity Field)這一概念。
平均速度是指「位移/時間間隔」的比值,本質上是對瞬時速度在時間軸上的積分。
僅基于這一定義,研究者推導出了平均速度與瞬時速度之間清晰且內在的數學關系,這為神經網絡訓練提供了理論依據。
在這一基本概念之上,直接訓練神經網絡,對平均速度場建模。
為此,研究者設計了新的損失函數,引導網絡去滿足平均速度與瞬時速度之間的內在關系,無需引入額外的啟發式方法。
由于存在明確定義的目標速度場,理論上最優解與網絡的具體結構無關,這種屬性有助于訓練過程更加穩健和穩定。
此外,新方法還能自然地將「無分類器引導」(Classifier-Free Guidance,CFG)融入目標速度場,在采樣階段使用引導時不會帶來額外的計算開銷。
詳細結果
在圖1和表2(左側)中,研究者將MeanFlow與現有的一步擴散/流模型進行了比較。
總體來看,MeanFlow在同類方法中表現顯著優越:
新模型在ImageNet 256×256上實現了3.43的FID分數,相比IMM的7.77,相對提升超過50%;
如果僅比較1-NFE(一次函數評估)的生成結果,MeanFlow相比此前的最優方法Shortcut(FID 10.60),相對提升接近70%。
這表明,MeanFlow在很大程度上縮小了一步與多步擴散/流模型之間的性能差距。
在2-NFE(兩次函數評估)設定下,新方法取得了2.20的FID分數(見表2左下角)。
這個結果已經可以媲美許多多步方法的最優基線。
它們都采用了XL/2級別的骨干網絡,且NFE達到250×2(見表2右側)。
這進一步表明,少步數的擴散/流模型已經具備挑戰多步模型的潛力。
此外,未來還能進一步提升性能。
圖5:1-NFE生成結果示例
在CIFAR-10數據集(32×32)上,研究人員進行了無條件生成實驗,結果列在表3中。
使用1-NFE采樣時,他們使用FID-50K分數作為性能指標。
所有方法均采用相同的U-Net架構(約5500萬參數)。
需要注意的是,其他所有對比方法均使用了EDM風格的預處理器(pre-conditioner),而新方法沒有使用任何預處理器。
在CIFAR-10這個數據集上,新方法在性能上與現有方法具有競爭力。
表3:CIFAR-10無條件生成結果
前身:流匹配
流匹配(Flow Matching,簡稱FM)是一種生成建模范式。
Flow Matching將「連續歸一化流」(Continuous Normalizing Flows,CNFs)與「擴散模型」(Diffusion Models,DMs)的一些關鍵思想相結合,從而緩解了這兩類方法各自存在的核心問題。
形式上,給定數據x~pdata(x)和先驗噪聲?~pprior(?),可以構造一條流動路徑,
其中t表示時間,a_t和b_t是預設的調度函數。
路徑的速度定義為
這個速度被稱為條件速度(conditional velocity)。參見圖2左側部分。
Flow Matching本質上是在對所有可能情況的期望進行建模,這種平均后的速度稱為邊緣速度(marginal velocity)(見圖2右側):
圖2:Flow Matching中的速度場示意圖。左圖:條件流(ConditionalFlows)。同一個z_t可能由不同的(x,?)組合生成,因此會對應不同的條件速度v_t。右圖:邊緣流(Marginal Flows)。通過對所有可能的條件速度進行邊緣化(求平均)得到邊緣速度場。這個邊緣速度場被作為訓練神經網絡時的「真實目標速度場」
圖例說明:灰點表示從先驗分布中采樣得到的樣本,紅點表示來自真實數據分布的樣本。
接著,學習由參數θ表示的神經網絡v_θ,來擬合這個邊緣速度場,其損失函數為:
但由于式(1)中的邊緣化過程難以直接計算,因此Flow Matching提出使用條件Flow Matching損失來代替:
其中目標速度v_t是條件速度。
可以證明,最小化上述兩個損失函數是等價的。
一旦得到了邊緣速度場v(z_t,t),就可以通過求解下面的常微分方程(ODE)來生成樣本:
初始值為z_1=?,上述微分方程的解可以寫成積分形式:
其中r表示另一個時間點。
在實際中,這個積分通常通過數值方法在離散時間步上進行近似。
值得注意的是,即便條件流被設計為「直線流動」(即所謂「校正流」),最終得到的邊緣速度場(公式(1))往往仍會誘導出彎曲的軌跡(見圖2的示意)。
這種軌跡的彎曲不僅僅是因為神經網絡的近似誤差,更是源于真實的邊緣速度場本身。
當對這些彎曲軌跡使用粗粒度的時間離散化時,數值ODE解法往往會產生較大的誤差,從而導致生成結果不準確。
MeanFlow模型
平均流(Mean Flows)的核心思想是:引入一個表示平均速度的新場(velocity field),而傳統Flow Matching所建模的是瞬時速度。
平均速度定義
平均速度被定義為兩個時間點t和r之間的位移(通過對瞬時速度積分獲得),再除以時間間隔。
形式上,平均速度u定義如下:
為了突出概念上的區別,統一用u表示平均速度,用v表示瞬時速度。
平均速度場u(z_t,r,t)同時依賴于起始時間r和終止時間t,如圖3所示。
圖3:平均速度場
需要注意的是,平均速度u本質上是瞬時速度v的泛函結果。
因此,平均速度場是由瞬時速度場決定的,并不依賴于任何神經網絡。
從概念上講,就像在Flow Matching中,瞬時速度v是訓練的「真實目標場」,在MeanFlow中,平均速度u則扮演著類似的角色,是學習所依據的「真實速度場」。
MeanFlow模型的最終目標是:用神經網絡近似平均速度場。
這樣做的優勢顯著:一旦平均速度被準確建模,就可以僅通過一次前向計算來近似整個流動路徑。
換句話說,這種方法非常適合一步或少步數的生成任務,因為它在推理階段不需要顯式計算時間積分——這是傳統建模瞬時速度方法所必須的步驟。
不過,在實踐中,直接使用公式(3)定義的平均速度作為訓練網絡的「真值」行不通,因為這要求在訓練時就對瞬時速度執行積分,計算成本高且不可行。
研究人員的關鍵見解是:可以對平均速度的定義公式進行數學變形,從而構造一個更易于訓練的優化目標,即使在只能訪問瞬時速度的前提下依然可行。
MeanFlow恒等式
為了得到適合訓練的形式,平均速度的定義公式(3)被重新改寫為:
接著,對這個等式的兩邊關于t求導(把r當作常數),然后運用函數積的求導法則和微積分基本定理,得到:
整理上式,即可得到核心的MeanFlow恒等式:
它刻畫了平均速度u和瞬時速度v之間的本質聯系。
需要說明的是,公式(6)與之前的積分公式(4)是等價的(詳見原文附錄B.3)。
在MeanFlow恒等式中,公式右側給出了可以作為訓練目標的形式,可以利用它構建損失函數,來訓練神經網絡預測u(z_t,r,t)。
為了構建這個損失函數,還需要進一步分解其中的時間導數項。
時間導數的計算
要計算公式(6)右側第二項全導數(total derivative),它可以用偏導數展開如下:
將導數關系帶入后得到:
這提供了另一種表達u和v關系的方式。
利用神經網絡自動微分,在訓練時高效計算時間導數項。
利用平均速度進行訓練
到目前為止,上述公式還沒有涉及任何網絡參數。現在引入可學習的模型u_θ,并希望它滿足MeanFlow恒等式(公式(6))。
研究者定義如下的損失函數來優化網絡參數:
其中,u_tgt是通過MeanFlow恒等式構造的訓練目標:
這個目標的幾個關鍵點如下:
訓練信號來自于瞬時速度v,不需要積分操作,因此相比平均速度定義式(3)更容易實現。
雖然公式中出現了對u的偏導數,但實際訓練中使用的是網絡輸出uθ的梯度(自動微分實現)。
使用了stop-gradient操作(記為sg):這是為了避免「二階反向傳播」,從而減小優化的計算負擔。
需要說明的是,即使在優化中進行了這些近似,只要u_θ最終能夠使損失為零,它就一定滿足MeanFlow恒等式,從而也滿足最初的平均速度定義。
條件速度替代邊緣速度
在公式(10)中的v(z_t,t)是Flow Matching中的邊緣速度(見圖2右),但它難以直接計算。
因此,借鑒Flow Matching已有的做法,使用條件速度(見圖2左)來替代:
這里vt=at′x+bt′?是條件速度,在默認設定下vt=??x。
論文鏈接:https://arxiv.org/abs/2210.02747
在算法1中,jvp操作(Jacobian-vector product)非常高效。
使用MeanFlow模型進行采樣非常簡單:只需將時間積分項替換為平均速度即可,偽代碼詳見算法2。
帶引導的MeanFlow
新方法能夠自然支持無分類器引導(Classifier-Free Guidance,CFG)。
與傳統做法在采樣階段直接應用CFG不同,研究者將CFG視為底層「真實速度場」的一部分屬性。
這種建模方式可以在保留CFG效果的同時,仍保持采樣時的1-NFE性能。
構建真實速度場
研究者定義新的帶引導的真實速度場vcfg:
這是一個類別條件場(class-conditional field)與無條件場(class-unconditional field)的線性組合。
其中,類別條件速度(即對給定類別c條件下的邊緣速度)、無條件邊緣速度,定義如下:
接下來,我們仿照MeanFlow的方式,為vcfg引入對應的平均速度。
根據MeanFlow恒等式(公式6),我們有:
我們再次強調,vcfg和ucfg都是理論上的真實速度場,與神經網絡參數無關。
此外,由公式(13)和MeanFlow恒等式導出:
這可以簡化計算。
帶引導的訓練方法
神經網絡ucfg,θ來擬合平均速度場,需要構造如下訓練目標:
其中目標值為:
這里的右側第一項是結合引導權重后的速度定義:
其中v_t是樣本條件速度,默認設定為vt=??x。
如果ω=1,即純類別條件引導,則損失函數退化為不含CFG的公式(9)。
stop-gradient操作用于阻斷目標對網絡參數的反向傳播,避免二階梯度計算。
此外,為了增強網絡對無類別輸入的泛化能力,以0%概率隨機丟棄類別條件。
單NFE下CFG采樣
在本方法中,網絡直接學習的是由引導速度vcfg所誘導的平均速度。
因此,在采樣階段,無需再進行線性組合計算,只需直接網絡調用即可完成一步采樣(見算法2)。
最終,新在保留CFG效果的同時,依然維持了理想的單步采樣性能(1-NFE),兼顧了效率與質量。
作者介紹
耿正陽(Zhengyang Geng)
耿正陽,卡內基梅隆大學(CMU)計算機科學博士生。
他熱衷于研究動態系統,致力于識別、理解并開發能夠自組織形成復雜系統的動態機制。
2020年,他畢業于四川大學,獲得計算機科學與技術學士學位。
他在北京大學、Meta等機構實習過多次。
參考資料:
https://arxiv.org/abs/2505.13447
https://mlg.eng.cam.ac.uk/blog/2024/01/20/flow-matching.html
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.