Generative Models IV

🐎 Elucidating the Design Space of Diffusion-Based Generative Models

如何进一步改进扩散模型,使其在图像质量、分布拟合质量、训练成本以及生成速度等方面取得更优表现,已成为当前研究中的核心问题。现有关于扩散模型的文献在理论层面极为密集,其关于采样日程、训练动力学、噪声水平参数化等内容的推导,通常都尽可能直接地建立在严密的理论框架之上。这种做法固然保证模型具有扎实的理论基础,但也可能带来一个问题:模型的可设计空间被遮蔽

  • 论文从实践而非纯理论视角重新审视这类模型,重点关注训练与采样阶段中实际可操作的对象和算法,而非过度依赖其背后的统计过程推导。其目的在于更清晰地揭示各组件之间的关系,并识别整体系统中的设计自由度。论文聚焦于这样一类模型:它们利用神经网络刻画训练数据在加入高斯噪声后、随噪声水平变化的边缘分布的 score,因此全文讨论均建立在去噪 score matching 框架之上
  • 论文围绕扩散模型的图像采样过程展开研究,识别效果最优的时间离散化方式,将高阶 Runge–Kutta 方法引入采样,评估不同采样调度策略,并分析采样中随机性的作用。这些改进显著减少图像生成所需的采样步数,且改进后的采样器可作为即插即用模块,直接应用于多种已有扩散模型
  • 论文关注 score 网络的训练机制。在仍采用 DDPM、NCSN 等常用架构的基础上,首次从原则层面对扩散模型中网络输入、输出及损失函数的预条件化进行系统分析,并总结出改善训练动态的最佳实践

🦨 Expressing diffusion models in a common framework

设数据分布为p_{\text{data}}(x),其标准差为\sigma_{\text{data}}。进一步考虑一族经平滑化后的分布p(x;\sigma),其构造方式为向原始数据中加入标准差为\sigma 的独立同分布高斯噪声。当\sigma_{\max} \gg \sigma_{\text{data}}时,p(x;\sigma_{\max})在实际中几乎无法与纯高斯噪声区分。扩散模型的基本思想是首先随机采样一个噪声图像x_0 \sim \mathcal{N}(0,\sigma_{\max}^2 I),随后按照噪声水平\sigma_0=\sigma_{\max} > \sigma_1 > \cdots > \sigma_N = 0逐步对其进行去噪,得到一系列图像x_i,并使得在每一个噪声水平下均满足x_i \sim p(x_i;\sigma_i)。因此,该过程的终点x_N便服从真实数据分布

Song 提出一种随机微分方程(SDE),用于刻画样本的时间演化过程,保证样本在各时刻均符合相应噪声水平下的边缘分布p(x;\sigma(t)) 。这使得上述过程可以通过一种随机求解器来实现,其在每一步迭代中既会去除噪声,也会重新注入噪声。与此同时,还可构造与之对应的概率流常微分方程(ODE),其中唯一的随机性来源于初始噪声图像x_0

概率流 ODE 在时间正向推进时连续增加图像中的噪声水平,而在时间反向推进时则连续降低噪声水平。为明确该 ODE 的具体形式,首先需要给定一个噪声调度函数\sigma(t),用于定义时刻t所对应的目标噪声水平。概率流 ODE 保持各噪声水平对应的边缘分布:若样本x_a \sim p(x;\sigma(t_a))从时刻t_a 演化至t_b(正向或反向),则有x_b \sim p(x;\sigma(t_b))

{ \mathrm{d}\boldsymbol{x} = - \frac{\mathrm{d}\sigma(t)}{\mathrm{d}t} \sigma(t) \nabla_{\boldsymbol{x}} \log p(\boldsymbol{x}; \sigma(t)) \, \mathrm{d}t }(1)

  • \nabla_x \log p(x;\sigma)即为 score function,它是一个向量场,在给定噪声水平下指向数据密度更高的方向。直观而言,该 ODE 的一个无穷小正向步会将样本从数据分布附近进一步推离,其速率取决于噪声水平的变化;等价地,一个反向步则会将样本推向数据分布

去噪 score matching:score function 具有一个重要性质,即它不依赖于底层密度函数p(x;\sigma)那个通常难以计算的归一化常数,因而更易于求值。若D(x;\sigma)表示一个去噪函数,并且对于每个\sigma都最小化如下期望L_2去噪误差

{ \mathbb{E}_{\boldsymbol{y} \sim p_{\text{data}}} \mathbb{E}_{\boldsymbol{n} \sim \mathcal{N}(\mathbf{0}, \sigma^2 \mathbf{I})} \| D(\boldsymbol{y} + \boldsymbol{n}; \sigma) - \boldsymbol{y} \|_2^2, \text{ then } \nabla_{\boldsymbol{x}} \log p(\boldsymbol{x}; \sigma) = (D(\boldsymbol{x}; \sigma) - \boldsymbol{x}) / \sigma^2 }(2)

  • 其中y为训练图像,n为噪声。从这一角度看,score function 实际上刻画样本x中噪声成分相对于信号成分的分离方式。扩散模型的关键观察在于,D(x;\sigma)可以由神经网络D_\theta(x;\sigma)来实现并进行训练。需要注意的是,D_\theta中还可能包含额外的预处理与后处理步骤

信号缩放调度通过改变概率密度影响 ODE 解轨迹。某些方法引入额外的缩放调度s(t)。这会改变随时间变化的概率密度,从而也会改变 ODE 的解轨迹

{ \mathrm{d}\boldsymbol{x} = \left[ \frac{1}{s(t)} \frac{\mathrm{d}s(t)}{\mathrm{d}t} \boldsymbol{x} - s(t)^2 \frac{\mathrm{d}\sigma(t)}{\mathrm{d}t} \sigma(t) \nabla_{\boldsymbol{x}} \log p \left( \frac{\boldsymbol{x}}{s(t)}; \sigma(t) \right) \right] \mathrm{d}t }(3)

数值积分通过离散时间步近似求解生成轨迹。ODE 的解可通过数值积分获得,即在离散时间区间上进行有限步推进。这一过程既需要选择积分方案,例如 Euler 方法或某种 Runge–Kutta 变体,也需要确定离散采样时刻\{t_0,t_1,\dots,t_N\}。许多已有工作采用 Euler 方法,但论文表明二阶求解器在计算代价与性能之间提供更优折中

🕊️ Improvements to deterministic sampling

作者提出的基本假设是:与采样过程相关的设计选择,在很大程度上独立于其他组件,例如网络结构和训练细节。换言之,D_\theta的训练方式不应决定\sigma(t)、s(t)和\{t_i\}的选取,反之亦然;从采样器的角度看,D_\theta只是一个黑箱

离散化与高阶积分器:对 ODE 的数值求解本质上是对真实解轨迹的近似。求解器在每一步都会引入截断误差,而这一误差会在N个步骤中不断累积。一般而言,局部误差相对于步长呈超线性缩放,因此增大N通常能够提升解的精度

常用的 Euler 方法是一阶 ODE 求解器,其局部误差相对于步长h为O(h^2)。更高阶的 Runge–Kutta 方法在误差缩放方面更有优势,但每一步需要多次评估D_\theta。通过大量实验,作者发现 Heun 二阶方法(亦称改进 Euler 法或梯形法)在截断误差与 NFE 之间提供极佳的折中;正如算法 1 所示,该方法在计算x_{i+1} 时引入额外的校正步骤,以补偿t_i到t_{i+1}之间dx/dt的变化。这一步校正以每步额外一次D_\theta评估为代价,将局部误差降低到O(h^3)。需要指出的是,当步进到\sigma=0时会导致除零问题,因此在该特殊情况下作者退回使用 Euler 方法

Algorithm 1 Deterministic Sampling Using Heun’s Second-Order Method

时间步\{t_i\}决定不同噪声水平下步长及截断误差的分布方式:随着\sigma的减小,步长应当单调减小,且步长无需针对每个样本单独调整。基于此,作者采用一种参数化方案,通过噪声水平序列\{\sigma_i\}定义时间步,即t_i=\sigma^{-1}(\sigma_i)​

\sigma_{i<N}=(Ai+B)^\rho(4)

  • 并选择常数 A 和 B,使得\sigma_0=\sigma_{\max},\sigma_{N-1}=\sigma_{\min},从而得到

\sigma_{i<N} = \left( \sigma_{\max}^{\frac{1}{\rho}} + \frac{i}{N-1} \left( \sigma_{\min}^{\frac{1}{\rho}} - \sigma_{\max}^{\frac{1}{\rho}} \right) \right)^\rho \quad \text{and} \quad \sigma_N = 0(5)

  • 其中\rho控制靠近\sigma_{\min}处的步长压缩程度,即在缩短低噪声区域步长的同时,允许高噪声区域使用更长步长。\rho=3几乎可以使每一步的截断误差相对均衡,但在图像采样任务中,\rho取 5 至 10 的范围效果更佳,这说明低噪声区域的误差对最终结果影响更大。基于实验,论文其余部分统一采用\rho=7

噪声与信号缩放调度通过轨迹曲率影响离散化误差。ODE 解轨迹的形状由\sigma(t)和s(t)两个函数共同决定。对这两个函数的选择,实际上也提供一条降低上述截断误差的途径,因为误差大小通常与dx/dt的曲率成正比。作者认为最优选择应为\sigma(t)=t与s(t)=1,这也是 DDIM 所采用的设置

dx/dt = \frac{x-D(x;t)}{t}(6)

  • 对于任意给定的x和t,只需执行一次到t=0的 Euler 步,就可得到去噪图像D_\theta(x;t)。因此,解轨迹的切向方向始终指向去噪器输出。这意味着该方向通常只会随噪声水平缓慢变化,也即解轨迹整体上接近线性

🐶 Stochastic sampling

确定性采样具有诸多优势,例如可以通过对 ODE 进行逆向求解,将真实图像映射为其对应的潜在表示。然而,与在每一步中向图像注入新噪声的随机采样相比,确定性采样往往会导致较差的生成质量。鉴于从理论上看,ODE 与 SDE 均可恢复相同的目标分布,那么随机性究竟发挥什么作用,便成为一个值得深入探讨的问题

Song 等提出的 SDE 可进一步推广为:由概率流 ODE 与一个时变的 Langevin diffusion SDE 共同构成的形式

{ \mathrm{d}\boldsymbol{x}_{\pm} = \underbrace{- \frac{\mathrm{d}\sigma(t)}{\mathrm{d}t} \sigma(t)\nabla_{\boldsymbol{x}} \log p(\boldsymbol{x}; \sigma(t)) \, \mathrm{d}t}_{\text{probability flow ODE}} \underbrace{ \pm\underbrace{\beta(t)\sigma(t)^2 \nabla_{\boldsymbol{x}} \log p(\boldsymbol{x}; \sigma(t)) \, \mathrm{d}t}_{\text{deterministic noise decay}} + \underbrace{\sqrt{2\beta(t)}\sigma(t) \, \mathrm{d}\omega_t}_{\text{noise injection}} }_{\text{Langevin diffusion SDE}} }(7)

  • 其中\omega_t表示标准 Wiener 过程。此时,dx^+与dx^-分别对应时间正向与反向演化的两个 SDE,它们之间通过 Anderson 的时间反演公式联系起来。进一步地,Langevin 项还可被理解为由一个基于 score 的确定性去噪项与一个随机噪声注入项所共同组成,而二者在净噪声水平上的贡献彼此抵消。因此,\beta(t)实际上表征旧噪声被新噪声替换的相对速率。当取\beta(t) = \frac{\mathrm{d}\sigma(t)}{\mathrm{d}t} / \sigma(t)时,便可恢复 Song 等人的原始 SDE;在该特例下,score 项在前向 SDE 中消失
  • 揭示随机性在实践中为何有助于提升性能:隐式的 Langevin diffusion 会在给定时刻主动推动样本向目标边缘分布靠近,从而不断纠正前面采样步骤中累积的误差。另一方面,以离散 SDE 求解器来近似 Langevin 项本身也会引入新的误差。已有研究表明,取非零的\beta(t)通常是有益的,但就作者所见,Song 等人对\beta(t)的隐式设定并不具有某种特殊的最优性质。因此,最优的随机性强度应通过经验实验加以确定

作者提出一种随机采样器,将前文的二阶确定性 ODE 积分器与显式的、类似 Langevin 的 churn 操作结合起来,即在每一步中交替执行加噪与去噪。在第i步,已知噪声水平t_i(=\sigma(t_i))下的样本x_i后,作者执行两个子步骤

  • 第一步,根据一个非负系数\gamma_i \ge 0,向样本中添加噪声,使其到达更高的噪声水平\hat{t}_i = t_i + \gamma_i t_i
  • 第二步,从加噪后的样本\hat{x}_i 出发,通过一次单步 ODE 反向求解,将其从\hat{t}_i推进到t_{i+1}。由此得到噪声水平为t_{i+1}的样本x_{i+1},并继续下一轮迭代
  • 作者特别强调,这并不是一个通用型 SDE 求解器,而是针对当前采样问题专门设计的采样程序。其正确性来源于两个子步骤的交替执行:二者各自都能保持正确的目标分布(ODE 步除截断误差外)
  • Euler–Maruyama 可以被理解为:首先添加噪声,然后执行一个 ODE 步;但这一 ODE 步并非从加噪后的中间状态出发,而是仍假定x与\sigma保持在该迭代开始时的初始状态。相比之下,在作者的方法中,算法 2 第 7 行用于评估D_\theta的参数对应的是噪声注入之后的状态,而类 Euler–Maruyama 的方法则会使用原始的(x_i,t_i),而不是(\hat{x}_i,\hat{t}_i)
Algorithm 2 Stochastic Sampling With σ(t) = t And s(t) = 1

🦔 Preconditioning and training

在监督学习训练神经网络时,已有许多被广泛认可的良好实践。例如,通常建议将输入与输出信号的幅值保持在相对稳定的范围内,如单位方差,并避免不同样本之间梯度幅值出现过大波动。若直接训练神经网络去建模D,显然并不理想。原因在于,输入x=y+n由干净信号y与噪声n\sim\mathcal N(0,\sigma^2I)共同构成,其幅值会随着噪声水平\sigma的变化而发生巨大改变。正因如此,现有方法通常并不直接将D_\theta表示为神经网络本身,而是转而训练另一个网络F_\theta,再由F_\theta构造出D_\theta

已有方法通过引入依赖于\sigma的归一化因子来处理输入缩放问题,并尝试通过训练F_\theta去预测经缩放后具有单位方差的噪声n来实现输出预条件化,随后再利用D_\theta(x;\sigma)=x-\sigma F_\theta(\cdot)恢复信号。然而,这种做法存在一个缺点

  • 当\sigma较大时,网络必须极为精细地调整其输出,才能恰好抵消已有噪声n,并使输出保持在正确尺度上;需要注意的是,网络产生的任何误差都会被\sigma放大。在这种情形下,直接预测期望输出D(x;\sigma)似乎会更容易

延续以往参数化方法中自适应混合信号与噪声的思想,提出利用一个依赖于\sigma的 skip connection 对神经网络进行预条件化,使其既可以估计y,也可以估计 n,或者估计二者之间的某种中间形式

D_\theta(x;\sigma)=c_{\text{skip}}(\sigma)x+c_{\text{out}}(\sigma)\, F_\theta\!\big(c_{\text{in}}(\sigma)x;\,c_{\text{noise}}(\sigma)\big)(8)

其中F_\theta是待训练的神经网络,c_{\text{skip}}(\sigma)用于调节 skip connection,c_{\text{in}}(\sigma)与c_{\text{out}}(\sigma)分别控制输入与输出的幅值缩放,而c_{\text{noise}}(\sigma)则负责将噪声水平\sigma映射为F_\theta的条件输入。整体训练损失如下

{ \mathbb{E}_{\sigma, \boldsymbol{y}, \boldsymbol{n}} \left[ \lambda(\sigma) \| D(\boldsymbol{y} + \boldsymbol{n}; \sigma) - \boldsymbol{y} \|_2^2 \right], \text{ where } \sigma \sim p_{\text{train}}, \boldsymbol{y} \sim p_{\text{data}}, \text{ and } \boldsymbol{n} \sim \mathcal{N}(\mathbf{0}, \sigma^2 \mathbf{I}) }(9)

给定噪声水平\sigma被采样到的概率由p_{\text{train}}(\sigma)决定,相应的损失权重由\lambda(\sigma)决定。进一步将该损失改写为关于原始网络输出F_\theta的形式

{ \mathbb{E}_{\sigma, \boldsymbol{y}, \boldsymbol{n}} \left[ \underbrace{\lambda(\sigma) c_{\text{out}}(\sigma)^2}_{\text{effective weight}} \bigg\| \underbrace{F_\theta(c_{\text{in}}(\sigma) \cdot (\boldsymbol{y} + \boldsymbol{n}); c_{\text{noise}}(\sigma))}_{\text{network output}} - \underbrace{\frac{1}{c_{\text{out}}(\sigma)} (\boldsymbol{y} - c_{\text{skip}}(\sigma) \cdot (\boldsymbol{y} + \boldsymbol{n}))}_{\text{effective training target}} \bigg\|_2^2 \right] }(10)

损失加权与噪声采样:预条件化网络F_\theta的有效损失权重为\lambda(\sigma)c_{\text{out}}(\sigma)^2。为使各噪声水平的有效权重均为 1,取\lambda(\sigma) = 1 / c_{\text{out}}(\sigma)^2 .噪声采样分布p_{\text{train}}(\sigma)采用对数正态分布,将训练集中于损失改善更明显的中等噪声区间

🦨 DPM-solver

扩散概率模型(DPMs)正逐渐成为一类性能强大的生成模型。DPMs 可以由离散时间随机过程或连续时间随机微分方程所定义,其核心思想是学习逐步去除施加在数据样本上的噪声。与广泛使用的生成对抗网络和变分自编码器相比,DPMs 不仅能够计算精确似然,而且在图像生成任务上往往能够获得更好的样本质量。现有的 DPM 快速采样方法大致可分为两类

  • 第一类包括知识蒸馏以及噪声水平或样本轨迹学习。这类方法在高效采样之前通常需要一个可能代价较高的训练阶段。此外,它们的适用性和灵活性也可能受到限制,因为将方法迁移到不同模型、数据集或不同采样步数设定下,往往需要付出额外努力
  • 第二类是免训练采样器,其能够以一种简单的即插即用方式适用于任意预训练 DPM。免训练采样器主要包括采用隐式或解析式的生成过程、先进的微分方程求解器,以及动态规划方法。然而,这些方法通常仍然需要约 50 次函数评估才能生成高质量样本,而这种质量大致才可与基础采样器在约 1000 次函数评估下生成的结果相当,因此依旧较为耗时

作者进一步将免训练采样器的效率提升到新的水平,使其能够在少步采样(few-step sampling)场景下生成高质量样本,即在大约 10 步顺序函数评估内完成采样

  • 将 DPM 采样视为求解相应扩散常微分方程的问题,并对扩散 ODE 的结构进行细致分析。扩散 ODE 具有一种半线性结构,即它由数据变量的线性函数与由神经网络参数化的非线性函数共同构成。此前的免训练采样方法并未充分利用这一结构,而是直接采用黑盒微分方程求解器
  • 为充分利用这一半线性结构,通过对解的线性部分进行解析计算,推导出扩散 ODE 解的精确表达,从而避免线性部分离散化所引入的误差。进一步地,通过变量替换,该解又可等价地化简为神经网络项的指数加权积分。由于这种积分形式具有特殊结构,因此可以借助指数积分器的数值方法进行高效近似

基于上述解的形式化表达,提出 DPM-Solver,通过对上述积分进行近似来高效求解扩散 ODE。构建 DPM-Solver 的一阶、二阶和三阶版本,并为其提供收敛阶保证。还进一步为 DPM-Solver 设计自适应步长调度策略。总体而言,DPM-Solver 既适用于连续时间 DPM,也适用于离散时间 DPM,同时还可用于带有分类器引导的条件采样

🕊️ Diffusion Probabilistic Models

前向过程与扩散 SDE:设x_0 \in \mathbb{R}^D是D维随机变量,其分布q_0(x_0)未知。扩散概率模型定义一个以前向时间演化过程\{x_t\}_{t\in[0,T]},其中T>0,且该过程以x_0为起点,使得对任意t\in[0,T],在给定x_0的条件下,x_t的条件分布满足

{ q_{0t}(\boldsymbol{x}_t|\boldsymbol{x}_0) = \mathcal{N}(\boldsymbol{x}_t|\alpha(t)\boldsymbol{x}_0, \sigma^2(t)\boldsymbol{I}) }(1)

其中\alpha(t),\sigma(t)\in\mathbb{R}_+是关于 t的可微函数,且其导数有界。为简便起见,分别记为\alpha_t,\sigma_t。\alpha_t与\sigma_t的具体选取被称为 DPM 的噪声调度。令q_t(x_t)表示x_t的边缘分布,DPM 通过适当设计噪声调度,使得q_T(x_T)\approx \mathcal{N}(x_T|0,\tilde{\sigma}^2I),\tilde{\sigma}>0,要求信噪比\alpha_t^2/\sigma_t^2关于t严格单调递减。进一步地,Kingma 等证明,对于任意t\in[0,T],如下随机微分方程与上式具有相同的转移分布q_{0t}(x_t|x_0)

{ \mathrm{d}\boldsymbol{x}_t = f(t)\boldsymbol{x}_t\mathrm{d}t + g(t)\mathrm{d}\boldsymbol{w}_t, \quad \boldsymbol{x}_0 \sim q_0(\boldsymbol{x}_0) }(2)

{ f(t) = \frac{\mathrm{d} \log \alpha_t}{\mathrm{d}t}, \quad g^2(t) = \frac{\mathrm{d}\sigma_t^2}{\mathrm{d}t} - 2\frac{\mathrm{d}\log \alpha_t}{\mathrm{d}t}\sigma_t^2 }(3)

在一定正则条件下,Song 等证明,所定义的前向过程存在一个等价的逆向过程,该过程从时刻T演化至 0,并以边缘分布q_T(x_T)为初始分布

{ \mathrm{d}\boldsymbol{x}_t = [f(t)\boldsymbol{x}_t - g^2(t)\nabla_{\boldsymbol{x}} \log q_t(\boldsymbol{x}_t)]\mathrm{d}t + g(t)\mathrm{d}\bar{\boldsymbol{w}}_t, \quad \boldsymbol{x}_T \sim q_T(\boldsymbol{x}_T) }(4)

其中\bar{w}_t 表示逆时间方向上的标准 Wiener 过程。唯一未知的项是任意时刻t对应的 score function  \nabla_x \log q_t(x_t).DPM 使用参数为\theta 的神经网络\epsilon_\theta(x_t,t)来估计缩放后的 score function,即-\sigma_t\nabla_x \log q_t(x_t)。参数\theta通过最小化如下目标函数进行优化

{
\begin{aligned}
\mathcal{L}(\theta; \omega(t)) &:= \frac{1}{2} \int_0^T \omega(t) \mathbb{E}_{q_t(\boldsymbol{x}_t)} \left[ \| \boldsymbol{\epsilon}_\theta(\boldsymbol{x}_t, t) + \sigma_t \nabla_{\boldsymbol{x}} \log q_t(\boldsymbol{x}_t) \|_2^2 \right] \mathrm{d}t \\
&= \frac{1}{2} \int_0^T \omega(t) \mathbb{E}_{q_0(\boldsymbol{x}_0)} \mathbb{E}_{q(\boldsymbol{\epsilon})} \left[ \| \boldsymbol{\epsilon}_\theta(\boldsymbol{x}_t, t) - \boldsymbol{\epsilon} \|_2^2 \right] \mathrm{d}t + C
\end{aligned}
}(5)

其中\omega(t)为加权函数,\epsilon\sim q(\epsilon)=\mathcal{N}(\epsilon|0,I),且x_t=\alpha_tx_0+\sigma_t\epsilon,常数C与\theta无关。由于\epsilon_\theta(x_t,t)也可以被理解为对加入到 x_t中的高斯噪声进行预测,通常被称噪声预测模型。鉴于\epsilon_\theta(x_t,t) 所对应的真实目标为-\sigma_t\nabla_x\log q_t(x_t),DPM 用-\epsilon_\theta(x_t,t)/\sigma_t替代 score function,从而定义出如下参数化逆过程(即扩散 SDE),其从T时刻演化至 0,并以x_T\sim \mathcal{N}(0,\tilde{\sigma}^2I)为初始条件

{ \mathrm{d}\boldsymbol{x}_t = \left[ f(t)\boldsymbol{x}_t + \frac{g^2(t)}{\sigma_t}\boldsymbol{\epsilon}_\theta(\boldsymbol{x}_t, t) \right] \mathrm{d}t + g(t)\mathrm{d}\bar{\boldsymbol{w}}_t, \quad \boldsymbol{x}_T \sim \mathcal{N}(\boldsymbol{0}, \tilde{\sigma}^2\boldsymbol{I}) }(6)

因此,可以通过数值求解器对扩散 SDE 从T到 0 进行离散化求解,从而生成样本。Song 等证明,DPM 的传统 ancestral sampling 方法可以被视为一个一阶 SDE 求解器。然而,这类一阶方法通常需要数百次甚至数千次函数评估才能收敛,从而导致极其缓慢的采样速度

扩散(概率流)ODE:在对 SDE 进行离散化时,步长大小受到 Wiener 过程随机性的限制。当步长较大(即步数较少)时,离散求解往往难以收敛,尤其是在高维空间中更是如此。为实现更快的采样,可以考虑与之对应概率流常微分方程,该 ODE 在每一时刻t上都与原 SDE 具有相同的边缘分布。对于 DPM,Song 等证明,对应的概率流 ODE 可写为

{ \frac{\mathrm{d}\boldsymbol{x}_t}{\mathrm{d}t} = f(t)\boldsymbol{x}_t - \frac{1}{2}g^2(t)\nabla_{\boldsymbol{x}} \log q_t(\boldsymbol{x}_t), \quad \boldsymbol{x}_T \sim q_T(\boldsymbol{x}_T) }(7)

其中x_t 的边缘分布同样为q_t(x_t)。进一步地,若用噪声预测模型替代 score function,则 Song 等定义如下参数化 ODE

{ \frac{\mathrm{d}\boldsymbol{x}_t}{\mathrm{d}t} = \boldsymbol{h}_\theta(\boldsymbol{x}_t, t) := f(t)\boldsymbol{x}_t + \frac{g^2(t)}{2\sigma_t}\boldsymbol{\epsilon}_\theta(\boldsymbol{x}_t, t), \quad \boldsymbol{x}_T \sim \mathcal{N}(\boldsymbol{0}, \tilde{\sigma}^2\boldsymbol{I}) }(8)

样本可以通过求解该 ODE 并将时间从 T推进到 0 来获得。与 SDE 相比,ODE 不包含随机项,因此允许采用更大的步长进行数值求解。此外,还可以借助高效的数值 ODE 求解器来进一步加速采样。Song 等采用 RK45 ODE 求解器来求解扩散 ODE,在 CIFAR-10 数据集上,仅需约 60 次函数评估,便可达到与对应的 1000 步 SDE 求解器相当的生成质量。然而,现有通用型 ODE 求解器在少步采样(约 10 步)场景下,仍难以生成令人满意的样本。目前在少步采样场景下,针对 DPM 的免训练采样器仍然较为缺乏,DPM 的采样效率问题仍然是一个亟待解决的关键挑战

🐇 Customized Fast Solvers for Diffusion ODEs

扩散 ODE 的线性部分可以精确计算,非线性部分则可转化为易于近似的积分。给定时刻s>0 处的初值x_s,考虑时刻t<s 处的解x_t的求解过程

利用扩散 ODE 的半线性结构,其解x_t的一部分实际上可以被精确计算。扩散 ODE 右端项由两部分构成:其中f(t)x_t 是关于x_t的线性函数,而另一项\frac{g^2(t)}{2\sigma_t}\epsilon_\theta(x_t,t)由于神经网络\epsilon_\theta(x_t,t)的存在,通常是x_t的非线性函数。这类 ODE 通常被称为半线性 ODE

  • 先前工作所采用的黑盒 ODE 求解器并未利用这种半线性结构,而是将整个h_\theta(x_t,t)视作输入统一处理,这会同时对线性项和非线性项引入离散化误差。对于半线性 ODE,其在时刻t的解可以通过常数变易公式被精确写为

{ \boldsymbol{x}_t = e^{\int_s^t f(\tau)\mathrm{d}\tau} \boldsymbol{x}_s + \int_s^t \left( e^{\int_\tau^t f(r)\mathrm{d}r} \frac{g^2(\tau)}{2\sigma_\tau} \boldsymbol{\epsilon}_\theta(\boldsymbol{x}_\tau, \tau) \right) \mathrm{d}\tau }(9)

  • 线性部分与非线性部分解耦。与黑盒 ODE 求解器不同,此时线性部分已被精确计算,从而消除线性项的近似误差。然而,非线性部分对应的积分仍然较为复杂,因为其中同时耦合与噪声调度有关的系数(即f(\tau), g(\tau), \sigma_\tau)以及复杂的神经网络\epsilon_\theta,因此仍难以直接近似

简化非线性积分,定义\lambda_t:=\log(\alpha_t/\sigma_t),即 \text{log-SNR}的一半。\lambda_t关于t是严格单调递减的。可以将g(t)重写为

{ g^2(t) = \frac{\mathrm{d}\sigma_t^2}{\mathrm{d}t} - 2\frac{\mathrm{d}\log \alpha_t}{\mathrm{d}t}\sigma_t^2 = 2\sigma_t^2\left(\frac{\mathrm{d}\log \sigma_t}{\mathrm{d}t} - \frac{\mathrm{d}\log \alpha_t}{\mathrm{d}t}\right) = -2\sigma_t^2\frac{\mathrm{d}\lambda_t}{\mathrm{d}t}. }(10)

  • 再结合f(t)=d\log\alpha_t/dt,进一步改写为

{ \boldsymbol{x}_t = \frac{\alpha_t}{\alpha_s}\boldsymbol{x}_s - \alpha_t \int_s^t \left(\frac{\mathrm{d}\lambda_\tau}{\mathrm{d}\tau}\right) \frac{\sigma_\tau}{\alpha_\tau} \boldsymbol{\epsilon}_\theta(\boldsymbol{x}_\tau, \tau)\mathrm{d}\tau. }(11)

  • 严格单调递减,记其反函数为 。令 ,并记 。将积分变量由时间换为 ,利用 ,即可得到扩散 ODE 的精确解
(12)
  • 线性部分由比例因子 精确计算,非线性部分则化为噪声预测模型的指数加权积分。DPM-Solver 对这一积分进行数值近似,进而构造不同阶数的求解器

对指数加权积分作泰勒展开,可以构造不同阶数的 DPM-Solver。给定时刻T处的初值x_T,以及一组从t_0=T递减到t_M=t_{\min}>0的M+1个时间步\{t_i\}_{i=0}^M。令 \tilde{x}_{t_0}=x_T为初始值。所提出的求解器通过M个迭代步骤,逐步计算一个序列\{\tilde{x}_{t_i}\}_{i=0}^M,以近似对应时间点\{t_i\}_{i=0}^M处的真实解。特别地,最后一个迭代值 \tilde{x}_{t_M}用于近似时刻 t_{\min} 处的真实解

  • 为减小\tilde{x}_{t_M}与时刻 t_{\min} 处真实解之间的误差,需要在每一步都尽可能减小各个\tilde{x}_{t_i}的近似误差。从前一个时刻t_{i-1}的近似值\tilde{x}_{t_{i-1}}出发,在时刻t_i 处的精确解x_{t_{i-1}\to t_i}可表示为

\boldsymbol{x}_{t_{i-1}\to t_i} = \frac{\alpha_{t_i}}{\alpha_{t_{i-1}}} \tilde{\boldsymbol{x}}_{t_{i-1}} - \alpha_{t_i} \int_{\lambda_{t_{i-1}}}^{\lambda_{t_i}} e^{-\lambda} \hat{\boldsymbol{\epsilon}}_\theta(\hat{\boldsymbol{x}}_\lambda, \lambda)\mathrm{d}\lambda.(13)

  • 关键在于近似\hat{\epsilon}_\theta从\lambda_{t_{i-1}}到\lambda_{t_i}区间上的指数加权积分。记h_i := \lambda_{t_i}-\lambda_{t_{i-1}},并记\hat{\epsilon}^{(n)}_\theta(\hat{x}_\lambda,\lambda) := \frac{d^n \hat{\epsilon}_\theta(\hat{x}_\lambda,\lambda)}{d\lambda^n}为\hat{\epsilon}_\theta(\hat{x}_\lambda,\lambda)关于\lambda的第n阶全导数。对于k\ge 1,在 \lambda_{t_{i-1}}处对\hat{\epsilon}_\theta(\hat{x}_\lambda,\lambda)作k-1阶泰勒展开,有

\hat{\boldsymbol{\epsilon}}_\theta(\hat{\boldsymbol{x}}_\lambda, \lambda) = \sum_{n=0}^{k-1} \frac{(\lambda - \lambda_{t_{i-1}})^n}{n!} \hat{\boldsymbol{\epsilon}}_\theta^{(n)}(\hat{\boldsymbol{x}}_{\lambda_{t_{i-1}}}, \lambda_{t_{i-1}}) + \mathcal{O}((\lambda - \lambda_{t_{i-1}})^k),(14)

  • 将上述泰勒展开代入,可得

\boldsymbol{x}_{t_{i-1}\to t_i} = \frac{\alpha_{t_i}}{\alpha_{t_{i-1}}} \tilde{\boldsymbol{x}}_{t_{i-1}} - \alpha_{t_i} \sum_{n=0}^{k-1} \hat{\boldsymbol{\epsilon}}_\theta^{(n)}(\hat{\boldsymbol{x}}_{\lambda_{t_{i-1}}}, \lambda_{t_{i-1}}) \int_{\lambda_{t_{i-1}}}^{\lambda_{t_i}} e^{-\lambda} \frac{(\lambda - \lambda_{t_{i-1}})^n}{n!} \mathrm{d}\lambda + \mathcal{O}(h_i^{k+1}),(15)

积分项\int e^{-\lambda}\frac{(\lambda-\lambda_{t_{i-1}})^n}{n!}\,d\lambda可以通过重复应用 n次分部积分解析求得。近似 x_{t_{i-1}\to t_i}只需要近似n\le k-1时的第n阶全导数\hat{\epsilon}^{(n)}_\theta(\hat{x}_\lambda,\lambda)。通过舍去误差项O(h_i^{k+1}),并利用 stiff order conditions 去近似前k-1阶全导数,便可以构造出扩散 ODE 的k阶求解器。将这类求解器统称为 DPM-Solver,而其特定阶数版本记为 DPM-Solver-k。以k=1的情形为例进行说明

\begin{aligned}
\boldsymbol{x}_{t_{i-1}\to t_i} &= \frac{\alpha_{t_i}}{\alpha_{t_{i-1}}} \tilde{\boldsymbol{x}}_{t_{i-1}} - \alpha_{t_i} \boldsymbol{\epsilon}_\theta(\tilde{\boldsymbol{x}}_{t_{i-1}}, t_{i-1}) \int_{\lambda_{t_{i-1}}}^{\lambda_{t_i}} e^{-\lambda}\mathrm{d}\lambda + \mathcal{O}(h_i^2) \\
&= \frac{\alpha_{t_i}}{\alpha_{t_{i-1}}} \tilde{\boldsymbol{x}}_{t_{i-1}} - \sigma_{t_i}(e^{h_i} - 1)\boldsymbol{\epsilon}_\theta(\tilde{\boldsymbol{x}}_{t_{i-1}}, t_{i-1}) + \mathcal{O}(h_i^2).
\end{aligned}(16)

DPM-Solver-1 沿递减时间步执行一阶更新。给定时刻T 处的初始值x_T,以及一组从t_0=T递减至t_M=t_{\min}>0的M+1个时间步\{t_i\}_{i=0}^M。从\tilde{x}_{t_0}=x_T开始,序列\{\tilde{x}_{t_i}\}_{i=1}^M按如下方式迭代计算

{ \tilde{\boldsymbol{x}}_{t_i} = \frac{\alpha_{t_i}}{\alpha_{t_{i-1}}} \tilde{\boldsymbol{x}}_{t_{i-1}} - \sigma_{t_i}(e^{h_i} - 1)\boldsymbol{\epsilon}_\theta(\tilde{\boldsymbol{x}}_{t_{i-1}}, t_{i-1}), \quad \text{where}\quad h_i = \lambda_{t_i} - \lambda_{t_{i-1}}. }(17)

  • 对于 k\ge 2 的情形,为近似泰勒展开中的前k项,通常需要在t与 s 之间引入额外的中间点

定理 3.2(DPM-Solver-k 是一个 k 阶求解器):设噪声预测模型\epsilon_\theta(x,t)具有足够的正则性,即其关于状态与时间的变化足够平滑,使得所涉及的泰勒展开与高阶导数近似成立;同时\epsilon_\theta(x,t)对状态变量 x满足 Lipschitz 连续条件。进一步离散求解时所采用的时间步长足够小,则对于 k=1,2,3,DPM-Solver-k 分别构成扩散 ODE 的一个 k 阶数值求解方法

{ \tilde{\boldsymbol{x}}_{t_M} - \boldsymbol{x}_{t_{\min}} = \mathcal{O}(h_{\max}^k), \text{ where } h_{max} = \max_{1 \le i \le M}(\lambda_{t_i} - \lambda_{t_{i-1}}). }(18)

当 \sigma(0)=0 时,\lambda(0) 发散,不能直接代入有限步长公式;实际求解取正的终止时间 t_{\min},并与同一终点的 ODE 精确解比较。需要零噪声输出时,再执行单独的最终去噪。求解器需要预先指定时间步序列\{t_i\}_{i=0}^M。作者提出两种时间步长调度方式

  • 其中一种是手工设定的步长方案,即在区间[\lambda_T,\lambda_{t_{\min}}]上进行均匀划分,也就是说:\lambda_{t_i}=\lambda_T+\frac{i}{M}(\lambda_{t_{\min}}-\lambda_T),\quad i=0,\dots,M.经验结果表明,当采用\lambda_{t_i}的均匀时间步时,DPM-Solver 在少步采样场景下已经能够生成质量相当不错的样本
  • 另一种方案是自适应步长算法。该方法通过组合不同阶数的 DPM-Solver,动态调节步长大小
  • 对于少步采样而言,需要充分利用全部可用函数评估次数。当 NFE 不能被 3 整除时,作者的策略是:优先尽可能多地使用 DPM-Solver-3,随后根据 NFE 除以 3 的余数,再额外补充一步 DPM-Solver-1 或 DPM-Solver-2

🦄 Comparison with Existing Fast Sampling Methods

DDIM 作为 DPM-Solver-1 的特例:去噪扩散隐式模型提出一种从 DPM 中快速采样的确定性方法。对于相邻的两个时间步t_{i-1}和t_i,假设在时刻t_{i-1}已有一个解\tilde{x}_{t_{i-1}},则 DDIM 从 t_{i-1}到 t_i的单步更新可写为

\tilde{x}_{t_i} = \frac{\alpha_{t_i}}{\alpha_{t_{i-1}}}\tilde{x}_{t_{i-1}} - \alpha_{t_i} \left( \frac{\sigma_{t_{i-1}}}{\alpha_{t_{i-1}}} - \frac{\sigma_{t_i}}{\alpha_{t_i}} \right) \epsilon_\theta(\tilde{x}_{t_{i-1}},t_{i-1}). (19)

尽管 DDIM 与 DPM-Solver-1 的出发视角完全不同,作者指出,二者的更新公式实际上是完全一致的。根据\lambda的定义,有

\frac{\sigma_{t_{i-1}}}{\alpha_{t_{i-1}}}=e^{-\lambda_{t_{i-1}}},\qquad \frac{\sigma_{t_i}}{\alpha_{t_i}}=e^{-\lambda_{t_i}}(20)

将这些关系以及h_i=\lambda_{t_i}-\lambda_{t_{i-1}}代入,便可精确得到 DPM-Solver-1 单步更新公式。因此,DDIM 实际上就是 DPM-Solver-1 的一个特例。不过,DPM-Solver 的优势在于,其基于扩散 ODE 的半线性结构进行建模,因此能够以一种更具原理性的方式推广到高阶求解器,并开展相应的收敛阶分析

与传统 Runge–Kutta 方法的比较:一种获得高阶求解器的直接思路,是将传统的显式 Runge–Kutta(RK)方法直接应用于扩散 ODE。RK 方法通常将解写为如下积分形式

x_t=x_s+\int_s^t h_\theta(x_\tau,\tau)\,d\tau = x_s+\int_s^t \left[ f(\tau)x_\tau+\frac{g^2(\tau)}{2\sigma_\tau}\epsilon_\theta(x_\tau,\tau) \right]d\tau(21)

并在区间[t,s]内选取若干中间时间点,通过组合这些时间点上对h_\theta的函数评估来近似整个积分。然而,显式 RK 方法的近似误差取决于 h_\theta本身,而h_\theta同时包含线性项 f(\tau)x_\tau与非线性的噪声预测模型\epsilon_\theta所带来的误差。问题在于,线性项的误差可能会被显著放大,因为线性项精确解中本身包含指数系数。已有大量经验研究表明,对于半线性 ODE,若直接采用显式 RK 方法并使用较大步长,往往会遭遇数值不稳定问题

面向 DPM 的训练型快速采样方法:需要额外训练或优化的快速采样方法主要包括:知识蒸馏、学习噪声水平或方差,以及学习噪声调度或样本轨迹。尽管渐进式蒸馏方法可以将快速采样压缩到 4 步以内,但它需要额外的训练成本,并且会损失原始 DPM 中的部分信息。例如,在蒸馏之后,噪声预测模型将不再能够在区间 [0,T]上的任意时间步都对噪声(即 score function)进行预测。免训练采样器能够完整保留原始模型中的全部信息,因此也更容易直接扩展到条件采样场景,例如通过结合原始模型和外部分类器来实现带有 classifier guidance 的条件采样

🦨 Consistency Models

一致性模型的核心目标,是在保留扩散模型迭代采样优势的基础上,实现高效的单步生成。具体而言,该模型既希望具备一步生成高质量样本的能力,又希望在必要时仍可通过增加采样步数来进一步提升生成质量,并继承扩散模型在零样本数据编辑任务中的灵活性

  • 基于概率流常微分方程的建模思想:该方法建立在连续时间扩散模型的概率流常微分方程之上。该 ODE 的轨迹能够将数据分布平滑地演化为易于处理的噪声分布。基于这一性质,作者提出学习一个映射函数,使其能够将任意时间步上的样本点映射回对应轨迹的起始点。该映射具有显著的自一致性特征,即同一条轨迹上的不同点都应映射到同一个初始点,因此这类模型被称为一致性模型
  • 兼顾单步高效生成与多步质量提升:一致性模型仅需一次网络前向传播,便可将随机噪声直接映射为数据样本,从而实现高效的单步生成。进一步地,通过在多个时间步上串联一致性模型的输出,还可以像扩散模型的迭代采样过程那样,以额外计算开销为代价换取更高的样本质量,并支持零样本数据编辑等任务

为训练一致性模型,作者提出两种基于自一致性约束的方法

  • 第一种方法依赖数值 ODE 求解器和预训练扩散模型,在概率流 ODE 轨迹上构造相邻点对,并通过约束模型在这些相邻点上的输出保持一致,实现将扩散模型有效蒸馏为一致性模型
  • 第二种方法则摆脱对预训练扩散模型的依赖,使一致性模型可以作为一种独立的生成模型直接进行训练
  • 基于这两种训练范式,一致性模型不仅能够继承扩散模型蒸馏带来的高质量生成能力,也由此发展为一类独立的生成模型框架。值得注意的是,这两种方法均不需要对抗训练,且对模型结构的限制较少,因此能够灵活适配多种神经网络架构

从方法效果上看,一致性模型在少步生成场景中展现出显著优势。作为蒸馏方法时,它能够实现高效的单步或少步采样;而作为独立生成模型训练时,即便不借助预训练扩散模型,也能够获得具有竞争力的生成性能。此外,一致性模型还具备良好的任务泛化能力,可广泛应用于多种零样本数据编辑任务,包括图像去噪、插值、图像修补、图像上色、超分辨率重建以及基于笔画引导的图像编辑等

🦫 Diffusion Models

一致性模型的提出,在理论上深受连续时间扩散模型框架的启发。在这一框架中,扩散模型首先通过连续注入高斯噪声,将真实数据分布逐步扩散为更简单、更易处理的噪声分布;随后再通过相应的反向过程,从噪声中恢复出数据样本。设p_{\text{data}}(x)表示真实数据分布,则其前向扩散过程可由如下随机微分方程描述

  • 其中,t\in[0,T]表示连续时间变量,T>0为预先设定的终止时刻;\mu(\cdot,\cdot)和\sigma(\cdot)分别表示漂移系数与扩散系数,\{w_t\}_{t\in[0,T]}表示标准布朗运动。将x_t的分布记为p_t(x),因此有p_0(x)=p_{\text{data}}(x)

该 SDE 的一个重要性质在于:存在一个常微分方程,Song 等将其称为概率流 ODE,其解轨迹在任意时刻t的边缘分布均服从p_t(x)

  • 其中,\nabla \log p_t(x)是分布 p_t(x)的分数函数;因此,扩散模型也常被称为基于分数的生成模型。通常,SDE 会被设计为使得p_T(x)接近某个易于处理的高斯分布\pi(x)。后续采用 Karras 的设定,即令\mu(x,t)=0,\sigma(t)=\sqrt{2t}。在该情况下,有

  • 在采样阶段,首先通过分数匹配训练一个分数模型s_\phi(x,t)\approx \nabla \log p_t(x),得到 PF ODE 的经验估计,其形式为

  • 上式被称为经验概率流 ODE。随后,从\pi=\mathcal{N}(0,T^2I)中采样得到 \hat{x}_T作为经验 PF ODE 的初始值,并利用任意数值 ODE 求解器(如 Euler 求解器和 Heun 求解器)沿时间反向求解,得到解轨迹\{\hat{x}_t\}_{t\in[0,T]}。最终得到的 \hat{x}_0可视为数据分布p_{\text{data}}(x)的近似样本。为避免数值不稳定性,通常不会将求解器真正积分到t=0,而是在t=\epsilon处停止,其中\epsilon为一个固定的较小正数,并将\hat{x}_\epsilon作为近似生成样本。按照 Karras 的设置,将图像像素值缩放到区间[-1,1],并设定T=80,\ \epsilon=0.002

扩散模型的一个主要瓶颈在于其采样速度较慢。显然,利用 ODE 求解器进行采样需要对分数模型s_\phi(x,t)进行多次迭代评估,这会带来较高的计算开销。现有的快速采样方法主要包括两类:一类是设计更高效的数值 ODE 求解器,另一类是采用蒸馏技术。然而,即便采用改进的 ODE 求解器,通常仍需要超过 10 次函数评估,才能生成具有竞争力的样本。而大多数蒸馏方法都依赖于在蒸馏之前预先从扩散模型中生成一个大规模样本数据集,而这一过程本身就具有较高的计算成本

🦥 Consistency Models

作者提出一致性模型,这是一类在设计上原生支持单步生成的新型模型,同时仍允许通过迭代生成在样本质量与计算开销之间进行权衡,并支持零样本数据编辑。一致性模型可以通过两种模式进行训练,即蒸馏模式和独立训练模式

  • 在前一种模式下,一致性模型将预训练扩散模型中的知识蒸馏到一个单步采样器中,在样本质量方面显著优于其他蒸馏方法,同时仍可支持零样本图像编辑应用
  • 在后一种模式下,一致性模型则在不依赖任何预训练扩散模型的情况下独立训练,从而使其成为一类独立的新型生成模型

一致性模型的定义:给定概率流 ODE 的一条解轨迹{\{x_t\}_{t\in[\epsilon,T]}},将一致性函数定义为f:(x_t,t)\mapsto x_\epsilon

  • 对于属于同一条概率流 ODE 轨迹的任意一对(x_t,t),其输出应保持一致,即对任意t,t'\in[\epsilon,T],都有f(x_t,t)=f(x_{t'},t')
  • 一致性模型f_\theta的目标是通过数据学习来逼近上述一致性函数f,其核心在于学习并显式施加这种自一致性约束

一致性模型的参数化:对于任意一致性函数f(\cdot,\cdot),都有f(x_\epsilon,\epsilon)=x_\epsilon,即当时间取边界值\epsilon时,函数f(\cdot,\epsilon) 应当退化为恒等映射。将这一约束称为边界条件。所有一致性模型都必须满足该边界条件,因为它在一致性模型的成功训练中起着关键作用。同时,这一边界条件也是一致性模型在架构层面最主要、最本质的约束

对于基于深度神经网络构建的一致性模型,作者讨论两种几乎无需额外代价即可实现该边界条件的参数化方式。设F_\theta(x,t) 是一个自由形式的深度神经网络,其输出维度与x相同。第一种方式是直接将一致性模型参数化为

第二种方式则利用跳跃连接进行参数化,其中c_{\text{skip}}(t)和c_{\text{out}}(t)是可微函数,并满足c_{\text{skip}}(\epsilon)=1,c_{\text{out}}(\epsilon)=0。在这种参数化方式下,只要F_\theta(x,t)、c_{\text{skip}}(t)和c_{\text{out}}(t)都是可微的,则一致性模型在t=\epsilon处也是可微的,而这一点对于连续时间一致性模型的训练至关重要

一致性模型的采样:在获得训练良好的一致性模型f_\theta(\cdot,\cdot)后,可以首先从初始分布中采样\hat{x}_T\sim \mathcal{N}(0,T^2I),然后通过一次一致性模型评估得到\hat{x}_\epsilon=f_\theta(\hat{x}_T,T).这一过程仅涉及一次前向传播,因此可以实现单步采样生成。更重要的是,一致性模型也可以通过交替执行去噪与噪声注入步骤,多次调用一致性模型,从而进一步提升样本质量

算法 1 对这种多步采样过程进行总结。该过程赋予模型一种灵活机制,使其能够在计算开销与样本质量之间进行权衡。在实际实现中,作者使用贪心算法来确定算法 1 中的时间点\tau_1,\tau_2,\ldots,\tau_{N-1}。具体而言,这些时间点是逐个选取的,并通过三分搜索来优化由算法 1 所生成样本的 FID

Algorithm 1 Multistep Consistency Sampling

🐳 Training Consistency Models via Distillation

首先介绍第一种一致性模型训练方法,即基于对预训练分数模型s_\phi(x,t)进行蒸馏的训练方式。讨论围绕经验概率流 ODE 展开,该方程是将分数模型s_\phi(x,t)代入概率流 ODE 后得到的

考虑将时间区间[\epsilon, T]离散为N-1个子区间,其边界点满足t_1=\epsilon < t_2 < \cdots < t_N=T.在实际实现中,遵循 Karras 的设定,按照如下公式确定这些时间边界

当N足够大时,可以通过数值 ODE 求解器执行一步离散更新,从x_{t_{n+1}}得到x_{t_n}的准确近似。将该估计记为\hat{x}^{\phi}_{t_n},并定义为

其中\Phi(\cdot,\cdot;\phi)表示应用于经验 PF ODE 的单步 ODE 求解器的更新函数。例如,当采用 Euler 求解器时,有\Phi(x,t;\phi)=-t\,s_\phi(x,t),相应地,其更新规则为

由于概率流 ODE 与 SDE 之间存在对应关系,可以通过以下方式对 ODE 轨迹上的分布进行采样:首先从真实数据分布中采样 x\sim p_{\text{data}},然后向其加入高斯噪声。具体而言,给定一个数据点x,可先从 SDE 的转移分布\mathcal{N}(x,t_{n+1}^2I)中采样得到x_{t_{n+1}},再通过数值 ODE 求解器执行一步离散更新,计算得到 \hat{x}^{\phi}_{t_n}。由此,便可高效生成概率流 ODE 轨迹上一对相邻点(\hat{x}^{\phi}_{t_n},x_{t_{n+1}})。随后,通过最小化一致性模型在这对点上的输出差异来进行训练。这一思想引出如下的一致性蒸馏损失

一致性蒸馏通过约束相邻轨迹点的预测一致来训练模型

  • 期望对 、 和 计算, 由式(10)得到
  • 为权重, 为目标网络参数;距离函数满足 ,且仅当 时取 0

作者考虑三种距离度量:平方\ell_2 距离d(x,y)=\|x-y\|_2^2、\ell_1距离d(x,y)=\|x-y\|_1,以及感知相似性指标 LPIPS。实验结果表明,在所有任务和数据集上,取 \lambda(t_n)\equiv 1都能够取得良好效果。在实际训练中,作者通过随机梯度下降优化模型参数\theta,同时利用指数滑动平均(EMA)更新\theta^-。具体而言,给定衰减率0\le \mu <1,在每一步优化之后执行如下更新

Algorithm 2 Consistency Distillation (CD)

Input: dataset , initial model parameter , learning rate , ODE solver , , , and

作者将f_{\theta^-}称为目标网络,将 f_\theta称为在线网络。实验中发现,相较于直接令\theta^-=\theta,EMA 更新和 stopgrad 操作能够显著稳定训练过程,并提升一致性模型的最终性能。下面,作者基于渐近分析给出一致性蒸馏的理论解释

定理 1:一致性蒸馏的误差界

记 , 为式(8)所定义经验概率流 ODE 的一致性函数

  • 假设 对输入满足一致的 Lipschitz 条件:存在 ,使任意 、 和 均满足
  • 假设 ODE 求解器在每个区间上的局部误差一致有界于 ,其中 ;若 ,则有

🦫 Training Consistency Models in Isolation

一致性模型的训练并不一定依赖于预训练扩散模型。这一点不同于现有的大多数扩散蒸馏方法,也正因此,使得一致性模型不仅是一种蒸馏框架,更成为一类独立的生成模型家族。回顾一致性蒸馏方法可知,该方法依赖于一个预训练分数模型s_\phi(x,t)来逼近真实分数函数\nabla \log p_t(x)。然而,作者指出,实际上可以完全摆脱这一预训练分数模型,而改为利用如下无偏估计量

  • 其中x\sim p_{\text{data}}且x_t\sim \mathcal{N}(x,t^2I).数据样本x在给定噪声样本x_t的后验条件分布下仍具有随机性;-\frac{x_t-x}{t^2}是\nabla \log p_t(x_t)的条件无偏估计。作者进一步表明,在采用 Euler 方法 作为 ODE 求解器、并考虑N\to\infty的极限情形时,这一无偏估计足以替代一致性蒸馏中原本所需的预训练扩散模型

定理 2:一致性训练与蒸馏的渐近关系

沿用 ,在下列条件下,一致性训练损失与蒸馏损失具有渐近关系

  • 与 二阶连续可微且二阶导数有界, 有界,并满足
  • 采用 Euler 求解器,且对所有 ,分数模型精确满足

其中,一致性训练目标定义为

  • 期望对 、 和 计算,两侧加噪输入使用同一个噪声样本
Algorithm 3 Consistency Training (CT)

Input: dataset , initial model parameter , learning rate , step schedule , EMA decay rate schedule , , and