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

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

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

在ImageNet 256x256基準(zhǔn)測(cè)試中，iMF-XL/2模型在 1-NFE（單步函數(shù)評(píng)估）中取得了1.72的FID成績(jī)，相較于原始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ù)測(cè)函數(shù)，回到標(biāo)準(zhǔn)的回歸問題

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)結(jié)果

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

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

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

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

iMF在2-NFE時(shí)的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ì)算機(jī)視覺，此前曾在清華叉院許華哲教授指導(dǎo)下研究機(jī)器人方向。

這篇論文部分的內(nèi)容由他們?cè)贛IT期間，于何愷明教授指導(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圖像實(shí)驗(yàn)室的高級(jí)首席科學(xué)家。他于2007加入 Adobe，并于2007–2010年間在華盛頓大學(xué)擔(dān)任博士后研究員。

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

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

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

就在最近的NeurIPS放榜中，何愷明參與的FastCNN還拿下了時(shí)間檢驗(yàn)獎(jiǎng)。

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

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

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

文章來源：量子位。