Generative Models III

🕊️ Denoising Diffusion Probabilistic Models

扩散概率模型定义:一个经过参数化的马尔可夫链,通过变分推断进行训练,使其在有限步采样后生成与真实数据分布一致的样本。该模型包含两个方向相反的过程

  • 其一是前向扩散过程,即一个固定的马尔可夫链,按照给定方差调度逐步向数据中加入高斯噪声,直到原始信号基本被破坏
  • 其二是反向生成过程,即学习一个从纯噪声逐步恢复样本的马尔可夫链,用来逆转前向加噪过程
  • 当每一步扩散只加入较小的高斯噪声时,反向过程的转移也可以设为条件高斯分布,因此能够采用非常简洁的神经网络参数化形式。训练时,模型优化负对数似然的变分上界。而在特定的噪声预测参数化下,这一目标又可写成与多噪声层级 denoising score matching 相似的形式,并使采样过程呈现出 annealed Langevin dynamics 的解释

DDPM 结合噪声预测与加权变分训练,实现逐步去噪生成

  1. 高质量生成能力验证:论文证明扩散模型能够生成与 GAN、自回归模型、流模型及 EBM 等主流方法相竞争的高质量样本,并在 CIFAR-10 上取得 9.46 的 IS 和 3.17 的 FID,确立扩散模型作为强生成范式的地位
  2. 扩散与 DSM/Langevin 的统一联系:论文建立扩散模型与 denoising score matching、annealed Langevin dynamics 之间的明确桥梁,并提出基于噪声预测与加权变分下界的实用训练方式
  3. 渐进式生成/压缩解释:论文提出扩散采样可以理解为一种由粗到细的渐进式解码过程,即先恢复全局结构,再补充局部细节;从压缩角度看,这相当于一种有损到无损的渐进编码,也可被视为对自回归解码的广义推广,从而解释样本观感强但似然不占优的现象

🐦 Background

1️⃣ 扩散模型本质上是一类潜变量生成模型,其概率形式表示为p_\theta(x_0):=\int p_\theta(x_{0:T})\,dx_{1:T},其中x_1,\ldots,x_T表示与观测数据x_0 \sim q(x_0)维度相同的一系列潜变量。模型的联合分布p_\theta(x_{0:T})通常称为反向生成过程,其被建模为一个带有可学习高斯转移核的马尔可夫链,并从各向同性高斯先验分布p(x_T)=\mathcal{N}(x_T;0,I)开始,逐步恢复出目标数据样本

p_\theta(\mathbf{x}_{0:T}) := p(\mathbf{x}_T) \prod_{t=1}^{T} p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t), \quad p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t) := \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))(1)

2️⃣ 扩散模型区别于其他潜变量模型的关键在于,其近似后验分布q(x_{1:T}\mid x_0)是预先固定的;该分布被称为前向过程或扩散过程。具体而言,它被设定为一个马尔可夫链,并按照给定的方差调度\beta_1,\ldots,\beta_T,逐步向原始数据中加入高斯噪声

q(\mathbf{x}_{1:T}|\mathbf{x}_0) := \prod_{t=1}^{T} q(\mathbf{x}_t|\mathbf{x}_{t-1}), \quad q(\mathbf{x}_t|\mathbf{x}_{t-1}) := \mathcal{N}(\mathbf{x}_t; \sqrt{1-\beta_t}\mathbf{x}_{t-1}, \beta_t\mathbf{I})(2)

3️⃣ 模型训练通过优化关于负对数似然的标准变分上界来进行,即把直接最小化负对数似然,换成最小化它的一个可计算上界

\mathbb{E}\left[-\log p_\theta(\mathbf{x}_0)\right] \leq \mathbb{E}_q\left[-\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}\right] = \mathbb{E}_q\left[-\log p(\mathbf{x}_T) - \sum_{t\geq 1}\log \frac{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_t|\mathbf{x}_{t-1})}\right] =: L(3)

\begin{align}
L &= \mathbb{E}_q\left[-\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}|\mathbf{x}_0)}\right] \\
&= \mathbb{E}_q\left[-\log p(\mathbf{x}_T) - \sum_{t\geq 1}\log \frac{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_t|\mathbf{x}_{t-1})}\right] \\
&= \mathbb{E}_q\left[-\log p(\mathbf{x}_T) - \sum_{t> 1}\log \frac{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_t|\mathbf{x}_{t-1})} - \log \frac{p_\theta(\mathbf{x}_0|\mathbf{x}_1)}{q(\mathbf{x}_1|\mathbf{x}_0)}\right] \\
&= \mathbb{E}_q\left[-\log p(\mathbf{x}_T) - \sum_{t> 1}\log \left[\frac{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)} \cdot \frac{q(\mathbf{x}_{t-1}|\mathbf{x}_0)}{q(\mathbf{x}_t|\mathbf{x}_0)}\right] - \log \frac{p_\theta(\mathbf{x}_0|\mathbf{x}_1)}{q(\mathbf{x}_1|\mathbf{x}_0)}\right] \\
&= \mathbb{E}_q\left[-\log \frac{p(\mathbf{x}_T)}{q(\mathbf{x}_T|\mathbf{x}_0)} - \sum_{t> 1}\log \frac{p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)}{q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0)} - \log p_\theta(\mathbf{x}_0|\mathbf{x}_1)\right] \\
&= \mathbb{E}_q\left[D_{\text{KL}}(q(\mathbf{x}_T|\mathbf{x}_0) \| p(\mathbf{x}_T)) + \sum_{t>1} D_{\text{KL}}(q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) \| p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t)) - \log p_\theta(\mathbf{x}_0|\mathbf{x}_1)\right]
\end{align}(4)

4️⃣ 前向过程中的方差参数\beta_t 可以通过重参数化方法进行学习,也可以作为超参数保持固定。反向过程的表达能力在一定程度上由p_\theta(x_{t-1}\mid x_t)中采用高斯条件分布这一选择所保证,因为当\beta_t足够小时,前向过程与反向过程具有相同的函数形式。前向过程的一个显著性质在于,它允许在任意时间步t以闭式形式对x_t 进行采样

\alpha_t := 1 - \beta_t \quad \text{and} \quad \bar{\alpha}_t := \prod_{s=1}^{t} \alpha_s.(5)


\boxed{q(\mathbf{x}_t | \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t}\mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I})}
(6)

5️⃣ 前向后验具有闭式高斯形式,所有 KL 散度项都变成高斯分布之间的比较,从而可以采用 Rao-Blackwellized 的方式,通过闭式表达式直接计算,而无需使用高方差的 Monte Carlo 估计

\mathbb{E}_q \left[ \underbrace{D_{\mathrm{KL}}(q(\mathbf{x}_T|\mathbf{x}_0) \parallel p(\mathbf{x}_T))}_{L_T} + \sum_{t>1} \underbrace{D_{\mathrm{KL}}(q(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1}|\mathbf{x}_t))}_{L_{t-1}} - \underbrace{\log p_\theta(\mathbf{x}_0|\mathbf{x}_1)}_{-L_0} \right](7)

q(x_{t-1} \mid x_t, x_0) = \mathcal{N}(x_{t-1}; \tilde{\mu}_t(x_t, x_0), \tilde{\beta}_t I)(8)

\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) := \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t}\mathbf{x}_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t}\mathbf{x}_t \quad \text{and} \quad \tilde{\beta}_t := \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t}\beta_t(9)

6️⃣ 前向后验q(x_{t-1}\mid x_t,x_0)的闭式表达式推导

  1. 由前向过程的闭式形式可知

q(x_t \mid x_{t-1}) = \mathcal{N}(x_t; \sqrt{\alpha_t}x_{t-1}, \beta_t I). q(x_{t-1} \mid x_0) = \mathcal{N}(x_{t-1}; \sqrt{\bar{\alpha}_{t-1}}x_0, (1 - \bar{\alpha}_{t-1})I).(10)

  1. 由 Bayes 公式写出后验

q(x_{t-1} \mid x_t, x_0) = \frac{q(x_t \mid x_{t-1}, x_0) q(x_{t-1} \mid x_0)}{q(x_t \mid x_0)} = \frac{q(x_t \mid x_{t-1}) q(x_{t-1} \mid x_0)}{q(x_t \mid x_0)}.(11)

q(x_{t-1} \mid x_t, x_0) \propto q(x_t \mid x_{t-1}) q(x_{t-1} \mid x_0).(12)

q(x_{t-1} \mid x_t, x_0) \propto \exp \left( -\frac{1}{2\beta_t} \| x_t - \sqrt{\alpha_t} x_{t-1} \|^2 - \frac{1}{2(1 - \bar{\alpha}_{t-1})} \| x_{t-1} - \sqrt{\bar{\alpha}_{t-1}} x_0 \|^2 \right).(13)

  1. 忽略与x_{t-1}无关的常数项

q(x_{t-1} \mid x_t, x_0) \propto \exp \left\{ -\frac{1}{2} \left[ \left( \frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}} \right) \|x_{t-1}\|^2 - 2 \left( \frac{\sqrt{\alpha_t}}{\beta_t} x_t + \frac{\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} x_0 \right)^\top x_{t-1} \right] \right\}.(14)

  1. 高斯分布\mathcal N(x;\mu,\Sigma)的指数项标准形式是-\frac12\left[x^\top \Sigma^{-1}x-2(\Sigma^{-1}\mu)^\top x\right]

\Sigma_t = \left( \frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}} \right)^{-1} I, \quad \Sigma_t = \frac{\beta_t(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} I, \quad \tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t.(15)

\tilde{\mu}_t = \frac{\beta_t(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \left( \frac{\sqrt{\alpha_t}}{\beta_t} x_t + \frac{\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} x_0 \right) = \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} x_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} x_0.(16)

🦣 Diffusion Models And Denoising Autoencoders

将扩散模型的训练目标写为负对数似然的变分上界,并进一步分解为若干时间步上的 KL 散度项

  • 对于t>1的中间各步,由于前向后验q(x_{t-1}\mid x_t,x_0)具有闭式高斯形式,相应的 KL 项都可以转化为反向过程均值之间的加权平方误差
  • 结合前向过程的重参数化表示,模型训练又可等价改写为:在给定带噪样本x_t和时间步t的条件下,直接预测其中所包含的噪声 \varepsilon
  • 对最后一步的离散像素解码p_\theta(x_0\mid x_1) 进行专门处理,以保证模型能够对离散图像数据给出合理的概率建模
  • 从完整的变分目标出发,进一步提出实践中更常用的简化训练目标,即随机采样时间步并最小化真实噪声与预测噪声之间的均方误差,从而得到更稳定且样本质量更优的训练方式

\boxed{ L \;=\; \mathbb{E}_q\Big[ \underbrace{D_{KL}(q(x_T|x_0)\|p(x_T))}_{L_T} \;+\;\sum_{t>1}\underbrace{D_{KL}(q(x_{t-1}|x_t,x_0)\|p_\theta(x_{t-1}|x_t))}_{L_{t-1}} \;-\;\underbrace{\log p_\theta(x_0|x_1)}_{-L_0} \Big] }(17)

1️⃣ 在实现里把前向噪声日程\beta_t 预先设定为随时间步变化的序列,不参与训练,因此L_T 在训练时是常数项,直接忽略

2️⃣ 把t>1 的 KL 变成均值回归

  1. 先固定反向方差,只学均值,并把协方差固定成各向同性常数, 把训练重点集中到如何预测均值\mu_\theta上。首先,将\Sigma_\theta(x_t, t) = \sigma_t^2 I 设为不参与训练的、随时间变化的常数。设定\sigma_t^2 = \beta_t \quad \text{or} \quad \sigma_t^2 = \tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t都得到相似的结果

p_{\theta}(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t)),(18)

  1. 把L_{t-1}写成均值回归,让网络输出的均值 \mu_\theta去拟合真实后验均值\tilde\mu_t

\boxed{q(x_{t-1} \mid x_t, x_0) = \mathcal{N}\left(x_{t-1}; \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} x_0 + \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} x_t, \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \beta_t I\right).}(19)

L_{t-1} =\mathbb E_q\Big[\frac{1}{2\sigma_t^2}\|\tilde\mu_t(x_t,x_0)-\mu_\theta(x_t,t)\|^2\Big]+C(20)

  1. 把x_t用噪声重参数化写出来,把预测后验均值转成预测噪声\epsilon的问题

x_t(x_0,\varepsilon)=\sqrt{\bar\alpha_t}\,x_0+\sqrt{1-\bar\alpha_t}\,\varepsilon,\quad \varepsilon\sim\mathcal N(0,I)(21)

\tilde{\mu}_t(x_t, x_0) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon \right).(22)

\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right).(23)

  1. 把L_{t-1}写成只依赖 (x_0,\varepsilon) 的形式, \mu_\theta 必须在给定 x_t的情况下预测一个包含\varepsilon 的量。不直接学 \mu_\theta,而是让网络学 \varepsilon_\theta(x_t,t)

\begin{aligned}
L_{t-1} - C &= \mathbb{E}_{\mathbf{x}_0, \boldsymbol{\epsilon}} \left[ \frac{1}{2\sigma_t^2} \left\| \tilde{\boldsymbol{\mu}}_t \left( \mathbf{x}_t(\mathbf{x}_0, \boldsymbol{\epsilon}), \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t(\mathbf{x}_0, \boldsymbol{\epsilon}) - \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}) \right) - \boldsymbol{\mu}_\theta(\mathbf{x}_t(\mathbf{x}_0, \boldsymbol{\epsilon}), t) \right\|^2 \right] \\
&= \mathbb{E}_{\mathbf{x}_0, \boldsymbol{\epsilon}} \left[ \frac{1}{2\sigma_t^2} \left\| \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t(\mathbf{x}_0, \boldsymbol{\epsilon}) - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon} \right) - \boldsymbol{\mu}_\theta(\mathbf{x}_t(\mathbf{x}_0, \boldsymbol{\epsilon}), t) \right\|^2 \right]
\end{aligned}(24)

  1. 训练目标直接变成多噪声尺度的去噪自编码器 / score matching。 优化一个类似去噪 score matching 的目标,与通过变分推断去拟合一个具有 Langevin dynamics 形式的有限时间采样链边缘分布,本质上是等价的

\boldsymbol{\mu}_{\theta}(\mathbf{x}_t, t) = \tilde{\boldsymbol{\mu}}_t \left( \mathbf{x}_t, \frac{1}{\sqrt{\bar{\alpha}_t}}(\mathbf{x}_t - \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}_{\theta}(\mathbf{x}_t,t)) \right) = \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_{\theta}(\mathbf{x}_t, t) \right)(25)

\mathbb{E}_{\mathbf{x}_0, \boldsymbol{\epsilon}} \left[ \frac{\beta_t^2}{2\sigma_t^2 \alpha_t (1 - \bar{\alpha}_t)} \left\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta \left( \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}, t \right) \right\|^2 \right](26)

  1. 模型将反向过程建模为高斯转移p_\theta(x_{t-1}\mid x_t)=\mathcal N\!\bigl(x_{t-1};\mu_\theta(x_t,t),\sigma_t^2 I\bigr),因此一旦网络\epsilon_\theta(x_t,t) 能够从当前带噪样本x_t中预测出噪声\epsilon,就可以先由噪声预测写出反向均值\mu_\theta(x_t,t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t-\frac{\beta_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t,t) \right)。从p_\theta(x_{t-1}\mid x_t)采样就等价于

x \sim \mathcal{N}(\mu, \sigma^2 I) \iff x = \mu + \sigma z, z \sim \mathcal{N}(0, I).(27)

x_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_\theta(x_t, t) \right) + \sigma_t z, \quad z \sim \mathcal{N}(0, I).(28)

3️⃣ 最后一步L_0 对离散图像进行建模:假设图像数据由\{0,1,\ldots,255\}中的整数像素组成,并将其线性缩放到区间[-1,1]。这样做可以保证反向过程中的神经网络从标准高斯先验p(x_T)出发时,始终处理尺度一致的输入。为获得离散形式的对数似然,将反向过程的最后一项设为一个由高斯分布\mathcal N\!\bigl(x;\mu_\theta(x_1,1),\sigma_1^2 I\bigr)导出的独立离散解码器,即

  • 前面t>1的各步都在连续空间里做高斯转移。但真实数据x_0本质是离散像素格点,直接用连续高斯密度\mathcal N(x_0;\mu_\theta,\sigma_1^2I)去当作离散概率,会出现离散似然不对齐的问题
  • 把最后一步定义成离散化高斯解码器:对每个维度(每个像素通道)把连续高斯在对应 bin 上的概率质量积分出来,并且假设各维独立

p_\theta(x_0\mid x_1)=\prod_{i=1}^D \int_{\delta_-(x_0^i)}^{\delta_+(x_0^i)} \mathcal N\!\big(x;\mu_\theta^i(x_1,1),\sigma_1^2\big)\,dx(29)

\delta_+(x)=\begin{cases}
\infty & \text{if}\quad x=1\\[3pt]
x+\frac{1}{255} & \text{if}\quad x<1
\end{cases}
\qquad
\delta_-(x)=\begin{cases}
-\infty & \text{if}\quad x=-1\\[3pt]
x-\frac{1}{255} & \text{if}\quad x>-1
\end{cases}(30)

  • 保证变分下界对应于离散数据的无损编码长度,而无需额外向数据中加入噪声,也不需要在对数似然中显式引入缩放操作的 Jacobian 项。在采样结束时,直接无噪声地输出 \mu_\theta(x_1,1)

4️⃣ 从样本质量和实现简洁性来看,采用更简单的训练目标更为有效,相应的简化目标为

L_{\text{simple}}(\theta) := \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\epsilon}} \left[ \left\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_\theta \left( \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}, t \right) \right\|^2 \right] (31)

  • 其中t 在1,\dots,T间均匀采样。这个目标本质上就是:随机选一个时间步,把原图加噪得到x_t,再让网络直接预测所加入的噪声\varepsilon
  • L_{\text{simple}}可以理解为对原始变分目标的一种重加权形式。这种简化实际减弱小t时刻损失项的权重。由于小t 对应的是仅含极少噪声的去噪任务,相对较容易;降低其权重后,模型就能把更多能力放在大t 时更困难的强噪声去噪任务上,从而得到更好的样本生成质量
Algorithm 1 Training
Algorithm 2 Sampling

🦨 Denoising Diffusion Implicit Models

去噪扩散概率模型与噪声条件 score 网络表明,即使不依赖对抗训练,也能够生成与生成对抗网络相当质量的样本。这类方法的基本思想是在不同高斯噪声水平下进行去噪学习,使模型从受扰动样本中恢复数据结构;在采样阶段,则从白噪声出发,经由逐步去噪的马尔可夫链生成目标样本。该生成过程通常可建立在 Langevin dynamics 之上,或被解释为前向扩散过程的逆过程。但其采样过程通常依赖多步马尔可夫链的逐步反演,因此生成效率仍然受限。为缓解这一问题,论文提出去噪扩散隐式模型,并围绕以下两个方面展开改进

  • 以提升采样效率为核心目标,提出去噪扩散隐式模型这一更高效的迭代式隐式概率生成框架。该方法在保持原有训练流程不变的前提下,实现对传统扩散采样方式的有效加速,因而兼具方法上的兼容性与实践上的高效性
  • 从生成过程的建模方式入手,对去噪扩散概率模型进行推广。传统方法将生成过程定义为特定马尔可夫扩散过程的逆过程,而论文进一步提出一类非马尔可夫扩散过程,并证明其仍可导出与原方法一致的训练目标。进一步地,这类非马尔可夫过程在特定条件下可退化为确定性生成过程,从而使模型能够以更少的采样步数生成高质量样本

去噪扩散隐式模型可被视为一类隐式概率生成模型,并与去噪扩散概率模型保持紧密联系,其关键在于二者共享相同的训练目标

  • 将原本马尔可夫形式的前向扩散过程推广为一类非马尔可夫扩散过程,并证明由此得到的变分训练目标具有共同的代理目标,而该代理目标恰与去噪扩散概率模型的训练目标一致。因此,无需重新训练网络,仅通过选择不同的非马尔可夫扩散过程及其对应的逆向生成链,便可在统一训练框架下构造出一族不同的生成模型
  • 进一步地,部分非马尔可夫扩散过程能够导出更短的生成链,从而以更少的采样步数完成生成,并在仅带来较小质量损失的情况下显著提升采样效率。特别是在特定条件下,这类过程还可对应于确定性的生成机制,从而赋予去噪扩散隐式模型隐式生成模型的性质
  • 相较于去噪扩散概率模型,去噪扩散隐式模型在加速采样时仍能保持较高的样本质量,并在快速采样场景下表现出更优性能。此外,该模型还具有一致性特征,即从相同初始潜变量出发、采用不同长度生成链得到的样本,在高层语义上仍具有较强相似性,因此更适合实现具有语义意义的潜空间插值

🦫 Background

去噪扩散概率模型是一类潜变量模型,核心思想是在观测变量x_0之外引入一系列与其处于相同样本空间\mathcal{X}中的潜变量x_1,\ldots,x_T,通过对联合分布的建模来刻画数据生成过程

p_\theta(\boldsymbol{x}_0) = \int p_\theta(\boldsymbol{x}_{0:T}) \mathrm{d}\boldsymbol{x}_{1:T}, \quad \text{where} \quad p_\theta(\boldsymbol{x}_{0:T}) := p_\theta(\boldsymbol{x}_T) \prod_{t=1}^T p_\theta^{(t)}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t)(1)

  • 模型参数\theta通过最大化数据对数似然的变分下界进行学习,即\max_{\theta} \mathbb{E}_{q(\boldsymbol{x}_0, \boldsymbol{x}_1, \dots, \boldsymbol{x}_T)} [\log p_\theta(\boldsymbol{x}_{0:T}) - \log q(\boldsymbol{x}_{1:T} | \boldsymbol{x}_0)]。其中q(x_{1:T}\mid x_0)表示潜变量的推断分布。不同于变分自编码器等典型潜变量模型,去噪扩散概率模型中的推断过程通常是预先设定且固定的,而非可学习的;同时,其潜变量通常具有较高维度

Ho 等人将潜变量的推断过程建模为一类由递减序列\alpha_{1:T}\in(0,1]^T参数化的高斯马尔可夫链

q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0) := \prod_{t=1}^T q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1}), \text{ where } q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1}) := \mathcal{N} \left( \sqrt{\frac{\alpha_t}{\alpha_{t-1}}} \boldsymbol{x}_{t-1}, \left( 1 - \frac{\alpha_t}{\alpha_{t-1}} \right) \boldsymbol{I} \right)(2)

  • 由于该过程以自回归方式将样本从x_0逐步演化至x_T,因此被称为前向过程。直观而言,前向过程不断向观测样本 x_0中注入噪声,使其逐步退化为近似纯噪声的状态
  • 与之对应,潜变量模型p_\theta(x_{0:T})则描述一个从x_T逐步生成至x_0的马尔可夫链,其目标在于逼近难以直接求解的逆条件分布q(x_{t-1}\mid x_t),因此被称为生成过程。从直观上看,生成过程正是对受噪样本逐步执行去噪,以恢复原始数据结构

前向扩散过程的一项关键性质是,给定x_0 时,x_t的条件分布具有闭式高斯形式,因此可表示为x_0与标准高斯噪声的线性组合。该性质揭示前向过程逐步加噪的本质,也简化训练与采样推导。当\alpha_T足够接近于零时,q(x_T\mid x_0)收敛到标准高斯分布,故通常设p_\theta(x_T)=\mathcal{N}(0,I)

q(x_t \mid x_0) = \mathcal{N}(x_t; \sqrt{\alpha_t} x_0, (1 - \alpha_t)I),(3)

x_t = \sqrt{\alpha_t} x_0 + \sqrt{1 - \alpha_t} \epsilon, \quad \epsilon \sim \mathcal{N}(0, I).(4)

当各步条件分布均被建模为具有可学习均值、固定方差的高斯分布时,变分目标可以进一步化简为加权去噪回归目标

L_\gamma(\theta) := \sum_{t=1}^T \gamma_t \mathbb{E}_{x_0 \sim q(x_0), \epsilon \sim \mathcal{N}(0,I)} \left[ \left\| \epsilon_\theta^{(t)} \left( \sqrt{\alpha_t} x_0 + \sqrt{1 - \alpha_t} \epsilon \right) - \epsilon \right\|_2^2 \right].(5)

  • 其中\theta=\{\theta^{(t)}\}_{t=1}^{T}表示一组按时间步索引的可训练函数,\epsilon_{\theta}^{(t)}: \mathcal{X} \to \mathcal{X}用于预测从原始样本构造时刻t的受扰动样本所用的累计噪声,\gamma=[\gamma_1,\ldots,\gamma_T]为由噪声调度\alpha_{1:T}决定的正权重系数。由此,模型训练可等价地理解为:在不同噪声水平下,让网络学习从受扰动样本中恢复对应的高斯噪声

🌱 Variational Inference For Non-Markovian Forward Processes

生成过程近似推断过程的逆过程,因此,减少采样迭代次数需要重新考虑推断过程的构造。DDPM 的代理目标L_\gamma仅依赖边缘分布q(x_t\mid x_0),而不直接依赖联合分布q(x_{1:T}\mid x_0)。具有相同边缘分布的联合推断分布并不唯一,因此可以将马尔可夫推断过程推广为非马尔可夫形式,并沿用同一代理训练目标

考虑由实向量\sigma \in \mathbb{R}_{\ge 0}^{T}索引的一族推断分布\mathcal{Q},定义为

\begin{aligned}
&q_\sigma(\boldsymbol{x}_{1:T}\mid\boldsymbol{x}_0)
:=q_\sigma(\boldsymbol{x}_T\mid\boldsymbol{x}_0)
\prod_{t=2}^{T}q_\sigma(\boldsymbol{x}_{t-1}\mid\boldsymbol{x}_t,\boldsymbol{x}_0),\\[5pt]
&\text{where}\quad q_\sigma(\boldsymbol{x}_T\mid\boldsymbol{x}_0)
=\mathcal{N}\!\left(\sqrt{\alpha_T}\boldsymbol{x}_0,(1-\alpha_T)\boldsymbol{I}\right),\\[5pt]
&q_\sigma(\boldsymbol{x}_{t-1}\mid\boldsymbol{x}_t,\boldsymbol{x}_0)
=\mathcal{N}\!\left(
\sqrt{\alpha_{t-1}}\boldsymbol{x}_0+
\sqrt{1-\alpha_{t-1}-\sigma_t^2}\,
\frac{\boldsymbol{x}_t-\sqrt{\alpha_t}\boldsymbol{x}_0}{\sqrt{1-\alpha_t}},
\sigma_t^2\boldsymbol{I}\right),\quad t>1
\end{aligned}(6)

均值函数的选择保证,对任意时间步t,都有 q_\sigma(\boldsymbol{x}_t\mid\boldsymbol{x}_0)=\mathcal{N}\!\left(\sqrt{\alpha_t}\boldsymbol{x}_0,(1-\alpha_t)\boldsymbol{I}\right),从而使联合推断分布具有所需的边缘分布。根据 Bayes 公式,前向过程为

q_\sigma(\boldsymbol{x}_t\mid\boldsymbol{x}_{t-1},\boldsymbol{x}_0)=\frac{q_\sigma(\boldsymbol{x}_{t-1}\mid\boldsymbol{x}_t,\boldsymbol{x}_0)\,q_\sigma(\boldsymbol{x}_t\mid\boldsymbol{x}_0)}{q_\sigma(\boldsymbol{x}_{t-1}\mid\boldsymbol{x}_0)}(7)

  • 该条件分布也是高斯分布。与 DDPM 的扩散过程不同,每个时刻的状态可以同时依赖x_{t-1}与x_0,因而前向过程不再具有马尔可夫性。 \sigma的大小控制前向过程的随机程度;当\sigma \to 0时,只要给定 x_0 和某一时刻的 x_t,x_{t-1} 就随之确定

接下来构造可训练的生成过程p_{\theta}(x_{0:T}),使每一步转移分布p_{\theta}^{(t)}(x_{t-1}\mid x_t) 利用已知的逆条件分布q_{\sigma}(x_{t-1}\mid x_t,x_0)。给定x_0 \sim q(x_0)和高斯噪声\epsilon \sim \mathcal{N}(0,I),时刻t的噪声样本可表示为x_t=\sqrt{\alpha_t}\,x_0+\sqrt{1-\alpha_t}\,\epsilon.模型\epsilon_\theta^{(t)}(x_t)在未知x_0 的情况下,仅根据x_t 预测噪声\epsilon 。重排上述关系,即得原始样本x_0的去噪估计

f_\theta^{(t)}(x_t) = \frac{x_t-\sqrt{1-\alpha_t}\,\epsilon_\theta^{(t)}(x_t)}{\sqrt{\alpha_t}}.(8)

将逆条件分布中的x_0替换为上述估计,即可构造反向生成过程。从标准高斯先验采样x_T后,每一步将 f_\theta^{(t)}(x_t)代入逆条件分布,得到从x_t到x_{t-1}的转移核。在t=1时加入协方差为\sigma_1^2 I的非退化高斯噪声,以保证生成分布具有全空间支撑。逐步转移分布为

p_\theta^{(t)}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) = \begin{cases} \mathcal{N}(f_\theta^{(1)}(\boldsymbol{x}_1), \sigma_1^2 \boldsymbol{I}) & \text{if}\quad t = 1 \\ q_\sigma(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, f_\theta^{(t)}(\boldsymbol{x}_t)) & \text{otherwise,} \end{cases}(9)

模型参数通过最小化以下变分目标进行学习

\begin{aligned}
J_\sigma(\epsilon_\theta) &:= \mathbb{E}_{\boldsymbol{x}_{0:T} \sim q_\sigma(\boldsymbol{x}_{0:T})} [\log q_\sigma(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0) - \log p_\theta(\boldsymbol{x}_{0:T})] \\
&= \mathbb{E}_{\boldsymbol{x}_{0:T} \sim q_\sigma(\boldsymbol{x}_{0:T})} \left[ \log q_\sigma(\boldsymbol{x}_T|\boldsymbol{x}_0) + \sum_{t=2}^T \log q_\sigma(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0) - \sum_{t=1}^T \log p_\theta^{(t)}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) - \log p_\theta(\boldsymbol{x}_T) \right]
\end{aligned}(10)

Theorem 1. 对任意\sigma>0,都存在权重向量\gamma\in\mathbb{R}_{>0}^{T}及与模型参数无关的常数C\in\mathbb{R},使得J_\sigma=L_\gamma+C

  • 若模型\epsilon_\theta^{(t)}在不同时间步t上不共享参数,则\epsilon_\theta的最优解不依赖于正权重\gamma,因为各时间步的损失可以独立优化,而正权重不会改变各项的最优解。这一性质既支持 DDPM 使用L_1 作为代理目标,也表明它可用于优化J_\sigma

🦋 Sampling From Generalized Generative Processes

以 为目标训练的模型,不仅适用于 DDPM 的马尔可夫推断过程,也适用于由 参数化的非马尔可夫前向过程。因此,可以直接使用预训练的 DDPM 模型,通过改变 选择符合采样需求的生成过程

Denoising Diffusion Implicit Models: 由前文定义的生成分布 ,对 ,可从 生成

(11)
  • 其中, 是与 独立的标准高斯噪声,并约定 ; 时使用前文定义的末端分布。不同的 对应不同的生成过程,但都使用同一个噪声预测模型 ,无需重新训练。当各步取 时,前向过程为马尔可夫过程,生成过程对应 DDPM

  • 当所有 时,除 外,前向过程在给定 和 后是确定性的;生成过程中的随机噪声项也随之消失。此时,模型通过固定映射将潜变量 转换为样本 ,形成隐式概率模型,称为去噪扩散隐式模型(DDIM)。它沿用 DDPM 的训练目标,而对应的前向过程不再是扩散过程

Accelerated Generation Processes: 上述前向过程包含 步,生成过程也相应需要执行 步采样。但只要边缘分布 保持不变,去噪目标 就不依赖具体的前向过程,因此可以构造长度小于 的前向过程,无需重新训练即可缩短生成链

  • 将前向过程定义在潜变量子集 上,其中 是 中长度为 的严格递增子序列。在这些变量上构造顺序前向过程,并保持边缘分布 。生成过程按 的顺序采样,该序列称为采样轨迹;当其长度远小于 时,所需的迭代次数显著减少

采样轨迹为 τ = [1, 3] 的加速生成图模型

  • 与前文相同的论证表明,仍可使用以 训练的模型,只需将采样更新中的相邻时间步改为所选轨迹中的相邻时间步。这一构造同时适用于 DDPM、DDIM 及上述其他生成过程,使训练步数与实际采样步数可以分别选择

Relevance To Neural ODEs: 将 DDIM 的确定性更新式重写为如下形式,可以看出它与求解常微分方程的 Euler 方法之间的关系

(12)
  • 重新参数化为 与 ,其中 是满足 的连续递增函数。这里的 表示重参数化后的连续噪声尺度,与前文控制采样随机性的 不同。上述更新可视为以下 ODE 的 Euler 离散化
(13)
  • 当离散步数足够多时,也可以沿相反方向从 积分至 ,将观测样本 编码为潜变量 ,从而获得与生成过程对应的潜在表示

在噪声预测模型达到最优时,上述 ODE 与 Variance-Exploding SDE 对应的 probability flow ODE 等价,但二者采用的离散采样更新不同。对后者按时间 使用 Euler 方法,得到

(14)
  • 当 与 足够接近时,两种更新近似一致;在较少采样步数下,离散化方式会带来差异

☃️ Proofs

Lemma 1. 设前向过程为关于x_0的非马尔可夫联合分布q_{\sigma}(x_{1:T}\mid x_0) = q_{\sigma}(x_T\mid x_0)\prod_{t=2}^{T} q_{\sigma}(x_{t-1}\mid x_t,x_0),其中

q_{\sigma}(x_{t-1}\mid x_t,x_0) = \mathcal N\!\left( \sqrt{\alpha_{t-1}}x_0 + \sqrt{1-\alpha_{t-1}-\sigma_t^2}\, \frac{x_t-\sqrt{\alpha_t}x_0}{\sqrt{1-\alpha_t}}, \, \sigma_t^2 I \right).(15)

在该构造下,任意时刻t的边缘分布仍满足q_{\sigma}(x_t\mid x_0) = \mathcal N\!\bigl(\sqrt{\alpha_t}x_0,\,(1-\alpha_t)I\bigr).因此,DDIM 改变的是前向演化路径的分解方式,而不是各时刻的边缘分布,这也保证其能够与 DDPM 共享相同的训练目标

  • 证明:对某个2\le t\le T,有q_{\sigma}(x_t\mid x_0)=\mathcal N\!\left(\sqrt{\alpha_t}x_0,(1-\alpha_t)I\right),若能进一步推出q_{\sigma}(x_{t-1}\mid x_0)=\mathcal N\!\left(\sqrt{\alpha_{t-1}}x_0,(1-\alpha_{t-1})I\right), 则由于基例t=T已经成立,因此可对t从T到 1 使用归纳法证明该结论
  • 首先,有q_\sigma(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0) := \int_{\boldsymbol{x}_t} q_\sigma(\boldsymbol{x}_t|\boldsymbol{x}_0)q_\sigma(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0)\mathrm{d}\boldsymbol{x}_t, 并且

q_\sigma(\boldsymbol{x}_t|\boldsymbol{x}_0) = \mathcal{N}(\sqrt{\alpha_t}\boldsymbol{x}_0, (1 - \alpha_t)\boldsymbol{I})(16)

q_\sigma(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0) = \mathcal{N}\left(\sqrt{\alpha_{t-1}}\boldsymbol{x}_0 + \sqrt{1 - \alpha_{t-1} - \sigma_t^2} \cdot \frac{\boldsymbol{x}_t - \sqrt{\alpha_t}\boldsymbol{x}_0}{\sqrt{1 - \alpha_t}}, \sigma_t^2 \boldsymbol{I}\right).​(17)

  • 由 Bishop (2006) 式 (2.115) 可知,q_\sigma(x_{t-1}\mid x_0)为高斯分布

\mathbb{E}[\boldsymbol{x}_{t-1}\mid\boldsymbol{x}_0] = \sqrt{\alpha_{t-1}}\boldsymbol{x}_0 + \sqrt{1 - \alpha_{t-1} - \sigma_t^2} \cdot \frac{\sqrt{\alpha_t}\boldsymbol{x}_0 - \sqrt{\alpha_t}\boldsymbol{x}_0}{\sqrt{1 - \alpha_t}} = \sqrt{\alpha_{t-1}}\boldsymbol{x}_0(18)

\mathrm{Cov}[\boldsymbol{x}_{t-1}\mid\boldsymbol{x}_0] = \sigma_t^2 \boldsymbol{I} + \frac{1 - \alpha_{t-1} - \sigma_t^2}{1 - \alpha_t}(1 - \alpha_t)\boldsymbol{I} = (1 - \alpha_{t-1})\boldsymbol{I}(19)

  • 因此,q_{\sigma}(x_{t-1}\mid x_0) = \mathcal N\!\left(\sqrt{\alpha_{t-1}}x_0,(1-\alpha_{t-1})I\right),从而归纳成立

Theorem 1. 对于使上述高斯转移有定义的任意\sigma \in \mathbb{R}_{>0}^{T},都存在\gamma \in \mathbb{R}_{>0}^{T}和常数C\in\mathbb{R},使得J_{\sigma}=L_{\gamma}+C

  • 证明:由J_{\sigma}的定义可得

\begin{aligned}
J_\sigma(\epsilon_\theta)
&:=\mathbb E_{q_\sigma(x_{0:T})}\Biggl[
 \log q_\sigma(x_T\mid x_0)-\log p(x_T)+\sum_{t=2}^T\log q_\sigma(x_{t-1}\mid x_t,x_0)
 -\sum_{t=1}^T\log p_\theta^{(t)}(x_{t-1}\mid x_t)\Biggr]\\
&\equiv\mathbb E_{q_\sigma(x_{0:T})}\Biggl[
 \sum_{t=2}^T D_{\mathrm{KL}}\!\left(
 q_\sigma(x_{t-1}\mid x_t,x_0)\,\|\,p_\theta^{(t)}(x_{t-1}\mid x_t)
 \right)-\log p_\theta^{(1)}(x_0\mid x_1)\Biggr].
\end{aligned}(20)

  • 对于t>1,有

\begin{aligned}
\kappa_t&:=\sqrt{\alpha_{t-1}}-\sqrt{\alpha_t}\,
 \sqrt{\frac{1-\alpha_{t-1}-\sigma_t^2}{1-\alpha_t}},\qquad t>1,\\
&\mathbb E_{x_0,x_t}D_{\mathrm{KL}}\!\left(
 q_\sigma(x_{t-1}\mid x_t,x_0)\,\|\,p_\theta^{(t)}(x_{t-1}\mid x_t)\right) =\frac{\kappa_t^2}{2\sigma_t^2}
 \mathbb E_{x_0,x_t}\left\|x_0-f_\theta^{(t)}(x_t)\right\|_2^2\\
&=\frac{\kappa_t^2(1-\alpha_t)}{2\sigma_t^2\alpha_t}
 \mathbb E_{x_0,\epsilon}\left\|\epsilon-\epsilon_\theta^{(t)}(x_t)\right\|_2^2,
 \quad x_t=\sqrt{\alpha_t}x_0+\sqrt{1-\alpha_t}\epsilon.
\end{aligned}(21)

  • 对于末端项x_0,有

\begin{aligned}
&\mathbb E_{x_0,x_1}\!\left[-\log p_\theta^{(1)}(x_0\mid x_1)\right] \equiv\frac{1}{2\sigma_1^2}\mathbb E_{x_0,x_1}
 \left\|x_0-f_\theta^{(1)}(x_1)\right\|_2^2\\
&=\frac{1-\alpha_1}{2\sigma_1^2\alpha_1}
 \mathbb E_{x_0,\epsilon}\left\|\epsilon-\epsilon_\theta^{(1)}(x_1)\right\|_2^2,
 \quad x_1=\sqrt{\alpha_1}x_0+\sqrt{1-\alpha_1}\epsilon.
\end{aligned}(22)

(23)

🦜 Score-Based Generative Modeling through Stochastic Differential Equations

从数据中构造噪声是容易的,而从噪声中恢复数据则构成生成建模问题

  • 论文提出一种随机微分方程(SDE)框架:其前向过程通过逐步注入噪声,将复杂数据分布平滑地转化为已知先验分布;相应地,逆向过程则通过逐步去噪,将先验分布重新映射回数据分布
  • 关键在于,该逆时间 SDE 仅依赖于随时间变化的梯度场,即受扰动数据分布的得分函数(score)。借助基于得分生成建模领域的最新进展,​利用神经网络对这些得分函数进行精确估计,并结合数值 SDE 求解器实现样本生成

已有两类成功的概率生成模型都遵循“逐步加噪—再逆向恢复”的基本思想,将这两类方法统一归为基于得分的生成模型

  • 基于 Langevin 动力学的得分匹配方法(SMLD)在每个噪声尺度上估计得分函数,并在生成阶段通过 Langevin 动力学沿着逐步减小的噪声尺度进行采样
  • 去噪扩散概率模型(DDPM)则学习一系列概率模型,以逐步逆转噪声破坏过程。对于连续状态空间而言,DDPM 的训练目标本质上也在各个噪声尺度上隐式估计得分函数

在此基础上,进一步提出一个以 SDE 为核心的统一理论框架,用以概括并推广上述方法

  • 与基于有限个离散噪声分布的建模方式不同,考虑的是一个随时间连续演化的分布族,其演化由预先设定、与具体数据点无关且不含可训练参数的前向 SDE 描述。通过逆转这一扩散过程,随机噪声可以被逐步塑造为数据样本
  • 更重要的是,只要已知各时刻边缘分布的得分函数,就能够由前向 SDE 推导出对应的逆时间 SDE。基于这一性质,通过训练一个时间依赖的神经网络来估计得分函数,并利用数值 SDE 求解器完成采样
  • 提出预测器—校正器(predictor-corrector)框架,以修正离散化逆时间 SDE 所带来的数值误差;同时推导出一个与 SDE 采样分布等价的神经常微分方程(ODE),从而支持精确似然计算并进一步提升采样效率

🐦‍🔥 Background

SMLD: 设p_{\sigma}(\tilde{\mathbf{x}}\mid \mathbf{x}) := \mathcal{N}(\tilde{\mathbf{x}};\mathbf{x},\sigma^2 \mathbf{I})为高斯扰动核,对应的扰动分布为p_{\sigma}(\tilde{\mathbf{x}}):=\int p_{\text{data}}(\mathbf{x})\,p_{\sigma}(\tilde{\mathbf{x}}\mid \mathbf{x})\,d\mathbf{x},其中p_{\text{data}}(\mathbf{x})表示数据分布。取递增的正噪声尺度序列\sigma_{\min}=\sigma_1<\sigma_2<\cdots<\sigma_N=\sigma_{\max}.选择足够小的⁡\sigma_{\min},使p_{\sigma_{\min}}(\mathbf{x}) \approx p_{\text{data}}(\mathbf{x}),选择足够大的⁡\sigma_{\max},使p_{\sigma_{\max}}(\mathbf{x}) \approx \mathcal{N}(\mathbf{x};\mathbf{0},\sigma_{\max}^2\mathbf{I})

  • 训练噪声条件得分网络 s_{\theta}(\mathbf{x},\sigma)时,对不同噪声尺度的去噪得分匹配损失加权求和

\theta^{*} = \arg\min_{\theta} \sum_{i=1}^{N} \sigma_i^2 \mathbb{E}_{p_{\text{data}}(\mathbf{x})} \mathbb{E}_{p_{\sigma_i}(\tilde{\mathbf{x}}\mid \mathbf{x})} \left[ \left\| s_{\theta}(\tilde{\mathbf{x}},\sigma_i) - \nabla_{\tilde{\mathbf{x}}}\log p_{\sigma_i}(\tilde{\mathbf{x}}\mid \mathbf{x}) \right\|_2^2 \right].(1)

  • 在数据充分、模型容量足够大的条件下,最优模型s_{\theta^*}(\mathbf{x},\sigma)几乎处处等于\nabla_{\mathbf{x}}\log p_{\sigma}(\mathbf{x}),其中\sigma \in \{\sigma_i\}_{i=1}^{N}。采样时,按噪声尺度从大到小,针对每个p_{\sigma_i}(\mathbf{x})运行 M 步 Langevin MCMC

\mathbf{x}_i^m = \mathbf{x}_i^{m-1} + \epsilon_i\, s_{\theta^*}(\mathbf{x}_i^{m-1},\sigma_i) + \sqrt{2\epsilon_i}\,\mathbf{z}_i^m, \qquad m=1,2,\cdots,M(2)

  • \epsilon_i>0为步长,\mathbf{z}_i^m为标准高斯噪声。按i=N,N-1,\cdots,1采样,取\mathbf{x}_N^0 \sim \mathcal{N}(\mathbf{x}\mid \mathbf{0},\sigma_{\max}^2\mathbf{I}),当i<N时,有\mathbf{x}_i^0=\mathbf{x}_{i+1}^M

DDPM: 选取正噪声尺度0<\beta_1,\beta_2,\cdots,\beta_N<1.对数据样本\mathbf{x}_0\sim p_{\text{data}}(\mathbf{x})构造离散马尔可夫链\{\mathbf{x}_0,\mathbf{x}_1,\cdots,\mathbf{x}_N\},其转移核为p(\mathbf{x}_i\mid \mathbf{x}_{i-1}) = \mathcal{N}\!\left( \mathbf{x}_i;\sqrt{1-\beta_i}\mathbf{x}_{i-1},\beta_i\mathbf{I} \right),由此可得p_{\alpha_i}(\mathbf{x}_i\mid \mathbf{x}_0) = \mathcal{N}\!\left( \mathbf{x}_i;\sqrt{\alpha_i}\mathbf{x}_0,(1-\alpha_i)\mathbf{I} \right),其中\alpha_i:=\prod_{j=1}^{i}(1-\beta_j).p_{\alpha_i}(\tilde{\mathbf{x}}) := \int p_{\text{data}}(\mathbf{x})\,p_{\alpha_i}(\tilde{\mathbf{x}}\mid \mathbf{x})\,d\mathbf{x}.

  • 预设噪声尺度,使终态\mathbf{x}_N近似服从标准高斯分布\mathcal{N}(\mathbf{0},\mathbf{I}),反向转移分布与 ELBO 加权训练目标分别为

p_{\theta}(\mathbf{x}_{i-1}\mid \mathbf{x}_i) = \mathcal{N}\!\left( \mathbf{x}_{i-1}; \frac{1}{\sqrt{1-\beta_i}}\bigl(\mathbf{x}_i+\beta_i s_{\theta}(\mathbf{x}_i,i)\bigr), \beta_i\mathbf{I} \right),(3)

\theta^* = \arg\min_{\theta} \sum_{i=1}^{N} (1-\alpha_i)\, \mathbb{E}_{p_{\text{data}}(\mathbf{x})} \mathbb{E}_{p_{\alpha_i}(\tilde{\mathbf{x}}\mid \mathbf{x})} \left[ \left\| s_{\theta}(\tilde{\mathbf{x}},i) - \nabla_{\tilde{\mathbf{x}}}\log p_{\alpha_i}(\tilde{\mathbf{x}}\mid \mathbf{x}) \right\|_2^2 \right](4)

  • 得到最优模型s_{\theta^*}(\mathbf{x},i)后,从\mathbf{x}_N\sim \mathcal{N}(\mathbf{0},\mathbf{I})出发,沿逆向马尔可夫链生成样本

\mathbf{x}_{i-1} = \frac{1}{\sqrt{1-\beta_i}} \bigl(\mathbf{x}_i+\beta_i s_{\theta^*}(\mathbf{x}_i,i)\bigr) +\sqrt{\beta_i}\mathbf{z}_i, \qquad i=N,N-1,\cdots,1(5)

  • 该过程称为祖先采样(ancestral sampling),等价于从图模型\prod_{i=1}^{N}p_{\theta}(\mathbf{x}_{i-1}\mid \mathbf{x}_i)中采样。训练目标同样是加权去噪得分匹配,因此最优模型s_{\theta^*}(\tilde{\mathbf{x}},i)估计扰动数据分布的得分函数\nabla_{\mathbf{x}}\log p_{\alpha_i}(\mathbf{x})。两种方法的权重均与相应扰动核的 score 均方范数成反比

\sigma_i^2 \propto \frac{1}{\mathbb{E} \left[ \| \nabla_{\tilde{\mathbf{x}}} \log p_{\sigma_i}(\tilde{\mathbf{x}} \mid \mathbf{x}) \|_2^2 \right]}, \quad (1 - \alpha_i) \propto \frac{1}{\mathbb{E} \left[ \| \nabla_{\tilde{\mathbf{x}}} \log p_{\alpha_i}(\tilde{\mathbf{x}} \mid \mathbf{x}) \|_2^2 \right]}.(6)

🦅 Score-Based Generative Modeling With SDEs

前述方法之所以能够取得良好生成效果,一个关键原因在于它们在多个噪声尺度下对数据进行扰动建模。基于这一观察,论文进一步将这一思想推广到连续情形,即引入无穷多个噪声尺度,使受扰动数据分布在噪声逐渐增强的过程中按照随机微分方程(SDE)连续演化

论文的目标是构造一个由连续时间变量t\in[0,T]索引的扩散过程x(t),使其在初始时刻满足x(0)\sim p_0,其中p_0为数据分布;在终止时刻满足x(T)\sim p_T,其中p_T为易于采样的先验分布。换言之,该过程旨在将数据分布连续地扩散到一个可处理的先验分布。为此,论文使用 Itô 随机微分方程对该扩散过程进行建模

dx=f(x,t)\,dt+g(t)\,dw,(7)

  • 其中,w为标准 Wiener 过程,f(x,t)为漂移项,g(t)为扩散系数。假设扩散系数为与状态x无关的标量函数,但这一设定可以自然推广到更一般的情形。在适当正则条件下,该 SDE 存在唯一强解。进一步地,论文记p_t(x)为时刻t的状态分布密度,并用p_{st}(x(t)\mid x(s))表示从时刻s到时刻t的转移核。在具体建模中,终态分布p_T通常被设定为不包含数据结构信息、但便于高效采样的简单先验分布,例如均值和方差固定的高斯分布。围绕这一目标,可以设计不同形式的 SDE,使数据分布在连续扩散过程中逐步演化为给定先验分布

Anderson(1982)的一个重要结论指出:扩散过程的时间反演仍然是一个扩散过程,其动力学由如下逆时间随机微分方程给出

d\mathbf{x} = \bigl[ f(\mathbf{x},t)-g(t)^2\nabla_{\mathbf{x}}\log p_t(\mathbf{x}) \bigr]dt + g(t)\,d\bar{\mathbf{w}}, (8)

  • 其中,\bar{\mathbf{w}}表示当时间从T向0反向流动时对应的标准 Wiener 过程,dt则表示一个无穷小的负时间步长。只要能够获得任意时刻边缘分布p_t(\mathbf{x})的得分函数\nabla_{\mathbf{x}}\log p_t(\mathbf{x}),便可以构造出对应的逆扩散过程,并通过数值模拟该逆过程,从先验分布p_T逐步生成服从数据分布p_0的样本

分布的得分函数可以通过在样本上训练基于得分的模型,并采用 score matching 方法进行估计。为估计\nabla_{\mathbf{x}}\log p_t(\mathbf{x}),论文训练一个时间相关的得分模型s_{\theta}(\mathbf{x},t),其目标函数可以看作对 SMLD 和 DDPM 的连续时间推广

\theta^* = \arg\min_{\theta} \mathbb{E}_{t} \left\{ \lambda(t)\, \mathbb{E}_{\mathbf{x}(0)} \mathbb{E}_{\mathbf{x}(t)\mid \mathbf{x}(0)} \left[ \left\| s_{\theta}(\mathbf{x}(t),t) - \nabla_{\mathbf{x}(t)} \log p_{0t}(\mathbf{x}(t)\mid \mathbf{x}(0)) \right\|_2^2 \right] \right\}. (9)

  • 其中,\lambda:[0,T]\to \mathbb{R}_{>0}为正的加权函数,t在区间[0,T]上均匀采样,\mathbf{x}(0)\sim p_0(\mathbf{x}),且\mathbf{x}(t)\sim p_{0t}(\mathbf{x}(t)\mid \mathbf{x}(0))。在数据充足且模型容量足够大的条件下,score matching 保证最优解s_{\theta^*}(\mathbf{x},t)在几乎所有的\mathbf{x}与t上都等于\nabla_{\mathbf{x}}\log p_t(\mathbf{x}).与 SMLD 和 DDPM 类似,通常可以取

\lambda(t)\propto \frac{1}{\mathbb{E}\!\left[\left\|\nabla_{\mathbf{x}(t)}\log p_{0t}(\mathbf{x}(t)\mid \mathbf{x}(0))\right\|_2^2\right]}.(10)

  • 为高效求解,通常需要已知转移核p_{0t}(\mathbf{x}(t)\mid \mathbf{x}(0)).当f(\cdot,t)为仿射函数时,该转移核总是高斯分布,其均值和方差通常具有闭式表达,因此可以利用标准方法直接求得。对于更一般的 SDE,论文指出,也可以通过求解 Kolmogorov 前向方程来获得p_{0t}(\mathbf{x}(t)\mid \mathbf{x}(0))。另一种可行方案是直接模拟 SDE,从条件分布p_{0t}(\mathbf{x}(t)\mid \mathbf{x}(0))中采样,并在训练时用 sliced score matching 替代去噪 score matching,从而绕开对\nabla_{\mathbf{x}(t)}\log p_{0t}(\mathbf{x}(t)\mid \mathbf{x}(0))的显式计算

🦊 Examples: VE, VP SDEs And Beyond

SMLD 与 DDPM 中使用的噪声扰动,都可以看作是两类不同随机微分方程(SDE)的离散化形式。对于 SMLD,当采用N个噪声尺度时,每个扰动核p_{\sigma_i}(\mathbf{x}\mid \mathbf{x}_0)都对应于如下马尔可夫链中变量\mathbf{x}_i的分布

\mathbf{x}_i = \mathbf{x}_{i-1} + \sqrt{\sigma_i^2-\sigma_{i-1}^2}\,\mathbf{z}_{i-1}, \qquad i=1,\cdots,N, (11)

  • 其中\mathbf{z}_{i-1} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})。当N\to\infty时,令t_i=i/N,\quad\Delta t=1/N,\quad\sigma_i=\sigma(t_i),以\sqrt{\Delta t}\,\mathbf z_{i-1}表示布朗增量\Delta\mathbf w_i,在通常的正则条件下,离散马尔可夫链\{\mathbf{x}_i\}_{i=1}^{N}收敛为连续随机过程\{x(t)\}_{t=0}^{1},其对应的 SDE 为

d\mathbf{x} = \sqrt{\frac{d[\sigma^2(t)]}{dt}}\,d\mathbf{w}.(12)

类似地,对于 DDPM 的扰动核\{p_{\alpha_i}(\mathbf{x}\mid \mathbf{x}_0)\}_{i=1}^{N},其离散马尔可夫链写为

\mathbf{x}_i = \sqrt{1-\beta_i}\,\mathbf{x}_{i-1} + \sqrt{\beta_i}\,\mathbf{z}_{i-1}, \qquad i=1,\cdots,N.(13)

  • 令N\to\infty,\quad\beta_i=\beta(t_i)/N,则在通常的正则条件下收敛到如下 SDE

d\mathbf{x} = -\frac{1}{2}\beta(t)\mathbf{x}\,dt + \sqrt{\beta(t)}\,d\mathbf{w}. (14)

因此,SMLD 与 DDPM 所采用的噪声扰动过程,本质上分别对应于两类 SDE 的离散近似。进一步地,论文指出,SMLD 所对应的过程在t\to\infty时方差会不断增大,因此称其为方差爆炸型 SDE(Variance Exploding, VE SDE);而 DDPM 在初始分布具有单位方差时,会诱导一个始终保持固定方差的过程,因此称其为方差保持型 SDE(Variance Preserving, VP SDE)。受 VP SDE 的启发,论文进一步提出一类在似然建模中表现更优的新型 SDE

d\mathbf{x} = -\frac{1}{2}\beta(t)\mathbf{x}\,dt + \sqrt{\beta(t)\Bigl(1-e^{-2\int_0^t\beta(s)\,ds}\Bigr)}\,d\mathbf{w}. (15)

在使用相同\beta(t)且从相同初始分布出发的条件下,所诱导随机过程的方差在任意中间时刻都被 VP SDE 对应的方差上界所约束。基于这一性质,论文称为 sub-VP SDE。最后,由于 VE、VP 与 sub-VP 三类 SDE 的漂移项都属于仿射形式,其扰动核p_{0t}(x(t)\mid x(0))均可写成高斯分布,并能够以闭式形式求得,这使得基于连续时间 score matching 训练可以高效实现

p_{0t}(\mathbf{x}(t) \mid \mathbf{x}(0)) = 
\begin{cases} 
\mathcal{N}(\mathbf{x}(t); \mathbf{x}(0), [\sigma^2(t) - \sigma^2(0)]\mathbf{I}), & \text{(VE SDE)} \\ 
\mathcal{N}(\mathbf{x}(t); \mathbf{x}(0)e^{-\frac{1}{2}\int_0^t \beta(s)\mathrm{d}s}, \mathbf{I} - \mathbf{I}e^{-\int_0^t \beta(s)\mathrm{d}s}) & \text{(VP SDE)} \\ 
\mathcal{N}(\mathbf{x}(t); \mathbf{x}(0)e^{-\frac{1}{2}\int_0^t \beta(s)\mathrm{d}s}, [1 - e^{-\int_0^t \beta(s)\mathrm{d}s}]^2\mathbf{I}) & \text{(sub-VP SDE)} 
\end{cases}.(16)

SMLD 与 DDPM 中使用的噪声扰动过程,分别对应于 VE SDE 与 VP SDE 的离散化形式。VE SDE 在t\to\infty时会诱导一个方差持续增大的随机过程;与之相对,VP SDE 则对应一个方差始终有界的过程。进一步地,当初始分布p(x_0)具有单位方差时,VP SDE 在所有t\in[0,\infty)上都保持单位方差不变。由于 VP SDE 的漂移项与扩散项均为仿射形式,论文借助 Särkkä 与 Solin 中的结果,得到其协方差随时间演化所满足的常微分方程

\frac{d\Sigma_{\mathrm{VP}}(t)}{dt} = \beta(t)\bigl(I-\Sigma_{\mathrm{VP}}(t)\bigr),(17)

  • 其中,\Sigma_{\mathrm{VP}}(t):=\mathrm{Cov}[x(t)]表示满足 VP SDE 的随机过程\{x(t)\}_{t\ge 0}在时刻t的协方差矩阵。对该方程求解可得

\Sigma_{\mathrm{VP}}(t) = I+e^{-\int_0^t \beta(s)\,ds}\bigl(\Sigma_{\mathrm{VP}}(0)-I\bigr). (18)

  • 只要初始协方差\Sigma_{\mathrm{VP}}(0)有界,则任意时刻的协方差\Sigma_{\mathrm{VP}}(t)也始终保持有界;特别地,当\Sigma_{\mathrm{VP}}(0)=I时,有\Sigma_{\mathrm{VP}}(t)\equiv I.即该过程在整个演化过程中始终保持单位方差。正是基于这一性质差异,论文将 DDPM 对应的过程称为方差保持型随机微分方程 VP SDE

🐇 Solving The Reverse SDE

在训练得到时间相关的得分模型s_\theta之后,便可以利用该模型构造逆时间 SDE,并通过数值方法对其进行模拟,从而生成服从p_0的样本。由于 SDE 的解析解通常难以直接获得,因此实际生成过程依赖于数值求解器对随机动力学轨迹进行近似。已有大量通用的 SDE 数值求解方法,例如 Euler–Maruyama 方法和随机 Runge–Kutta 方法等,它们本质上对应于对随机动力学过程的不同离散化方式。因此,这些通用数值方法都可以直接应用于逆时间 SDE 的采样过程​

进一步地,论文指出,DDPM 中使用的祖先采样实际上可以看作逆时间 VP SDE 的一种特殊离散化形式。然而,对于新的 SDE,直接推导其对应的祖先采样规则往往并不容易。为解决这一问题,论文提出逆扩散采样器:其核心思想是在离散化逆时间 SDE 时,采用与前向过程相同的离散化方式,从而可以在已知前向离散形式的基础上直接构造出对应的逆向采样规则

由于已经获得基于得分的模型s_{\theta^*}(\mathbf{x},t)\approx \nabla_{\mathbf{x}}\log p_t(\mathbf{x}),因此可以进一步引入基于得分的 MCMC 方法,例如 Langevin MCMC 或 HMC,直接对各时刻的边缘分布p_t进行采样,并据此对数值 SDE 求解器给出的结果进行校正

具体来说,在每一个时间步,数值 SDE 求解器首先给出下一时刻样本的估计,这一过程起到预测器(predictor)的作用;随后,基于得分的 MCMC 方法对该估计样本的边缘分布进行修正,从而起到校正器(corrector)的作用。这一思想与求解方程组时常用的预测器—校正器方法相类似,因此论文将这种混合采样算法称为预测器—校正器(Predictor-Corrector, PC)采样器。从更一般的角度看,PC 采样器实际上统一并推广 SMLD 与 DDPM 的原始采样方式:对于前者,预测器可以看作恒等映射,而校正器对应于退火 Langevin 动力学;对于后者,预测器对应祖先采样,而校正器则退化为恒等映射

Algorithm 2 PC Sampling (VE SDE)
Algorithm 3 PC Sampling (VP SDE)

基于得分的模型还提供另一种求解逆时间 SDE 的数值途径。对于任意扩散过程,都存在一个与之对应的确定性过程,其轨迹虽然不再具有随机性,但在各个时刻仍与原 SDE 共享相同的边缘概率密度\{p_t(\mathbf{x})\}_{t=0}^{T}。这一确定性过程满足如下常微分方程

d\mathbf{x} = \left[ f(\mathbf{x},t)-\frac{1}{2}g(t)^2\nabla_{\mathbf{x}}\log p_t(\mathbf{x}) \right]dt,(19)

该方程在已知得分函数后即可由原始 SDE 唯一确定。论文称为概率流常微分方程(probability flow ODE)。当其中的得分函数由时间相关的得分模型近似,而该模型通常取为神经网络时,这一过程便构成神经 ODE 的一个具体实例

在此基础上,论文进一步指出,借助神经 ODE 中的瞬时变量替换公式,可以上式所定义的密度进行显式计算,从而实现对任意输入数据的精确似然估计。 除似然计算外,概率流 ODE 还使得潜在表示操控成为可能。通过上式进行积分,可以将任意数据点\mathbf{x}(0)编码到潜空间中的\mathbf{x}(T);相应地,再对对应的逆时间 ODE 进行积分,则可以完成解码过程。与其他可逆模型,如神经 ODE 与正规化流类似,这种机制使得模型能够对潜在表示进行操作,从而支持图像插值、温度缩放以及更一般的图像编辑任务

更进一步地,论文强调,该编码过程具有唯一可识别性。在固定前向 SDE、数据充分、模型容量足够且优化充分的条件下,精确 score 确定相应的概率流 ODE;在解唯一的条件下,同一起点产生确定的编码轨迹。概率流 ODE 与前向 SDE 在各时刻具有相同的边缘分布,但二者的样本轨迹并不相同

最后,在采样方面,论文表明也可以通过求解上式从不同的终态条件\mathbf{x}(T)\sim p_T出发生成\mathbf{x}(0)\sim p_0的样本。采用固定离散化策略时,概率流 ODE 已能够生成具有竞争力的样本,尤其是在与校正器结合使用时效果更佳。若进一步采用黑盒 ODE 求解器,则不仅能够生成高质量样本,还能够显式控制采样精度与计算效率之间的权衡

🍄 SDEs In The Wild

VE SDE、VP SDE 与 sub-VP SDE 的具体实例化形式,说明它们如何分别对应 SMLD、DDPM 及论文实验中使用的 sub-VP 模型

对于 SMLD,噪声尺度\{\sigma_i\}_{i=1}^N通常取为几何序列,连续极限下有

\sigma(t)=\sigma_{\min}\left(\frac{\sigma_{\max}}{\sigma_{\min}}\right)^t,\quad t\in(0,1](20)

  • 对应的 VE SDE 为

dx=\sigma_{\min}\left(\frac{\sigma_{\max}}{\sigma_{\min}}\right)^t \sqrt{2\log\frac{\sigma_{\max}}{\sigma_{\min}}}\,dw,\quad t\in(0,1], (21)

  • 从给定初值出发、不额外注入初始噪声时,该连续 SDE 的精确扰动核为

p_{0t}(x(t)\mid x(0)) = \mathcal N\!\left( x(t);x(0),\, \sigma_{\min}^2\left[\left(\frac{\sigma_{\max}}{\sigma_{\min}}\right)^{2t}-1\right]I \right),\quad t\in[0,1].(22)

  • 这里的条件方差在t=0时必须为零。原论文实现直接按几何尺度扰动数据,其方差还包含初始噪声尺度的平方,与上述精确 SDE 转移核的约定不同。为避开零方差端点,实际计算通常将时间区间限制为[\epsilon,1]

对于 DDPM,离散噪声参数\{\beta_i\}通常取为等差序列。在连续极限下

\beta(t)=\bar\beta_{\min}+t(\bar\beta_{\max}-\bar\beta_{\min}),\qquad t\in[0,1],(23)

  • 由此得到具体的 VP SDE

dx = -\frac12\bigl(\bar\beta_{\min}+t(\bar\beta_{\max}-\bar\beta_{\min})\bigr)x\,dt + \sqrt{\bar\beta_{\min}+t(\bar\beta_{\max}-\bar\beta_{\min})}\,dw, \qquad t\in[0,1],(24)

  • 其扰动核为

p_{0t}(x(t)\mid x(0)) = \mathcal N\!\left( x(t); e^{-\frac14 t^2(\bar\beta_{\max}-\bar\beta_{\min})-\frac12 t\bar\beta_{\min}}x(0),\, I-Ie^{-\frac12 t^2(\bar\beta_{\max}-\bar\beta_{\min})-t\bar\beta_{\min}} \right),\qquad t\in[0,1].(25)

  • VP SDE 的扰动核在初始时刻也退化为给定初值;由于t\to0 时方差趋近于 0,训练和采样时同样会将时间限制在[\epsilon,1]

对于 sub-VP SDE,论文采用与 VP SDE 相同的\beta(t),它与 VP SDE 具有相同的均值衰减形式,但方差更受控制,因此通常在似然建模和样本质量上都更有优势。其扰动核进一步改写为

p_{0t}(x(t)\mid x(0)) = \mathcal N\!\left( x(t); e^{-\frac14 t^2(\bar\beta_{\max}-\bar\beta_{\min})-\frac12 t\bar\beta_{\min}}x(0),\, \left[ 1-e^{-\frac12 t^2(\bar\beta_{\max}-\bar\beta_{\min})-t\bar\beta_{\min}} \right]^2 I \right),\quad t\in[0,1]. (26)