当前位置:首页 > 科技资讯 > 正文

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破

自今年五月MeanFlow (MF) 问世后,何恺明领导的研究小组近日正式推出了其增强版本——

Improved MeanFlow (iMF),这一新模型成功解决了原始MF在训练稳定性、指导灵活性以及架构效率三大核心领域的挑战。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第1张

通过将训练目标重构为更稳健的瞬时速度损失,并引入灵活的无分类器指导(CFG)以及高效的上下文内条件机制,模型性能得到了显著提升。

在ImageNet 256x256基准测试中,iMF-XL/2模型在单步函数评估(1-NFE)中取得了1.72的FID分数,相比原始MF提高了50%,这表明从头训练的单步生成模型可以达到与多步扩散模型相竞争的水平。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第2张

MeanFlow的第一作者耿正阳继续主导工作,值得注意的是共同第一作者Yiyang Lu目前仍是清华大学姚班的大二学生,而何恺明教授也位列作者名单。

其他贡献者包括:Adobe研究员Zongze WuEli Shechtman,以及卡内基梅隆大学机器学习系主任Zico Kolter

重构预测函数,转化为标准回归问题

iMF(Improved MeanFlow)的核心改进在于重构预测函数,将训练过程转变为标准回归问题

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第3张

在原始MeanFlow (MF)(左图所示)中,直接最小化平均速度的损失。其中,Utgt是根据MeanFlow恒等式和条件速度e-x推导的目标平均速度。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第4张

这里的关键问题是,推导出的目标Utgt包含了网络自身预测输出的导数项,这种“目标自依赖”结构导致优化过程极不稳定、方差过大。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第5张

基于此,iMF从瞬时速度的角度构建损失,使训练变得稳定。

值得注意的是,网络输出仍然是平均速度,而训练损失则转换为瞬时速度损失,以实现稳定、标准的回归训练。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第6张

它首先将输入简化为单一的含噪数据z,并在内部巧妙修改预测函数的计算方式。

具体来说,iMF在计算复合预测函数V(代表对瞬时速度的预测)时,雅可比向量积(JVP)项所需的切向量输入不再依赖外部的e-x,而是由网络自身预测的边缘速度提供。

通过这一系列步骤,iMF成功移除了复合预测函数V对目标近似值e-x的依赖。此时,iMF将损失函数的目标设定为稳定的条件速度e-x。

最终,iMF 将训练流程转换为一个稳定、标准的回归问题,为平均速度的学习提供了坚实的优化基础。

除了改良训练目标外,iMF还通过以下两大突破,全面提升了MeanFlow框架的实用性和效率:

灵活的无分类器指导(CFG)

原始MeanFlow框架的一个主要局限是:为了支持单步生成,无分类器指导(CFG)的指导尺度在训练时必须固定,这极大限制了推理时通过调整尺度优化图像质量或多样性的能力。

iMF通过将指导尺度内化为可学习条件来解决此问题。

具体来说,iMF直接将指导尺度作为输入条件提供给网络。

在训练阶段,模型会从偏向较小值的幂分布中随机采样不同指导尺度。这种处理使网络能够适应并学习不同指导强度下的平均速度场,从而在推理时释放CFG的全部灵活性。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第7张

此外,iMF 还将这种灵活的条件作用扩展到支持CFG区间,进一步增强了模型对样本多样性的控制。

高效的上下文内条件作用(In-context Conditioning)架构

原始MF依赖于参数量巨大的adaLN-zero机制来处理多种异构条件(如时间步、类别标签和指导尺度)。

当条件数量增多时,简单对所有条件嵌入进行求和并交给adaLN-zero处理,会变得效率低下且参数冗余。

iMF引入了改进的上下文内条件作用来解决此问题。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第8张

它的创新点在于:将所有条件(包括时间步、类别以及 CFG 因子等)编码成多个可学习的Token(而非单一向量),并将这些条件Token直接沿序列轴与图像潜在空间的Token进行拼接,然后一起输入到 Transformer 块中进行联合处理。

这一架构调整带来的最大益处是:iMF可以彻底移除参数量巨大的adaLN-zero模块。

这使得iMF在性能提升的同时,模型尺寸得到了大幅优化,例如 iMF-Base 模型尺寸减小了约1/3(从 133M 降至 89M),极大提升了模型的效率和设计灵活性。

实验结果

iMF在最具挑战性的ImageNet 256x256上的1-NFE中展示了卓越的性能。

iMF-XL/2在1-NFE下的FID达到了1.72,将单步生成模型的性能推到了新高度。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第9张

iMF从头开始训练的性能甚至优于许多从预训练多步模型中蒸馏而来的快进模型,证明了 iMF 框架在基础训练上的优越性。

下图在ImageNet 256x256上进行1-NFE(单步函数评估)生成的结果。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第10张

iMF在2-NFE时的FID达到1.54,将单步模型与多步扩散模型(FID约1.4-1.7)的差距进一步缩小。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第11张

One more thing

如前所述,iMF 第一作者延续前作Mean Flow(已入选 NeurIPS 2025 Oral)的核心班底——耿正阳。

他本科毕业于四川大学,目前在卡内基梅隆大学攻读博士学位,师从Zico Kolter教授。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第12张

共同第一作者为清华姚班大二学生Yiyang Lu,现于麻省理工学院跟随何恺明教授研究计算机视觉,此前曾在清华交叉信息研究院许华哲教授指导下研究机器人方向。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第13张

这篇论文的部分内容由他们在MIT期间,于何恺明教授指导下完成。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第14张

此外,论文的其他作者还包括:Adobe研究员Zongze Wu、Eli Shechtman,卡内基梅隆大学机器学习系主任J. Zico Kolter以及何恺明教授。

其中,Zongze Wu本科毕业于同济大学,并在耶路撒冷希伯来大学获得博士学位,他目前在Adobe旧金山研究院担任研究科学家。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第15张

同样,Eli Shechtman也来自Adobe,他是Adobe Research图像实验室的高级首席科学家。他于2007年加入 Adobe,并于2007–2010年间在华盛顿大学担任博士后研究员。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第16张

J. Zico Kolter是论文第一作者耿正阳的导师,他是卡内基梅隆大学计算机科学学院教授,并担任机器学习系主任。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第17张

论文的尾作者是著名的机器学习科学家何恺明教授,他目前是MIT的终身副教授。

他最著名的工作是ResNet,是21世纪被引用次数最多的论文之一。

何恺明团队发布Improved MeanFlow (iMF),单步生成模型性能实现重大突破 Improved  单步生成模型 图像生成 人工智能研究 第18张

就在最近的NeurIPS放榜中,何恺明参与的FastCNN还获得了时间检验奖。

参考链接:

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

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

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