Denoising diffusion implicit models 阅读笔记2

Denoising diffusion probabilistic models (DDPMs)从马尔科夫链中采样生成样本，需要迭代多次，速度较慢。Denoising diffusion implicit models (DDIMs)的提出是为了在复用DDPM训练的网络的前提下，加速采样过程。
加速采样的基本思路是，原本的生成过程是从$[T,\cdots ,1]$的序列逐步采样，加速时考虑从子序列$\{{\tau }_1,\dots ,{\tau }_S\},{\tau }_1>{\tau }_2>\cdots >{\tau }_S\in [1,T]$采样，通过跳步的方式减少采样的步数。比如DDPM网络原始训练包含1000步，但是采样时可以只从1000步中均匀的选出50步，用这50步采样出图像。
DDPM和DDIM都可以跳步采样，这个作者在实验中也进行了证明。DDIM的贡献主要是在复用DDPM训练的网络的前提下，提出了一种可以调节方差的生成形式，在步数较少的时候使用小的方差生成效果好。

DDIM论文中的符号和DDPM论文不同，本笔记中采用DDPM论文的符号。

引子

DDPM的优化目标是
$$\begin{array}{rl}{L}_{\text{VLB}}& ={\mathbb{E}}_{q({\mathbf{x}}_{0:T})}\left[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}\right]\\ & ={\mathbb{E}}_q[\underset{L_T}{D_{\text{KL}}(q({\mathbf{x}}_T|{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_T))}+\sum _{t=2}^T\underset{L_{t-1}}{D_{\text{KL}}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))}\underset{L_0}{-\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]\end{array}$$
其中$L_{t-1}={\mathbb{E}}_q[D_{\text{KL}}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))]$是DDPM网络优化的目标项。
注意$L_{t-1}$只和边际分布$q(x_t|x_0)$有关，而不是联合分布$q(x_{1:T}|x_0)$。所以，我们可以定义一个更为灵活的推理过程，只要它的边际分布$q(x_t|x_0)$和DDPM一致，就可以复用DDPM优化的网络。

非马尔科夫的前向过程

DDPM中推理分布$q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)$（推理分布是从可观测变量${\mathbf{x}}_0$推理隐变量${\mathbf{x}}_{1:T}$的分布）是固定的马尔科夫链（DDPM中要求$q(x_t|x_{t-1},x_0)=q(x_t|x_{t-1}):=\mathcal{N}(\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)$）。现在放宽限制，不要求前向过程是马尔科夫的，也就是对$q(x_t|x_{t-1})$不做形式要求。

作者定义由实向量$\sigma \in {\mathbb{R}}_{\ge 0}^T$索引的推理分布族Q：
$$\begin{array}{r}{q}_{\sigma }(x_{1:T}|x_0):=q_{\sigma }(x_T|x_0)\prod _{t=2}^T{q}_{\sigma }(x_{t-1}|x_t,x_0)\end{array}$$只要求满足边际分布和DDPM一样，即$q_{\sigma }(x_t|x_0)=\mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I)$。
通过待定系数法（参考[1]）可以得到形式更自由的生成过程：
$$q_{\sigma }(x_{t-1}|x_t,x_0):=\mathcal{N}(\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}\cdot \frac{x_t-\sqrt{{\bar{\alpha }}_t}{x}_0}{\sqrt{1-{\bar{\alpha }}_t}},{\sigma }_t^{2}I)$$
对应的前向过程也是高斯分布，但前向过程变成了非马尔科夫的，因为每一步都依赖${\mathbf{x}}_0$
$$q_{\sigma }(x_t|x_{t-1},x_0)=\frac{q_{\sigma }(x_{t-1}|x_t,x_0)q_{\sigma }(x_t|x_0)}{q_{\sigma }(x_{t-1}|x_0)}$$
如下图所示，DDPM的推理过程是非马尔科夫的。

注意，DDIM构造的推理分布和DDPM不同，但和DDPM优化相同的优化目标。

反向生成过程

根据上面的推理过程，定义需要学习的生成过程$p_{\theta }({\mathbf{x}}_{0:T})$，该过程利用$q_{\sigma }(x_{t-1}|x_t,x_0)$。
直观地说，给定$x_t$，我们首先预测对应的$x_0$，然后使用我们定义的反向条件分布$q_{\sigma }(x_{t-1}|x_t,x_0)$获得$x_{t-1}$。
预测对应的$x_0$如下：
$${\hat{x}}_0=f_{\theta }^{(t)}(x_t):=\frac{x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }^{(t)}(x_t)}{\sqrt{{\bar{\alpha }}_t}}$$使用预测的${\hat{x}}_0$通过$q_{\sigma }(x_{t-1}|x_t,x_0)$获得$x_{t-1}$如下：
$${\hat{x}}_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}{\hat{x}}_0+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}\cdot \frac{x_t-\sqrt{{\bar{\alpha }}_t}{\hat{x}}_0}{\sqrt{1-{\bar{\alpha }}_t}}+{\sigma }_{t}z,z\sim \mathcal{N}(0,I)$$写成使用${ϵ}_{\theta }^{(t)}(x_t)$的形式：
$${\hat{x}}_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}\frac{x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }^{(t)}(x_t)}{\sqrt{{\bar{\alpha }}_t}}+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}\cdot {ϵ}_{\theta }^{(t)}(x_t)+{\sigma }_{t}z,z\sim \mathcal{N}(0,I)$$
选择不同的$\sigma$值会导致不同的生成过程，但它们使用相同的${ϵ}_{\theta }$模型。

DDPM和DDIM对比

$$\begin{array}{rl}\text{DDPM}:& q(x_{t-1}|x_t,x_0)=\mathcal{N}(\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0+\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t,\frac{{\beta }_t(1-{\bar{\alpha }}_{t-1})}{(1-{\bar{\alpha }}_t)}I)\\ \text{DDIM}:& q_{\sigma }(x_{t-1}|x_t,x_0):=\mathcal{N}(\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}\cdot \frac{x_t-\sqrt{{\bar{\alpha }}_t}{x}_0}{\sqrt{1-{\bar{\alpha }}_t}},{\sigma }_t^{2}I)\end{array}$$当${\sigma }_t=\sqrt{(1-{\bar{\alpha }}_{t-1})/(1-{\bar{\alpha }}_t)}\sqrt{(1-{\bar{\alpha }}_t/{\bar{\alpha }}_{t-1})}=\sqrt{\frac{{\beta }_t(1-{\bar{\alpha }}_{t-1})}{(1-{\bar{\alpha }}_t)}}$时，DDPM的$q(x_{t-1}|x_t,x_0)$和DDIM的$q_{\sigma }(x_{t-1}|x_t,x_0)$是一样的，生成过程变成和DDPM是一样的。另外，DDIM的前向过程变成了马尔科夫的。

当${\sigma }_t=0$时，随机噪声前的系数是0，${\mathbf{x}}_0$和${\mathbf{x}}_T$之间的关系是固定的，这属于隐概率模型（implicit probabilistic model）。
这时生成过程的每一步变为$$\begin{gathered} {\hat{x}}_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}\frac{x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }^{(t)}(x_t)}{\sqrt{{\bar{\alpha }}_t}}+\sqrt{1-{\bar{\alpha }}_{t-1}}\cdot {ϵ}_{\theta }^{(t)}(x_t) \\ =\frac{1}{\sqrt{{\alpha }_t}}(x_t-(\sqrt{1-{\bar{\alpha }}_t}-\sqrt{{\alpha }_t}\sqrt{1-{\bar{\alpha }}_{t-1}}){ϵ}_{\theta }^{(t)}(x_t)) \end{gathered}$$对比DDPM生成过程的每一步
$${\hat{x}}_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }^{(t)}(x_t))+{\sigma }_{t}z$$

[2]中给出了下面的分解：
$$\begin{array}{rl}{\mathbf{x}}_{t-1}& =\sqrt{{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_{t-1}}{\mathbf{ϵ}}_{t-1}\\ & =\sqrt{{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{\mathbf{ϵ}}_t+{\sigma }_t\mathbf{ϵ}\\ & =\sqrt{{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}\frac{{\mathbf{x}}_t-\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0}{\sqrt{1-{\bar{\alpha }}_t}}+{\sigma }_t\mathbf{ϵ}\end{array}$$分解的依据是不相关高斯随机变量$\mathcal{N}(0,{\sigma }_1^2\mathbf{I})$和$\mathcal{N}(0,{\sigma }_2^2\mathbf{I})$之和的分布是$\mathcal{N}(0,({\sigma }_1^2+{\sigma }_2^2)\mathbf{I})$。

在DDPM的逆向生成过程的每一步中，也可以认为是先估计${\hat{x}}_0$，再求${\hat{x}}_{t-1}$。
DDIM与DDPM的主要区别是DDIM构造了一种更自由的过程，通过$\sigma$改变了方差的大小，同时改变均值，使其依然符合DDPM的边际分布$q(x_t|x_0)$。

参考资料

[1] https://kxz18.github.io/2022/06/21/DDIM/
[2] https://lilianweng.github.io/posts/2021-07-11-diffusion-models/