DDIM
扩散模型之DDIM
Author: [小小将]
Link: [https://zhuanlan.zhihu.com/p/565698027]
“What I cannot create, I do not understand.” – Richard Feynman
上一篇文章https://zhuanlan.zhihu.com/p/563661713介绍了经典扩散模型DDPM的原理和实现,对于扩散模型来说,一个最大的缺点是需要设置较长的扩散步数才能得到好的效果,这导致了生成样本的速度较慢,比如扩散步数为1000的话,那么生成一个样本就要模型推理1000次。这篇文章我们将介绍另外一种扩散模型DDIM(https://arxiv.org/abs/2010.02502),DDIM和DDPM有相同的训练目标,但是它不再限制扩散过程必须是一个马尔卡夫链,这使得DDIM可以采用更小的采样步数来加速生成过程,DDIM的另外是一个特点是从一个随机噪音生成样本的过程是一个确定的过程(中间没有加入随机噪音)。
DDIM原理
在介绍DDIM之前,先来回顾一下DDPM。在DDPM中,扩散过程(前向过程)定义为一个马尔卡夫链:
$q(\mathbf{x}{1:T} \vert \mathbf{x}0) = \prod^T{t=1} q(\mathbf{x}t \vert \mathbf{x}{t-1}) \quad q(\mathbf{x}t \vert \mathbf{x}{t-1}) = \mathcal{N}(\mathbf{x}t; \sqrt{\frac{\alpha_t}{\alpha{t-1}}} \mathbf{x}{t-1}, \Big(1-\frac{\alpha_t}{\alpha_{t-1}}\Big)\mathbf{I}) \$
注意,在DDIM的论文中,$\alpha_t$ 实是DDPM论文中的$\bar{\alpha}_t$,那么DDPM论文中的前向过程$\beta_t$ 为:
$\beta_t = \Big(1-\frac{\alpha_t}{\alpha_{t-1}}\Big) \$
扩散过程的一个重要特性是可以直接用$\mathbf{x}_0$ 对任意的$\mathbf{x}_t$ 行采样:
$q(\mathbf{x}_t \vert \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{{\alpha}_t} \mathbf{x}_0, (1 - {\alpha}_t)\mathbf{I}) \$
而DDPM的反向过程也定义为一个马尔卡夫链:
$p_\theta(\mathbf{x}{0:T}) = p(\mathbf{x}T) \prod^T{t=1} p\theta(\mathbf{x}{t-1} \vert \mathbf{x}t) \quad p\theta(\mathbf{x}{t-1} \vert \mathbf{x}t) = \mathcal{N}(\mathbf{x}{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}t, t), \boldsymbol{\Sigma}\theta(\mathbf{x}_t, t))\$
这里用神经网络$p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}t)$ 拟合真实的分布$q(\mathbf{x}{t-1} \vert \mathbf{x}_t)$。DDPM的前向过程和反向过程如下所示:
我们近一步发现后验分布$q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)$ 一个可获取的高斯分布:
$q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}0) = \mathcal{N}(\mathbf{x}{t-1}; \color{blue}{\tilde{\boldsymbol{\mu}}}(\mathbf{x}_t, \mathbf{x}_0), \color{red}{\tilde{\beta}_t} \mathbf{I}) \$
其中这个高斯分布的方差是定值,而均值是一个依赖$\mathbf{x}_0$ $\mathbf{x}_t$ 组合函数:
$\tilde{\boldsymbol{\mu}}_t (\mathbf{x}t, \mathbf{x}0) = \frac{\sqrt{\alpha_t}(1 - {\alpha}{t-1})}{\sqrt{\alpha{t-1}}(1 - {\alpha}_t)} \mathbf{x}t + \frac{\sqrt{{\alpha}{t-1}}\beta_t}{1 - {\alpha}_t} \mathbf{x}_0\$
然后我们基于变分法得到如下的优化目标:
$\begin{aligned} L &= \mathbb{E}{q(\mathbf{x}{1:T}\vert \mathbf{x}{0})} \Big[ \log\frac{q(\mathbf{x}{1:T}\vert\mathbf{x}0)}{p\theta(\mathbf{x}{0:T})} \Big] \ &= \underbrace{D\text{KL}(q(\mathbf{x}T \vert \mathbf{x}0) \parallel p\theta(\mathbf{x}T))}{L_T} + \sum{t=2}^T \underbrace{\mathbb{E}{q(\mathbf{x}{t}\vert \mathbf{x}{0})}\Big[D\text{KL}(q(\mathbf{x}{t-1} \vert \mathbf{x}t, \mathbf{x}0) \parallel p\theta(\mathbf{x}{t-1} \vert\mathbf{x}t))\Big]}{L{t-1}} -\underbrace{\mathbb{E}{q(\mathbf{x}{1}\vert \mathbf{x}{0})}\log p\theta(\mathbf{x}_0 \vert \mathbf{x}1)}{L_0} \end{aligned}\$
根据两个高斯分布的KL公式,我们近一步得到:
$L_{t-1}=\mathbb{E}{q(\mathbf{x}{t}\vert \mathbf{x}_{0})}\Big[ \frac{1}{2{\sigma_t^2}}|\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}0) - {\boldsymbol{\mu}\theta(\mathbf{x}_t, t)} |^2\Big] \$
根据扩散过程的特性,我们通过重参数化可以近一步简化上述目标:
$L_{t-1}=\mathbb{E}{\mathbf{x}{0},\mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\Big[ \frac{\beta_t^2}{2{\sigma_t^2}\alpha_t(1-\bar{\alpha}t)}| \mathbf{\epsilon}- \mathbf{\epsilon}\theta\big(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{\epsilon}, t\big)|^2\Big] \$
如果去掉系数,那么就能得到更简化的优化目标:
$L_{t-1}^{\text{simple}}=\mathbb{E}{\mathbf{x}{0},\mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\Big[ | \mathbf{\epsilon}- \mathbf{\epsilon}_\theta\big(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{\epsilon}, t\big)|^2\Big] \$
仔细分析DDPM的优化目标会发现,DDPM其实仅仅依赖边缘分布$q(\mathbf{x}_t \vert \mathbf{x}0)$,而并不是直接作用在联合分布$q(\mathbf{x}{1:T} \vert \mathbf{x}0)$。这带来的一个启示是:DDPM这个隐变量模型可以有很多推理分布来选择,只要推理分布满足边缘分布条件(扩散过程的特性)即可,而且这些推理过程并不一定要是马尔卡夫链。但值得注意的一个点是,我们要得到DDPM的优化目标,还需要知道分布$q(\mathbf{x}{t-1} \vert \mathbf{x}_t, \mathbf{x}0)$,之前我们在根据贝叶斯公式推导这个分布时是知道分布$q(\mathbf{x}t \vert \mathbf{x}{t-1})$ ,而且依赖了前向过程的马尔卡夫链特性。如果要解除对前向过程的依赖,那么我们就需要直接定义这个分布$q(\mathbf{x}{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)$。 基于上述分析,DDIM论文中将推理分布定义为:
$q_{\sigma}(\mathbf{x}{1:T} \vert \mathbf{x}0) = q{\sigma}(\mathbf{x}{T} \vert \mathbf{x}0)\prod^T{t=2} q_{\sigma}(\mathbf{x}{t-1} \vert \mathbf{x}{t},\mathbf{x}_{0}) \$
这里要同时满足$q_{\sigma}(\mathbf{x}_{T} \vert \mathbf{x}0)=\mathcal{N}(\sqrt{\alpha_T}\mathbf{x}{0},(1-{\alpha_T})\mathbf{I})$ 及对于所有的$t\ge2$ :
$q_\sigma(\mathbf{x}_{t-1} \vert \mathbf{x}t, \mathbf{x}0) = \mathcal{N}(\mathbf{x}{t-1}; \sqrt{\alpha{t-1}}\mathbf{x}0 + \sqrt{1 - \alpha{t-1} - \sigma_t^2} \frac{\mathbf{x}_t - \sqrt{\alpha_t}\mathbf{x}_0}{\sqrt{1 - \alpha_t}}, \sigma_t^2 \mathbf{I})\$
这里的方差$\sigma_t^2$ 一个实数,不同的设置就是不一样的分布,所以$q_{\sigma}(\mathbf{x}{1:T} \vert \mathbf{x}0)$ 实是一系列的推理分布。可以看到这里分布$q\sigma(\mathbf{x}{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)$ 均值也定义为一个依赖$\mathbf{x}0$ $\mathbf{x}t$ 组合函数,之所以定义为这样的形式,是因为根据$q{\sigma}(\mathbf{x}{T} \vert \mathbf{x}_0)$,我们可以通过数学归纳法证明,对于所有的$t$ 满足:
$q_{\sigma}(\mathbf{x}_t \vert \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{{\alpha}_t} \mathbf{x}_0, (1 - {\alpha}_t)\mathbf{I}) \$
这部分的证明见DDIM论文的附录部分,另外博客https://kexue.fm/archives/9181也从待定系数法来证明了分布$q_\sigma(\mathbf{x}_{t-1} \vert \mathbf{x}t, \mathbf{x}0)$ 构造的形式。 可以看到这里定义的推理分布$q{\sigma}(\mathbf{x}{1:T} \vert \mathbf{x}0)$ 没有直接定义前向过程,但这里满足了我们前面要讨论的两个条件:边缘分布$q{\sigma}(\mathbf{x}_t \vert \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{{\alpha}t} \mathbf{x}0, (1 - {\alpha}t)\mathbf{I})$,同时已知后验分布$q\sigma(\mathbf{x}{t-1} \vert \mathbf{x}t, \mathbf{x}0)$。同样地,我们可以按照和DDPM的一样的方式去推导优化目标,最终也会得到同样的$L^{\text{simple}}$(虽然VLB的系数不同,论文3.2部分也证明了这个结论)。 论文也给出了一个前向过程是非马尔可夫链的示例,如下图所示,这里前向过程是$q\sigma(\mathbf{x}{t} \vert \mathbf{x}{t-1}, \mathbf{x}_0)$,由于生成$\mathbf{x}t$ 仅依赖$\mathbf{x}{t-1}$,而且依赖$\mathbf{x}_0$,所以是一个非马尔可夫链:
注意,这里只是一个前向过程的示例,而实际上我们上述定义的推理分布并不需要前向过程就可以得到和DDPM一样的优化目标。与DDPM一样,这里也是用神经网络$\mathbf{\epsilon}\theta$ 预测噪音,那么根据$q\sigma(\mathbf{x}_{t-1} \vert \mathbf{x}t, \mathbf{x}0)$ 形式,在生成阶段,我们可以用如下公式来从$\mathbf{x}{t}$ 成$\mathbf{x}{t-1}$:
$\mathbf{x}{t-1} = \sqrt{\alpha{t-1}}\Big(\underbrace{\frac{\mathbf{x}t-\sqrt{1-\alpha{t}}\mathbf{\epsilon}\theta(\mathbf{x}t, t)}{\sqrt{\alpha{t}}}}{\text{predicted}\ \mathbf{x}0}\Big) + \underbrace{\sqrt{1 - \alpha{t-1} - \sigma_t^2} \cdot \mathbf{\epsilon}_\theta(\mathbf{x}t, t)}{\text{direction pointing to }\ \mathbf{x}t} + \underbrace{\sigma_t\epsilon_t}{\text {random noise}}\$
这里将生成过程分成三个部分:一是由预测的$\mathbf{x}_0$ 产生的,二是由指向$\mathbf{x}_t$ 部分,三是随机噪音(这里$\epsilon_t$ 与$\mathbf{x}_t$ 关的噪音)。论文将$\sigma_t^2$ 一步定义为:
$\sigma_t^2 = \eta \cdot \tilde{\beta}t=\eta \cdot\sqrt{(1-\alpha{t-1})/(1-\alpha_{t})}\sqrt{(1-\alpha_{t}/\alpha_{t-1})} \$
这里考虑两种情况,一是$\eta=1$,此时$\sigma_t^2 = \tilde{\beta}_t$,此时生成过程就和DDPM一样了。另外一种情况是$\eta=0$,这个时候生成过程就没有随机噪音了,是一个确定性的过程,论文将这种情况下的模型称为DDIM(denoising diffusion implicit model),一旦最初的随机噪音$\mathbf{x}_T$ 定了,那么DDIM的样本生成就变成了确定的过程。
上面我们终于得到了DDIM模型,那么我们现在来看如何来加速生成过程。虽然DDIM和DDPM的训练过程一样,但是我们前面已经说了,DDIM并没有明确前向过程,这意味着我们可以定义一个更短的步数的前向过程。具体地,这里我们从原始的序列$[1, …, T]$ 样一个长度为$S$ 子序列$[\tau_1,…,\tau_S]$,我们将$\mathbf{x}{\tau {1}},…,\mathbf{x}{\tau{S}}$ 前向过程定义为一个马尔卡夫链,并且它们满足:$q(\mathbf{x}{\tau{i}} \vert \mathbf{x}0) = \mathcal{N}(\mathbf{x}t; \sqrt{{\alpha}{\tau{i}}} \mathbf{x}0, (1 - {\alpha}{\tau_{i}})\mathbf{I})$。下图展示了一个具体的示例:
那么生成过程也可以用这个子序列的反向马尔卡夫链来替代,由于$S$ 以设置比原来的步数$L$ 小,那么就可以加速生成过程。这里的生成过程变成:
$\mathbf{x}{\tau{i-1}} = \sqrt{\alpha_{\tau_{i-1}}}\Big(\frac{\mathbf{x}{\tau{i}}-\sqrt{1-\alpha_{\tau_{i}}}\mathbf{\epsilon}\theta(\mathbf{x}{\tau_{i}}, \tau_{i})}{\sqrt{\alpha_{\tau_{i}}}}\Big) + \sqrt{1 - \alpha_{\tau_{i-1}} - \sigma_{\tau_{i}}^2} \cdot \mathbf{\epsilon}\theta(\mathbf{x}{\tau_{i}}, \tau_{i})+\sigma_{\tau_{i}}\epsilon\$
其实上述的加速,我们是将前向过程按如下方式进行了分解:
$q_{\sigma,\tau}(\mathbf{x}{1:T} \vert \mathbf{x}0) = q{\sigma,\tau}(\mathbf{x}{T} \vert \mathbf{x}0)\prod^S{i=1} q_{\sigma}(\mathbf{x}{\tau{i-1}} \vert \mathbf{x}{\tau{i}},\mathbf{x}{0})\prod{t\in \bar\tau}q_{\sigma,\tau}(\mathbf{x}{t}\vert \mathbf{x}{0}) \$
其中$\bar\tau={1,…,T}\backslash \tau$。这包含了两个图:其中一个就是由${\mathbf{x}{\tau_i}}{i=1}^S$ 成的马尔可夫链,另外一个是剩余的变量${\mathbf{x}{t}}{t\in \bar\tau}$ 成的星状图。同时生成过程,我们也只用马尔可夫链的那部分来生成:
$p_\theta(\mathbf{x}{0:T}) = p(\mathbf{x}T) \underbrace {\prod^S{i=1} p\theta(\mathbf{x}{\tau{i-1}} \vert \mathbf{x}{\tau{i}})}{\text{use to produce sample}} \times \underbrace {\prod{t\in \bar\tau}p_\theta(\mathbf{x}{0} \vert \mathbf{x}{t})}_{\text {only for VLB}} \$
论文共设计了两种方法来采样子序列,分别是:
- Linear:采用线性的序列$\tau_i=\lfloor ci\rfloor$;
- Quadratic:采样二次方的序列$\tau_i=\lfloor ci^2\rfloor$;
这里的$c$ 一个定值,它的设定使得$\tau_{-1}$ 接近$T$。论文中只对CIFAR10数据集采用Quadratic序列,其它数据集均采用Linear序列。
实验结果
下表为不同的$\eta$ 以及不同采样步数下的对比结果,可以看到DDIM($\eta=0$)在较短的步数下就能得到比较好的效果,媲美DDPM($\eta=1$)的生成效果。如果$S$ 置为50,那么相比原来的生成过程就可以加速20倍。

代码实现
DDIM和DDPM的训练过程一样,所以可以直接在DDPM的基础上加一个新的生成方法(这里主要参考了https://github.com/ermongroup/ddim以及https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_ddim.py),具体代码如下所示:
1 | |
这里以MNIST数据集为例,训练的扩散步数为500,直接采用DDPM(即推理500次)生成的样本如下所示:
同样的模型,我们采用DDIM来加速生成过程,这里DDIM的采样步数为50,其生成的样本质量和500步的DDPM相当:
完整的代码示例见https://github.com/xiaohu2015/nngen。
其它:重建和插值
在DDIM论文中,还额外讨论了两个小点的内容**:重建和插值**。所谓重建是指的首先用原始图像求逆得到对应的噪音然后再进行生成的过程;而插值是指的对两个随机噪音进行插值从而得到融合两种噪音的图像。 首先是重建,对于DDIM,其$\eta=0$,这个时候从$\mathbf{x}{t}$ 成$\mathbf{x}{t-1}$ 更新公式就变为:
$\mathbf{x}{t-1} = \sqrt{\alpha{t-1}}\Big(\frac{\mathbf{x}t-\sqrt{1-\alpha{t}}\mathbf{\epsilon}\theta(\mathbf{x}t, t)}{\sqrt{\alpha{t}}}\Big)+ \sqrt{1 - \alpha{t-1}} \cdot \mathbf{\epsilon}_\theta(\mathbf{x}_t, t)\$
我们进一步对上述公式进行变换可得到:
$\frac{\mathbf{x}{t-1}}{\sqrt{\alpha{t-1}}} = \frac{\mathbf{x}t}{\sqrt{\alpha{t}}} + \Big(\sqrt{\frac{1-\alpha_{t-1}}{\alpha_{t-1}}} - \sqrt{\frac{1-\alpha_{t}}{\alpha_{t}}} \Big)\mathbf{\epsilon}_\theta(\mathbf{x}_t, t) \$
当$T$ 够大时,以上公式其实可以看成用欧拉法来求解一个常微分方程(ODE,ordinary differential equation):
$\frac{\mathbf{x}{t-\Delta t}}{\sqrt{\alpha{t-\Delta t}}} = \frac{\mathbf{x}t}{\sqrt{\alpha{t}}} + \Big(\sqrt{\frac{1-\alpha_{t-\Delta t}}{\alpha_{t-\Delta t}}} - \sqrt{\frac{1-\alpha_{t}}{\alpha_{t}}} \Big)\mathbf{\epsilon}_\theta(\mathbf{x}_t, t) \$
这里令$\sigma=\sqrt{1-\alpha}/\sqrt{\alpha}$,$\bar{\mathbf{x}}=\mathbf{x}/\sqrt{\alpha}$,它们都是关于$t$ 函数,这样对应的ODE就是:
$\text{d}\bar{\mathbf{x}}(t)=\mathbf{\epsilon}_\theta(\frac{\bar{\mathbf{x}}(t)}{\sqrt{\sigma^2+1}}, t)\text{d}\sigma(t) \$
看成ODE后,我们可以利用如下公式对生成过程进行逆操作:
$\frac{\mathbf{x}{t+1}}{\sqrt{\alpha{t+1}}} = \frac{\mathbf{x}t}{\sqrt{\alpha{t}}} + \Big(\sqrt{\frac{1-\alpha_{t+1}}{\alpha_{t+1}}} - \sqrt{\frac{1-\alpha_{t}}{\alpha_{t}}} \Big)\mathbf{\epsilon}_\theta(\mathbf{x}_t, t) \$
这意味着,我们可以由一个原始图像$\mathbf{x}_0$ 到对应的随机噪音$\mathbf{x}_T$,然后我们再用$\mathbf{x}_T$ 行生成就可以重建原始图像$\mathbf{x}_0$(具体的代码实现见https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py#L524-L560)。论文也通过在CIFAR10测试集上的实验来证明当步数足够时,这种方式可以得到较低的重建误差:
第二个插值,对于DDIM,两个不同的随机噪音会产生不同的图像,但是如果我们对这两个随机噪音进行插值生成新的$\mathbf{x}_T$,那么将生成融合的图像。这里采用的插值方法是球面线性插值( spherical linear interpolation):
$\mathbf{x}_T^{(\alpha)}=\frac{\sin ((1-\alpha)\theta)}{\sin (\theta)}\mathbf{x}_T^{(0)}+\frac{\sin (\alpha\theta)}{\sin (\theta)}\mathbf{x}_T^{(1)} \quad \theta = \arccos\Big(\frac{(\mathbf{x}_T^{(0)})^\text{T}\mathbf{x}_T^{(1)}}{| \mathbf{x}_T^{(0)}||\mathbf{x}_T^{(1)}|}\Big) \$
这里的参数$\alpha\in[0,1]$ 制插值系数,具体的代码实现见https://github.com/ermongroup/ddim/blob/main/runners/diffusion.py#L296-L334。下图展示了一些具体的插值效果:
DDIM的重建和插值也在文本转图像模型DALLE-2中使用,不过这里插值的是扩散模型的条件CLIP image embedding,详情见论文**https://arxiv.org/abs/2204.06125**。

小结
如果从直观上看,DDIM的加速方式非常简单,直接采样一个子序列,其实论文https://arxiv.org/abs/2102.09672也采用了类似的方式来加速。另外DDIM和其它扩散模型的一个较大的区别是其生成过程是确定性的。
参考
- https://arxiv.org/abs/2010.02502
- https://github.com/ermongroup/ddim
- https://github.com/openai/improved-diffusion
- https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_ddim.py
- https://github.com/CompVis/latent-diffusion/blob/main/ldm/models/diffusion/ddim.py
- https://kexue.fm/archives/9181












