Diffusion Model（3）：训练目标以及训练过程

观看本文之前建议先观看以下两篇文章：

• Diffusion Model（1）：预备知识
• Diffusion Model（2）：前向扩散过程和逆向降噪过程

Training Loss

训练目标

​ 首先回顾一下我们的问题，我们在逆向降噪过程中由于没办法得到$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$，因此我们定义了一个 需要学习的模型模型 $p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$来对其进行近似，并且在训练阶段我们可以利用后验$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)$来对$p_{\theta }$进行优化。

​ 那么现在的问题是我们如何$p_{\theta }$优化得到理想的${\mathbf{\mu }}_{\theta }$和${\mathbf{\Sigma }}_{\theta }$？类似于 VAE ，我们可以最小化在真实数据期望下，模型预测分布的负对数似然，即最小化预测$p_{\mathrm{data}}=q({\mathbf{x}}_0)$和$p_{\theta }({\mathbf{x}}_0)$的交叉熵：
$$\begin{array}{c}\mathcal{L}={\mathbb{E}}_{{\mathbf{x}}_0\sim q\left({\mathbf{x}}_0\right)}\left[-\log p_{\theta }\left({\mathbf{x}}_0\right)\right]\end{array}$$
​ 但是，我们没法得到$p_{\theta }({\mathbf{x}}_0)$的表达式，因此公式1的交叉熵是没法计算的。那么可以借助公式Diffusion Model（2）：前向扩散过程和逆向降噪过程
2-6 进行一些数学推导。将公式1中的$p_{\theta }({\mathbf{x}}_0)$转化为已知的项：
$$\begin{array}{c}\begin{array}{rl}\mathcal{L}& =-{\mathbb{E}}_{q\left({\mathbf{x}}_0\right)}\log p_{\theta }\left({\mathbf{x}}_0\right)\\ & =-{\mathbb{E}}_{q\left({\mathbf{x}}_0\right)}\log \left(\int p_{\theta }\left({\mathbf{x}}_{0:T}\right)d{\mathbf{x}}_{1:T}\right)\\ & =-{\mathbb{E}}_{q\left({\mathbf{x}}_0\right)}\log \left(\int q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right)\frac{p_{\theta }\left({\mathbf{x}}_{0:T}\right)}{q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right)}d{\mathbf{x}}_{1:T}\right)\\ & =-{\mathbb{E}}_{q\left({\mathbf{x}}_0\right)}\log \left({\mathbb{E}}_{q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right)}\frac{p_{\theta }\left({\mathbf{x}}_{0:T}\right)}{q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right)}\right)\\ & \le -{\mathbb{E}}_{q\left({\mathbf{x}}_{0:T}\right)}\log \frac{p_{\theta }\left({\mathbf{x}}_{0:T}\right)}{q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right)}\\ & ={\mathbb{E}}_{q\left({\mathbf{x}}_{0:T}\right)}\left[\log \frac{q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{0:T}\right)}\right]={\mathcal{L}}_{\mathrm{VLB}}\end{array}\end{array}$$
​ 上式中$q({\mathbf{x}}_0)$是真实的数据分布，而$p_{\theta }({\mathbf{x}}_0)$是模型，从第四行到第五行使用了Jensen不等式$\log \mathbb{E}[f(x)]\le \mathbb{E}[\log f(x)]$并结合了对$q({\mathbf{x}}_0)$的期望和对$q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)$的期望。

​ 为了最小化这个损失，结合公式2可以将其转化为最小化其上界${\mathcal{L}}_{\mathrm{VLB}}$：
$$\begin{array}{c}\begin{array}{l}{\mathcal{L}}_{\mathrm{VLB}}={\mathbb{E}}_{q\left({\mathbf{x}}_{0:T}\right)}\left[\log \frac{q\left({\mathbf{x}}_{1:T}|{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{0:T}\right)}\right]\\ ={\mathbb{E}}_q\left[\log \frac{{\prod }_{t=1}^{T}q\left({\mathbf{x}}_t|{\mathbf{x}}_{t-1}\right)}{p_{\theta }\left({\mathbf{x}}_T\right){\prod }_{t=1}^T{p}_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}\right]\\ ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+{\sum }_{t=1}^T\log \frac{q\left({\mathbf{x}}_t|{\mathbf{x}}_{t-1}\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}\right]\\ ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+{\sum }_{t=2}^T\log \frac{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}+\log \frac{q\left({\mathbf{x}}_1|{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)}\right]\\ ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+{\sum }_{t=2}^T\log \left(\frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}\cdot \frac{q({\mathbf{x}}_t|{\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}\right)+\log \frac{q\left({\mathbf{x}}_1|{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)}\right]\\ ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+{\sum }_{t=2}^T\log \frac{q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}+{\sum }_{t=2}^T\log \frac{q\left({\mathbf{x}}_t|{\mathbf{x}}_0\right)}{q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_0\right)}+\log \frac{q\left({\mathbf{x}}_1|{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)}\right]\\ ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+{\sum }_{t=2}^T\log \frac{q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}+\log \frac{q\left({\mathbf{x}}_T|{\mathbf{x}}_0\right)}{q\left({\mathbf{x}}_1|{\mathbf{x}}_0\right)}+\log \frac{q\left({\mathbf{x}}_1|{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)}\right]\\ ={\mathbb{E}}_q\left[\log \frac{q\left({\mathbf{x}}_T|{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_T\right)}+{\sum }_{t=2}^T\log \frac{q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)}-\log p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)\right]\\ ={\mathbb{E}}_q[\underset{L_0}{\underbrace{-\log p_{\theta }\left({\mathbf{x}}_0|{\mathbf{x}}_1\right)}}]+{\sum }_{t=2}^T\underset{L_{t-1}}{\underbrace{D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right)\right)}}+\underset{L_T}{\underbrace{D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_T|{\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_T\right)\right)}}\end{array}\end{array}$$
​ 上述式子中：

• 从第三行到第四行，是将t=1的情况与总的求和拆开
• 从第四行到第五行，使用了前向过程的马尔科夫过程结合贝叶斯公式Diffusion Model（1）：预备知识1-8

$$\begin{array}{c}\begin{array}{rl}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})& =q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)=\frac{q({\mathbf{x}}_t,{\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}=q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\cdot \frac{q({\mathbf{x}}_t|{\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}\end{array}\end{array}$$

• 从第六行到第七行，将对数和转化为了乘积的形式，然后消去相同的分子和分母
• 第七行到第八行，首先将最后两项的求和转化为乘积然后消除$q({\mathbf{x}}_1|{\mathbf{x}}_0)$，然后将第一项放在分母上，将原来的分母变为减法
• 从第八行到第九行，使用了KL散度的公式 $D_{KL}(q(\mathbf{x})||p(\mathbf{x}))={\mathbb{E}}_q[\log \frac{q(\mathbf{x})}{p(\mathbf{x})}]$

​ 我们可以得到${\mathcal{L}}_{VLB}$实际上是由一个熵（$L_0$）以及多个KL散度（$(L_t,t\in 1,2,3,...,T)$）构成。其中最后一项$L_T$中的${\mathbf{x}}_t$和${\mathbf{x}}_0$分别是数据分布和先验分布，都是固定的，因此它是一个常数，在最小化时可以忽略。转而去研究$L_0和L_t,t\in 1,2,3,...,T-1$。

$L_t$的计算

​ 首先来考虑公式2中比较复杂的$L_t$。

​ 我们的模型$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mu }_{\theta }({\mathbf{x}}_t,t),{\Sigma }_{\theta }({\mathbf{x}}_t,t)\right)$。

​ 根据公式Diffusion Model（2）：前向扩散过程和逆向降噪过程
2-11，我们知道,对于平均值${\mu }_{\theta }({\mathbf{x}}_t,t)$在给定${\mathbf{x}}_t$和$t$的情况下，我们期望它接近于${\tilde{\mu }}_t({\mathbf{x}}_t,{\mathbf{x}}_0)=\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{(1-{\bar{\alpha }}_t)}}{z}_t)$。

​ 因此我们可以通过重参数化，通过学习高斯噪声$z_{\theta }({\mathbf{x}}_t,t)$来接近于$z_t$：
$$\begin{array}{c}{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right)=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{z}}_{\theta }\left({\mathbf{x}}_t,t\right)\right)\end{array}$$
​ 这个过程可以归结为我们的模型实际上学会了估计真正的逆向过程中的噪声。

​ 对于方差${\Sigma }_{\theta }({\mathbf{x}}_t,t)$的处理，DDPM (Ho et al 2020)将其设置成了${\sigma }_t^2\mathbf{I}$,其中${\sigma }_t^2={\beta }_t$或者${\sigma }_t^2={\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$。当然此项也可以通过模型学习，此处只是为了简化，并不是唯一的。

​ 现在，结合KL散度的公式Diffusion Model（1）：预备知识1-10我们可以写出模型$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mu }_{\theta }({\mathbf{x}}_t,t)=\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{z}_{\theta }({\mathbf{x}}_t,t)),{\Sigma }_{\theta }({\mathbf{x}}_t,t)={\sigma }_t^2\mathbf{I})$关于$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};{\tilde{\mu }}_t({\mathbf{x}}_t,{\mathbf{x}}_0)=\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{z}_t),{\tilde{\Sigma }}_t={\bar{\beta }}_t\mathbf{I})$的KL散度：
$$\begin{array}{c}\begin{array}{l}{L}_t={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{z}}\left[\frac{1}{2{\Vert {\mathbf{\Sigma }}_{\theta }\left({\mathbf{x}}_t,t\right)\Vert }_2^2}{\Vert {\tilde{\mathbf{\mu }}}_t\left({\mathbf{x}}_t,{\mathbf{x}}_0\right)-{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right)\Vert }^2\right]\\ ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{z}}\left[\frac{1}{2{\Vert {\mathbf{\Sigma }}_{\theta }\Vert }_2^2}{\Vert \frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{z}}_t\right)-\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{z}}_{\theta }\left({\mathbf{x}}_t,t\right)\right)\Vert }^2\right]\\ ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{z}}\left[\frac{{\beta }_t^2}{2{\alpha }_t\left(1-{\bar{\alpha }}_t\right){\Vert {\mathbf{\Sigma }}_{\theta }\Vert }_2^2}{\Vert {\mathbf{z}}_t-{\mathbf{z}}_{\theta }\left({\mathbf{x}}_t,t\right)\Vert }^2\right]\\ ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{z}}\left[\frac{{\beta }_t^2}{2{\alpha }_t\left(1-{\bar{\alpha }}_t\right){\Vert {\mathbf{\Sigma }}_{\theta }\Vert }_2^2}{\Vert {\mathbf{z}}_t-{\mathbf{z}}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}{\mathbf{z}}_t,t\right)\Vert }^2\right]\end{array}\end{array}$$
​ 我们可以发现，$L_t$的训练目标实际上是在使用MSE最小化两个高斯噪声$z_t$和$z_{\theta }({\mathbf{x}}_t,t)$。

​ 在训练中发现上述带有加权（$\frac{{\beta }_t^2}{2{\alpha }_t\left(1-{\bar{\alpha }}_t\right){\Vert {\mathbf{\Sigma }}_{\theta }\Vert }_2^2}$）的MSE Loss不太稳定，因此DDPM (Ho et al 2020)使用了不带权重项的简化损失。
$$\begin{gathered} \begin{array}{c}{L}_t^{\text{simple }}={\mathbb{E}}_{{\mathbf{x}}_0,{\mathbf{z}}_t}\left[{\Vert {\mathbf{z}}_t-{\mathbf{z}}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}{\mathbf{z}}_t,t\right)\Vert }^2\right] \\ {L}_{\text{simple }}=L_t^{\text{simple }}+C\end{array} \end{gathered}$$
​ 其中的$C$是一项不依赖于$\theta$的常数。

$L_0$的计算

​ 然后来考虑$L_0$的计算。

​ 已知$L_0=-{\mathbb{E}}_{{\mathbf{x}}_0,{\mathbf{x}}_1}\log (p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1))$，而$p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)=\mathcal{N}({\mu }_{\theta }({\mathbf{x}}_1),1,{\sigma }_1^2\mathbf{I})$。因此$L_0$实际上是一个多元高斯分布的负对数似然期望，即其熵。多元高斯分布的熵仅与其协方差有关，即$L_0$仅与${\sigma }_1^2\mathbf{I}$有关，$L_0$是一个常数。

​ 然而，论文DDPM指出，一般而言，${\mathbf{x}}_0$的分布实际上是离散的，而不是连续的。比如图片数据，像素值取值必须是整数，归一化到$[-1,1]$后，依然是离散的点。Diffusion前向过程的第一步实际上是为离散数据添加噪声。那么，逆Diffusion的最后一步，即从${\mathbf{x}}_1$到${\mathbf{x}}_0$，也不能被简单地看作从$\mathcal{N}({\mathbf{\mu }}_{\theta }({\mathbf{x}}_1,1),{\sigma }_1^2\mathbf{I})$中采样，而是在从$\mathcal{N}({\mathbf{\mu }}_{\theta }({\mathbf{x}}_1,1),{\sigma }_1^2\mathbf{I})$采样的基础上再加上离散化操作。$L_0$也不再是一个常数，而是一个与${\mu }_{\theta }({\mathbf{x}}_1,1)$相关的积分，其具体表达式可以参考DDPM (Ho et al 2020)的Sec3.3。在忽略${\sigma }_1^2$和边缘效应后，$L_0$的取值可以被$\mathcal{N}({\mathbf{\mu }}_{\theta }({\mathbf{x}}_1,1),{\sigma }_1^2\mathbf{I})$的密度函数与离散时的分块大小(bin width)相乘所拟合。

​ 另外值得一提的是，逆Diffusion的最后一步，DDPM直接取${\mu }_{\theta }({\mathbf{x}}_1,1)$作为${\mathbf{x}}_0$。

The overall training and sampling algorithms

训练过程

​ 训练时，分别从$q({\mathbf{x}}_0)$、$Uniform(1,...,T)$、$\mathcal{N}(0,\text{I})$中采样得到${\mathbf{x}}_0$，$t$和$ϵ$（这里的$ϵ$就是前面说的$z$），利用公式Diffusion Model（2）：前向扩散过程和逆向降噪过程
2-4计算得到${\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{z}$，将${\mathbf{x}}_t$和$t$送入网络，预测得到一个噪声。结合公式7最小化预测噪声和真实采样的$ϵ$之间的距离。重复这一过程直到网络收敛。

采样过程

​ 采样时，需要从${\mathbf{x}}_T$一步一步的变回${\mathbf{x}}_0$，其中的每一步都包含三个操作：

• 将${\mathbf{x}}_t$和$t$送入网络，预测得到噪声$ϵ$
• 利用估计的噪声$ϵ$和${\mathbf{x}}_t$，计算${\mu }_{\theta }=\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ)$
• 如果$t>1$，需要从$\mathcal{N}({\mu }_{\theta },{\sigma }_t^2\mathbf{I})$中采样得到${\mathbf{x}}_{t-1}$，利用重参数化技巧，可以将采样过程转换为首先采样$z\in \mathcal{N}(0,\text{I})$，然后计算$x_{t-1}={\mu }_{\theta }+{\sigma }_{t}z$。如果$t=1$，直接令${\mathbf{x}}_0={\mu }_{\theta }$

网络中的参数选择

​ 在前向扩散过程中，需要确定的超参数有${\beta }_t$以及总的步数$T$。在前向扩散过程中，其值通常会增加。

• 在 DDPM (Ho et al 2020)中，前向扩散过程中的方差${\beta }_t$被设置为从${\beta }_1=10^{-4}$到${\beta }_T=0.02$线性增加。与 $[-1,1] $之间的归一化图像像素值相比，它们相对较小。
• 在Nichol & Dhariwal (2021) 中，提出了可以使用基于余弦的方差表。调度函数的选择可以是任意的，只要它在训练过程的中间提供一个接近线性的下降，以及在$t=0$和$t=T$附近的细微变化。

$${\beta }_t=\mathrm{clip}\left(1-\frac{{\bar{\alpha }}_t}{{\bar{\alpha }}_{t-1}},0.999\right){\bar{\alpha }}_t=\frac{f(t)}{f(0)}\text{ where }f(t)=\cos \left(\frac{t/T+s}{1+s}\cdot \frac{\pi }{2}\right)$$

​ 注意此图的纵坐标是${\bar{\alpha }}_t=1-{\beta }_t$而我们讨论的是${\beta }_t$的一个取值。

​ 对于逆向降噪过程中，我们需要选择一个网络来预测一个噪声来从${\mathbf{x}}_t$得到${\mathbf{x}}_{t-1}$（参照采样的过程）。在这个过程中唯一的要求是我们需要保证输入和输出的维度是一样的。比如输入是$1\times 256\times 256$那么输出也需要是$1\times 256\times 256$。因此此处可以选择U-Net。

References:

• 什么是Diffusion模型？
  
• What are Diffusion Models?
  
• Diffusion Models for Deep Generative Learning