diffusion model（二）—— DDIM技术小结

系列阅读

• diffusion model（一）DDPM技术小结 (denoising diffusion probabilistic)
• diffusion model（二）—— DDIM技术小结
• diffusion model（三）—— classifier guided diffusion model
• diffusion model（四）文生图diffusion model（classifier-free guided）
• diffusion model（五）stable diffusion底层原理（latent diffusion model， LDM

论文地址：Denoising Diffusion Implicit Models
github地址：https://github.com/ermongroup/ddim

背景

去噪扩散概率模型 (DDPM¹) 在没有对抗训练的情况下实现了高质量的图像生成，但其采样过程依赖马尔可夫假设，需要较多的时间步才能得到较好的生成效果。本文提出的DDIM(denoising diffusion implicit models ²)是更有效的迭代隐式概率模型，其训练过程与 DDPM 相同，但相比DDPM，采样过程快 10 到 50 倍。

DDPM为什么慢

从DDPM中我们知道，其扩散过程(前向过程、或加噪过程，forward process)被定义为一个马尔可夫过程
$$\begin{gathered} q(x_{1:T}|x_0):=\prod _{t=1}^{T}q(x_t|x_{t-1}), \\ \text{其中:}q(x_t|x_{t-1}):=\mathcal{N}(x_t;\sqrt{{\alpha }_t}{x}_{t-1},(1-{\alpha }_t)\text{I})\text{\hspace{1em}(1)} \end{gathered}$$
通过这样设置，前向过程有一个很好的性质，可以通过$x_0$得到任意时刻$x_t$的分布，而无需烦琐的链式计算。
$$q(x_t|x_0):=\int q(x_{1:t}|x_0)d{x}_{1:(t-1)}=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\text{I})\text{\hspace{1em}(2)}$$
其去噪过程(也有叫逆向过程，reverse process)也是一个马尔可夫过程。
$$\begin{gathered} p_{\theta }(x_0)=\int p_{\theta }(x_{0:T})d{x}_{1:T} \\ \text{其中:}{p}_{\theta }(x_{0:T}):=p_{\theta }(x_T)\prod _{t=1}^T{p}_{\theta }^{(t)}(x_{t-1}|x_t) \\ {p}_{\theta }^{(t)}(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\sigma }_t\text{I})\text{\hspace{1em}(3)} \end{gathered}$$
从式(2)可以看出，当$t$足够大时，$q(x_t|x_0)$对所有$x_0$都收敛于标准高斯分布。因此DDPM在去噪过程定义$p_{\theta }(x_T):=\mathcal{N}(0,\text{I})$，并且采用一个较大的采样时间步$T$。

在对$p_{\theta }^{(t)}(x_{t-1}|x_t)$的推导中，DDPM用到了一阶马尔可夫假设，使得$p(x_t|x_{t-1},x_0)=p(x_t|x_{t-1})$。因此重建的步长非常长，导致速度慢。

DDIM如何加速采样

下面我们来看DDIM如何解决生成速度慢这个问题。DDPM速度慢的本质原因是对马尔可夫假设的依赖，导致重建需要较多的步长。那么不用一阶马尔可夫假设，有没有另一种方法推导出$p(x_{t-1}|x_t,x_0)$。前面我们提到，DDPM的前向过程主要是用于构造数据集，来训练噪声预测模型。在实际优化过程中，其实并没有用到前向过程的马尔可夫假设，只用到了推理分布$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\text{I})$ ，因此如果DDIM也能满足这个推理分布的话，那么直接就可以复用DDPM的噪声预测模型了

归纳一下DDIM想要做的

• DDIM希望构建一个采样分布$p(x_{t-1}|x_t,x_0)$，这个采样分布不依赖一阶马尔科夫假设
• DDIM希望维持这个前向推理分布$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\text{I})$，

⚠️文本的符号约定与论文不同，本文的${\overline{\alpha }}_t$是论文中的${\alpha }_t$。${\overline{\alpha }}_t=\prod _{i=1}^t{\alpha }_i$

采样分布$p(x_{t-1}|x_t,x_0)$的求解

下面我们看DDIM是如何推导出采样分布$p(x_{t-1}|x_t,x_0)$

根据DDPM的结果参考，采样分布$p(x_{t-1}|x_t,x_0)$是一个高斯分布，并且均值是$x_0,x_t$的线性函数，方差是一个与时间步有关的函数，这里不妨假设DDIM的采样分布也是一个高斯分布，并且均值也是$x_0,x_t$的线性函数，方差是时间步$t$的函数，即
$$p(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\lambda x_0+k{x}_t,{\sigma }_t^2\text{I})\text{\hspace{1em}(4)}$$
现在我们已经有了采样分布，并且该采样分布不依赖一阶马尔可夫假设。这个采样分布有3个自由变量$\lambda ,k,{\sigma }_t$，理论上有无穷多个解。但DDIM想要维持与DDPM一致的推理分布$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\text{I})$。这样问题就转化成，找到一组解${\lambda }^{\ast },k^{\ast },{\sigma }_t^{\ast }$，使得DDIM的推理分布满足$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\text{I})$。下面我们来看DDIM要满足上述推理分布需要什么条件

当$t=1$时，DDIM可以满足$q(x_1|x_0)=\mathcal{N}(x_1;\sqrt{{\overline{\alpha }}_1}{x}_0,(1-{\overline{\alpha }}_1)\text{I})$

假设当$t=t$时，DDIM也满足$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\text{I})$

那么根据数学归纳法,只需保证$q(x_{t-1}|x_0)=\mathcal{N}(x_{t-1};\sqrt{{\overline{\alpha }}_{t-1}}{x}_0,(1-{\overline{\alpha }}_{t-1})\text{I})$,就能对所有时间步都能满足上述推理分布。

到这里，问题进一步转化为：已知 $q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\text{I})$ , 找到$p(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\lambda x_0+k{x}_t,{\sigma }_t^2\text{I})$的一组解${\lambda }^{\ast },k^{\ast },{\sigma }_t^{\ast }$，使得$q(x_{t-1}|x_0)=\mathcal{N}(x_{t-1};\sqrt{{\overline{\alpha }}_{t-1}}{x}_0,(1-{\overline{\alpha }}_{t-1})\text{I})$，即求解
$${\int }_{x_t}p(x_{t-1}|x_t,x_0)q(x_t|x_0)d{x}_t=q(x_{t-1}|x_0)\text{\hspace{1em}(5)}$$
我们可以根据待定系数法进行求解

根据$p(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\lambda x_0+k{x}_t,{\sigma }_t^2\text{I})$，可采样$x_{t-1}$, ${\mathbf{ϵ}}_{t-1}^{\prime }\sim \mathcal{N}(0,\text{I})$
$$x_{t-1}=\lambda x_0+k{x}_t+\sigma {\mathbf{ϵ}}_{t-1}^{\prime }\text{\hspace{1em}(6)}$$
根据$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\text{I})$，可采样$x_t$, ${\mathbf{ϵ}}_t^{\prime }\sim \mathcal{N}(0,\text{I})$
$$x_t=\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{(1-{\overline{\alpha }}_t)}{\mathbf{ϵ}}_t^{\prime }\text{\hspace{1em}(7)}$$
带入上式有
$$\begin{array}{rl}{x}_{t-1}& =\lambda x_0+k(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{(1-{\overline{\alpha }}_t)}{\mathbf{ϵ}}_t^{\prime })+{\sigma }_t{\mathbf{ϵ}}_{t-1}^{\prime }\\ & =(\lambda +k\sqrt{{\overline{\alpha }}_t})x_0+\underset{\sqrt{(k^2(1-{\overline{\alpha }}_t)+{\sigma }_t^2)}{\overline{\mathbf{ϵ}}}_{t-1}\text{正太分布可加性}}{\underbrace{k\sqrt{(1-{\overline{\alpha }}_t)}{\mathbf{ϵ}}_t^{\prime }+{\sigma }_t{\mathbf{ϵ}}_{t-1}^{\prime }}}\end{array}\text{\hspace{1em}(8)}$$
式中${\overline{\mathbf{ϵ}}}_{t-1}\sim \mathcal{N}(0,\text{I})$

根据$q(x_{t-1}|x_0)=\mathcal{N}(x_{t-1};\sqrt{{\overline{\alpha }}_{t-1}}{x}_0,(1-{\overline{\alpha }}_{t-1})\text{I})$，采样$x_{t-1}$, ${\mathbf{ϵ}}_{t-1}\sim \mathcal{N}(0,\text{I})$
$$x_{t-1}=\sqrt{{\overline{\alpha }}_{t-1}}{x}_0+\sqrt{(1-{\overline{\alpha }}_{t-1})}{\mathbf{ϵ}}_{t-1}\text{\hspace{1em}(9)}$$
为了满足式5

上述方程有3个未知数，两个方程，显然有无穷组解能够满足式5，DDIM中将${\sigma }_t$看作可变变量（该数值的大小可视为采样过程的随机性有多大）。

综上所得，我们得到DDIM的采样分布
$$\begin{array}{rl}p(x_{t-1}|x_t,x_0)& =\mathcal{N}\left(x_{t-1};\left(\sqrt{{\overline{\alpha }}_{t-1}}-\sqrt{{\overline{\alpha }}_t}\sqrt{\frac{1-{\overline{\alpha }}_{t-1}-{\sigma }_t^2}{1-{\overline{\alpha }}_t}}\right)x_0+\left(\sqrt{\frac{1-{\overline{\alpha }}_{t-1}-{\sigma }_t^2}{1-{\overline{\alpha }}_t}}\right)x_t,{\sigma }_t^2\text{I}\right)\\ & =\mathcal{N}\left(x_{t-1};\sqrt{{\overline{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\overline{\alpha }}_{t-1}-{\sigma }_t^2}\frac{x_t-\sqrt{{\overline{\alpha }}_t}{x}_0}{\sqrt{1-{\overline{\alpha }}_t}},{\sigma }_t^2\text{I}\right)\end{array}\text{\hspace{1em}(11)}$$
从上式可见可以找到一簇解，使得$p(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\lambda x_0+k{x}_t,{\sigma }_t^2\text{I})$,且满足推理分布$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\overline{\alpha }}_t}{x}_0,(1-{\overline{\alpha }}_t)\text{I})$。不同的${\sigma }_t$对应不同生成过程。由于前向过程没变，故可以直接用DDPM训练的噪声预测模型。其采样过程如下：
$$\begin{array}{rl}{x}_{t-1}& =\sqrt{{\overline{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\overline{\alpha }}_{t-1}-{\sigma }_t^2}\frac{x_t-\sqrt{{\overline{\alpha }}_t}{x}_0}{\sqrt{1-{\overline{\alpha }}_t}}+{\sigma }_t\mathbf{ϵ}\end{array}\text{\hspace{1em}(12)}$$
式中
$$x_0=\frac{1}{\sqrt{{\overline{\alpha }}_t}}(x_t-\sqrt{1-{\overline{\alpha }}_t}{\mathbf{ϵ}}_{\theta }(x_t,t))\text{\hspace{1em}(13)}$$
整理得

有两个特殊的case值得注意

1. 当${\sigma }_t=\sqrt{(1-{\overline{\alpha }}_{t-1})/(1-{\overline{\alpha }}_t)}\sqrt{1-{\overline{\alpha }}_t/{\overline{\alpha }}_{t-1}}$时，此时的生成过程与DDPM一致
2. 当${\sigma }_t=0$时，此时采样过程中添加的随机噪声项为0，当给定$z=x_t\sim \mathcal{N}(0,\text{I})$时，采样过程是确定的，此时的生成模型是一个隐概率模型(implicit probabilstic model)³。作者将此时diffusion model称之为denoising diffusion implicit model（DDIM）。此时的采样递推公式

$$\begin{array}{rl}{x}_{t-1}& =\sqrt{{\overline{\alpha }}_{t-1}}\underset{预测的x_0}{\underbrace{\frac{x_t-\sqrt{1-{\overline{\alpha }}_t}{\mathbf{ϵ}}_{\theta }(x_t,t)}{\sqrt{{\overline{\alpha }}_t}}}}+\underset{x_0的方向}{\underbrace{\sqrt{1-{\overline{\alpha }}_{t-1}}{\mathbf{ϵ}}_{\theta }(x_t,t)}}\end{array}\text{\hspace{1em}(15)}$$

DDIM是如何加速采样的

前面我们提到DDPM在推导采样分布$p(x_{t-1}|x_t,x_0)$时用到了马尔可夫假设$p(x_t|x_{t-1},x_0)=p(x_t|x_{t-1})$使得在采样过程中需要从$t=T$依次采样到 $t=1$。在DDPM中往往需要较大的$T$才能得到较好的结果(论文中$T=1000$)，这导致需要T次采样步长。DDIM加速采样的思路很简单，假定原本的生成序列为$L=[T,T-1,\cdots ,1]$,其长度为$\dim (L)=T$,我们可以从生成序列$L$中构建一个子序列$\tau =[{\tau }_s,{\tau }_{s-1},\cdots ,{\tau }_1]$，其长度$\dim (\tau )=S$, $S\ll T$。在生成过程，按照构造序列$\tau$进行采样。此时的采样递推序列为
$$\begin{array}{rl}{x}_{{\tau }_{s-1}}& =\sqrt{{\overline{\alpha }}_{{\tau }_{s-1}}}\frac{x_{{\tau }_s}-\sqrt{1-{\overline{\alpha }}_{{\tau }_s}}{\mathbf{ϵ}}_{\theta }(x_{{\tau }_s},t_{{\tau }_s})}{\sqrt{{\overline{\alpha }}_{{\tau }_s}}}+\sqrt{1-{\overline{\alpha }}_{{\tau }_{s-1}}}{\mathbf{ϵ}}_{\theta }(x_{{\tau }_s},t_{{\tau }_s})\end{array}\text{\hspace{1em}(16)}$$
论文给出了不同的采样步长的生成效果。可以看到DDIM在较小采样步长时就能达到较好的生成效果。如CIFAR10 $S=50$就达到了$S=1000$的90%的效果，与之相对DDPM只能达到10%左右的FID效果。可见DDPM在推导采样分布中用了马尔可夫假设的确限制了它的采样间隔。但不考虑生成质量

可能有同学会有疑问，不按推理步骤按序预测噪声会不会有问题。答案是不会。
从下图可以看到，在训练时，是从时间序列$[1,2,\cdots ,T]$进行采样时间步，再从高斯分布采样噪声，直接从$x_0$出发计算$x_t$，随后传入模型，根据预测噪声与实际噪声计算损失迭代模型。我们可以看到在这个训练pipeline中并没有限制时序信息。当噪声预测模型充分训练时DDIM的采样序列$\tau$也是在训练过程中被充分训练的。（论文附录$C.1$给出了更为详尽的解释推导，感兴趣的同学可以看论文）。

DDIM区别于DDPM两个重要的特性

采样一致性(sample consistency)

根据式16我们知道DDIM将${\sigma }_t$设置为0，这让采样过程是确定的，只受$x_T$影响。作者发现，当给定$x_T$，不同的的采样时间序列$\tau$所生成图片都很相近，$x_T$似乎可以视作生成图片的隐编码信息。

有个小trick，我们在实际的生成中可以先设置较小的采样步长进行生成，若生成的图片是我们想要的，则用较大的步长重新生成高质量的图片。

语义插值效应(sementic interpolation effect)

即然$x_T$可能是生成图片的隐空间编码，那么它是否具备其它隐概率模型(如GAN⁴)所观察到的语义插值效应呢?下面来看怎么做：

首先从高斯分布采样两个随机变量$x_T^{(0)},x_T^{(1)}$，并用他们做图像生成得到下图最左侧与最右侧的结果。随后用球面线性插值方法（spherical linear interpolation，Slerp）[^4, 7]对$x_T^{(1)},x_T^{(2)}$他们进行插值，得到一系列中间结果
$$x_T^{(\alpha )}=\frac{\sin ((1-\alpha )\theta )}{\sin (\theta )}{x}_T^{(0)}+\frac{\sin (\alpha \theta )}{\sin (\theta )}{x}_T^{(1)}\text{\hspace{1em}(17)}$$
其中 $\theta =\arccos \left(\frac{(x_T^{(0)}{)}^T{x}_T^{(1)}}{\parallel x_T^{(0)}\parallel \parallel x_T^{(1)}\parallel }\right)$

参考文献

---

1. Ho J, Jain A, Abbeel P. Denoising diffusion probabilistic models[J]. Advances in Neural Information Processing Systems, 2020, 33: 6840-6851. ↩︎

2. Song J, Meng C, Ermon S. Denoising diffusion implicit models[J]. arXiv preprint arXiv:2010.02502, 2020. ↩︎

3. Shakir Mohamed and Balaji Lakshminarayanan. Learning in implicit generative models. arXiv preprint arXiv:1610.03483, October 2016. ↩︎

4. Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural information processing systems, pp. 2672–2680, 2014. ↩︎