DDIM详解

DDIM详解

参考：https://www.bilibili.com/video/BV1VP411u71p/

虽然 DDIM 现在主要用于加速采样，但他的实际意义远不止于此。本文将首先回顾 DDPM 的训练和采样过程，再讨论 DDPM 与 DDIM 的关系，然后推导 DDIM 的采样公式，最后给出几个不同的理解 DDIM 的角度。

DDPM回顾

DDPM 实际是建模两个分布：diffusion 过程的分布 $q(x_t|x_0)$ 、$q(x_t|x_{t-1})$ 和 reverse 过程的分布 $q(x_{t-1}|x_t)$。

diffusion 过程：
$$\begin{gathered} q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I) \\ {x}_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,ϵ\sim \mathcal{N}(0,I) \end{gathered}$$
reverse 过程：
$$q(x_{t-1}|x_t)=q(x_{t-1}|x_t,x_0)=\frac{q(x_t|x_{t-1})q(x_{t-1}|x_0)}{q(x_t|x_0)}=\mathcal{N}(x_{t-1};\mu (x_t,x_0),{\sigma }_t^{2}I)$$
其中：
$$\begin{gathered} \mu (x_t,x_0)=\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 \\ {\sigma }_t^2=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t \end{gathered}$$
之前的文章已经详细介绍过推导过程，这里仅列出结论。

强调一个点：注意 reverse 过程中的 $x_0$ 是未知的，因此这里每一步的 $x_0$ 实际都是该步对 $x_0$ 的估计值，最终的 $x_0$ 相当于是每一步估计 $x_0$ 的加权和。DDPM 的 reverse 过程的每一步可以理解为做了两件事情：

1. 根据 $x_t$ 和 UNet 预测出的噪声 ${ϵ}_{\theta }(x_t,t)$ 估计当前步的 ${\hat{x}}_{0|t}$：${\hat{x}}_{0|t}=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }(x_t,t))$
2. 计算出 $x_{t-1}$：$x_{t-1}=\mu (x_t,{\hat{x}}_{0|t})+{\sigma }_tϵ,ϵ\sim \mathcal{N}(0,I)$

DDIM理解

在 DDPM 的采样公式 $q(x_{t-1}|x_t,x_0)$ 中 ，当我们已知 $x_0$ 时，当然可以通过扩散公式来计算 $x_{t-1}$ 。但在采样时，我们显然并不知道真实的 $x_0$，因此，我们使用 $x_t$ 来计算 $x_0$ 的估计值 $x_{0|t}$ ，然后计算 $x_{t-1}$ ：
$$\begin{array}{rl}q(x_{t-1}|x_t,x_0)& =\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}}ϵ,ϵ\sim \mathcal{N}(0,I)\\ & =\sqrt{{\bar{\alpha }}_{t-1}}{\hat{x}}_{0|t}+\sqrt{1-{\bar{\alpha }}_{t-1}}ϵ\\ & =\sqrt{{\bar{\alpha }}_{t-1}}\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }(x_t,t))+\sqrt{1-{\bar{\alpha }}_{t-1}}ϵ\end{array}$$
这与我们上面复习 DDPM 时的思路相同。

对于这个式子：
$$q(x_{t-1}|x_t,x_0)=\sqrt{{\bar{\alpha }}_{t-1}}{\hat{x}}_{0|t}+\sqrt{1-{\bar{\alpha }}_{t-1}}ϵ$$
我们知道 $ϵ$ 是采样自标准高斯分布，我们可以将其替换为模型估计的噪声 ${ϵ}_{\theta }(x_t,t)$：
$$q(x_{t-1}|x_t,x_0)=\sqrt{{\bar{\alpha }}_{t-1}}{\hat{x}}_{0|t}+\sqrt{1-{\bar{\alpha }}_{t-1}}{ϵ}_{\theta }(x_t,t)$$
同时考虑式 (5,6)，我们可以将这两种噪声都考虑进来，则有：
$$q(x_{t-1}|x_t,x_0)=\sqrt{{\bar{\alpha }}_{t-1}}{\hat{x}}_{0|t}+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{ϵ}_{\theta }+{\sigma }_tϵ$$
注意这里要保证加的两个噪声的方差和之前一样，即 ${\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}}^2+{\sigma }_t^2={\sqrt{1-{\bar{\alpha }}_{t-1}}}^2$ ，明显上式中是满足的。

观察上述推导过程，其实有一个 DDPM 中的条件我们是一直没有用到的。即：$q(x_t|x_{t-1})$ 。 这件事情其实是 DDIM 的核心，即丢弃掉 $q(x_t|x_{t-1})$ 这个条件。这样一来，我们的采样公式就不依赖于两步是相邻的，从而可以实现跳步的采样，即 $q(x_s|x_k,x_0)$ （因此 DDIM 可以用来加速采样）。我们有更通用的 DDIM 的采样公式：
$$x_s=q(x_s|x_k,x_0)=\sqrt{{\bar{\alpha }}_s}{\hat{x}}_{0|s}+\sqrt{1-{\bar{\alpha }}_s-{\sigma }_k^2}{ϵ}_{\theta }+{\sigma }_kϵ$$
在这个公式中，我们丢弃了 $q(x_t|x_{t-1})$ 这个约束条件，实现了从任意步 $k$ 到任意步 $s$ 的跳步采样。在 DDPM 采样时，必须一步一步进行，而使用 DDIM 采样，则可以自己定义任意的采样步数与步长：
$$\begin{gathered} \text{DDPM Samping}:T=1000,999,998,...,2,1,0 \\ \text{DDIM Samping}:T=1000,888,666,...,123,0 \end{gathered}$$

DDIM 丢掉了 DDPM 中的一个条件 $q(x_t|x_{t-1})$ ，可以进行跳步采样，实际上是一种更一般的形式。也就是说 DDIM 加一个条件约束 $q(x_t|x_{t-1})$ ，即可推出 DDPM，DDPM 去掉该条件，即可得到更一般的 DDIM。

DDIM推导

接下来我们来正式地推导一下 DDIM。

首先整理一下我们推到的目标：给定 $q(x_t|x_0)$ 和 $q(x_{t-1}|x_0)$ ，不能用 $q(x_t|x_{t-1})$ ，求 $q(x_{t-1}|x_t,x_0)$ 。

这里假设 $x_{t-1}$ 是 $x_t$ 和 $x_0$ 的线性组合，记其系数分别为 $m_t$ 和 $n_t$ ，即有：
$$x_{t-1}=m_t{x}_t+n_t{x}_0+{\sigma }_t{ϵ}_1$$
又知道：
$$\begin{gathered} x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_2 \\ {x}_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}}{ϵ}_3 \end{gathered}$$
这里用 ${ϵ}_{123}$ 的下标来区分对高斯分布的不同采样。(11) 代入 (10)，有：
$$\begin{array}{rl}{x}_{t-1}& =m_t(\sqrt{{\bar{\alpha }}_t}{x}_t+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_2)+n_t{x}_0+{\sigma }_t{ϵ}_1\\ & =(m_t\sqrt{{\bar{\alpha }}_t}+n_t)x_0+m_t\sqrt{1-{\bar{\alpha }}_t}{ϵ}_2+{\sigma }_t{ϵ}_1\end{array}$$
从而：
$$\{\begin{array}{l}{m}_t\sqrt{{\bar{\alpha }}_t}+n_t=\sqrt{{\bar{\alpha }}_{t-1}}\\ m_t^2(1-{\alpha }_t)+{\alpha }_t^2=1-{\bar{\alpha }}_{t-1}\end{array}$$
立即可以计算出 $m_t$ 和 $n_t$：
$$\begin{gathered} m_t=\sqrt{\frac{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{1-{\bar{\alpha }}_t}} \\ {n}_t=\sqrt{{\bar{\alpha }}_{t-1}}-\sqrt{\frac{{\bar{\alpha }}_t}{1-{\bar{\alpha }}_t}(1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2)} \end{gathered}$$

代回到式 (10)，有：
$$\begin{array}{rl}{x}_{t-1}& =\sqrt{\frac{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{1-{\bar{\alpha }}_t}}{x}_t+(\sqrt{{\bar{\alpha }}_{t-1}}-\sqrt{\frac{{\bar{\alpha }}_t}{1-{\bar{\alpha }}_t}(1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2)})x_0+{\sigma }_tϵ\\ & =\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}(\frac{1}{\sqrt{1-{\bar{\alpha }}_t}}{x}_t-\frac{\sqrt{{\bar{\alpha }}_t}}{\sqrt{1-{\bar{\alpha }}_t}}{x}_0)+{\sigma }_tϵ\\ & =\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}\frac{x_t-\sqrt{{\bar{\alpha }}_t}{x}_0}{\sqrt{1-{\bar{\alpha }}_t}}+{\sigma }_tϵ\end{array}$$
代换成模型的预测值：
$$x_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}{\hat{x}}_{0|t}+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{ϵ}_{\theta }(x_t,t)+{\sigma }_tϵ$$
至此，推导出的结果式 (16) 就与式 (8) 完全一致了。即：我们使用 $q(x_t|x_0)$ 和 $q(x_{t-1}|x_0)$ 两个条件，没有用 $q(x_t|x_{t-1})$ ，求出了 $q(x_{t-1}|x_t,x_0)$ 。DDPM 去掉条件 $q(x_t|x_{t-1})$ ，得到了更一般的 DDIM 的采样公式。而如果我们再加上该条件，就可以把方差 ${\sigma }_t$ 的值也确定下来，就得到了 DDPM 的采样公式。

理解DDIM的三个视角

训练目标的角度

DDPM 推导了那么多公式，但是在网络训练的时候是很简单的。就是先定义一系列超参数 ${\bar{\alpha }}_{\mathrm{1...}T}$ ，均匀采样一个时间步 $t$，从训练集中采样真实图片，按照 $\bar{\alpha }$ 和 $t$ 计算出噪声加到图片上，得到 $x_t$，训练一个 UNet 网络根据 $x_t$ 和 $t$ 来预测出图中噪声。也就是说，DDPM 训练的就是一个 UNet 网络在不同的噪声步条件 $t$ 下对应的去噪能力。

而 DDIM 在采样时可以跳步，也就是说 UNet 只需要具备对某几个时间步的去噪能力就行了。即：DDIM 的训练目标是 DDPM 的子集。

反过来说，如果一个 DDPM 已经完整地训练完成了，UNet 具备 $\mathrm{1...}T$ 所有时间步的去噪能力，那么就肯定包含了 DDIM 所需要的能力。因此，一个训练好的 DDPM 网络，可以直接使用 DDIM 采样算法进行加速采样。

图示理解的角度

$$x_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}{\hat{x}}_{0|t}+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{ϵ}_{\theta }(x_t,t)+{\sigma }_tϵ$$

观察 DDIM 的采样公式，其中一共有三项：

1. 第一项 $\sqrt{{\bar{\alpha }}_{t-1}}{\hat{x}}_{0|t}$ 的意义是当前步对 $x_0$ 的估计，乘一个系数（图中蓝色部分）
2. 第二项 $\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{ϵ}_{\theta }(x_t,t)$ 则相当于按照原来的方向又往回走了一段距离（图中绿色部分）
3. 第三项 ${\sigma }_tϵ$ 则表示在当前到达的位置又加了一个小扰动（图中黄色部分）

经过这样一步一步，最终逼近 $x_0$。

这又有两种特殊情况：

特殊情况一：${\sigma }_t=0$ 。此时相当于第三项没有了，即没有了小扰动。此时采样的过程就没有随机性了，初始点 $x_T$ 就已经决定了最终结果 $x_0$ 。这样的好处就是 $x_T$ 可以被看做是一个隐变量，类似 GAN Inversion 那中编辑、插值的方法就都可以搞起来了。

特殊情况二：$\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}=0$ 。此时相当于第二项没有了，这样采样过程会比较震荡。所以说，第二项的存在，能够使得采样的过程更加平滑。

常微分方程的角度

考虑 ${\sigma }_t=0$ 的情况：
$$\begin{array}{rl}{x}_{t-1}& =\sqrt{{\bar{\alpha }}_{t-1}}{\hat{x}}_{0|t}+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{ϵ}_{\theta }(x_t,t)\\ & =\sqrt{{\bar{\alpha }}_{t-1}}\cdot \frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }(x_t,t))+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{ϵ}_{\theta }(x_t,t)\\ \frac{x_{t-1}}{\sqrt{{\bar{\alpha }}_{t-1}}}& =\frac{x_t}{\sqrt{{\bar{\alpha }}_t}}-\sqrt{\frac{1-{\bar{\alpha }}_t}{{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)+\sqrt{\frac{1-{\bar{\alpha }}_{t-1}}{{\bar{\alpha }}_{t-1}}}{ϵ}_{\theta }(x_t,t)\\ \frac{x_{t-1}}{\sqrt{{\bar{\alpha }}_{t-1}}}-\frac{x_t}{\sqrt{{\bar{\alpha }}_t}}& =(\sqrt{\frac{1-{\bar{\alpha }}_{t-1}}{{\bar{\alpha }}_{t-1}}}-\sqrt{\frac{1-{\bar{\alpha }}_t}{{\bar{\alpha }}_t}}){ϵ}_{\theta }(x_t,t)\end{array}$$
可以看到，这已经是一个差分的形式。

记 自变量为 $s\in [0,1]$，$a=\sqrt{\bar{\alpha }}$ ，$\sigma =\sqrt{\frac{1-\bar{\alpha }}{\bar{\alpha }}}$ ，则有：
$$\frac{d}{ds}(\frac{x(s)}{a(s)})=\frac{d}{ds}\sigma (s){ϵ}_{\theta }(x(s),t(s))$$
这就把离散的形式写成连续的形式。

那么，扩散模型采样的过程就相当于：给定 $x(1)\sim \mathcal{N}(0,I)$，求 $x(0)$ 。这样一来，扩散模型的采样过程就相当于求解这个常微分方程。从而，很多加速求解常微分方程的方法，都可以用来加速扩散模型的采样过程。