diffusion model(二)：DDIM技术小结 (denoising diffusion implicit model)

DDIM

论文地址：Denoising Diffusion Implicit Models

github地址：https://github.com/ermongroup/ddim
个人博客地址：http://myhz0606.com/article/ddim

背景

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

DDPM为什么慢

从DDPM中我们知道，其扩散过程(前向过程、或加噪过程，forward process)被定义为一个马尔可夫过程

\[q(x_{1:T}|x_0):= \prod \limits_{t=1}^Tq(x_t|x_{t-1}) ,\\ \text{其中:} \space q(x_t|x_{t-1}):=\mathcal{N}(x_t; \sqrt{\alpha_t}x_{t-1}, (1 - \alpha_t) \textbf{I}) \tag{1} \]

通过这样设置，前向过程有一个很好的性质，可以通过\(x_0\)得到任意时刻\(x_t\)的分布，而无需烦琐的链式计算。

\[q(x_t|x_0) := \int q(x_{1:t}|x_0)dx_{1:(t-1)} = \mathcal{N}(x_t; \sqrt{\overline{\alpha}_{t}}x_0,(1 - \overline{\alpha}_{t})\textbf{I}) \tag{2} \]

其去噪过程(也有叫逆向过程，reverse process)也是一个马尔可夫过程。

\[p_\theta(x_{0})= \int p_\theta(x_{0:T})dx_{1:T} \\ \text{其中:} \space 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 \textbf{I}) \tag{3} \]

从式(2)可以看出，当\(t\)足够大时，\(q(x_t|x_0)\)对所有\(x_0\)都收敛于标准高斯分布。因此DDPM在去噪过程定义\(p_{\theta}(x_T):=\mathcal{N}(0, \textbf{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})\textbf{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})\textbf{I})，\)

⚠️文本的符号约定与论文不同，本文的\(\overline{\alpha}_t\)是论文中的\(\alpha_t\)。\(\overline{\alpha}_t=\prod \limits_{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 + kx_t, \sigma^2_t\textbf{I}) \tag{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})\textbf{I})\)。这样问题就转化成，找到一组解\(\lambda^*,k^*, \sigma ^* _t\)，使得DDIM的推理分布满足\(q(x_t|x_0)=\mathcal{N}(x_t; \sqrt{\overline{\alpha}_{t}}x_0,(1 - \overline{\alpha}_{t})\textbf{I})\)。下面我们来看DDIM要满足上述推理分布需要什么条件

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

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

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

到这里，问题进一步转化为：已知 \(q(x_t|x_0)=\mathcal{N}(x_t; \sqrt{\overline{\alpha}_{t}}x_0,(1 - \overline{\alpha}_{t})\textbf{I}) ,\) 找到\(p(x_{t-1}|x_{t}, x_0)=\mathcal{N}(x_{t-1}; \lambda x_0 + kx_t, \sigma^2_t\textbf{I})\)的一组解\(\lambda^*,k^*, \sigma ^*_t\)，使得\(q(x_{t-1}|x_0)=\mathcal{N}(x_{t-1}; \sqrt{\overline{\alpha}_{t-1}}x_0,(1 - \overline{\alpha}_{t-1})\textbf{I})\)，即求解

\[\int_{x_t} p(x_{t-1}|x_t, x_0) q(x_t|x_0)dx_t = q(x_{t-1}|x_0) \tag{5} \]

我们可以根据待定系数法进行求解

根据\(p(x_{t-1}|x_{t}, x_0)=\mathcal{N}(x_{t-1}; \lambda x_0 + kx_t, \sigma^2_t\textbf{I})\)，可采样\(x_{t-1}, \boldsymbol{\epsilon}_{t-1}' \sim \mathcal{N}(0, \textbf{I})\)

\[x_{t-1} = \lambda x_0 + kx_t + \sigma \boldsymbol{\epsilon}_{t-1}' \tag{6} \]

根据\(q(x_t|x_0)=\mathcal{N}(x_t; \sqrt{\overline{\alpha}_{t}}x_0,(1 - \overline{\alpha}_{t})\textbf{I})\)，可采样\(x_t, \boldsymbol{\epsilon}'_{t} \sim \mathcal{N}(0, \textbf{I})\)

\[x_{t} = \sqrt{\overline{\alpha}_{t}}x_0 + \sqrt{(1 - \overline{\alpha}_{t})} \boldsymbol{\epsilon}_t' \tag{7} \]

带入上式有

\[\begin{aligned} x_{t-1} &= \lambda x_0 + k(\sqrt{\overline{\alpha}_{t}}x_0 + \sqrt{(1 - \overline{\alpha}_{t})}\boldsymbol{\epsilon}_t') + \sigma_t \boldsymbol{\epsilon}_{t-1}' \\ &= (\lambda + k\sqrt{\overline{\alpha}_{t}}) x_0 + \underbrace{ k \sqrt{(1 - \overline{\alpha}_{t})} \boldsymbol{\epsilon}_t' + \sigma_t \boldsymbol{\epsilon}_{t-1}'}_{ \sqrt{(k^2(1 - \overline{\alpha}_{t}) + \sigma^2_t)} \overline{\boldsymbol{\epsilon}} _{t-1} \space \text{正太分布可加性} } \end{aligned} \tag{8} \]

式中\(\overline{\boldsymbol{\epsilon}} _{t-1} \sim \mathcal{N}(0, \textbf{I})\)

根据\(q(x_{t-1}|x_0)=\mathcal{N}(x_{t-1}; \sqrt{\overline{\alpha}_{t-1}}x_0,(1 - \overline{\alpha}_{t-1})\textbf{I})\)，采样\(x_{t-1}, \boldsymbol{\epsilon}_{t-1} \sim \mathcal{N}(0, \textbf{I})\)

\[x_{t-1} = \sqrt{\overline{\alpha}_{t-1}}x_0 + \sqrt{(1 - \overline{\alpha}_{t-1})} \boldsymbol{\epsilon}_{t-1} \tag{9} \]

为了满足式5

\[\begin{equation} \left\{ \begin{array}{lr} \lambda + k\sqrt{\overline{\alpha}_{t}} = \sqrt{\overline{\alpha}_{t-1}} \\ k^2(1 - \overline{\alpha}_{t}) + \sigma^2_t = 1 - \overline{\alpha}_{t-1} \end{array} \right. \Rightarrow \begin{pmatrix} \lambda^* \\ k^* \\ \sigma ^*_t \end{pmatrix} = \begin{pmatrix} \sqrt{\overline{\alpha}_{t-1}} - \sqrt{\overline{\alpha}_{t}} \sqrt{ \frac{1 - \overline{\alpha}_{t-1} - \sigma^2_{t} } {1 - \overline{\alpha} _t} } \\ \sqrt{ \frac{1 - \overline{\alpha}_{t-1} - \sigma^2_{t} } {1 - \overline{\alpha} _t} } \\ \sigma_t \end{pmatrix}\tag{10} \end{equation} \]

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

综上所得，我们得到DDIM的采样分布

\[\begin{aligned} 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^2_{t} } {1 - \overline{\alpha} _t} } \right)x_0 + \left(\sqrt{ \frac{1 - \overline{\alpha}_{t-1} - \sigma^2_{t} } {1 - \overline{\alpha} _t} } \right)x_t, \sigma^2_t\textbf{I}\right) \\ &= \mathcal{N}\left(x_{t-1}; \sqrt{\overline{\alpha}_{t-1}}x_0 + \sqrt{1 - \overline{\alpha}_{t-1} - \sigma^2_{t}} \frac{x_t - \sqrt{\overline{\alpha}_{t}}x_0 }{\sqrt{1 - \overline{\alpha} _t}}, \sigma^2_t\textbf{I} \right) \end{aligned} \tag{11} \]

从上式可见可以找到一簇解，使得\(p(x_{t-1}|x_{t}, x_0)=\mathcal{N}(x_{t-1}; \lambda x_0 + kx_t, \sigma^2_t\textbf{I})\),且满足推理分布\(q(x_t|x_0)=\mathcal{N}(x_t; \sqrt{\overline{\alpha}_{t}}x_0,(1 - \overline{\alpha}_{t})\textbf{I})\)。不同的\(\sigma_t\)对应不同生成过程。由于前向过程没变，故可以直接用DDPM训练的噪声预测模型。其采样过程如下：

\[\begin{aligned} x_{t-1} &= \sqrt{\overline{\alpha}_{t-1}}x_0 + \sqrt{1 - \overline{\alpha}_{t-1} - \sigma^2_{t}} \frac{x_t - \sqrt{\overline{\alpha}_{t}}x_0 }{\sqrt{1 - \overline{\alpha} _t}} + \sigma_t \boldsymbol{\epsilon} \\ \end{aligned}\tag{12} \]

式中

\[x_0 = \frac{1}{\sqrt{\overline{\alpha}_{t}}}(x_t - \sqrt{1 - \overline{\alpha}_{t}}\boldsymbol{\epsilon}_\theta(x_t, t))\tag{13} \]

整理得

\[\begin{aligned} x_{t-1} &= \sqrt{\overline{\alpha}_{t-1}} (\frac{1}{\sqrt{\overline{\alpha}_{t}}}(x_t - \sqrt{1 - \overline{\alpha}_{t}}\boldsymbol{\epsilon}_\theta(x_t, t))) + \sqrt{1 - \overline{\alpha}_{t-1} - \sigma^2_{t}} \frac{x_t - \sqrt{\overline{\alpha}_{t}}( \frac{1}{\sqrt{\overline{\alpha}_{t}}}(x_t - \sqrt{1 - \overline{\alpha}_{t}}\boldsymbol{\epsilon}_\theta(x_t, t))) }{\sqrt{1 - \overline{\alpha} _t}} + \sigma_t \boldsymbol{\epsilon} \\ &= \sqrt{\overline{\alpha}_{t-1}} \underbrace{ \frac{x_t - \sqrt{1 - \overline{\alpha}_{t}}\boldsymbol{\epsilon}_\theta(x_t, t)} {\sqrt{\overline{\alpha}_{t}}} }_{预测的x_0} + \underbrace{ \sqrt{1 - \overline{\alpha}_{t-1} - \sigma^2_{t}} \boldsymbol{\epsilon}_\theta(x_t, t) }_{x_0的方向} + \underbrace{\sigma_t \boldsymbol{\epsilon}}_{随机噪声项} \end{aligned}\tag{14} \]

有两个特殊的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, \textbf{I})\)时，采样过程是确定的，此时的生成模型是一个隐概率模型(implicit probabilstic model)[1]。作者将此时diffusion model称之为denoising diffusion implicit model（DDIM）。此时的采样递推公式

\[\begin{aligned} x_{t-1} &= \sqrt{\overline{\alpha}_{t-1}} \underbrace{ \frac{x_t - \sqrt{1 - \overline{\alpha}_{t}}\boldsymbol{\epsilon}_\theta(x_t, t)} {\sqrt{\overline{\alpha}_{t}}} }_{预测的x_0} + \underbrace{ \sqrt{1 - \overline{\alpha}_{t-1} } \boldsymbol{\epsilon}_\theta(x_t, t) }_{x_0的方向} \end{aligned}\tag{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]\),其长度为\(\mathrm{dim}(L)=T\),我们可以从生成序列L中构建一个子序列\(\tau=[\tau_s, \tau_{s-1}, \cdots, \tau_1]\)，其长度\(\mathrm{dim}(\tau)=S, S \ll T\)。在生成过程，按照构造序列\(\tau\)进行采样。此时的采样递推序列为

\[\begin{aligned} x_{\tau_{s-1}} &= \sqrt{\overline{\alpha}_{\tau_{s-1}}} \frac{x_{\tau_{s}} - \sqrt{1 - \overline{\alpha}_{\tau_{s}}}\boldsymbol{\epsilon}_\theta(x_{\tau_{s}}, t_{\tau_{s}})} {\sqrt{\overline{\alpha}_{\tau_{s}}}} + \sqrt{1 - \overline{\alpha}_{\tau_{s-1}} } \boldsymbol{\epsilon}_\theta(x_{\tau_{s}}, t_{\tau_{s}}) \end{aligned}\tag{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\)可能是生成图片的隐空间编码，那么它是否具备其它隐概率模型(如GAN2)所观察到的语义插值效应呢?下面来看怎么做：

首先从高斯分布采样两个随机变量\(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)}\tag{17} \]

其中 \(\theta = \arccos \left( \frac{ (x_T^{(0)})^T x_T^{(1)}} {\parallel x_T^{(0)} \parallel \parallel x_T^{(1)} \parallel} \right)\)

参考文献

[1]Shakir Mohamed and Balaji Lakshminarayanan. Learning in implicit generative models. arXiv preprint arXiv:1610.03483, October 2016.

[2]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.

[3]https://kexue.fm/archives/9181

[4]Ken Shoemake. Animating rotation with quaternion curves. In Proceedings of the 12th annual conference on Computer graphics and interactive techniques, pp. 245–254, 1985.

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

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

[7]https://en.wikipedia.org/wiki/Slerp