论文阅读——mean-flow
本文最后更新于 2026年7月8日 20:41:18
本文已发在小红书上 # Mean-Flow
论文标题：Mean Flows for One-step Generative Modeling
作者：Zhengyang Geng, Mingyang Deng, Xingjian Bai, J. Zico Kolter, Kaiming He
GitHub链接：https://github.com/Gsunshine/meanflow
一、概述
本文提出了单步生成模型的一种严谨有效的新方法，即Mean Flow。它的核心思想是使用平均速度(average velocity)来刻画流场(flow fields)，而非flow matching 方法中的瞬时速度(instantaneous velocity)。Mean flow独立性强，自成一体，无需预训练、蒸馏或课程学习。
二、背景知识
1、Diffusion & Flow Matching：
Diffusion模型设计前向加噪和反向去噪两个过程。其中，前向过程会逐步向图片中加入噪声，反向过程则训练一个神经网络从噪声中逐步恢复数据，从而学习到从噪声分布到数据分布的逆过程。
Flow Matching在Diffusion理论的基础上更进一步，直接让神经网络建模速度场(velocity fields)。这种方法不需要估计对数概率密度的梯度（即 score function），从而降低了计算难度。
具体来说，给定初始数据\(x\)（在这里就是真实图像）和噪声\(\epsilon\)，可以构造一条随着时间\(t\)变化的路径\(z_t = a_t x + b_t \epsilon\)。其中\(z_t\)表示数据在时间\(t\)下的状态，\(\alpha_t\)和\(\beta_t\)是预先定义好的时间调度函数(schedules)。等式两边同时对t求导，便可以得到条件速度(conditional velocity)的表达式。由于\(\alpha_t,\beta_t\)已知，\(z_t\)的演化只受到\(x\)和\(\epsilon\)的影响，这个\((x,\epsilon)\)便是影响路径演化的“条件”。
这两张图展现了Flow Matching中的速度场。其中黑色点是噪声，红色点是真实图像。
我们常用线性调度函数\(\alpha_t=1-t, \beta_t=t\)学习速度场，这时，条件速度\(v_t=\epsilon-x\)，对应左图。
为了消除随机噪声\(\epsilon\)带来的随机性，我们对条件速度在\(z_t\)上进行边缘化(marginalization)，得到边际速度(marginal velocity)：
\[ v(z_t, t) \triangleq \mathbb{E}_{p_t(v_t \mid z_t)}[v_t] \]
对应右图。
但是，这种边缘化操作往往是不可行的，因为需要根据复杂的联合分布\(p(z_t,v_t)\)进行积分，难以计算，所以我们通过Conditional Flow Matching 来绕过边缘化，直接建模条件概率\(p(v_t|z_t)\)。实验证明，这种方法在避免复杂的积分操作的同时，也能有效建模真实的演化路径。
2、现有的少步生成模型：
出于实际应用和理论研究的角度，减少采样步数是很重要的发展方向。早期的方法主要是蒸馏学习，将预训练好的多步扩散模型蒸馏成few-step模型。与此同时，一致性模型(Consistency Model)通过引入一致性约束，迫使模型在相同采样路径上的不同时间步都能输出清晰的图像。随着研究不断深入，利用两个时间点进行建模的方法，如Flow Map, Shortcut Models和Inductive Moment Matching也被发明出来。
三、核心思想
①平均速度
\[ u(z_t, r, t) \triangleq \frac{1}{t-r} \int_r^t v(z_{\tau}, \tau) d\tau \]
这种定义方式使得\(u\)天然具有边界条件（\(\lim_{r \to t} u = v\)）和一致性（积分的可加性），并可通过设置初始时间\(r=0\)，结束时间\(t=1\)的方式实现一步生图。
②Mean Flow恒等式
Mean Flow恒等式：推导出平均速度与瞬时速度之间的关系，为
\[ \underbrace{u(z_t, r, t)}_{\text{average vel.}} = \underbrace{v(z_t, t)}_{\text{instant. vel.}} - \underbrace{(t-r) \frac{d}{dt}u(z_t, r, t)}_{\text{time derivative}} \]
③训练与采样
\[ \mathcal{L}(\theta) = \mathbb{E}\left\|u_{\theta}(z_t, r, t) - \text{sg}(u_{\text{tgt}})\right\|_2^2 \\ \text{where} \quad u_{\text{tgt}} = v(z_t, t) - (t-r)\left(v(z_t, t)\partial_z u_{\theta} + \partial_t u_{\theta}\right), \]
想要得到真正的ground truth就要在高维空间中进行积分操作，难以算出，因此采用神经网络计算出的\(\partial_z u_{\theta}\)和\(\partial_t u_{\theta}\)来近似得到。同时，采用\(sg()\)(stop-gradient)操作，在计算时将\(u_{tgt}\)视作常数，避免两次反向传播和更高阶的优化计算。
进一步，我们用conditional velocity来替换\(v(z_t,t)\)，避免复杂的期望计算，可以得到：
\[ u_{\text{tgt}} = v_t - (t - r) \left( v_t \partial_z u_\theta + \partial_t u_\theta \right). \]
在采样时，仅需设置起止时间点便可以得到推理结果。
④融入CFG(classifier-free guidance)
通过修改训练目标的方式，使得模型天然具有CFG属性，从而实现1-NFE（单次网络函数评估）。
通过推导，可以得到CFG情况下\(u_{cfg}\)和\(v_{cfg}\)的直接关系：
\[ 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) \]
对应地，修改损失函数
\[ \begin{align*}\mathcal{L}(\theta) &= \mathbb{E}\big\|u^{\text{cfg}}_{\theta}(z_t, r, t \mid c) - \text{sg}(u_{\text{tgt}})\big\|_2^2, \\\text{where} \quad u_{\text{tgt}} &= \tilde{v}_t - (t-r)(\tilde{v}_t\partial_z u^{\text{cfg}}_{\theta} + \partial_t u^{\text{cfg}}_{\theta}).\end{align*} \]
这里的\(v_t\)是线性插值得到的。取\(\omega\)=1时，此式退化为无CFG时的式子。
\[ \tilde{v}_t \triangleq \omega v_t + (1 - \omega) u^{\text{cfg}}_{\theta}(z_t, t, t) \]
采样时同样仅需设置两个起止时间点便可以得到推理结果。
四、实验，模型表现与总结
1、消融实验
作者通过实验，通过控制\(\mathbf{r} \boldsymbol{\ne} \mathbf{t}\)样本的占比，调整JVP的计算方式，证明了Mean Flow的有效性。实验中也发现，不同的时间步编码方式、采样方式，损失函数和CFG幅度同样会影响最终结果，并证明Mean Flow具有良好的可扩展性。
2、模型表现
在ImageNet上，Mean Flow最高达到了3.43的FID，在1-NFE的情况下相比sota模型提升了70%，并缩小了和多步推理模型之间的差距。
3、总结
Mean Flow找到了粗粒度层面上描述底层量的方法，使计算机得以从宏观上找到一定的规律，提升推理速度和质量。
五、个人评价
Mean Flow最原创的部分就在于通过平均速度来刻画图像生成过程，这种宏观视角为少步图像生成提供了一种全新的思路。不仅如此，作者还提出了严谨的数学推导，提出了MeanFlow Identity，架起了连接平均速度场与瞬时速度场之间的桥梁。然而，这个方法仍然存在一些问题：
①训练不稳定性
这篇论文基于JAX框架，针对TPU有着更好的表现。但绝大部分研究者使用的都是GPU。在GitHub上的非官方PyTorch实现版本中，许多研究者都反应了结果无法复现和训练不稳定的问题。除去JAX框架的使用，Mean Flow的损失函数使用了stop-gradient方法。尽管这是非常流行的做法，但仍然会带来潜在的梯度不平衡，增加训练不稳定的可能性。
②理论发展不完全
尽管Mean Flow是平均速度场的开山之作，在理论推导上仍然存在可以提高的地方。MeanFlow Identity中仍然混用了平均速度与瞬时速度，并没有摆脱微分形式、充分地利用平均速度场的概念。有学者在此之上进一步提出了Interval Splitting Consistency的概念，其对应的恒等式是：
\[ (t-r)u(z_t, r, t) = (s-r)u(z_s, r, s) + (t-s)u(z_t, s, t) \]
这一概念是Mean Flow的一般化形式，并得到了更好的实验结果。