繼今年5月提出MeanFlow (MF) 之后，何愷明團(tuán)隊于近日推出了最新的改進(jìn)版本——

Improved MeanFlow (iMF)，iMF成功解決了原始MF在訓(xùn)練穩(wěn)定性、指導(dǎo)靈活性和架構(gòu)效率上的三大核心問題。

其通過將訓(xùn)練目標(biāo)重新表述為更穩(wěn)定的瞬時速度損失，同時引入靈活的無分類器指導(dǎo)（CFG）和高效的上下文內(nèi)條件作用，大幅提升了模型性能。

在ImageNet 256x256基準(zhǔn)測試中，iMF-XL/2模型在 1-NFE（單步函數(shù)評估）中取得了1.72的FID成績，相較于原始MF提升了50%，證明了從頭開始訓(xùn)練的單步生成模型可以達(dá)到與多步擴(kuò)散模型相媲美的結(jié)果。

MeanFlow一作耿正陽依舊，值得注意的是共同一作的Yiyang Lu目前還是大二學(xué)生——來自清華姚班，而何愷明也在最后署了名。

其他合作者包括：Adobe研究員Zongze Wu、Eli Shechtman，及CMU機(jī)器學(xué)習(xí)系主任Zico Kolter。

重構(gòu)預(yù)測函數(shù)，回到標(biāo)準(zhǔn)的回歸問題

iMF (Improved MeanFlow) 的核心改進(jìn)是通過重構(gòu)預(yù)測函數(shù)，將訓(xùn)練過程轉(zhuǎn)換為一個標(biāo)準(zhǔn)的回歸問題。

在原始的MeanFlow (MF) （上圖左）中，其直接最小化平均速度的損失。其中，Utgt是根據(jù)MeanFlow恒等式和條件速度e-x推導(dǎo)出來的目標(biāo)平均速度。

這里的問題在于，推導(dǎo)出來的目標(biāo)Utgt包含網(wǎng)絡(luò)自身預(yù)測輸出的導(dǎo)數(shù)項，而這種“目標(biāo)自依賴”的結(jié)構(gòu)使得優(yōu)化極不穩(wěn)定、方差極大。

基于此，iMF從瞬時速度的角度去構(gòu)建損失，使整個訓(xùn)練就變得穩(wěn)定。

值得注意的是，網(wǎng)絡(luò)輸出仍然是平均速度，而訓(xùn)練損失則變成了瞬時速度損失，以獲得穩(wěn)定的、標(biāo)準(zhǔn)的回歸訓(xùn)練。

它首先將輸入簡化為單一的含噪數(shù)據(jù)z，并在內(nèi)部巧妙地修改了預(yù)測函數(shù)的計算方式。

具體來說，iMF讓用于計算復(fù)合預(yù)測函數(shù)V（代表對瞬時速度的預(yù)測）中，雅可比向量積（JVP）項所需的切向量輸入不再是外部的e-x，而是由網(wǎng)絡(luò)自身預(yù)測的邊緣速度。

通過這一系列步驟，iMF成功移除了復(fù)合預(yù)測函數(shù)V對目標(biāo)近似值e-x的依賴。此時，iMF再將損失函數(shù)的目標(biāo)設(shè)定為穩(wěn)定的條件速度e-x。

最終，iMF 成功將訓(xùn)練流程轉(zhuǎn)換成了一個穩(wěn)定的、標(biāo)準(zhǔn)的回歸問題，為平均速度的學(xué)習(xí)提供了堅實的優(yōu)化基礎(chǔ)。

除了對訓(xùn)練目標(biāo)進(jìn)行改良外，iMF還通過以下兩大突破，全面提升了MeanFlow框架的實用性和效率：

靈活的無分類器指導(dǎo)（CFG）。

原始MeanFlow框架的一大局限是：為了支持單步生成，無分類器指導(dǎo)（CFG）的指導(dǎo)尺度在訓(xùn)練時必須被固定，這極大地限制了在推理時通過調(diào)整尺度來優(yōu)化圖像質(zhì)量或多樣性的能力。

iMF通過將指導(dǎo)尺度內(nèi)化為一個可學(xué)習(xí)的條件來解決此問題。

具體來說，iMF直接將指導(dǎo)尺度作為一個輸入條件提供給網(wǎng)絡(luò)。

在訓(xùn)練階段，模型會從一個偏向較小值的冪分布中隨機(jī)采樣不同的指導(dǎo)尺度。這種處理方式使得網(wǎng)絡(luò)能夠適應(yīng)并學(xué)習(xí)不同指導(dǎo)強(qiáng)度下的平均速度場，從而在推理時解鎖了CFG的全部靈活性。

此外，iMF 還將這種靈活的條件作用擴(kuò)展到支持CFG區(qū)間，進(jìn)一步增強(qiáng)了模型對樣本多樣性的控制。

高效的上下文內(nèi)條件作用（In-context Conditioning）架構(gòu)

原始MF依賴于參數(shù)量巨大的adaLN-zero機(jī)制來處理多種異構(gòu)條件（如時間步、類別標(biāo)簽和指導(dǎo)尺度）。

當(dāng)條件數(shù)量增多時，簡單地對所有條件嵌入進(jìn)行求和并交給adaLN-zero處理，會變得效率低下且參數(shù)冗余。

iMF引入了改進(jìn)的上下文內(nèi)條件作用來解決此問題。

它的創(chuàng)新點在于：它將所有條件（包括時間步、類別以及 CFG 因子等）編碼成多個可學(xué)習(xí)的Token（而非單一向量），并將這些條件Token直接沿序列軸與圖像潛在空間的Token進(jìn)行拼接，然后一起輸入到 Transformer 塊中進(jìn)行聯(lián)合處理。

這一架構(gòu)調(diào)整帶來的最大益處是：iMF可以徹底移除參數(shù)量巨大的adaLN-zero模塊。

這使得iMF在性能提升的同時，模型尺寸得到了大幅優(yōu)化，例如 iMF-Base 模型尺寸減小了約1/3（從 133M 降至 89M），極大地提升了模型的效率和設(shè)計靈活性。

實驗結(jié)果

iMF在最具挑戰(zhàn)性的ImageNet 256x256上的1-NFE中展示了卓越的性能。

iMF-XL/2在1-NFE下的FID達(dá)到了1.72，將單步生成模型的性能推到了一個新的高度。

iMF從頭開始訓(xùn)練的性能甚至優(yōu)于許多從預(yù)訓(xùn)練多步模型中蒸餾而來的快進(jìn)模型，證明了 iMF 框架在基礎(chǔ)訓(xùn)練上的優(yōu)越性。

下圖在ImageNet 256x256上進(jìn)行1-NFE（單步函數(shù)評估）生成的結(jié)果。

iMF在2-NFE時的FID達(dá)到1.54，將單步模型與多步擴(kuò)散模型（FID約1.4-1.7）的差距進(jìn)一步縮小。

One more thing

如前文所述，IMF 一作延續(xù)前作Mean Flow（已入選 NeurIPS 2025 Oral）的核心班底——耿正陽。

他本科畢業(yè)于四川大學(xué)，目前在CMU攻讀博士，師從Zico Kolter教授。

共一作者為清華姚班大二學(xué)生Yiyang Lu，現(xiàn)于MIT跟隨何愷明教授研究計算機(jī)視覺，此前曾在清華叉院許華哲教授指導(dǎo)下研究機(jī)器人方向。

這篇論文部分的內(nèi)容由他們在MIT期間，于何愷明教授指導(dǎo)下完成。

此外，論文的其他作者還包括：Adobe研究員Zongze Wu、Eli Shechtman，CMU機(jī)器學(xué)習(xí)系主任J. Zico Kolter以及何愷明教授。

其中，Zongze Wu本科畢業(yè)于同濟(jì)大學(xué)，并在Hebrew University of Jerusalem獲得博士學(xué)位，他目前在Adobe舊金山研究院擔(dān)任研究科學(xué)家，

同樣的，Eli Shechtman也同樣來自Adobe，他是Adobe Research圖像實驗室的高級首席科學(xué)家。他于2007加入 Adobe，并于2007–2010年間在華盛頓大學(xué)擔(dān)任博士后研究員。

J. Zico Kolter是論文一作耿正陽的導(dǎo)師，他是CMU計算機(jī)科學(xué)學(xué)院教授，并擔(dān)任機(jī)器學(xué)習(xí)系主任。

論文的尾作則是著名的機(jī)器學(xué)習(xí)科學(xué)家何愷明教授，他目前是MIT的終身副教授。

他最出名的共工作是ResNet，是21世紀(jì)被引用次數(shù)最多的論文。

就在最近的NeurIPS放榜中，何愷明參與的FastCNN還拿下了時間檢驗獎。

[1]https://arxiv.org/pdf/2505.13447

[2]https://gsunshine.github.io/

[3]https://arxiv.org/pdf/2512.02012

文章來源：量子位。