扩散模型Diffusion model | DDIM

扩散模型Diffusion model | DDIM

论文原文：Denoising Diffusion Implicit Models

有关DDPM的解释可以参考我的上一篇博客：扩散模型Diffusion model | DDPM

DDPM的贝叶斯解释

直接根据贝叶斯定理我们有
$$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\frac{p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})p({\mathbf{x}}_{t-1})}{p({\mathbf{x}}_t)}$$
但是$p(x_{t-1}),p(x_t)$难以直接计算，因而转向计算
$$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\frac{p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})p({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{p({\mathbf{x}}_t|{\mathbf{x}}_0)}$$
代入各自的表达式得到：（本文中与原文定义不同的是：${\beta }_t=\sqrt{1-{\alpha }_t^2}$，原文是${\alpha }_t=1-{\beta }_t$）
$$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}\left({\mathbf{x}}_{t-1};\frac{{\alpha }_t{\bar{\beta }}_{t-1}^2}{{\bar{\beta }}_t^2}{\mathbf{x}}_t+\frac{{\bar{\alpha }}_{t-1}{\beta }_t^2}{{\bar{\beta }}_t^2}{\mathbf{x}}_0,\frac{{\bar{\beta }}_{t-1}^2{\beta }_t^2}{{\bar{\beta }}_t^2}\mathbf{I}\right)$$

用$\bar{\mathbf{\mu }}({\mathbf{x}}_t)$来预估${\mathbf{x}}_0$损失$\Vert {\mathbf{x}}_0-\bar{\mathbf{\mu }}({\mathbf{x}}_t){\Vert }^2$，就可以消去$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$中的${\mathbf{x}}_0$，使得它只依赖于${\mathbf{x}}_t$了。

实际上这也就是个去噪的过程，对应DDPM的Denoising过程。

而由$p({\mathbf{x}}_t|{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_t;{\bar{\alpha }}_t{\mathbf{x}}_0,{\bar{\beta }}_t^2\mathbf{I})$可以推出${\mathbf{x}}_0=\frac{1}{{\bar{\alpha }}_t}\left({\mathbf{x}}_t-{\bar{\beta }}_t\mathbf{\epsilon }\right)$，于是我们可以构造
$$\bar{\mathbf{\mu }}(x_t)=\frac{1}{{\bar{\alpha }}_t}({\mathbf{x}}_t-{\bar{\beta }}_t{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t))$$

代回即可得到DDPM的损失函数。

用$\bar{\mathbf{\mu }}(x_t)$来预估$x_0$不会太准，它仅仅起到了一个前瞻性的预估作用，然后只用$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$来推进一小步（类似于梯度下降中，找到最速下降方向，但是仅仅下降一小步的预估-修正思想）。

DDIM

DDPM的推导思路为：
$$x_t\stackrel{model}{\to }{ϵ}_{\theta }(x_t,t)\stackrel{P(x_t|x_0)\to P(x_0|x_t,{ϵ}_{\theta })}{\to }{\hat{x}}_0(x_t,{ϵ}_{\theta })\stackrel{\text{推导}}{\to }\mu (x_t,{\hat{x}}_0),{\beta }_t\stackrel{P(x_{t-1}|x_t,x_0)}{\to }{\hat{x}}_{t-1}$$
在上述推导中，可以看到

1. 损失函数只依赖于$p({\mathbf{x}}_t|{\mathbf{x}}_0)$；
2. 采样过程只依赖于$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$。

因此由于对马尔可夫假设的依赖，导致DDPM的重建过程需要较多的步长。

但是实际上推理与$p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})$好像并没有关系（仅仅是马尔可夫的约束），如果想要加速这个重建过程，可以考虑解除对前向过程$p({\mathbf{x}}_t|{\mathbf{x}}_{t-1})$马尔可夫特性的依赖，直接定义分布$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$，原来的前向过程即可变化为：
$$q_{\sigma }({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)=q_{\sigma }({\mathbf{x}}_T|{\mathbf{x}}_0)\prod _{t=2}^T{q}_{\sigma }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$$

但为了保持前向过程与DDPM等价，需要满足边缘分布条件
$$\int p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)p({\mathbf{x}}_t|{\mathbf{x}}_0)d{\mathbf{x}}_t=p({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)$$

使用待定系数法，即可得到（具体参见苏老师的博客：DDIM = 高观点DDPM）
$$p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}\left({\mathbf{x}}_{t-1};\frac{\sqrt{{\bar{\beta }}_{t-1}^2-{\sigma }_t^2}}{{\bar{\beta }}_t}{\mathbf{x}}_t+\left({\bar{\alpha }}_{t-1}-\frac{{\bar{\alpha }}_t\sqrt{{\beta }_{t-1}^2-{\sigma }_t^2}}{{\bar{\beta }}_t}\right){\mathbf{x}}_0,{\sigma }_t^2\mathbf{I}\right)$$

那么接下来的过程就与DDPM相同了…

$$\begin{array}{rl}p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)& \approx p({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0=\bar{\mathbf{\mu }}({\mathbf{x}}_t))\\ & =\mathcal{N}\left({\mathbf{x}}_{t-1};\frac{1}{{\alpha }_t}\left({\mathbf{x}}_t-\left({\bar{\beta }}_t-{\alpha }_t\sqrt{{\bar{\beta }}_{t-1}^2-{\sigma }_t^2}\right){\mathbf{ϵ}}_{\mathbf{\theta }}({\mathbf{x}}_t,t)\right),{\sigma }_t^2\mathbf{I}\right)\end{array}$$

此式带有一个自由参数${\sigma }_t$，和DDPM相比训练过程没有变化，但生成过程却有一个可变动的参数${\sigma }_t$，不同${\sigma }_t$的采样过程会呈现出不同的特点。

而当${\sigma }_t=\frac{{\bar{\beta }}_{t-1}{\beta }_t}{{\bar{\beta }}_t}$时，DDIM与DDPM等价；当${\sigma }_t=0$时，此时从$x_t$到$x_{t-1}$是一个确定性变换，也就是论文特指的部分。

DDPM的训练结果实质上包含了它的任意子序列参数的训练结果。

原文推导

如果将条件改回${\alpha }_t=1-{\beta }_t$

则$p({\mathbf{x}}_t|{\mathbf{x}}_0)$可以形式化为：
$$x_0=\frac{{\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }^{(t)}\left({\mathbf{x}}_t\right)}{\sqrt{{\bar{\alpha }}_t}}$$

而反向条件概率$q_{\sigma }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)$为：
$$q_{\sigma }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)=\mathcal{N}\left({\mathbf{x}}_{t-1};\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^2\mathbf{I}\right)$$

采样过程即为：
$$x_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}\underset{\text{predicied }{x}_0}{\underbrace{\left(\frac{x_t-\sqrt{1-{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_t}}\frac{{ϵ}_{\theta }^{(t)}(x_t)}{\sqrt{{\bar{\alpha }}_t}}\right)}}+\underset{\text{direction pointing to }{x}_t}{\underbrace{\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}\cdot {ϵ}_{\theta }^{(t)}(x_t)}}+\underset{\text{random noise}}{\underbrace{{\sigma }_t\cdot {ϵ}_t}}$$

而可以将${\sigma }_t$进一步定义为：
$${\sigma }_t=\eta \sqrt{(1-{\bar{\alpha }}_{t-1})/(1-{\bar{\alpha }}_t)}\sqrt{1-{\bar{\alpha }}_t/{\bar{\alpha }}_{t-1}}$$

如果$\eta =0$，生成过程就没有随机噪音了，是一个确定性的过程，论文将这种情况下的模型称为DDIM（denoising diffusion implicit model）；而如果$\eta =1$，该前向过程变成了马尔科夫链，模型为DDPM。

重建与插值

对于$\eta =0$时的DDIM，实质上就是将任意正态噪声向量变换为图片的一个确定性变换，此时$x_t$生成$x_{t-1}$的更新公式就变为：
$${\mathbf{x}}_{t-1}=\sqrt{{\alpha }_{t-1}}\left(\frac{{\mathbf{x}}_t-\sqrt{1-{\alpha }_t}{ϵ}_{\theta }({\mathbf{x}}_t,t)}{\sqrt{{\alpha }_t}}\right)+\sqrt{1-{\alpha }_{t-1}}\cdot {ϵ}_{\theta }({\mathbf{x}}_t,t)$$
对上式作等价变换可以得到：
$$\frac{{\mathbf{x}}_{t-1}}{\sqrt{{\alpha }_{t-1}}}=\frac{{\mathbf{x}}_t}{\sqrt{{\alpha }_t}}+(\sqrt{\frac{1-{\alpha }_{t-1}}{{\alpha }_{t-1}}}-\sqrt{\frac{1-{\alpha }_t}{{\alpha }_t}}){ϵ}_{\theta }({\mathbf{x}}_t,t)$$

当$T$足够大，或者说${\alpha }_t$与${\alpha }_{t-1}$足够小时，我们可以将上式视为某个常微分方程ODE的差分形式。
$$\frac{{\mathbf{x}}_{t-\Delta t}}{\sqrt{{\alpha }_{t-\Delta t}}}=\frac{{\mathbf{x}}_t}{\sqrt{{\alpha }_t}}+(\sqrt{\frac{1-{\alpha }_{t-\Delta t}}{{\alpha }_{t-\Delta t}}}-\sqrt{\frac{1-{\alpha }_t}{{\alpha }_t}}){ϵ}_{\theta }({\mathbf{x}}_t,t)$$
这里令$\sigma =\sqrt{1-\alpha }/\sqrt{\alpha },\bar{\mathbf{x}}=\mathbf{x}/\sqrt{\alpha }$，它们都是关于$t$的函数，这样对应的ODE就是：
$$\mathrm{d}\bar{\mathbf{x}}(t)={ϵ}_{\theta }(\frac{\bar{\mathbf{x}}(t)}{\sqrt{{\sigma }^2+1}},t)\mathrm{d}\sigma (t)$$
那么可以由一个原始图像$x_0$得到对应的随机噪音$x_T$，然后我们再用$x_T$进行生成就可以重建原始图像$x_0$，可以得到较低的重建误差
$$\frac{{\mathbf{x}}_{t+1}}{\sqrt{{\alpha }_{t+1}}}=\frac{{\mathbf{x}}_t}{\sqrt{{\alpha }_t}}+(\sqrt{\frac{1-{\alpha }_{t+1}}{{\alpha }_{t+1}}}-\sqrt{\frac{1-{\alpha }_t}{{\alpha }_t}}){ϵ}_{\theta }({\mathbf{x}}_t,t)$$

也就是说，将生成过程等同于求解常微分方程后，可以借助常微分方程的数值解法，为生成过程的加速提供更丰富多样的手段。

$\eta =0$时，DDIM，所以跟GAN类似，我们可以对这两个随机噪音进行插值生成新的$x_T$，那么将生成融合的图像。这里采用的插值方法是球面线性插值（ spherical linear interpolation）,参数$\alpha$控制插值系数：
$${\mathbf{x}}_T^{(\alpha )}=\frac{\sin ((1-\alpha )\theta )}{\sin (\theta )}{\mathbf{x}}_T^{(0)}+\frac{\sin (\alpha \theta )}{\sin (\theta )}{\mathbf{x}}_T^{(1)}\theta =\arccos \left(\frac{({\mathbf{x}}_T^{(0)}{)}^{\mathrm{T}}{\mathbf{x}}_T^{(1)}}{\Vert {\mathbf{x}}_T^{(0)}\Vert \Vert {\mathbf{x}}_T^{(1)}\Vert }\right)$$