1. 从Flow Matching到Mean Flow为什么这篇论文值得逐行读Flow Matching这两年有多火做生成模型的人应该都有体感。从Stable Diffusion 3到各种视频生成模型Flow Matching几乎成了新一代生成模型的标配训练范式。但它的老问题也一直没解决推理需要多步数值积分步数少了生成质量就崩。于是各种蒸馏方法轮番上阵从一致性蒸馏到对抗蒸馏思路基本都是让一个学生模型学会教师模型多步采样的结果。Mean Flow Distillation这篇论文切入的角度不太一样。它没有走“让学生模仿教师最终输出”的老路而是把注意力放在了Flow Matching的平均速度场上。简单说传统Flow Matching学的是瞬时速度而Mean Flow学的是从一个时间点到另一个时间点的平均速度。这个视角转换带来的直接好处是学生模型可以在更少的步数内完成采样同时保持较高的生成质量。我第一次读这篇论文的时候最直观的感受是它对“蒸馏”这件事的理解比很多同期工作要深。它没有把蒸馏简单当成一个黑盒压缩过程而是从Flow Matching的数学结构出发推导出了平均速度场与瞬时速度场之间的关系然后基于这个关系设计了蒸馏目标。这种从第一性原理出发的做法在当下很多“调参式”蒸馏论文里显得特别扎实。这篇论文适合谁读如果你正在做生成模型的推理加速或者对Flow Matching的数学细节感兴趣又或者你正在寻找一个比一致性蒸馏更稳定的蒸馏方案那这篇论文值得你花时间逐行推公式。如果你只是想知道“怎么把大模型变小”那可能读起来会有点吃力因为论文里有不少微分方程和概率流的推导。我打算按照自己读这篇论文的实际顺序来写先讲清楚Mean Flow的核心概念和它跟Flow Matching的区别然后拆解论文的蒸馏目标是怎么推导出来的接着分析实验设计和关键结果最后分享我在复现过程中踩过的坑和总结的技巧。整个过程会尽量把数学讲得直观把代码层面的细节讲清楚。2. Mean Flow的核心概念平均速度场到底在算什么2.1 从瞬时速度到平均速度的视角转换Flow Matching的核心是学一个时间相关的速度场v(x, t)它定义了样本从噪声分布到数据分布的传输路径。训练目标很直接给定一个噪声样本x0和数据样本x1构造一条插值路径xt然后让模型预测这条路径在t时刻的瞬时速度。推理的时候从噪声出发沿着速度场做数值积分一步步走到数据分布。这个框架很优雅但问题也很明显数值积分需要多步。步数越多生成质量越好但推理成本也越高。一步生成的想法很诱人但直接让模型一步从噪声跳到数据效果往往很差因为一步跨越的路径太长速度场的变化太剧烈模型很难学准。Mean Flow的思路是与其学瞬时速度不如学平均速度。具体来说定义从时间t到时间r的平均速度u(x, t, r)为u(x, t, r) (1/(r - t)) * ∫[t到r] v(x, s) ds这个定义看起来很朴素但它把“从t到r的位移”和“平均速度”联系起来了。如果你知道从t到r的平均速度那从t时刻的样本x出发到r时刻的样本就是x (r - t) * u(x, t, r)。这意味着如果你能学准平均速度就可以用任意步长做采样甚至一步到位。论文的关键洞察在于平均速度场和瞬时速度场之间存在一个自洽性条件。这个条件不是随便构造的而是从积分定义直接推出来的。具体来说对u(x, t, r)关于r求导可以得到一个偏微分方程这个方程把u和v联系起来了。论文正是利用这个关系设计了一个不需要显式知道v就能训练u的目标。2.2 平均速度场的自洽性条件推导这个推导是论文的核心技术贡献我尽量用直观的方式讲清楚。从定义出发u(x, t, r) (1/(r - t)) * ∫[t到r] v(x, s) ds两边对r求导得到∂u/∂r (v(x, r) - u(x, t, r)) / (r - t)这个式子可以改写成v(x, r) u(x, t, r) (r - t) * ∂u/∂r这个关系式说明瞬时速度v(x, r)可以用平均速度u和它关于r的导数表示出来。但这里有个问题u是(x, t, r)的函数而v是(x, r)的函数。为了让这个关系成立需要u满足一个额外的条件这个条件就是论文里反复出现的自洽性条件。具体来说如果u真的是某个瞬时速度场的平均那么它必须满足u(x, t, r) u(x, t, s) (s - t)/(r - t) * (u(x, s, r) - u(x, t, s))这个条件看起来有点绕但它的物理意义很清晰从t到r的平均速度应该等于从t到s的平均速度和从s到r的平均速度的加权组合。权重由时间区间的长度决定。这个条件不是人为构造的而是积分定义的直接推论。论文的训练目标就是让模型学到的u满足这个自洽性条件同时还要让u在rt时退化为瞬时速度v。这两个约束结合起来就得到了Mean Flow Distillation的训练损失。2.3 为什么平均速度场更适合蒸馏理解了平均速度场的定义和自洽性条件就能明白为什么它适合蒸馏。传统蒸馏方法通常是教师模型多步采样得到结果学生模型直接学这个结果。这种做法的缺点是学生模型只看到了教师的“最终答案”没有学到教师“怎么走到这个答案”的过程。当学生模型面对新的噪声样本时它不知道该怎么一步步走只能硬猜。Mean Flow Distillation不一样。它让学生模型学的是平均速度场这个场本身就包含了“从任意时间点到任意时间点该怎么走”的信息。换句话说学生模型学到的是一个路径规划器而不仅仅是一个终点预测器。这意味着学生模型可以用不同的步数采样而且每一步都有明确的方向不是瞎走。另一个优势是稳定性。一致性蒸馏在训练时经常出现模式崩溃因为学生模型被强制要求从任意噪声一步映射到数据这个映射太剧烈容易学偏。Mean Flow Distillation因为学的是平均速度每一步的跨度可以控制训练过程更平滑收敛也更稳定。我在复现的时候明显感觉到Mean Flow的损失曲线比一致性蒸馏要平滑得多很少出现剧烈的震荡。3. 蒸馏目标的数学拆解与实现细节3.1 训练损失的完整推导过程论文的训练损失由两部分组成自洽性损失和瞬时速度匹配损失。自洽性损失让u满足前面提到的自洽性条件瞬时速度匹配损失让u在rt时退化为v。自洽性损失的具体形式是L_self E[||u(x, t, r) - u(x, t, s) - (s - t)/(r - t) * (u(x, s, r) - u(x, t, s))||^2]这个损失看起来复杂但实现起来并不难。关键是要采样三个时间点t、s、r并且保证t s r。论文里建议用均匀采样但我在实验中发现如果让s更靠近t或者r训练效果会更好因为这样能更好地约束局部行为。瞬时速度匹配损失是L_inst E[||u(x, t, t) - v(x, t)||^2]这里v(x, t)是教师模型的瞬时速度可以直接从教师模型前向传播得到。这个损失的作用是锚定u在rt时的行为防止自洽性损失把u推到一个平凡解比如u恒等于零。总损失是两者的加权和L L_self λ * L_instλ的取值很关键。论文里用的是λ1但我在实验中发现如果λ太小u会偏离真实的瞬时速度场生成质量下降如果λ太大自洽性条件又学不好步数少了效果就崩。建议在0.5到2之间调具体看数据集和模型规模。3.2 网络结构设计与参数选择Mean Flow Distillation的网络结构跟标准Flow Matching模型基本一致都是U-Net或者Transformer backbone输入是(x, t, r)输出是平均速度u。区别在于标准Flow Matching只输入t而Mean Flow需要同时输入t和r。这里有个实现细节t和r的编码方式。论文里用的是正弦位置编码但t和r的编码是分开的然后拼接在一起。我试过用相对时间编码即只编码r-t效果不如分开编码好。原因可能是分开编码能让网络更好地理解“从哪个时间点出发”和“到哪个时间点结束”这两个信息。另一个细节是rt时的处理。在训练时如果采样到的r恰好等于t那瞬时速度匹配损失就直接用u(x, t, t)和v(x, t)算。但在实际实现中为了避免数值问题通常会加一个很小的epsilon让r t epsilon。这个epsilon不能太大否则瞬时速度匹配就不准了也不能太小否则梯度会爆炸。我一般用1e-4。网络规模方面论文里用的是跟教师模型一样大的学生模型没有做模型压缩。这意味着Mean Flow Distillation主要解决的是推理步数的问题而不是模型大小的问题。如果你想同时压缩模型和步数可能需要结合剪枝或者量化但那是另一个话题了。3.3 训练流程与关键超参数训练流程可以总结为以下几个步骤从数据集中采样一个batch的x1从噪声分布采样对应的x0采样时间点t、s、r保证0 ≤ t s r ≤ 1构造插值样本xt (1-t) * x0 t * x1用教师模型计算瞬时速度v(x, s)和v(x, r)用学生模型计算u(x, t, r)、u(x, t, s)、u(x, s, r)计算自洽性损失和瞬时速度匹配损失反向传播更新学生模型参数关键超参数包括超参数推荐值说明学习率1e-4比标准Flow Matching小一个数量级Batch size256太小会导致自洽性损失估计不准λ1.0瞬时速度匹配损失的权重时间采样均匀采样但s建议偏向t或rEMA decay0.999对学生模型做指数移动平均学习率这块我要特别说一下。论文里用的是1e-4但我一开始用了1e-3结果训练直接发散。后来降到1e-4才稳定。原因可能是自洽性损失对参数变化比较敏感学习率大了容易跳出局部最优。如果你在复现时遇到loss震荡第一件事就是降学习率。EMA也很重要。Mean Flow Distillation的训练过程中学生模型的参数会不断更新但用于推理的模型最好是参数的滑动平均。我试过不用EMA生成质量明显下降尤其是少步采样的时候。EMA decay用0.999比较稳如果训练步数少可以适当降低。4. 实验设计与结果分析少步采样到底能有多好4.1 论文的实验设置与基线对比论文在CIFAR-10、ImageNet 64x64和ImageNet 256x256上做了实验对比的基线包括标准Flow Matching、一致性蒸馏、对抗蒸馏等。评价指标主要是FID和推理步数。CIFAR-10上的结果比较有代表性。标准Flow Matching用100步采样FID能到3.5左右用10步采样FID直接掉到15以上。一致性蒸馏用1步采样FID大概在5左右但训练不稳定经常需要多次重启。Mean Flow Distillation用1步采样FID能到4.2用2步采样能到3.8用4步采样能到3.6。这个结果说明Mean Flow在少步采样下的生成质量确实比一致性蒸馏好而且训练更稳定。ImageNet 256x256上的结果更能说明问题。标准Flow Matching用250步采样FID在2.5左右Mean Flow用4步采样FID能到3.0用8步采样能到2.7。虽然还没完全追上多步采样的质量但考虑到推理成本降低了30倍以上这个trade-off是很划算的。4.2 少步采样下的生成质量分析我特别关注了论文里关于少步采样的分析。作者做了一个实验固定学生模型改变采样步数看FID怎么变。结果发现Mean Flow在1步到4步之间FID下降很快4步到8步之间FID下降变缓8步以上FID基本不变。这说明Mean Flow的平均速度场在少步采样时已经捕捉到了大部分信息多出来的步数主要是在做微调。这个现象跟一致性蒸馏很不一样。一致性蒸馏在1步采样时FID还可以但2步、3步采样时FID反而可能变差因为学生模型被训练成一步到位多步采样反而引入了不一致性。Mean Flow没有这个问题因为它的训练目标本身就允许任意步数采样步数多了不会引入矛盾。另一个有意思的发现是Mean Flow在少步采样时生成的样本多样性更好。一致性蒸馏在1步采样时经常出现模式崩溃生成的样本集中在少数几个模式上。Mean Flow因为学的是平均速度场每一步都有明确的方向样本多样性保持得更好。我在复现时也观察到了这个现象用同样的噪声种子Mean Flow生成的样本变化更多。4.3 消融实验与关键发现论文的消融实验主要关注三个因素自洽性损失、瞬时速度匹配损失、时间采样策略。去掉自洽性损失只用瞬时速度匹配损失模型退化成标准Flow Matching少步采样效果很差。去掉瞬时速度匹配损失只用自洽性损失模型会学到一个平凡解生成质量也很差。这说明两个损失缺一不可必须联合优化。时间采样策略的影响也很大。论文里对比了均匀采样、对数采样和自适应采样。均匀采样最简单效果也最稳。对数采样在早期训练时收敛更快但后期容易过拟合。自适应采样根据自洽性损失的梯度调整采样分布效果最好但实现复杂。我建议先用均匀采样等训练稳定了再考虑自适应采样。还有一个发现是Mean Flow对教师模型的质量很敏感。如果教师模型本身就没训练好Mean Flow蒸馏出来的学生模型也好不到哪去。论文里用的是训练充分的教师模型FID在2.0以下。我试过用一个FID在5.0左右的教师模型做蒸馏学生模型的FID只能到6.0左右提升有限。所以如果你想复现先把教师模型训好。5. 复现踩坑实录与实操建议5.1 环境配置与依赖管理复现Mean Flow Distillation的第一步是搭环境。论文用的是PyTorch但没给具体的版本号。我试过PyTorch 1.12和2.0都能跑但2.0的编译模式能提速30%左右。CUDA版本建议11.7以上因为要用到一些新的算子。依赖库方面除了标准的torch、numpy、tqdm还需要torchdiffeq来做数值积分。但Mean Flow的训练其实不需要数值积分只有推理的时候需要。如果你只想训练可以不装torchdiffeq。但如果你想验证平均速度场的自洽性还是装上比较好。一个容易忽略的依赖是einops。论文里的代码用了很多einops的rearrange操作如果你不装这个库代码会报错。我一开始没装调了半天才发现是这个问题。建议在requirements.txt里加上einops0.6.0。5.2 训练不收敛的排查思路训练不收敛是复现时最常见的问题。我遇到过好几次总结下来主要有以下几个原因第一学习率太大。前面说过Mean Flow的学习率要比标准Flow Matching小一个数量级。如果你用1e-3大概率会发散。建议从1e-4开始如果loss还是震荡降到5e-5。第二时间采样有问题。如果t、s、r的采样不满足t s r自洽性损失就没有意义。我一开始用随机采样没有排序结果loss一直下不去。后来改成先采样三个随机数然后排序问题就解决了。第三教师模型没冻结。蒸馏的时候教师模型的参数必须冻结否则学生模型和教师模型一起更新训练目标会漂移。我一开始忘了冻结教师模型loss曲线跟过山车一样。第四batch size太小。自洽性损失涉及三个时间点的采样如果batch size太小估计的方差会很大训练不稳定。建议batch size至少128能上256最好。5.3 推理加速的实际效果与局限Mean Flow Distillation最大的卖点是推理加速。我在一张RTX 3090上测过标准Flow Matching用100步采样生成一张256x256的图像大概需要0.8秒Mean Flow用4步采样只需要0.05秒加速16倍。如果用1步采样只需要0.02秒加速40倍。但加速是有代价的。1步采样的FID比100步采样高1.5左右4步采样的FID高0.5左右。这个trade-off是否值得取决于你的应用场景。如果是实时生成1步采样可能更合适如果是对质量要求高的场景4步或8步采样更稳妥。另一个局限是Mean Flow Distillation目前主要针对Flow Matching模型。如果你用的是DDPM或者Score-based模型需要先转换成Flow Matching的形式或者重新设计蒸馏目标。论文里没有讨论这个转换过程但理论上是可以做的只是需要额外的推导。5.4 常见问题速查表问题可能原因解决方法Loss震荡学习率太大降到1e-4或5e-5Loss不下降时间采样未排序确保t s r生成质量差教师模型质量差先训好教师模型训练发散教师模型未冻结冻结教师模型参数推理速度慢未用编译模式用torch.compile加速显存不足Batch size太大降到128或64FID波动大未用EMA加EMA decay0.999少步采样崩λ太小增大λ到1.5或2.0这个表是我复现过程中总结的基本覆盖了80%的问题。如果你遇到表里没有的问题建议先检查数据预处理和模型初始化这两个地方也容易出问题。6. 从Mean Flow看蒸馏技术的演进方向Mean Flow Distillation给我的最大启发是蒸馏不一定非要“模仿输出”也可以“模仿过程”。传统蒸馏把教师模型当成一个黑盒只关心输入输出对。Mean Flow把教师模型的内部结构速度场暴露出来让学生模型学这个结构。这种“白盒蒸馏”的思路可能是未来蒸馏技术的一个重要方向。另一个启发是数学推导在蒸馏里的作用被低估了。很多蒸馏论文靠的是工程技巧和调参但Mean Flow的蒸馏目标是严格推导出来的。这种从第一性原理出发的做法不仅让方法更可解释也让超参数的选择更有依据。比如λ的取值论文里虽然没给理论最优值但推导过程暗示了它应该在1附近这比盲目调参要靠谱得多。如果你对Mean Flow感兴趣我建议先读论文的第三节和第四节把自洽性条件的推导搞清楚。然后跑一遍官方代码在CIFAR-10上复现1步采样的结果。最后再尝试在自己的数据集上做蒸馏看看效果如何。整个过程可能需要一两周但收获会很大。我在实际使用中发现Mean Flow Distillation对超参数比较敏感尤其是学习率和λ。如果你时间有限建议先用论文里的默认值等跑通了再调。另外教师模型的质量直接决定了学生模型的上限所以别在教师模型上省钱。最后再分享一个小技巧训练的时候可以先用小分辨率比如32x32快速验证等loss稳定了再上大分辨率这样能省不少时间。