一步生成ImageNet FID 1.29!斯坦福用Wasserstein 梯度流重写模型

训练时让分布沿最优传输的 “下山方向” 走，推理时只需一次网络前向。W-Flow 把多步演化压进静态生成器，在 ImageNet 256×256 上刷新一步生成指标。

• 论文标题：One-Step Generative Modeling via Wasserstein Gradient Flows
• 论文链接：https://arxiv.org/abs/2605.11755

扩散模型在图像生成领域有着广泛的应用，但它的代价也很清楚：采样常常需要几十到数百步。GAN 倒是一步出图，却经常要在训练稳定性上付出代价。有没有一种路线，训练像扩散模型 / 流匹配一样有清晰的分布演化，推理却像 GAN 一样只走一步？

斯坦福大学的新论文W-Flow给出了一个颇有野心的答案：不要手工设计一步生成器的漂移场，也不要依赖多步蒸馏，而是直接让生成分布沿Wasserstein 梯度流下降，再把这条训练时的演化压缩到一个静态神经网络生成器里。最终，模型推理只需一步。论文目前已发布在 arXiv, 论文共同第一作者为斯坦福大学博士生 Jiaqi Han 和 Puheng Li。论文代码和权重已开源。

图 1（论文 Figure 1）：左侧为 W-Flow-L/2 在 ImageNet-256×256 上的一步采样样本；右侧为 FID 与有效采样计算量对比。

论文最吸睛的结果是：W-Flow-XL/2在 ImageNet 256×256 class-conditional 生成上达到1.29 FID，W-Flow-L/2 达到1.35 FID；论文还称，相比相近 FID 的多步扩散模型，采样速度约提升100 倍

这不只是把采样器调快一点，而是在问一个更本质的问题：一步生成器的训练轨迹，应该由什么原则来定义？

3 个看点先读

1.从头训练的一步生成器：不是把已有扩散模型蒸馏成一步，而是直接设计生成器的训练动力学。

2.Sinkhorn 散度 + Wasserstein 梯度流：用全局最优传输结构更新粒子，替代局部启发式漂移。

3.结果和稳定性同时提升：ImageNet FID 1.29；FFHQ 域迁移和小众模式覆盖也更好。

01 为什么一步生成难？

多步扩散模型 / 流匹配 把简单分布到数据分布的复杂运输拆成很多小步，每一步都相对容易学，但推理慢。一步生成器则试图一次把噪声映射到数据，推理快，但训练信号往往难以稳定。

GAN 使用判别器诱导的梯度，强大但容易振荡；漂移模型用吸引 / 排斥的粒子场，效果很强，但论文认为这种局部规则仍偏启发式，可能留下不可预测收敛和模式坍缩的风险。

问题不在于 “能不能一步生成”，而在于：训练时每个样本、每一批样本应该往哪里移动？W-Flow 的回答是：让整个分布沿一个可解释的能量景观下坡。

02 W-Flow：先规定分布怎么走，再让网络学会这条路

在 W-Flow 中，生成建模被看作一个映射：把参考分布中的样本映射到目标数据分布。训练时，作者并不是把这个映射拆成推理时要执行的长链，而是用一系列局部更新在训练中规定分布演化；随后用一个神经网络学习这些更新，把多步训练轨迹压缩成推理时的一步映射。

这个演化被写成 Wasserstein 梯度流：当前生成分布以某个能量函数的最陡下降方向移动。选择能量函数为当前生成分布与真实数据分布之间的差异，就得到一条往真实数据分布下降的运动轨迹。

图 2（论文 Figure 2 (a)）：W-Flow 概念图。

图 3（论文 Figure 2 (b) 的动图）：训练轨迹在 Sinkhorn 散度景观上的可视化。

直观地看，W-Flow 并不是给每个生成样本找一个最近的真实样本，而是在一个批次上计算当前生成分布和真实分布之间应该如何运输。

03 为什么是 Sinkhorn 散度？

能量函数有很多选择。MMD 的核相互作用在生成分布和真实分布距离很远时可能梯度变弱；KL 需要估计生成分布的 score，实际实现中会引入额外偏差和不稳定。

W-Flow 选择Sinkhorn 散度，这是一个带熵正则的最优传输差异度量：它可以通过 Sinkhorn 迭代在 mini-batch 上高效估计，同时保留全局质量约束。

论文中，Sinkhorn 诱导的速度场可以写成一个非常有解释力的形式：

V (x) = 生成到真实的运输投影 − 生成到生成的自运输投影

前一项把生成粒子推向真实数据；后一项是自传输修正，防止粒子只按局部吸引塌缩。

所以 W-Flow 的速度场由两张全局运输计划共同决定，而不是由最近邻或单样本匹配决定。

04 训练实际怎么跑？两张运输计划，一个回归目标

每次训练，W-Flow 抽一批噪声生成样本 xi=fθ(zi)，抽一批真实样本 yj，再额外抽一批生成样本 x′j 做自传输。然后计算两次 Sinkhorn 重心投影：生成到真实、生成到生成。二者相减就是粒子的速度 V。

接下来，作者把停止梯度后的 xi+ηV 作为目标，让当前网络回归到这个更新后的位置。也就是说，训练最终变成一个简单的回归问题；不需要像 Sinkhorn GAN 那样穿过 Sinkhorn 迭代反传梯度。

关键细节是双批次自传输。若用同一批生成样本做自匹配，每个粒子都能和自己零成本匹配，容易产生自传输偏差。作者用第二批生成样本来消除这个偏差。

图 4（论文 Figure 3）：单批次与双批次自传输估计在二维高斯分布学习中的差异。

05 ImageNet 主结果：1-NFE 的 1.29 FID

在 ImageNet 256×256 条件生成任务中，W-Flow 直接从头训练，不依赖蒸馏。核心数字如下：W-Flow-B/2：1.52 FID；W-Flow-L/2：1.35 FID；W-Flow-XL/2：1.29 FID；采样步数全部为 1

图 5（论文 Table 3）：ImageNet 256×256 条件生成对比。

这个结果很有意思：它不仅压过漂移模型（B/2 1.75，L/2 1.54），也超过论文表中多种单步 扩散 / 流方法。更重要的是，W-Flow-B/2 仅用 133M 生成器参数，就超过了更大的 Drifting-L/2。

与需要大量采样步数的扩散模型 / 流方法相比，W-Flow 的优势在于采样时间极短。对低延迟生成、在线交互式创作和高吞吐部署来说，这个差异直接转化为响应速度和算力成本。

06 消融：不是某个 trick，而是动力学真的在起作用

在控制实验中，Sinkhorn 散度的 FID 为7.29，优于 MMD 的、KL 的以及同配置下 漂移模型的8.46；速度场引导（下文介绍）又把 FID 推到7.08

双批次估计也不是小优化。论文报告，单批次自传输让 FID 增加到 17.57；加入双批次估计后回到 7.08，并优于对角线掩码的 7.45。

图 6（论文 Table 2）：能量函数、CFG、最优传输代价与双批次估计的消融结果。

这些消融共同支持一个结论：W-Flow 的提升并不只是大模型或调参，而来自更合理的传输动力学。

07 CFG：一步模型也能吃到引导增益

论文还把无分类器引导注入到速度场中，而不是简单改写目标分布。这样做的直观好处是，条件与非条件的差别体现在 “速度” 上，几何解释更自然。

图 7（论文 Figure 4）：CFG 规模变化下的 FID / IS 曲线与样本视觉变化。

在 CFG 规模增加时，样本的类别细节更清晰，FID / IS 的权衡优于漂移模型。这说明 W-Flow 并不是只能靠不加引导的裸模型取胜，而能兼容现代扩散模型常用的质量增强机制。

08 不只是出图快：还更会保留小众模式

如果只看 ImageNet FID，W-Flow 像是一个 “快采样” 工作。但论文更想强调的是：全局最优传输动力学带来的分布覆盖能力。

在域迁移中，W-Flow 可以把源分布直接设成另一个数据域。例如从中老年人的脸到青年人的脸。因为优化的是两个分布之间的 Wasserstein 梯度流，它不必从高斯噪声出发，也可以学习 “同一语义空间内” 的一步迁移。

图 8（论文 Figure 5 (a) 的动图）：二维椭圆到圆的域迁移。

图 9（论文 Figure 5 (b)&(c)）：FFHQ 中老年到青年人脸的一步年龄迁移。

在二维椭圆分布到圆分布迁移的任务中，漂移模型的轨迹更散、更不稳定；W-Flow 的粒子轨迹更协调、路径更短。在 FFHQ 年龄迁移中，W-Flow 的潜空间 L2 距离更短，身份保持也更好。

更重要的是模式覆盖。作者构造了极端不平衡分布：95% 中老年人脸 + 5% 小孩脸。漂移模型几乎丢掉少数的小孩脸部分，而 W-Flow 仍能捕捉到该部分。这也呼应了 Sinkhorn 全局质量约束的作用。

图 10（论文 Figure 6）：不平衡目标分布下的模式覆盖。

09 工作意义

从工程角度看，W-Flow 一步生成器带来强指标：高质量、低延迟、可从头训练。

从方法论角度看，它把问题从 “怎么蒸馏一个多步模型” 换成 “如何设计生成分布的训练动力学”。这可能比单个 ImageNet 结果更重要，因为它给一步生成模型提供了一个原则化框架。

从应用角度看，一步生成器如果能保持质量，意味着实时交互式创作、端侧生成、低成本数据增强、快速域迁移都有更大想象空间。

10 局限与未来工作

论文也没有把问题全部解决。实验主要集中在 ImageNet-256 和 FFHQ；高分辨率、文字到图片、视频和多模态生成仍是未来工作。大规模实验依赖预训练 特征编码器 / 自编码器；理论部分证明的是粒子动力学在一定假设下的一致性，还不能完全解释有限网络、有限算力下的优化。

但这并不削弱 W-Flow 的信号：一步生成不一定只能靠 GAN 对抗或扩散蒸馏，也可以靠 Wasserstein 空间里原理驱动的动力学

结语

一句话总结：W-Flow 不是把扩散模型的采样步数硬砍掉，而是在训练阶段为生成分布设计了一条可计算、可解释的 “下坡路”，再把它压进一个一步生成器。

生成模型的下一场速度战，可能不只发生在采样器，也发生在训练动力学本身。