Denoising Diffusion Implicit Models (DDIM)

因为扩散模型的正向过程可以实现跳步，其马尔科夫性并不重要，因此DDIM尝试在反向过程中也舍弃马尔科夫性的约束。在引入一个可控参数${\sigma }_t$后，反向过程的迭代被整理为预测$x_0$后重新进行正向扩散的过程，扩散时使用的噪声为新初始化噪声与上一步预测噪声的加权，在特定$\sigma$选择下，可以导出DDPM的采样过程和新的确定性的采样过程，并且使得反向过程同样支持跳步。

方法概述

舍弃马尔科夫性

在DDPM中，反向过程的损失函数是根据马尔科夫性推导得出的，而虽然正向过程也是马尔科夫链，但是因为高斯分布的可加性，其完全可以舍弃马尔科夫性，实现从第 0 步直接跳到第 t 步。

而DDIM的核心思想就是在反向过程时也舍弃马尔科夫性的约束，使得采样的过程也可以实现跳步。

注意，DDIM的$p$和$q$符号也与本文相反，在“理论推导”小节需要注意与原文对照。

采样公式

在DDPM中，我们先是使用了一种没有严格理论支撑但容易理解的方法导出了一个采样迭代公式：
$$x_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }\left(x_t,t\right)\right)+\sqrt{1-{\bar{\alpha }}_{t-1}}z$$
然后又在后续理论推导出中给出了DDPM原版的采样迭代公式：
$$x_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)\right)+{\tilde{\beta }}_{t}z$$
而这两种都是DDIM这个工作提出的采样迭代公式的两种特殊情况：
$$x_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}\left(\frac{x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }(x_t,t)}{\sqrt{{\bar{\alpha }}_t}}\right)+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}\cdot {ϵ}_{\theta }(x_t,t)+{\sigma }_{t}z$$
带入${\sigma }_t=\sqrt{1-{\bar{\alpha }}_{t-1}}$和${\sigma }_t=\sqrt{(1-{\bar{\alpha }}_{t-1})/(1-{\bar{\alpha }}_t)}\sqrt{1-{\alpha }_t}$可以分别得到以上两个迭代过程。为了便于分析和与DDIM原文对齐，这里将上式的${\bar{\alpha }}_t$表示为${\alpha }_t$，上式的${\alpha }_t$则对应${\alpha }_t/{\alpha }_{t-1}$，将$t$移到符号上标，$z$替换为${ϵ}_t$：
$$x_{t-1}=\sqrt{{\alpha }_{t-1}}\left(\frac{x_t-\sqrt{1-{\alpha }_t}{ϵ}_{\theta }^{(t)}(x_t)}{\sqrt{{\alpha }_t}}\right)+\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\cdot {ϵ}_{\theta }^{(t)}(x_t)+{\sigma }_t{ϵ}_t$$
此时可以得到DDIM最重要的特殊情况，即设置${\sigma }_t=0$，此时上式等于：
$$x_{t-1}=\sqrt{{\alpha }_{t-1}}\left(\frac{x_t-\sqrt{1-{\alpha }_t}{ϵ}_{\theta }^{(t)}(x_t)}{\sqrt{{\alpha }_t}}\right)+\sqrt{1-{\alpha }_{t-1}}\cdot {ϵ}_{\theta }^{(t)}(x_t)=\sqrt{{\alpha }_{t-1}}{x}_{\theta }^{(0)}(x_t)+\sqrt{1-{\alpha }_{t-1}}\cdot {ϵ}_{\theta }^{(t)}(x_t)$$
可以发现，此时DDIM在迭代过程中没有随机性，即每一步反向过程迭代的时候并没有采样新的噪声，而是直接使用前向过程$${\mathbf{x}}_t=\sqrt{{\alpha }_t}{\mathbf{x}}_0+\sqrt{1-{\alpha }_t}ϵ$$在估计的干净图像$x_{\theta }^{(0)}(x_t)$上回加估计的噪声${ϵ}_{\theta }^{(t)}(x_t)$，可以看作一个迭代细化的过程。当然，DDIM采样并不是完全没有随机性，因为迭代前需要采样$x_T\sim \mathcal{N}(0,\mathbf{I})$。

加速采样

对于DDPM，如果要贯彻推导时使用的马尔科夫性，那么在采样时需要使用训练时设置的$T$步进行采样，而在绝大部分扩散模型中，都设置$T=1000$，这导致DDPM生成图像需要执行1000次去噪网络${ϵ}_{\theta }$，而${ϵ}_{\theta }$是一个参数量和计算量都很大的UNet（甚至是DiT），导致耗时不可接受，因此减少采样步数是必不可少的。
一个简单的方法是，在训练时依旧使用1000步 t ∈ { 0 , 1 , 2 , . . . , 999 } t\in \{0, 1, 2, ..., 999\} t∈{0,1,2,...,999}，但在采样时只使用timestep的子集，比如 t ∈ { 249 , 499 , 749 , 999 } t\in \{249, 499, 749, 999\} t∈{249,499,749,999}，并索引训练时计算好的${\alpha }_t$列表，实验证明，这种跳步方法在50步左右依旧可以生成清晰的图像。DDIM将timestep的子序列定义为 { τ 1 , … , τ S } \{\tau_{1},\ldots,\tau_{S}\} {τ1​,…,τS​}，并一般设置为等间隔。此时生成过程为：
$$z=x_T\to x_{{\tau }_S}\to x_{{\tau }_{S-1}}\to \cdots \to x_{{\tau }_1}\to x_0=x$$
此外，DDIM将${\sigma }_t$定义为：
$${\sigma }_{{\tau }_i}(\eta )=\eta \sqrt{(1-{\alpha }_{{\tau }_{i-1}})/(1-{\alpha }_{{\tau }_i})}\sqrt{1-{\alpha }_{{\tau }_i}/{\alpha }_{{\tau }_{i-1}}}$$
其中$\eta \in [0,1]$，$\eta =0$时对应确定性的反向迭代过程，$\eta =1$对应DDPM。如果不特指$\eta$的值，DDIM默认采用确定性采样，一般采样步数在20~30步，并且一般质量比DDPM更好。

理论推导

反向过程的非马尔可夫推断

回顾DDPM的简化训练目标：
$$L={\mathbb{E}}_{t,x_0,ϵ}\left[{\Vert ϵ-{ϵ}_{\theta }^{(t)}\left(\sqrt{{\alpha }_t}{x}_0+\sqrt{1-{\alpha }_t}ϵ\right)\Vert }^2\right]$$
该损失函数关于 $t$的每一项，都只与$x_0$以及$x_t$相关，而且根据$x_0=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\sqrt{1-{\alpha }_t}ϵ\right)$，上式等价于每一步都估计干净图像$x_0$，只是在 $t$越大时估计越不准确。因此，只需要在采样到达 $t$步时$x_t$包含的噪声的方差与训练时正向过程添加的噪声方差一致，那么训练和推理就不会有明显的gap。对于正向过程，正向过程扩散到$x_{t-1}$时$p(x_{t-1}|x_0)=\mathcal{N}(x_{t-1};\sqrt{{\alpha }_{t-1}}{x}_0,(1-{\alpha }_{t-1})\mathbf{I})$，即方差为$1-{\alpha }_{t-1}$。因为我们在 $t$步估计出${\hat{x}}_0^{(t)}=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\sqrt{1-{\alpha }_t}{ϵ}_{\theta }^{(t)}({\hat{x}}_t)\right)$，迭代采样过程中会替代正向过程的$x_0$，因此我们只需要在${\hat{x}}_0^{(t)}$上添加方差为$1-{\alpha }_{t-1}$即可满足反向过程的状态转移，而无需考虑状态转移是否满足马尔科夫性。因此可以引入一个可调的辅助变量${\sigma }_t$：
$${ϵ}^{(t)}=\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\cdot {ϵ}_{\theta }^{(t)}(x_t)+{\sigma }_t{ϵ}_t$$
很明显，${\left(\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\right)}^2+{\sigma }_t^2=1-{\alpha }_{t-1}$，满足训练时噪声方差的设定。

边缘分布

DDIM原文并不是直接从噪声方差的训练-推断一致的角度进行推导的，而是从边缘分布$p_{\sigma }(x_t)$的一致性考虑。因为$x_0\sim p_{data}$不受参数$\sigma$影响，则只需证明参数$\sigma$不影响条件分布$p_{\sigma }(x_t|x_0)$，即：
$$p_{\sigma }(x_t|x_0)=\mathcal{N}(\sqrt{{\alpha }_t}{x}_0,(1-{\alpha }_t)I)$$
接下来开始证明：考虑DDPM的训练过程，其最小化$p(x_{t-1}|x_t,x_0)$和$q_{\theta }(x_{t-1}|x_t)$的KL散度，训练完成后可以使用$q_{\theta }(x_{t-1}|x_t)$替代$p(x_{t-1}|x_t,x_0)$。当前，我们的反向过程迭代定义为：
$$x_{t-1}=\sqrt{{\alpha }_{t-1}}\left(\frac{x_t-\sqrt{1-{\alpha }_t}{ϵ}_{\theta }^{(t)}(x_t)}{\sqrt{{\alpha }_t}}\right)+\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\cdot {ϵ}_{\theta }^{(t)}(x_t)+{\sigma }_t{ϵ}_t$$
其中${\hat{x}}_0=\frac{x_t-\sqrt{1-{\alpha }_t}{ϵ}_{\theta }^{(t)}(x_t)}{\sqrt{{\alpha }_t}}$可近似$x_0$，那么：
$$p_{\sigma }(x_{t-1}|x_t,x_0)=\mathcal{N}\left(\sqrt{{\alpha }_{t-1}}{x}_0+\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\cdot \frac{x_t-\sqrt{{\alpha }_t}{x}_0}{\sqrt{1-{\alpha }_t}},{\sigma }_t^2\mathbf{I}\right)$$
上式可表示为$$x_{t-1}=\sqrt{{\alpha }_{t-1}}{x}_0+\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\cdot \frac{x_t-\sqrt{{\alpha }_t}{x}_0}{\sqrt{1-{\alpha }_t}}+{\sigma }_t{ϵ}_t$$。考虑$$p_{\sigma }(x_t|x_0)=\mathcal{N}(\sqrt{{\alpha }_t}{x}_0,(1-{\alpha }_t)I)=\sqrt{{\alpha }_t}{x}_0+\sqrt{1-{\alpha }_t}ϵ,ϵ\sim \mathcal{N}(0,\mathbf{I})$$，则$p_{\sigma }(x_{t-1}|x_0)$的均值和方差：
$$\begin{array}{rl}{\mu }_{t-1}& =\sqrt{{\alpha }_{t-1}}{x}_0+\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\cdot \frac{\sqrt{{\alpha }_t}{x}_0-\sqrt{{\alpha }_t}{x}_0}{\sqrt{1-{\alpha }_t}}=\sqrt{{\alpha }_{t-1}}{x}_0\\ {\sigma }_{t-1}& ={\sigma }_t^2\mathbf{I}+\frac{1-{\alpha }_{t-1}-{\sigma }_t^2}{1-{\alpha }_t}(1-{\alpha }_t)\mathbf{I}=(1-{\alpha }_{t-1})\mathbf{I}\end{array}$$
即$$p_{\sigma }(x_{t-1}|x_0)=\mathcal{N}(\sqrt{{\alpha }_{t-1}}{x}_0,(1-{\alpha }_{t-1})I)$$，与$$p_{\sigma }(x_t|x_0)=\mathcal{N}(\sqrt{{\alpha }_t}{x}_0,(1-{\alpha }_t)I)$$一致。
因此，只要$q_{\theta }(x_{t-1}|x_t)$可以很好地近似$p(x_{t-1}|x_t,x_0)$，那么DDIM的转移过程与DDPM共享相同的边缘分布。DDIM原文使用与DDPM相似的变分推导来更严谨地导出了该结论，有兴趣可以阅读DDIM附录。

加速采样

当设定采样子序列为 { τ 1 , … , τ S } \{\tau_{1},\ldots,\tau_{S}\} {τ1​,…,τS​}时，对应的前向过程为：
$$p_{\tau }(x_{1:T}|x_0)=p_{\tau }(x_{{\tau }_S}|x_0)\prod _{i=1}^S{p}_{\tau }(x_{{\tau }_{i-1}}|x_{{\tau }_i},x_0)\prod _{t\in \bar{\tau }}{p}_{\tau }(x_t|x_0)$$
对应的反向过程为：
$$q_{\theta }(x_{0:T}):=q_{\theta }(x_T)\prod _{i=1}^S{q}_{\theta }^{({\tau }_i)}(x_{{\tau }_{i-1}}|x_{{\tau }_i})\times \prod _{t\in \bar{\tau }}{q}_{\theta }^{(t)}(x_0|x_t)$$
上式乘号的右半部分实际对应训练时的去噪损失，它与反向过程的迭代过程本身没有关系，当在完整的timestep序列上训练后，这部分得到良好定义。左半部分的单项则是一个与$\theta$无直接关联的项：
$$q_{\theta }^{({\tau }_i)}(x_{{\tau }_{i-1}}|x_{{\tau }_i})=p_{\tau }(x_{{\tau }_{i-1}}|x_{{\tau }_i},{\hat{x}}_0^{({\tau }_i)}(x_{{\tau }_{i-1}}))$$
因此只要在完整timestep上训练，那么在任意子序列上推断都是受到良好约束的。