何凱明團隊新作：均值流用于一步生成建模 原創
摘要
我們提出了一種原則性且高效的一步生成模型框架。我們引入了平均速度的概念來表征流場，這與流匹配方法中建模的瞬時速度形成對比。我們推導出了平均速度與瞬時速度之間的明確關系，并用其指導神經網絡訓練。我們的方法，稱為均值流模型，是自包含的，無需預訓練、蒸餾或課程學習。均值流模型展示了強大的實證性能：在ImageNet 256×256上從頭訓練，僅用一次函數評估（1-NFE）即可達到3.43的Fréchet Inception Distance（FID），顯著優于之前最先進的一步擴散/流模型。我們的研究大幅縮小了一步擴散/流模型與其多步前驅之間的差距，我們希望它能激勵未來研究重新審視這些強大模型的基礎。
圖1：ImageNet 256×256上的一步生成
我們的均值流（MF）模型在生成質量上顯著優于之前最先進的一步擴散/流方法。此處，iCT [43]、Shortcut [13]和我們的MF均為1-NFE生成，而IMM的1步結果 [52]涉及2-NFE指導。詳細數據見表2。所示圖像由我們的1-NFE模型生成。
1、引言
生成建模的目標是將先驗分布轉換為數據分布。流匹配 [28,2,30] 為構建將一個分布映射到另一個分布的流路徑提供了一個直觀且概念簡單的框架。與擴散模型 [42,44,19] 密切相關，流匹配側重于指導模型訓練的速度場。自提出以來，流匹配已在現代生成建模中得到廣泛應用 [11,33,35]。
流匹配和擴散模型在生成過程中都需要迭代采樣。最近的研究對少步 —— 尤其是一步前饋 —— 生成模型給予了極大關注。一致性模型 [46,43,15,31] 率先在這一方向上為沿同一路徑采樣的輸入引入了網絡輸出的一致性約束。盡管取得了令人鼓舞的結果，但一致性約束是作為網絡行為的屬性強加的，而指導學習的潛在真實場的屬性仍不明確。因此，訓練可能不穩定，需要精心設計的 “離散化課程”[46,43,15] 來逐步約束時間域。
在這項工作中，我們提出了一種名為 MeanFlow 的原理性強且有效的一步生成框架。核心思想是引入一個表示平均速度的新真實場，以區別于流匹配中通常建模的瞬時速度。平均速度定義為位移與時間間隔的比率，其中位移由瞬時速度的時間積分給出。僅從這個定義出發，我們推導出平均速度與瞬時速度之間明確的內在關系，這自然成為指導網絡訓練的原理性基礎。
（我們在這項研究里搞了個叫MeanFlow的靠譜框架，專門解決“一步生成圖像”的問題。核心思路是換了個底層概念：以前的方法算的是“瞬時速度”（類似汽車某一秒的速度表讀數），得一步步調整才能生成圖像；而我們改用“平均速度”（類似一段路的平均速度），它等于“總位移（路程）÷時間”，其中總位移是把瞬時速度在時間里一點點加起來得到的。就憑這個定義，我們找出了平均速度和瞬時速度之間的數學關系，拿這個關系當“指導手冊”來訓練神經網絡，讓模型學一次就能直接生成圖，不用像以前那樣反復算了。）
基于這一基本概念，我們訓練神經網絡直接對平均速度場進行建模。我們引入一個損失函數，促使網絡滿足平均速度與瞬時速度之間的內在關系，無需額外的一致性啟發式方法。真實目標場的存在確保了最優解原則上獨立于特定網絡，這在實踐中可以帶來更穩健和穩定的訓練。我們進一步表明，我們的框架可以自然地將無分類器引導（CFG）[18] 納入目標場，在使用引導時采樣時不會產生額外成本。
我們的 MeanFlow 模型在一步生成建模中表現出強大的實驗性能。在 ImageNet 256×256 [7] 上，我們的方法使用 1-NFE 生成實現了 3.43 的 FID 分數。這一結果以 50% 到 70% 的相對優勢顯著優于同類中之前的最先進方法（圖 1）。此外，我們的方法是一個自包含的生成模型：它完全從頭開始訓練，無需任何預訓練、蒸餾或課程學習。我們的研究在很大程度上縮小了一步擴散 / 流模型與其多步前身之間的差距，希望能啟發未來的工作重新考慮這些強大模型的基礎。
2、相關工作
擴散與流匹配
在過去十年中，擴散模型 [42, 44, 19, 45] 發展成為生成模型的非常成功的框架。這些模型逐步向干凈數據添加噪聲，并訓練神經網絡來逆轉這一過程。這一過程涉及求解隨機微分方程（SDE），隨后被重新表述為概率流常微分方程（ODE） [45, 22]。流匹配方法 [28, 2, 30] 通過建模定義分布之間流路徑的速度場擴展了這一框架。流匹配也可以看作是連續時間歸一化流 [36] 的一種形式。
少步擴散/流模型
從實際和理論角度來看，減少采樣步驟已成為重要考量。一種方法是將預訓練的多步擴散模型蒸餾為少步模型，例如 [39, 14, 41] 或分數蒸餾 [32, 50, 53]。早期對少步模型訓練的探索 [46] 建立在基于蒸餾的方法演變之上。同時，一致性模型 [46] 作為獨立生成模型開發，不需要蒸餾。這些模型對不同時間步的網絡輸出施加一致性約束，鼓勵它們在軌跡上產生相同的端點。已研究了多種一致性模型和訓練策略 [46, 43, 15, 31, 49]。
在最近的工作中，幾種方法側重于刻畫基于兩個時間相關變量的擴散 / 流相關量。在 [3] 中，流圖被定義為兩個時間步之間的流的積分，并開發了幾種形式的匹配損失用于學習。與我們的方法所基于的平均速度相比，流圖對應于位移。Shortcut 模型 [13] 除了流匹配之外，還引入了自一致性損失函數，捕捉不同離散時間間隔的流之間的關系。歸納矩匹配 [52] 對不同時間步的隨機插值的自一致性進行建模。
3、背景：流匹配中的速度場
圖2：流匹配中的速度場 [28]
左圖：條件流 [28]。給定的 z_t 可以由不同的 (x, \epsilon) 對產生，導致不同的條件速度 v_t。右圖：邊際流 [28]，通過對所有可能的條件速度進行邊際化獲得。邊際速度場作為網絡訓練的底層真實場。圖示中的所有速度本質上是瞬時速度。插圖遵循 [12]。（灰色點：來自先驗的樣本；紅色點：來自數據的樣本。）
由于給定的 z_t 和其 v_t 可以由不同的 x 和 \epsilon 產生，流匹配本質上建模所有可能性的期望，稱為邊際速度 [28]（圖2右圖）：
給定邊際速度場 v(z_t, t)，通過求解 z_t 的常微分方程生成樣本：
4、均值流模型
4.1 均值流
我們方法的核心是引入一個表示平均速度的新場，而流匹配中建模的速度表示瞬時速度。
平均速度
我們定義平均速度為兩個時間步 t 和 r 之間的位移（通過積分獲得）除以時間間隔。形式上，平均速度 u 為：
圖3：平均速度 u(z, r, t) 的場
最左圖：瞬時速度 v 決定路徑的切線方向，而平均速度 u(z, r, t)（方程（3）定義）通常與 v 不對齊。平均速度與位移對齊，位移為 (t-r) u(z, r, t)。右三子圖：u(z, r, t) 的場依賴于 r 和 t，此處顯示 t=0.5, 0.7, 1.0 的情況。
均值流模型的最終目標是用神經網絡 u_\theta(z_t, r, t) 逼近平均速度。假設我們能夠精確逼近這一量，我們可以通過單次評估 u_\theta(\epsilon, 0, 1) 來逼近整個流路徑。換句話說，正如我們將通過實驗證明的，這種方法更適合單步或少步生成，因為它在推理時無需顯式逼近時間積分，而建模瞬時速度時需要這樣做。然而，直接使用方程（3）定義的平均速度作為訓練網絡的真實值是不可行的，因為它需要在訓練期間計算積分。我們的關鍵洞察是將平均速度的定義方程進行操作，構建一個適合訓練的優化目標，即使只能訪問瞬時速度。
均值流恒等式
為了得到適合訓練的公式，我們重寫方程（3）為：
現在我們對 t 進行微分，將 r 視為獨立于 t。這導致：
其中左側的操作使用了乘積法則，右側使用了微積分基本定理。重新排列項，我們得到身份：
我們將此方程稱為“均值流身份”，它描述了 v 和 u 之間的關系。易于證明方程（6）和方程（4）是等價的（見附錄B.3）。
方程（6）的右側為 u(z_t, r, t) 提供了“目標”形式，我們將利用它來構建訓練神經網絡的損失函數。為了作為合適的目標，我們必須進一步分解時間導數項，接下來討論。
計算時間導數
使用平均速度進行訓練
到目前為止，公式獨立于任何網絡參數化。我們現在引入一個模型來學習 u。形式上，我們參數化一個網絡 u_\theta 并鼓勵它滿足均值流身份（方程（6））。具體來說，我們最小化以下目標：
項 utgt 作為有效的回歸目標，由方程（6）驅動。此目標僅使用瞬時速度 v 作為真實信號，無需計算積分。雖然目標應涉及 u 的導數（即 ? u），但它們被替換為參數化的對應項（即 \partial u_\theta，?uθ）。在損失函數中，對目標 u_{\text{tgt}} 應用停止梯度（sg）操作，遵循常見實踐 [46, 43, 15, 31, 13]：在我們的案例中，它消除了通過雅可比向量積進行“雙重反向傳播”的需要，從而避免了高階優化。盡管采用了這些優化實踐，如果 u_\theta 實現零損失，易于證明它將滿足均值流身份（方程（6）），從而滿足原始定義（方程（3））。
方程（10）中的速度 v(z_t, t) 是流匹配中的邊際速度 [28]（見圖2右圖）。我們遵循 [28] 將其替換為條件速度（圖2左圖）。因此，目標為：
采樣
使用均值流模型進行采樣只需將時間積分替換為平均速度：
與先前工作的關系
雖然與先前的一步生成模型 [46, 43, 15, 31, 49, 23, 13, 52] 相關，但我們的方法提供了一個更原則性的框架。我們方法的核心是兩個底層場 v 和 u 之間的函數關系，這自然導致 u 必須滿足的均值流身份（方程（6））。此身份不依賴于神經網絡的引入。相比之下，先前工作通常依賴于施加在神經網絡行為上的額外一致性約束。一致性模型 [46, 43, 15, 31] 專注于錨定在數據側的路徑：在我們的符號中，這對應于對任何 t 固定 r \equiv 0。因此，一致性模型僅條件于單個時間變量，與我們的不同。另一方面，Shortcut [13] 和 IMM [52] 模型條件于兩個時間變量：它們引入了額外的雙時間自一致性約束。相比之下，我們的方法僅由平均速度的定義驅動，用于訓練的均值流身份（方程（6））從該定義自然推導，無需額外假設。
4.2 帶指導的均值流
我們的方法自然支持無分類器指導（CFG）[18]。我們不是在采樣時簡單應用 CFG（這會使 NFE 加倍），而是將 CFG 視為底層真實場的屬性。這種公式允許我們在采樣時保持 1-NFE 行為，同時享受 CFG 的好處。
真實場
我們構建一個新的真實場 v^{\text{cfg}}：
這是類條件場和類無條件場的線性組合：
\[
v(z_t, t \mid \mathbf{c}) \triangleq \mathbb{E}_{p_t(v_t \mid z_t, \mathbf{c})}[v_t] \quad \text{和} \quad v(z_t, t) \triangleq \mathbb{E}_{\mathbf{c}}[v(z_t, t \mid \mathbf{c})],
\]
其中 v_t 是條件速度 [28]（更精確地說，在此上下文中是樣本條件速度）。遵循均值流的精神，我們引入與 v^{\text{cfg}} 對應的平均速度 u^{\text{cfg}}。根據均值流身份（方程（6）），u^{\text{cfg}} 滿足：
\[
u^{\text{cfg}}(z_t, r, t \mid \mathbf{c}) = v^{\text{cfg}}(z_t, t \mid \mathbf{c}) - (t-r) \fracowocgsc{dt} u^{\text{cfg}}(z_t, r, t \mid \mathbf{c}),
\]
同樣，v^{\text{cfg}} 和 u^{\text{cfg}} 是獨立于神經網絡的底層真實場。此處，v^{\text{cfg}} 如方程（13）定義，可以重寫為：
\[
v^{\text{cfg}}(z_t, t \mid \mathbf{c}) = \omega v(z_t, t \mid \mathbf{c}) + (1-\omega) u^{\text{cfg}}(z_t, t, t),
\]
我們利用關系：v(z_t, t) = v^{\text{cfg}}(z_t, t)，以及 v^{\text{cfg}}(z_t, t) = u^{\text{cfg}}(z_t, t, t)。
帶指導的訓練
通過方程（15）和方程（16），我們構建網絡及其學習目標。我們直接通過函數 u_\theta^{\text{cfg}} 參數化 u^{\text{cfg}}。基于方程（15），我們得到目標：
其中 u_{\text{tgt}} = \bar{v}_t - (t-r) (\bar{v}_t \partial_z u_\theta^{\text{cfg}} + \partial_t u_\theta^{\text{cfg}})。
此公式類似于方程（9），唯一區別在于它具有修改后的 \bar{v}_t：
\[
\bar{v}_t \triangleq \omega v_t + (1-\omega) u_\theta^{\text{cfg}}(z_t, t, t),
\]
這由方程（16）驅動：方程（16）中的項 v(z_t, t \mid \mathbf{c})，即邊際速度，被替換為（樣本）條件速度 v_t，遵循 [28]。如果 \omega=1，此損失函數退化為無 CFG 情況（方程（9））。
為了使方程（17）中的網絡 u_\theta^{\text{cfg}} 暴露于類無條件輸入，我們以 10% 的概率丟棄類條件，遵循 [18]。出于類似動機，我們還可以使方程（19）中的 u_\theta^{\text{cfg}}(z_t, t, t) 暴露于類無條件和類條件版本：細節在附錄 B.1 中。
帶 CFG 的單 NFE 采樣
4.3 設計決策
損失度量
5、實驗
實驗設置
我們在 ImageNet [7] 上以 256×256 分辨率進行主要實驗。我們在 50K 張生成圖像上評估 Fréchet Inception Distance（FID）[17]。我們檢查函數評估次數（NFE），并默認研究 1-NFE 生成。遵循 [34, 13, 52]，我們在預訓練變分自編碼器（VAE）分詞器的潛在空間上實現我們的模型。對于 256×256 圖像，分詞器產生 32×32×4 的潛在空間，這是模型的輸入。我們的模型全部從頭訓練。實現細節在附錄 A 中。
在我們的消融研究中，我們使用 ViT-B/4 架構（即“基礎”大小，補丁大小為4）[9]，如 [34] 開發，訓練 80 個 epoch（400K 次迭代）。作為參考，[34] 中的 DiT-B/4 具有 68.4 FID，[33] 中的 SiT-B/4（我們的重現）具有 58.9 FID，兩者均使用 250-NFE 采樣。
5.1 消融研究
表1：ImageNet 256×256 數據集上 1 - NFE 生成的消融實驗 說明：評估指標為 FID - 50K。默認配置用灰色標注：使用 B/4 骨干網絡，從頭開始訓練 80 個 epoch。
我們在表1中研究模型屬性，分析如下：
從流匹配到均值流
JVP 計算
條件于 (r, t)
時間采樣器
損失度量
據報告 [43]，損失度量的選擇顯著影響少步/單步生成的性能。我們在表1e中研究這一方面。我們的損失度量通過自適應損失加權 [15] 實現，冪為 p（4.3節）。表1e 顯示，p=1 實現最佳結果，而 p=0.5（類似于 [43] 中的偽 Huber 損失）也具有競爭力。標準平方 L2 損失（此處 p=0）相較于其他設置表現不佳，但仍能產生有意義的結果，與 [43] 的觀察一致。
指導尺度
表1f 報告了帶 CFG 的結果。與多步生成中的觀察一致 [34]，CFG 在我們的 1-NFE 設置中也顯著提高了生成質量。我們強調，我們的 CFG 公式（4.2節）天然支持 1-NFE 采樣。
可擴展性
圖4：均值流模型在 ImageNet 256×256 上的可擴展性。
報告了 1-NFE 生成 FID。所有模型從頭訓練。應用 CFG，同時保持 1-NFE 采樣行為。我們的方法在模型大小方面表現出良好的可擴展性。
圖4：MeanFlow在不同模型規模和訓練時長下的1-NFE FID結果。該圖展示了使用不同模型規模和訓練時長訓練的MeanFlow模型在ImageNet 256×256上的1-NFE FID性能。每個點代表一個單獨訓練的模型。
5.2 與先前工作的比較
ImageNet 256×256 比較
在圖1中，我們與先前的一步擴散/流模型進行比較，總結在表2（左）中。總體而言，均值流在同類方法中大幅優于先前方法：它實現了 3.43 FID，相較于 IMM 的一步結果 7.77 [52] 有超過 50% 的相對改進；如果僅比較 1-NFE（不僅僅是一步）生成，均值流相較于之前最先進的 10.60（Shortcut [13]）有近 70% 的相對改進。我們的方法大幅縮小了一步與多步擴散/流模型之間的差距。
在 2-NFE 生成中，我們的方法實現了 2.20 FID（表2，左下）。此結果與領先的多步擴散/流模型基準相當，即 DiT [34]（FID 2.27）和 SiT [33]（FID 2.15），兩者均具有 250×2 NFE（表2，右），在相同的 XL/2 骨干下。我們的結果表明，少步擴散/流模型可以與多步前驅競爭。正交改進，如 REPA [51]，是適用的，留待未來工作。
值得注意的是，我們的方法是自包含的，完全從頭訓練。它在不使用任何預訓練、蒸餾或 [43, 15, 31] 中采用的課程學習的情況下實現了強大結果。
表2：ImageNet-256×256上的類條件生成
所有條目在適用時均報告 CFG。左：從頭訓練的 1-NFE 和 2-NFE 擴散/流模型。右：作為參考的其他生成模型家族。在兩表中，“×2”表示 CFG 每采樣步引入 2 NFE。我們的均值流模型均訓練 240 個 epoch，除了“MeanFlow-XL+”訓練更多 epoch 并使用為長期訓練選擇的配置，詳見附錄。?：iCT [43] 結果由 [52] 報告。
CIFAR-10 比較
我們在表3中報告了 CIFAR-10 [25]（32×32）上的無條件生成結果。報告了 50K FID，1-NFE 采樣。所有條目使用相同的 U-net [38]，由 [44] 開發（55M），直接應用于像素空間。所有其他競爭者使用 EDM 風格預處理器 [22]，而我們沒有預處理器。實現細節在附錄中。在此數據集上，我們的方法與先前方法具有競爭力。
表3：CIFAR-10 無條件生成
6、結論
我們提出了均值流，一個原則性且高效的一步生成框架。廣義上講，本工作考慮的場景與物理中的多尺度模擬問題相關，可能涉及空間或時間上的多種尺度、長度和分辨率。進行數值模擬本質上受到計算機分辨尺度范圍能力的限制。我們的公式涉及在較粗糙的粒度級別描述底層量，這是物理中許多重要應用的共同主題。我們希望我們的工作將連接生成模型、模擬和相關領域的動態系統研究。
本文轉載自???AIRoobt?? ，作者：Zhengyang Geng等
