ICLR 2026:解构一步扩散生成模型 ESC,探寻本质与优化方向

本文提出ESC方法,改进一步扩散模型,降低训练方差,提升生成质量,ImageNet 256x256上实现FID 2.53的SOTA性能。

原文标题:ICLR 2026 | ESC — 解构一步生成,厘清细节,探寻本质

原文作者:数据派THU

冷月清谈:

本文深入探讨了一步扩散生成模型的设计原理,旨在解决现有方法理论复杂、实现细节耦合的问题。研究人员提出了一个统一的设计框架,揭示了离散和连续时间方法背后的共通原则:用单步参数化预测逼近两步扩散构建的目标。通过对关键模块(如路径选择、时间采样器)的细粒度分析,发现损失函数中速度场监督带来的方差是影响训练稳定性和生成质量的关键因素。针对此问题,研究人员提出了ExplicitShortCut (ESC) 变体,通过插入理想速度、渐进时间采样器等方式,显著提升了训练稳定性和生成质量。实验结果表明,ESC 在 ImageNet 256x256 上实现了 SOTA 的一步生成性能,为未来高效的一步扩散生成模型设计提供了关键见解。

怜星夜思:

1、文章中提到,线性路径在一步生成中通常比余弦路径更具竞争力,这与传统认知有所不同。你认为除了文章中提到的“凸运输代价”外,还有哪些因素可能导致这一现象?
2、ESC通过引入“即插即用的理想速度”来降低监督方差,但这种方法依赖于mini-batch内的样本。那么,当batch size很小或者数据集类别分布不均匀时,这种方法的效果会受到影响吗?如果是,有什么改进方案吗?
3、文章最后提到了一步生成模型面临的挑战,例如依赖CFG和超参数空间庞大等。你认为未来一步生成模型发展的主要方向会是什么?

原文内容

图片
来源:机器之心
本文约5000字,建议阅读10分钟
本文对现有的一步扩散生成方法进行了改进,并提出了 ESC(ExplicitShortCut) 变体。


  • 论文标题: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) 则展示实际学习情形,此时目标偏离原始轨迹,从而导致模型预测也相应发生漂移。

扩散模型的本质是在边缘速度场   下模拟 PF-ODE 从时间 1 到 0 的轨迹,其中   代表数据样本,  代表噪音,一步扩散模型的目标是可以学到路径上任意时刻同一采样路径上的样本对直接的映射。理想情况下,应该从   采样到  ,再用真实采样的轨迹对样本对   为模型提供监督;然而,一旦   被采样,在不用多步迭代模拟路径的情况下,  是难以计算的。


总结得到,实际上现有模型的解决方案是通过中间状态  ,构建两步映射的目标:即利用网络自身或条件速度场进行从   到  ,再到   而构建的   的样本对,再训练网络通过一步映射来匹配它。如 Consistency Training Model,是在   的条件下,通过条件速度将   映射到  ,再用网络将   到   而构建的步骤学习目标;而 MeanFlow 和 sCT 是其时间   的特殊形式,区别是 MeanFlow 中   可以随机采样,而 sCT 中   固定为 0。具体而言,表 1 展示了主流的一步生成模型以及其具体的参数化形式、时间采样器选择、两步映射的目标等设计组件:

表 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 随机化

根据前面讨论,Consistency Model 往往将 r 固定在 0,而 MeanFlow 中的 r 是可以随机采样的,因此本文探寻了终点时间的采样对结果的影响。图 2 (a),(b),(c) 中 sCT-Linear 与 MeanFlow 的对比,得到的结论更偏向一种 “阶段性取舍”。固定 r=0 的做法能够带来更快的早期收敛,因为这时模型持续接收到对   的显式监督,学习任务更接近一个去噪问题,优化上更直接,也更容易快速进入较好的局部区域。相对地,随机化时间 r 虽然在训练初期未必占优,但它能够让模型接触到更丰富的中间路径信号,从而更好地学习整体的捷径模式。

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

如何改进现有难题?

基于上述理论和实证分析,我们采用连续时间线性路径基线(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的网络架构上,进行了实验消融,说明几种改进的结合可以稳定地提升模型的生成能力。

训练技巧

FID50k

Baseline

MeanFlow 基线 (CFG)

6.09

+ A1

+ Plug-in velocity (p=1.0)

6.01

+ A2

+ Plug-in velocity (p=0.5)

5.98

+ B1

+ Plug-in velocity (p=0.5) + Class-consistent mini-batching

5.96

 

+ C

+ Gradual time sampler

5.99

+ D

+ Other techniques

5.95

Baseline+B1+C+D 

ESC 

5.77


大规模实验评估结果

我们在 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) 性能提升在更大的骨干网络上更明显,表明降低方差对于模型稳定性的提升,随着模型容量增长,效果越来越明显

方法

参数量

NFE

FID50k

iCT

675M

1

34.24

SCD

675M

1

10.60

IMM

675M

1×2

7.77

MeanFlow

676M

1

3.43

MeanFlow

676M

2

2.93

ESC (类别一致批次处理

676M

1

2.85

ESC+ (两倍训练时长)

676M

1

2.53


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(包括扩散模型、多智能体等)及其在聚变、水下具身智能、虚拟细胞领域的应用,长期招募希望在以上领域做出有影响力工作的同学(博士生、博后、助理研究员、实习生)。

编辑:文婧



关于我们

数据派THU作为数据科学类公众号,背靠清华大学大数据研究中心,分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识,努力建设数据人才聚集平台、打造中国大数据最强集团军。




新浪微博:@数据派THU

微信视频号:数据派THU

今日头条:数据派THU


我觉得未来的突破口可能在以下几个方面:一是如何设计更有效的网络结构,例如 Transformer 在计算机视觉领域的应用,或许能为生成模型带来新的灵感;二是如何解决一步生成模型对计算资源的高需求,降低训练成本;三是如何提高生成模型的可控性,让用户能够更加精确地控制生成图像的内容。

我比较关注生成模型的理论研究。目前很多模型都是基于经验和实验结果,缺乏扎实的理论基础。未来,如果能够建立更加完善的理论框架,或许就能更好地理解生成模型的本质,从而设计出更加高效、更加强大的模型。而且,说不定能像GAN那样,出现更多变体。

我更倾向于实践角度的考虑。线性路径可能更容易调整超参数,例如扩散的速率,方便研究人员进行更精细的控制和优化,从而找到更好的模型配置。而且线性本身计算复杂度也更低,训练效率更高。

这个问题很有意思!除了文章提到的凸运输代价更低,我个人觉得线性路径可能在数学形式上更简单,对于神经网络来说,更容易学习和逼近这种简单的映射关系。毕竟一步生成本身就是希望一步到位,过于复杂的路径可能反而适得其反。

Batch size 影响的不仅是理想速度的估计,还会影响梯度下降的效率。过大的 batch size 会导致梯度更新的频率降低,模型收敛速度减慢,训练时间延长。因此,不能简单地认为 batch size 越大越好,需要根据具体情况进行调整。

可以考虑一些梯度累积的技巧,在显存不足的情况下模拟大 batch size 的效果。另外,一些自适应的优化器,例如 LAMB,对大 batch size 的训练也更友好,可以尝试使用。

我倒是觉得可以往多模态生成方向发展,比如根据文字描述生成图像、根据图像生成音乐等等。这样可以更好地利用不同模态的信息,创造出更加丰富多彩的内容。AIGC嘛,就是要够 fancy!

理论上,更大的 batch size 可以提供更精确的理想速度估计,从而进一步降低方差。但要注意,Batch size 越大,对硬件的要求越高,显存消耗也会快速增长。而且,过大的 batch size 可能会导致模型陷入局部最优解,泛化能力下降。所以,需要在计算资源、模型性能和泛化能力之间找到一个平衡点。

从信息论的角度看,线性路径可能信息损耗更少。扩散过程本质上是逐步引入噪声,线性路径的噪声引入方式可能更加均匀,避免了某些频率的信息过早丢失,有利于后续的生成过程。

减轻显存开销方面,可以尝试使用更高效的计算框架,或者探索近似计算方法,比如低精度量化。摆脱对CFG的依赖,可能需要更强的条件编码器,或者探索无条件生成的方式。总的来说,一步扩散模型还有很大的优化空间,未来可期!

我认为Plug-in velocity是最关键的。一步扩散模型的核心痛点是目标构建中的高方差,Plug-in velocity直接降低了监督信号的方差,这相当于从源头上提高了训练的稳定性。其他方法可能只是锦上添花,而 Plug-in velocity 是雪中送炭。

我觉得可以从两个方面入手:一是模型结构上的创新,比如设计更高效的神经网络架构,减少参数量和计算复杂度;二是训练方法上的改进,比如探索新的损失函数,或者利用自监督学习等技术,降低对大量标注数据的依赖。一步扩散模型想要真正落地,需要解决实际应用中的问题。

mini-batch 的 plug-in velocity 方法,我觉得最大的局限性在于它对 batch size 的依赖。如果 batch size 太小,那这个“理想速度”可能并不那么理想,仍然会受到单个样本的影响,方差可能降不下来。而且,不同类别的数据在同一个 mini-batch 中混合,可能会引入额外的偏差,尤其是对于类别数量很多的数据集。

未来,可以考虑使用一些更精细的方差估计方法,比如 moving average 或者 exponential moving average,来平滑速度场的估计。另外,也许可以借鉴GAN中的discriminator思想,训练一个专门用来估计速度场方差的神经网络,然后用这个网络来指导训练过程。