📊 statistics

One-Step Generative Modeling via Wasserstein Gradient Flows

本文介绍了 W-Flow，这是一种单步生成建模框架，它将 Wasserstein 梯度流压缩至单次神经网络推理中，从而实现了 1.29 的 FID 指标和比多步扩散模型快约 100 倍的采样速度，在 ImageNet 生成任务上达到了最先进水平。

原作者： Jiaqi Han, Puheng Li, Qiushan Guo, Renyuan Xu, Stefano Ermon, Emmanuel J. Candès

发布于 2026-05-13

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jiaqi Han, Puheng Li, Qiushan Guo, Renyuan Xu, Stefano Ermon, Emmanuel J. Candès

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你想教一个机器人画出完美的猫。

旧方法（扩散模型）：
大多数现代 AI 绘画工具的工作原理，就像雕塑家从一块石头上凿下碎片，或者摄影师先拍一张模糊的照片然后慢慢使其清晰。它们从随机噪声开始，经过数百个微小的步骤，逐渐将噪声转化为猫。

问题所在： 这就像去商店时，不是径直走过去，而是采取 100 个微小、犹豫的步伐。要得到一张图片，这需要很长时间并消耗大量能量（计算能力）。

新方法（W-Flow）：
这篇论文介绍了W-Flow，一种教导机器人一步跨越就能画出完美猫的方法。它不再需要走 100 步，而是瞬间学会从“噪声”到“猫”的完美捷径。

它是如何工作的？“河流”类比

要理解 W-Flow 如何学会这条捷径，想象有两组人：

人群（目标）： 一群人站在一个完美的圆圈里（代表真实数据，如真实的猫照片）。
流浪者（生成器）： 一群人随机散落在田野中（代表 AI 当前混乱的猜测）。

目标是将流浪者移动，使它们形成与人群相同的完美圆圈，但 AI 需要学会一条规则，能够一次性完成这一过程。

1. “能量”地图（坡度）
作者设想流浪者与人群之间的空间是一片丘陵地带。“人群”位于山谷的最底部（最低能量点）。流浪者则位于山坡的某处。

规则： 如果你是一个流浪者，你想尽可能快地滚下山坡到达人群所在的位置。
创新点： 以前的方法使用“启发式”（一种猜测）来判断哪边是下坡。有时它们猜错了，导致卡住，或者将人们向错误的方向推得太猛。
W-Flow 的窍门： 它们使用一种名为Sinkhorn 散度的数学工具。你可以将其想象为一个超精准的 GPS，它为流浪者群体中的每一个人计算出确切的最陡下坡路径，同时考虑整个群体的协同移动，而不仅仅是个体。

2. “双批次”安全网
在计算流浪者应如何移动时，存在一个棘手的问题：如果你要求一个人远离他们自己的群体，他们可能会不小心试图远离自己，这没有意义。

解决方案： 论文使用了一种巧妙的技巧，称为**“双批次”**策略。想象将流浪者分成两条独立的队伍。你根据 B 队计算 A 队应如何移动，反之亦然。这防止了它们因自身的倒影而感到困惑，并确保它们平稳地向目标移动而不会卡住。

3. 压缩旅程
这是神奇的一步：

首先，AI 多次模拟这种“滚下山坡”的过程（就像观看流浪者慢慢形成圆圈的影片）。
然后，它训练一个神经网络（“生成器”）来记忆整部影片。
结果： 一旦训练完成，网络就不再需要观看影片了。它知道起点和终点，因此可以一步直接从“随机噪声”跳到“完美的猫”。

为什么这很重要？

速度： 论文声称，这种方法比旧的多步方法快约100 倍。如果旧方法制作一张图片需要 10 秒，那么新方法只需不到一秒的时间。
质量： 尽管只有一步，生成的图片质量极高。在著名的 ImageNet 测试中，它们取得了1.29的分数（FID），这是一步生成器的新纪录。这意味着图片看起来几乎与真实照片无法区分。
稳定性： 由于数学基础建立在坚实的原则（Wasserstein 梯度流）之上，而非猜测，AI 不太可能“崩溃”（即只学会画一种猫而忽略其他所有类型）。它能更好地覆盖数据的所有不同“模式”。

总结

将W-Flow想象成教学生解数学题。

旧方法： 老师向学生一步步展示解题过程，学生反复练习这些步骤，直到能够独立完成。
W-Flow： 老师向学生展示解题的逻辑（梯度流），让他们练习这种逻辑，然后要求他们立即写出最终答案。学生将这条“捷径”学得非常透彻，以至于不再需要展示解题过程。

论文证明，通过使用这种特定的数学“指南针”（Sinkhorn 散度）来指导训练，你可以构建出一个既快如闪电又高度准确的生成器。

技术摘要：基于 Wasserstein 梯度流的单步生成建模

问题陈述

现代生成模型，特别是扩散模型和基于流的模型，虽然在样本质量上达到了最先进水平，但由于依赖多步迭代采样，面临着高昂的计算成本和延迟。尽管存在单步或少步蒸馏方法，但它们往往继承了原本为迭代生成设计的路径中的训练信号，或者依赖于缺乏原则性收敛保证且易导致模式崩溃的启发式动力学（例如 Drifting Models）。核心挑战在于设计一种训练动力学，使生成器能够学习从参考分布到目标数据分布的直接单步映射，同时确保稳定的收敛性和全局模式覆盖。

方法论：W-Flow

作者提出了W-Flow，这是一个将训练动力学的设计与生成器的参数化解耦的框架。其核心思想是利用**Wasserstein 梯度流（WGF）**规定模型分布向目标分布演化的过程，然后训练一个静态神经网络，将这种多步演化压缩为单次推理步骤。

1. 基于 Wasserstein 梯度流的训练动力学

W-Flow 不直接学习去噪步骤序列或速度场，而是将训练轨迹定义为概率测度空间中的连续时间演化。

能量泛函：演化由能量泛函 $F(q) = D(q \| p)$ 的最速下降驱动，其中 $q$ 是当前模型分布， $p$ 是目标数据分布。
速度场：连续时间速度场 $V_t$ 源自能量泛函的一阶变分： $V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$ 。
离散化：连续流使用显式欧拉方案进行离散化。在每个训练步 $k$ ，模型分布 $q^{(k)}$ 通过局部传输映射 $g^{(k)}(x) = x + \eta V^{(k)}(x)$ 进行更新，其中 $\eta$ 是步长。

2. 基于 Sinkhorn 散度的实例化

为了确保速度场在高维空间中具有计算可行性且几何行为良好，作者使用Sinkhorn 散度（ $S_\epsilon$ ，一种熵正则化的最优传输（OT）度量）来实例化能量泛函。

速度场公式：诱导的速度场由两个重心投影的差给出：
$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$
其中， $T^\epsilon_{q,p}$ 将质量从生成分布 $q$ 传输到目标分布 $p$ ，而 $T^\epsilon_{q,q}$ 代表一个自传输项，用于强制与当前分布的一致性。
全局协调：与局部核交互（如 MMD 或 Drifting Models 中）不同，Sinkhorn 速度场源自全局 OT 计划，确保粒子在整个批次中的移动是协调的，从而防止模式崩溃。
双批次估计器：为了避免粒子以零成本匹配自身的自匹配伪影，作者采用双批次策略。他们使用来自当前分布的独立第二批次样本来计算自传输项 $T^\epsilon_{q,q}$ ，而不是对单个批次成本矩阵的对角线进行掩码处理。

3. 生成器训练

生成器 $f_\theta$ 被训练以模仿由 WGF 定义的离散粒子更新。

目标：损失函数最小化生成器输出与源自 Sinkhorn 更新的“目标”粒子位置之间的平方距离：
$\mathcal{L}(\theta) = \mathbb{E}_{z \sim p_{ref}} \| f_\theta(z) - \text{sg}(f_\theta(z) + \eta V^\epsilon(\hat{q}_\theta, \hat{p})) \|^2$
其中 $\text{sg}(\cdot)$ 表示停止梯度算子，将更新后的粒子位置视为固定的回归目标。
无分类器引导（CFG）：该框架通过直接修改速度场来支持 CFG。作者没有线性混合分布，而是提出了速度引导，即添加一个与条件速度场和无条件速度场之差成正比的项。这种方法被证明在使用 KL 散度时能恢复标准的指数倾斜目标分布，并为 Sinkhorn 散度提供了几何一致的扩展。

主要贡献

原则性单步框架：W-Flow 通过将训练动力学建立在 Wasserstein 梯度流之上，为单步生成建模确立了新范式，提供了启发式方法（如 Drifting Models）所缺乏的清晰变分解释。
基于 Sinkhorn 的动力学：利用 Sinkhorn 散度实例化 WGF，提供了一种高效、基于最优传输的更新规则，能够捕捉全局分布差异，并相比基于核或局部交互的方法改善了模式覆盖。
理论收敛性：作者证明了在适当的正则性假设下，W-Flow 的有限样本粒子动力学收敛于连续时间分布动力学，为学习过程提供了理论依据。
双批次估计器：引入了一种新颖的双批次策略来估计自传输项，消除了单批次估计器和对角线掩码启发式方法中固有的自匹配偏差。

实验结果

作者在 ImageNet 256×256 和 FFHQ 上评估了 W-Flow。

最先进水平性能：在 ImageNet 的单步类条件生成任务中，W-Flow 取得了新的最先进水平 FID 分数，XL 规模为 1.29，L 规模为 1.35。这显著优于之前的单步方法（如 Drifting Models、MeanFlow），并且尽管仅使用单次函数评估（NFE），其表现仍超越了许多多步扩散模型（如 LightningDiT、SiT）。
采样效率：与具有相似 FID 分数的多步扩散模型相比，W-Flow 的采样速度快约 100 倍（例如，在单张 H100 GPU 上，W-Flow 为 77.88 张图像/秒，而 SiT-XL/2 为 0.93 张图像/秒）。
模式覆盖：在不平衡分布的实验（例如 95% 老年面孔，5% 儿童面孔）中，W-Flow 成功捕捉到了少数模式，而 Drifting Models 则表现出模式崩溃，完全忽略了少数类。
收敛速度：W-Flow 在训练期间收敛速度显著更快，仅需 384 个 epoch 即可达到 Drifting Models 训练 1280 个 epoch 的性能。
领域迁移：该方法在领域迁移任务（例如 FFHQ 上的年龄转换）中表现出强大的能力，与启发式基线相比，产生了更短且更连贯的传输轨迹。

意义与主张

该论文声称，Wasserstein 梯度流为快速、高保真度的生成建模提供了原则性且有效的基础。通过将训练动力学的设计与生成器参数化分离，W-Flow 证明了单步生成器可以被训练以达到与复杂多步迭代过程相当的性能。作者强调，最优传输动力学所提供的全局协调是稳定训练并确保全面模式覆盖的关键，解决了现有单步方法的关键局限性。结果表明，W-Flow 可作为实时创意工具和仿真等需要低延迟、高质量生成应用的扩展性替代方案。