Part3: Dive into DDPM

背景

整个系列有相对完整的公式推导，若正文中有涉及到的省略部分，皆额外整理在Part4，并会在正文中会指明具体位置。

• Part1: Overview of Diffusion Process
• Part2: DDPM as Example of Variational Inference
• Part3: Dive into DDPM
• Part4: Appendix

在Part2基于$\text{Variational Inference}$，找到原目标函数$-\ln p_{\theta }(x_0)$的上界$L$，定义如下：

$$\begin{array}{}\text{(1)}& \begin{array}{rl}L:=& {\mathbb{E}}_q[-\log \frac{p\left({\mathbf{x}}_T\right)}{q({\mathbf{x}}_T\mid {\mathbf{x}}_0)}-\sum _{t>1}\log \frac{p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)}{q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)}-\log p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)]\\ =& {\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{\mathrm{KL}}(q({\mathbf{x}}_T\mid {\mathbf{x}}_0)\Vert p\left({\mathbf{x}}_T\right))}}+\sum _{t>1}\underset{L_{t-1}}{\underbrace{D_{\mathrm{KL}}(q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t))}}\underset{L_0}{\underbrace{-\log p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}}]\end{array}\end{array}$$

沿着论文的思路对$L$继续精简，得到最终在代码层面实现的损失函数$L_{simple}$。同样的，补充的推导见Part4；“扩散过程”的梗概介绍见Part1。

简化过程

不难看出$L$中的每一项皆为KL散度。回顾forward process与reverse process两个阶段的定义，马尔可夫链的状态转移皆服从高斯分布，如下所示：

$$\begin{array}{}\text{(2)}& \begin{array}{rl}q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})& :=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})\\ p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)& :=\mathcal{N}({\mathbf{x}}_{t-1};{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))\end{array}\end{array}$$

同时，经过推导（见Part4推导二），易知：

$$\begin{array}{}\text{(3)}& \begin{array}{rl}q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)& =\mathcal{N}({\mathbf{x}}_{t-1};{\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0),{\tilde{\beta }}_t\mathbf{I})\\ \text{ where }{\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)& :=\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}{\mathbf{x}}_0+\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{\mathbf{x}}_t,{\tilde{\beta }}_t:=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}{\beta }_t\end{array}\end{array}$$

故KL散度比较皆发生在两个Gaussian间。

$L_t$的简化

可以看到，$(1)$式中的$L_T$代表前向扩散过程，与待求解的参数项$\theta$无关，因此可被忽略：

$$\arg \underset{\theta }{min}(L)⟺\arg \underset{\theta }{min}\left({\mathbb{E}}_q\left[\sum _{t>1}\underset{L_{t-1}}{\underbrace{D_{\mathrm{KL}}(q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t))}}\underset{L_0}{\underbrace{-\log p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}}\right]\right)$$

注：在2015年提出diffusion框架的论文，前向扩散过程中的${\beta }_t$是可以被学习的参数，故此处可视作DDPM的第一处简化。

$L_{t-1}$的简化

对于反向扩散过程的分布$p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$，共涉及到两组参数${\mathit{\mu }}_{\theta }$与${\mathbf{\Sigma }}_{\theta }$，DDPM的第二处简化是定义$\mathbf{\Sigma }$为常数${\sigma }_t^2$，在计算中使用${\beta }_t$或${\tilde{\beta }}_t$代替，故$p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t),{\sigma }_t^2\mathbf{I})$

基于Part4两个高斯的KL散度，对于$L_{t-1}$，有：

$$\begin{array}{}\text{(4)}& L_{t-1}={\mathbb{E}}_q\left[\frac{1}{2{\sigma }_t^2}{\Vert {\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)-{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)\Vert }^2\right]+C\end{array}$$

其中$C$是个常数项。

仔细观察$(4)$不难发现，想要目标函数最小化，则${\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)$与${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)$间的“距离必须要近”。也就是说，深度网络通过训练，使得${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)$趋近于${\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)$。为了使训练更加简单，尝试对$(4)$式改写。

在Foward Process中，${\mathbf{x}}_t$可由${\mathbf{x}}_0$与$ϵ$表示（见Part4推导一），不妨将${\mathbf{x}}_t$记作${\mathbf{x}}_t({\mathbf{x}}_0,ϵ)$，故${\mathbf{x}}_0$可以展开表示为${\mathbf{x}}_t({\mathbf{x}}_0,ϵ)$与$ϵ$的差：

$$\begin{array}{}\text{(5)}& {\mathbf{x}}_t=\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t}ϵ\Rightarrow {\mathbf{x}}_0=\frac{1}{\sqrt{{\overline{\alpha }}_t}}({\mathbf{x}}_t-\sqrt{1-{\overline{\alpha }}_t}ϵ)\end{array}$$

又因为$(3)$式，故有：

$$\begin{array}{}\text{(6-1)}& \begin{array}{rl}{\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)& =\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}{\mathbf{x}}_0+\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{\mathbf{x}}_t\\ & =\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}\ast \frac{{\mathbf{x}}_t-\sqrt{1-{\overline{\alpha }}_t}ϵ}{\sqrt{{\overline{\alpha }}_t}}+\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{\mathbf{x}}_t\\ & =\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}ϵ)\end{array}\end{array}$$

前文提到，要优化$(4)$式，则必然有：${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)\to {\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)$。其中，${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)$是深度网络的输出（预测）结果，${\mathbf{x}}_t$与$t$作为模型的输入参数。

由$\text{(6-1)}$可知${\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)$能展开为${\mathbf{x}}_t$与$ϵ$的表达，${\mathbf{x}}_t$已知，那不妨令原本要预测${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)$的深度网络直接预测${\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t)$，变换前后依然等价。即

$$\begin{array}{}\text{(6-2)}& \begin{array}{r}{\mathit{\mu }}_{{\theta }^{\prime }}({\mathbf{x}}_t,t)⟺\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t))\end{array}\end{array}$$

此处以${\theta }^{\prime }$与$\theta$对变换前后的深度网络参数进行区分，故$\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t))$需要无限趋近于${\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)$。

将$\text{(6-1)}$与$\text{(6-2)}$代入$(4)$式，有：

$$\begin{array}{}\text{(7)}& \begin{array}{rl}& L_{t-1}-C^{\prime }\\ =& {\mathbb{E}}_q\left[\frac{1}{2{\sigma }_t^2}{\Vert {\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)-{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)\Vert }^2\right]\\ ⟺& {\mathbb{E}}_{{\mathbf{x}}_0,\mathit{ϵ}}\left[\frac{1}{2{\sigma }_t^2}{\Vert \frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t({\mathbf{x}}_0,\mathit{ϵ})-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}\mathit{ϵ})-\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t({\mathbf{x}}_0,\mathit{ϵ})-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\mathit{ϵ}}_{\mathit{\theta }}({\mathbf{x}}_t,t))\Vert }^2\right]\\ =& {\mathbb{E}}_{{\mathbf{x}}_0,\mathit{ϵ}}\left[\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\overline{\alpha }}_t)}{\Vert \mathit{ϵ}-{\mathit{ϵ}}_{\theta }(\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t}\mathit{ϵ},t)\Vert }^2\right]\end{array}\end{array}$$

对比$(4)$与$(7)$，不难发现参数$\theta$作用的对象发生变化。在$(4)$中，$\theta$的参数化对象为高斯分布的均值$\mathit{\mu }$；而在$(7)$中，$\theta$的参数化对象转移到$\mathit{ϵ}$。实际上，不仅可以参数化$\mathit{\mu }$和$\mathit{ϵ}$，也可以参数化${\mathbf{x}}_0$，只需要对$(5)$中表示的主体进行变换即可。

并且，重新审视$\text{(6-2)}$，该式与Part1中的采样算法联系上了。上述目标函数的设定及推理，皆是为了获取反向过程的分布$p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t):=\mathcal{N}({\mathbf{x}}_{t-1};{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))$。

通过公式$\text{(6-2)}$，按照反向过程相邻状态间的图像转换服从高斯分布的定义，反向过程中知晓${\mathbf{x}}_t$与$t$后，通过深度网络预测出$\mathit{ϵ}$，再基于此求出${\mathit{\mu }}_t$，结合自定义的${\mathit{\sigma }}_t$，可采样得到${\mathbf{x}}_{t-1}$，便实现反向过程的一次“降噪”。

$L_0$的简化

这一项对应着信息由隐变量转变回${\mathbf{x}}_{\mathbf{0}}$，故而需要特殊考虑。

真实图片中各个像素由0到255的数值组成，在处理时通常将所有像素值归一化到区间[-1,1]。论文中将该项对应的优化目标定义为：

$$\begin{array}{}\text{(8)}& \begin{array}{rl}{p}_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)& =\prod _{i=1}^D{\int }_{{\delta }_{-}\left(x_0^i\right)}^{{\delta }_{+}\left(x_0^i\right)}\mathcal{N}(x;{\mu }_{\theta }^i({\mathbf{x}}_1,1),{\sigma }_1^2)dx\\ {\delta }_{+}(x)& =\{\begin{array}{ll}\mathrm{\infty }& \text{ if }x=1\\ x+\frac{1}{255}& \text{ if }x<1\end{array}{\delta }_{-}(x)=\{\begin{array}{ll}-\mathrm{\infty }& \text{ if }x=-1\\ x-\frac{1}{255}& \text{ if }x>-1\end{array}\end{array}\end{array}$$

其中，积分项是为了与图片真实像素的离散特性保持一致，$D$为像素点的个数。

该项的优化目标是：对于输入图片$x_0$的所有像素位置，使得基于神经网络产生的高斯分布在该位置的采样结果，与$x_0$对应位置的真实值相差不大。

直接文字阐述并不好理解，下方是对于单个位置的具体实例，截图来自视频：

当前有一张真实的图片$x_0$，对应上图内靠左边的图片，经过缩放后，在位置$i$的值为$x_0^i=\frac{10}{255}$；

并且，中间图片表示在$x_1^i$（此时还处于有噪声状态），经过神经网络模型，预测出该位置的值是服从均值为$\frac{11}{255}$的高斯分布${\mathcal{N}}^{\mathcal{1}}$；

在左下角画出该${\mathcal{N}}^{\mathcal{1}}$的概率密度曲线，此时积分的上下界为$(\frac{9}{255},\frac{11}{255})$，从图上可以直观地看出积分对应的阴影面积相对来说比较大。故基于此采样得到的${\widehat{x_0}}^i$与输入图片$x_0^i$接近的置信度很高。在训练时反映出来的是，神经网络在该位置的预测表现对$(8)$式即Loss的贡献程度较低；

但如果神经网络预测出该位置的值服从服从均值为$\frac{105}{255}$的高斯分布${\mathcal{N}}^{\mathcal{2}}$，此时概率密度曲线整体会往右平移，$(\frac{9}{255},\frac{11}{255})$区域属于长尾位置，显然积分结果比较小，从侧面来说，基于此采样得到的${\widehat{x_0}}^{i_{\prime }}$与输入图片$x_0^i$接近的置信度很低，在训练时对Loss的贡献程度高，在反向传播时的梯度也大。

实际代码实现中，该项被省略，这是第三处简化。

简化的损失函数

回顾$(1)$式，目前只剩下以$L_{t-1}$为主体的求和部分，如下所示：

$$\begin{array}{}\text{(9)}& \begin{array}{rl}\arg \underset{\theta }{min}(L)& ⟺\arg \underset{\theta }{min}{\mathbb{E}}_q\left[\sum _{t>1}\underset{L_{t-1}}{\underbrace{D_{\mathrm{KL}}(q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t))}}\right]\\ & ⟺\arg \underset{\theta }{min}{\mathbb{E}}_{{\mathbf{x}}_0,\mathit{ϵ}}\left[\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\overline{\alpha }}_t)}{\Vert \mathit{ϵ}-{\mathit{ϵ}}_{\theta }(\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t}\mathit{ϵ},t)\Vert }^2\right]\end{array}\end{array}$$

对于$(9)$式，DDPM的第四处简化在于省略了均方差损失项的权重，最终的损失函数$L_{simple}$为：

$$\begin{array}{r}{L}_{\text{simple }}(\theta ):={\mathbb{E}}_{t,{\mathbf{x}}_0,\mathit{ϵ}}\left[{\Vert \mathit{ϵ}-{\mathit{ϵ}}_{\theta }(\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t}\mathit{ϵ},t)\Vert }^2\right]\end{array}$$

总结

回顾本文，DDPM在损失函数上做了很多简化，对于代码侧的实现非常友好。同时，论文作者也给出实验对比，验证简化并不会使得结果变差，有些简化（比如设置reverse过程中的$\mathrm{\Sigma }$为非参数项）甚至取得大幅度的提升效果。

Reference

• Denoising Diffusion Probabilistic Models