论文标题：On the Design of One-step Diffusion via Shortcutting Flow Paths论文地址: https://openreview.net/forum?id=k6q8rRYVQR代码开源: https://github.com/EDAPINENUT/ExplicitShortCut/项目主页: https://edapinenut.github.io/explicitshortcut-project-page/

前言

近期，基于捷径化概率流路径（shortcut probability flow trajectory）并从头训练的一步扩散生成模型，展现出强大的实证有效性。然而，这类方法的提出通常建立在较为复杂的理论推导之上，并且往往与具体实现细节高度耦合。这带来一个直接的问题：究竟哪些设计是方法成立的本质要素，哪些又只是可以灵活替换的实现组件。

对于研究人员，这种复杂性进一步体现在多个层面：一方面是深奥的理论背景，繁复的方法形式推导、对应的学习目标设计，另一方面，是应用实现上的技术细节，包括时间采样器、课程学习方法等。面对这些内容，研究人员往往很难迅速抓住不同方法之间真正共享的核心思想，也因此缺少一种足够直观的设计范式，来理解这类一步生成模型到底 “为什么这样设计”。

为了解决这个问题，西湖大学研究人员首先从一个更实用的角度出发，尝试为这类 shortcut 模型提出一个统一的设计框架。其指出，无论是离散时间还是连续时间的一步生成方法，其背后都遵循着同一个基本原则：用单步参数化预测去逼近一个两步扩散构建的目标。在此基础上，该工作进一步给出了这一设计范式成立的一般性理论说明。这样的统一视角，使得原本耦合在一起的生成方法模块能够被拆解开来，以供研究人员更清晰地理解各个组件之间是如何协同作用的，以及在整体方法设计中究竟还保留了多少可调整的自由度。

得益于这一框架，本文对当前具有代表性的离散时间与连续时间一步扩散生成方法进行更细粒度的组件级分析，包括路径选择、时间采样器等关键模块。通过这一分析，本文进一步定位了影响训练 — 推理稳定性以及逼近误差的一个关键损失函数设计中，学习目标里的速度场监督带来的方差。基于这一发现，本文对现有的一步扩散生成方法进行了改进，并提出了 ESC（ExplicitShortCut） 变体。该方法在训练稳定性和生成质量上都取得了进一步提升。在 ImageNet 256×256 上，基于 SiT-XL/2 架构，ESC 在 1-NFE 设置下达到 FID50k 2.53，超越了此前所有从头训练的捷径扩散模型。

如何从头训练一个一步扩散模型？

图 1 离散时间与连续时间的一步扩散模型在理想学习与实际学习中的物理图景 (a) 展示了从 N (0,1) 到高斯混合分布的边缘速度场。(b)(c) 对应理想学习情形，此时 xr 采样自 PF-ODE 的同一条轨迹，因此构造的学习目标可作为正确监督信号；(d)(e) 则展示实际学习情形，此时目标偏离原始轨迹，从而导致模型预测也相应发生漂移。

表 1 代表性的一步生成模型中，具体的组件选择

如何选择设计组件？

图 2 不同一步生成模型在训练过程中 FID50k 的比较：(a) 表示在 CIFAR-10 上的无条件生成（Uncond.）；(b) 表示在 CIFAR-10 上的类别条件生成（Cond.）；(c) 表示在 ImageNet-256×256 上的 classifier-free guidance（CFG）训练。

Q1: 路径选择：线性 vs 余弦？

一般而言，线性路径扩散如 Rectified Flow、DiT 等工作在理论上更加简洁，实践上更方便以使用各种训练和采样技巧，而余弦路径扩散如 EDM 等工作往往被认为像素空间生成中通常因固定方差的随机过程而被视为更稳定。

本文在此探寻路径选择对模型的影响。根据图 2 (a),(b) 中 sCT (-Cosine) 以及 sCT-Linear 的实验结果表明，在线性路径与余弦路径的对比中，线性路径在一步生成的设定下通常更有竞争力。一种合理的解释是：线性路径作为条件路径时，对应的边缘速度场具有更低的凸运输代价，因此轨迹弯曲程度更小。这意味着，由两步采样映射构造出来的监督目标更接近理想情形，从而让 “一步去逼近两步” 这件事本身更容易成立。换句话说，在一步生成的问题设定里，线性路径并不只是实现更方便，而是更符合这一类模型的近似需求。

Q2: 损失类型：离散 vs 连续？

前面讨论过，当 s→t 的时候，损失函数就会将一步扩散推至连续时间模型。本文在此探寻连续与离散模型的优劣势。在相同训练设置和统一代码实现下，根据图 2 (a), (b), (c) 中的 MeanFlow、sCT 这样的连续时间上的一步扩散模型与 CT、SCD、IMM 等离散时间方法的对比，可以得到连续时间一步扩散模型整体上明显优于离散时间变体。

进一步，本文给出了推理误差分析，其核心结论是：对于线性路径，连续时间方法的误差上界相比离散时间方法少了一项与局部 Lipschitz 常数、时间步长以及监督方差耦合的额外误差项；而这项额外项很可能会直接放大训练不稳定性和一步推理误差。也就是说，连续时间方法更强，并不只是因为 “参数化更平滑”，而是因为它在目标逼近误差的结构上更占优。这也解释了为什么同样是一步生成扩散，基于连续时间的模型往往表现得更稳、更准。

Q3: 时间采样：固定终点时间 vs 随机化

换句话说，固定终止时间更像是在优化 “早期易学性”，而随机化终止时间则是在优化 “最终生成与表达能力”。

如何改进现有难题？

基于上述理论和实证分析，我们采用连续时间线性路径基线（MeanFlow with SiT-B/2），并通过三个方向进行改进：

1. 即插即用的理想速度（Plug-in velocity）

第一个改进方向是用 plug-in velocity 替代直接使用单样本对应的条件速度作为监督信号。其动机来自前面的误差分析：在训练时，真正难以处理的不是目标形式本身，而是引入条件速度而带来的高方差。本文指出，如果能够使用基于整个经验分布构造的 “理想速度”，那么速度项的方差可以显著下降；但这需要对整个训练集求和，在大规模数据集上几乎不可行。为此，作者采用在一个 mini-batch 上的 plug-in velocity，将一个批次内多个样本的条件速度以不同加权的形式组合起来，从而在仅引入较小偏差的代价下，将监督方差从依赖单样本的水平降到与批次大小相关的更低水平。这种替换直接带来更稳定的训练监督和更低的一步推理误差，是对核心瓶颈的直接修正。

2. 渐进时间采样器（Gradual time sampler）

第二个改进方向是设计一个渐进式时间采样器，让训练从更容易的目标开始，再平滑过渡到完整的 MeanFlow 式时间采样。这个设计直接回应了前面对 “固定终止时间还是随机化” 的分析：固定 r=0 的监督更像去噪任务，早期更容易优化，而覆盖完整区间的随机时间采样虽然更难，却更有助于模型学习全局 “捷径模式”。基于这一点，本文采用了一个 渐进式课程策略：在训练初期，以较高概率采样 r=0，让模型先获得更稳定、更直接的监督；随着训练推进，再逐步衰减这一概率，最终完全过渡到 MeanFlow 的采样方式。这样做的好处是同时兼顾了前期收敛速度和后期全局建模能力，本质上是在易学性与最终性能之间建立了一条更平滑的训练路径。

3. 优化技巧 (Other techniques)

第三个方向则是吸收已有连续时间一步扩散方法中已经验证有效的训练技巧，并将其系统整合到新的框架中。包括一些已经在原始工作中验证过的优化手段，例如 variational adaptive loss weighting，进一步提升优化过程的稳定性、缓解训练早期的脆弱性，并改善不同监督项之间的平衡。此外，在 classifier-free guidance 的训练场景下，本文额外引入类别一致的批次处理技术（class-consistent mini-batching）：通过同类采样避免 mini-batch 平均过程过度稀释类别信号，以达到大规模训练中真正 “跑得稳、训得动、效果好” 的工程化方案。

综合上述提出的技术，本文在SiT-B/2的网络架构上，进行了实验消融，说明几种改进的结合可以稳定地提升模型的生成能力。

大规模实验评估结果

我们在 ImageNet-256x256 的潜空间中使用 SiT-XL/2（约 676M 参数）评估 ESC。遵循 MeanFlow 在分类器自由引导下的训练协议，ESC 从头训练 240 个 epoch（约 120 万次迭代），ESC+ 延长至 480 个 epoch（约 240 万次迭代）。

在 ImageNet-256x256 的 1-NFE 设置下，ESC 将 FID 从 3.43 提升至 2.85，ESC+ 进一步提升至 2.53。这一结果不仅大幅超越了此前所有从头训练的捷径模型，甚至超过了 MeanFlow 的两步结果（2-NFE，FID50k 2.93）。

此外，我们还可以从中观察到：

1)类别一致的批次处理可以提高收敛速度。

2)即插即用的理想速度几乎不增加计算开销 (千分之九)，同时提高稳定性。

3) 性能提升在更大的骨干网络上更明显，表明降低方差对于模型稳定性的提升，随着模型容量增长，效果越来越明显

ESC 在 ImageNet-256x256 上展现出卓越的一步生成能力（SiT-XL/2 架构），且能够生成高质量、多样化的图像。（FID:2.85）

ExplicitShortCut (ESC)通过统一的设计空间分析，系统性地理解了一步扩散模型的关键设计选择。通过改进目标构建和降低监督方差，ESC 在 ImageNet-256x256 上实现了 FID 2.53 的 SOTA 一步生成性能，也为之后设计高效的一步扩散生成提供了关键见解。

然而，目前的从头训练的一步生成模型仍然面临着极大的挑战，比如在连续情形下 jvp 函数的使用以及其导致的显存开销、对 CFG 的依赖导致设计不够灵活、超参数空间庞大等，即便如此，一步生成的扩散模型仍然是未来的研究热门，正如何凯明在 CVPR 上的演讲所说：

“生成模型正处在一个类似「AlexNet 之前」的阶段：从多步推理，走向端到端的一步生成。” 一步扩散不是终点，它是下一个 AIGC 时代的开始。

作者及团队简介

本文第一作者为林海涛，来自西湖大学，共同一作为胡佩炎，来自中科院数学所。其研究方向关注生成模型，一步扩散与蒸馏，离散扩散模型等，以第一作者在 ICML、ICLR、NeurIPS 等人工智能顶级会议或期刊上发表论文超过十篇；本文的通讯作者为西湖大学李子青、吴泰霖，以及深势科技柯国霖。李子青实验室关注人工智能及生命科学交叉领域。吴泰霖实验室关注生成式 AI（包括扩散模型、多智能体等）及其在聚变、水下具身智能、虚拟细胞领域的应用，长期招募希望在以上领域做出有影响力工作的同学（博士生、博后、助理研究员、实习生）。