扩散模型的训练与采样算法

训练目标的推导

需要使得去噪过程所产生的${\mathbf{x}}^{(i)}$的总体出现概率最大，先不考虑第几个样本，省略上标，即最大化$p(\mathbf{x}|{\theta }_{1:T})$，也等价于最大化$\log \left[p(\mathbf{x}|{\theta }_{1:T})\right]$。直接最大化该式是无从下手的，考虑寻找该式的一个置信下界ELBO：

$$\begin{array}{rl}\log \left[p(\mathbf{x}|{\theta }_{1:T})\right]& =\log \left[p(\mathbf{x},{\mathbf{z}}_{1:T}|{\theta }_{1:T})d{\mathbf{z}}_{1:T}\right]\\ & =\log \left[\int q({\mathbf{z}}_{1:T}|\mathbf{x})\frac{p(\mathbf{x},{\mathbf{z}}_{1:T}|{\theta }_{1:T})}{q({\mathbf{z}}_{1:T}|\mathbf{x})}d{\mathbf{z}}_{1:T}\right]\\ & \ge \int q({\mathbf{z}}_{1:T}|\mathbf{x})\log \left[\frac{p({\mathbf{z}}_0,{\mathbf{z}}_{1:T}|{\theta }_{1:T})}{q({\mathbf{z}}_{1:T}|\mathbf{x})}\right]d{\mathbf{z}}_{1:T}\end{array}$$

而其中，

$$\begin{array}{rl}\log \left[\frac{p(\mathbf{x},{\mathbf{z}}_{1:T}|{\theta }_{1:T})}{q({\mathbf{z}}_{1:T}|\mathbf{x})}\right]& \begin{array}{r}=\log \lfloor \frac{p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1){\prod }_{t=2}^{T}p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)\cdot p({\mathbf{z}}_T)}{q({\mathbf{z}}_1|\mathbf{x}){\prod }_{t=2}^{T}q({\mathbf{z}}_t|{\mathbf{z}}_{t-1})}\rfloor \end{array}\\ & \begin{array}{r}=\log \left[\frac{p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)}{q({\mathbf{z}}_1|\mathbf{x})}\right]+\log \left[\frac{{\prod }_{t=2}^{T}p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)}{{\prod }_{t=2}^{T}q({\mathbf{z}}_t|{\mathbf{z}}_{t-1})}\right]+\log [p({\mathbf{z}}_T)]\end{array}\end{array}$$

由于扩散过程的马尔科夫链性质

$$\begin{array}{c}q({\mathbf{z}}_t|{\mathbf{z}}_{t-1})=q({\mathbf{z}}_t|{\mathbf{z}}_{t-1},\mathbf{x})=\frac{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})q({\mathbf{z}}_t|\mathbf{x})}{q({\mathbf{z}}_{t-1}|\mathbf{x})}\end{array}$$

所以该展开式可以继续简化为：

$$\begin{array}{rl}\log \left[\frac{p(\mathbf{x},{\mathbf{z}}_{1:T}|{\theta }_{1:T})}{q({\mathbf{z}}_{1:T}|\mathbf{x})}\right]& =\log \left[\frac{p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)}{q({\mathbf{z}}_1|\mathbf{x})}\right]+\log \left[\frac{{\prod }_{t=2}^{T}p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)\cdot q({\mathbf{z}}_0|\mathbf{x})}{{\prod }_{t=2}^{T}p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})\cdot q({\mathbf{z}}_T|\mathbf{x})}\right]+\log [p({\mathbf{z}}_T)]\\ & =\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]+\log \left[\frac{{\prod }_{t=2}^{T}p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)}{{\prod }_{t=2}^{T}p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})}\right]+\log \left[\frac{p({\mathbf{z}}_T)}{q({\mathbf{z}}_T|\mathbf{x})}\right]\\ & \approx \log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]+\sum _{t=2}^T\log \left[\frac{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)}{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})}\right]\end{array}$$

第二行到第三行中，$p({\mathbf{z}}_T)$为标准的高斯分布，而$q({\mathbf{z}}_T|\mathbf{x})$近似为标准的高斯分布，故$\log \left[\frac{p({\mathbf{z}}_T)}{q({\mathbf{z}}_T|\mathbf{x})}\right]\approx \log 1=0$。

至此，可以完整地写出

$$\begin{array}{rl}\log \left[p(\mathbf{x}|{\theta }_{1:T})\right]& \ge \log \left[\int q({\mathbf{z}}_{1:T}|\mathbf{x})\frac{p(\mathbf{x},{\mathbf{z}}_{1:T}|{\theta }_{1:T})}{q({\mathbf{z}}_{1:T}|\mathbf{x})}d{\mathbf{z}}_{1:T}\right]\\ & \approx \int q({\mathbf{z}}_{1:T}|\mathbf{x})\left(\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]+\sum _{t=2}^T\log \left[\frac{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)}{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})}\right]\right)d{\mathbf{z}}_{1:T}\\ & =\int q({\mathbf{z}}_{1:T}|\mathbf{x})\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]d{\mathbf{z}}_{1:T}+\int q({\mathbf{z}}_{1:T}|\mathbf{x})\sum _{t=2}^T\log \left[\frac{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)}{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})}\right]d{\mathbf{z}}_{1:T}\\ & =\int q({\mathbf{z}}_1|\mathbf{x})\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]d{\mathbf{z}}_1+\sum _{t=2}^T\int q({\mathbf{z}}_{1:T}|\mathbf{x})\log \left[\frac{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)}{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})}\right]d{\mathbf{z}}_{1:T}\\ & =E_{q({\mathbf{z}}_1|{\mathbf{z}}_0)}\left[\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]\right]+\sum _{t=2}^T\iint q({\mathbf{z}}_{t-1},{\mathbf{z}}_t|\mathbf{x})\log \left[\frac{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)}{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})}\right]d{\mathbf{z}}_{t-1}d{\mathbf{z}}_t\\ & =E_{q({\mathbf{z}}_1|{\mathbf{z}}_0)}\left[\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]\right]+\sum _{t=2}^T\iint q({\mathbf{z}}_t|\mathbf{x})p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})\log \left[\frac{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)}{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})}\right]d{\mathbf{z}}_{t-1}d{\mathbf{z}}_t\\ & =E_{q({\mathbf{z}}_1|{\mathbf{z}}_0)}\left[\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]\right]+\sum _{t=2}^T\int q({\mathbf{z}}_t|\mathbf{x})\left(\int p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})\log \left[\frac{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)}{p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})}\right]d{\mathbf{z}}_{t-1}\right)d{\mathbf{z}}_t\\ & =E_{q({\mathbf{z}}_1|{\mathbf{z}}_0)}\left[\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]\right]-\sum _{t=2}^T\int q({\mathbf{z}}_t|\mathbf{x})\cdot D_{KL}\left[p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})||p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)\right]d{\mathbf{z}}_t\\ & =E_{q({\mathbf{z}}_1|{\mathbf{z}}_0)}\left[\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]\right]-\sum _{t=2}^T{E}_{q({\mathbf{z}}_t|\mathbf{x})}\left[D_{KL}\left[p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})||p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)\right]\right]\\ & =E_{q({\mathbf{z}}_1|\mathbf{x})}\left[\log \left[N(f_1({\mathbf{z}}_1,{\theta }_1),{\sigma }_1^2\mathbf{I})\right]\right]\\ & -\sum _{t=2}^T{E}_{q({\mathbf{z}}_t|\mathbf{x})}\left[\frac{1}{2{\sigma }_t^2}{\Vert \frac{(1-{\alpha }_{t-1})}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{\mathbf{z}}_t+\frac{\sqrt{{\alpha }_{t-1}}{\beta }_t}{1-{\alpha }_t}\mathbf{x}-f_t[{\mathbf{z}}_t,{\theta }_t]\Vert }^2+C\right]\\ & \approx \log \left[N(f_1({\mathbf{z}}_1^{\ast },{\theta }_1),{\sigma }_1^2\mathbf{I})\right]-\sum _{t=2}^T\frac{1}{2{\sigma }_t^2}{\Vert \frac{(1-{\alpha }_{t-1})}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{\mathbf{z}}_t^{\ast }+\frac{\sqrt{{\alpha }_{t-1}}{\beta }_t}{1-{\alpha }_t}\mathbf{x}-f_t[{\mathbf{z}}_t^{\ast },{\theta }_t]\Vert }^2-C\end{array}$$

其中，

$$\begin{array}{c}\log \left[p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)\right]=\log \left[N(f_1({\mathbf{z}}_1,{\theta }_1),{\sigma }_1^2\mathbf{I})\right]\end{array}$$

$$\begin{array}{c}{D}_{KL}\left[p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})||p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)\right]=\frac{1}{2{\sigma }_t^2}{\Vert \frac{(1-{\alpha }_{t-1})}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{\mathbf{z}}_t+\frac{\sqrt{{\alpha }_{t-1}}{\beta }_t}{1-{\alpha }_t}\mathbf{x}-f_t[{\mathbf{z}}_t,{\theta }_t]\Vert }^2+C\end{array}$$

两个均值已使用蒙特卡洛方法近似，${\mathbf{z}}_1^{\ast }$和${\mathbf{z}}_t^{\ast }$是分别从$q({\mathbf{z}}_1|\mathbf{x})$和$q({\mathbf{z}}_t|\mathbf{x})$中采样出来的某个样本，实际优化时常数$C$不用考虑。

最终的目标是$\max {\prod }_{i=1}^{n}p({\mathbf{x}}^{(i)}|{\theta }_{1:T})$，等价于$\max {\sum }_{i=1}^n\log \left[p({\mathbf{x}}^{(i)}|{\theta }_{1:T})\right]$，其中$n$为训练样本的总数。

$$\begin{array}{rl}{\bar{\theta }}_{1:T}& =\arg \min -\sum _{i=1}^n\log \left[p({\mathbf{x}}^{(i)}|{\theta }_{1:T})\right]\\ & =\arg \min \\ & -\sum _{i=1}^n\left[\underset{①重建损失}{\underbrace{\log \left[N(f_t({\mathbf{z}}_1^{(i)},{\theta }_1),{\sigma }_1^2\mathbf{I})\right]}}-\sum _{t=2}^T\frac{1}{2{\sigma }_t^2}{\Vert \underset{②p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})的均值}{\underbrace{\frac{(1-{\alpha }_{t-1})}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{\mathbf{z}}_t^{(i)}+\frac{\sqrt{{\alpha }_{t-1}}{\beta }_t}{1-{\alpha }_t}{\mathbf{x}}^{(i)}}}-\underset{③神经网络估计的均值}{\underbrace{f_t[{\mathbf{z}}_t^{(i)},{\theta }_t]}}\Vert }^2\right]\end{array}$$

①重建损失：实际上是近似已知${\mathbf{z}}_1$的情况下$\mathbf{x}$的概率密度函数。

②$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})$的均值。

③神经网络估计的均值。

该目标函数实际上是在做两件事：一是使得最终$\mathbf{x}$的出现概率最大；二是使得解码过程中神经网络估计的均值尽可能逼近$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})$的均值。

由于${\mathbf{z}}_t=\sqrt{{\alpha }_t}\mathbf{x}+\sqrt{1-{\alpha }_t}{\mathbf{ϵ}}_t,t=1,2,\cdots ,T$，将$\mathbf{x}$替换为${\mathbf{z}}_t$的表达式：

$$\begin{array}{c}\mathbf{x}=\frac{1}{\sqrt{{\alpha }_t}}{\mathbf{z}}_t-\frac{\sqrt{1-{\alpha }_t}}{\sqrt{{\alpha }_t}}{\mathbf{ϵ}}_t\end{array}$$

将该表达式代入目标函数，并利用$\frac{\sqrt{{\alpha }_{t-1}}}{\sqrt{{\alpha }_t}}=\frac{1}{\sqrt{1-{\beta }_t}}$，可继续简化为

$$-\sum _{i=1}^n\left[\underset{①重建损失}{\underbrace{\log \left[N(f_t({\mathbf{z}}_1^{(i)},{\theta }_1),{\sigma }_1^2\mathbf{I})\right]}}-\sum _{t=2}^T\frac{1}{2{\sigma }_t^2}{\Vert \underset{②p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\mathbf{z}}_0)的均值}{\underbrace{\frac{1}{\sqrt{1-{\beta }_t}}{\mathbf{z}}_t^{(i)}-\frac{{\beta }_t}{\sqrt{(1-{\beta }_t)(1-{\alpha }_t)}}{\mathbf{ϵ}}_t^{(i)}}}-\underset{③神经网络预测的均值}{\underbrace{f_t[{\mathbf{z}}_t^{(i)},{\theta }_t]}}\Vert }^2\right]$$

其中③由${\mathbf{z}}_t$和$g_t[{\mathbf{z}}_t,{\theta }_t]$计算得到：

$$\begin{array}{c}{f}_t[{\mathbf{z}}_t,{\theta }_t]=\frac{1}{\sqrt{1-{\beta }_t}}{\mathbf{z}}_t-\frac{{\beta }_t}{\sqrt{(1-{\beta }_t)(1-{\alpha }_t)}}{g}_t[{\mathbf{z}}_t,{\theta }_t]\end{array}$$

其中$g_t$是用于估计噪声${\mathbf{ϵ}}_t$的神经网络，而${\mathbf{z}}_t$在采样阶段是已知的。参数为了方便统一仍记为${\theta }_t$。将该表达式代入，目标函数变为：

$$\begin{array}{r}{\bar{\theta }}_{1:T}=\arg \min -\sum _{i=1}^n\left[\underset{①重建损失}{\underbrace{\log \left[N(f_1({\mathbf{z}}_1^{(i)},{\theta }_1),{\sigma }_1^2\mathbf{I})\right]}}-\sum _{t=2}^T\frac{{\beta }_t^2}{2{\sigma }_t^2(1-{\beta }_t)(1-{\alpha }_t)}{\Vert \underset{②神经网络预测的噪声}{\underbrace{g_t[{\mathbf{z}}_t^{(i)},{\theta }_t]}}-\underset{③对样本{\mathbf{x}}^{(i)}在第t步添加的噪声}{\underbrace{{\mathbf{ϵ}}_t^{(i)}}}\Vert }^2\right]\end{array}$$

对于上式中①，由多元高斯分布的定义可以算出，$\Sigma ={\sigma }_1^{2}I$，$|\Sigma {|}^{1/2}$为一个常数，${\Sigma }^{-1}=\frac{1}{{\sigma }_1^2}\mathbf{I}$，故可以写为：

$$\begin{array}{r}N(f_1({\mathbf{z}}_1^{(i)},{\theta }_1),{\sigma }_1^2\mathbf{I})=-\log \left[(2\pi {)}^{n/2}|\Sigma {|}^{1/2}\right]-\frac{1}{2{\sigma }_1^2}{\Vert {\mathbf{x}}^{(i)}-f_1({\mathbf{z}}_1^{(i)},{\theta }_1)\Vert }^2\end{array}$$

而由（47）（48），可知

$$\begin{array}{rl}\frac{1}{2{\sigma }_1^2}{\Vert {\mathbf{x}}^{(i)}-f_1({\mathbf{z}}_1^{(i)},{\theta }_1)\Vert }^2& =\frac{1}{2{\sigma }_1^2}{\Vert \frac{1}{\sqrt{{\alpha }_1}}{\mathbf{z}}_1^{(i)}-\frac{\sqrt{1-{\alpha }_1}}{\sqrt{{\alpha }_1}}{\mathbf{ϵ}}_1^{(i)}-\frac{1}{\sqrt{1-{\beta }_1}}{\mathbf{z}}_1^{(i)}+\frac{{\beta }_1}{\sqrt{(1-{\beta }_1)(1-{\alpha }_1)}}{g}_1[{\mathbf{z}}_1^{(i)},{\theta }_1]\Vert }^2\\ & =\frac{{\beta }_1^2}{2{\sigma }_1^2(1-{\beta }_1)(1-{\alpha }_1)}{\Vert g_1[{\mathbf{z}}_1^{(i)},{\theta }_1]-{\mathbf{ϵ}}_1^{(i)}\Vert }^2\end{array}$$

至此，目标可以进一步简化为

$$\begin{array}{c}{\bar{\theta }}_{1:T}=\arg \min \sum _{i=1}^n\sum _{t=1}^T\frac{{\beta }_t^2}{2{\sigma }_t^2(1-{\beta }_t)(1-{\alpha }_t)}{\Vert g_t[{\mathbf{z}}_t^{(i)},{\theta }_t]-{\mathbf{ϵ}}_t^{(i)}\Vert }^2\end{array}$$

其中$-\log \left[(2\pi {)}^{n/2}|\Sigma {|}^{1/2}\right]$为一个常数，已从目标中省去。另外，在实际的实验中发现，优化目标的系数并不重要，可以在训练中设置为1以简化目标。

训练

根据编码器和训练目标，可以得到训练算法如下：

对所有的观测数据${\mathbf{x}}^{(i)}$，${\mathbf{z}}_0^{(i)}={\mathbf{x}}^{(i)},i=1,\cdots ,n$，loss=0，循环执行：

——循环$t=1,2,\cdots ,T$，执行：

————计算${\mathbf{z}}_t^{(i)}$：${\mathbf{z}}_t^{(i)}=\sqrt{{\alpha }_t}{\mathbf{x}}^{(i)}+\sqrt{1-{\alpha }_t}{\mathbf{ϵ}}_t^{(i)}$。其中${\mathbf{ϵ}}^{(i)}\sim N(0,\mathbf{I})$。

————训练$g_t[{\mathbf{z}}_t^{(i)},{\theta }_t]$，并累积损失loss+=$\frac{{\beta }_t^2}{2{\sigma }_t^2(1-{\beta }_t)(1-{\alpha }_t)}{\Vert g_t[{\mathbf{z}}_t^{(i)},{\theta }_t]-{\mathbf{ϵ}}_t^{(i)}\Vert }^2$

执行反向传播和梯度下降，更新${\theta }_t$。

采样

根据解码器，可以得到采样算法：

从$N(0,\mathbf{I})$采样出${\mathbf{z}}_T$

循环$t=T-1,T-2,\cdots ,1$，执行：

——将${\mathbf{z}}_t$输入神经网络$g_t$，获取$g_t[{\mathbf{z}}_t,{\theta }_t]$

——估计$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)$的均值$\mathbf{\mu }$：$\mathbf{\mu }=\frac{1}{\sqrt{1-{\beta }_t}}{\mathbf{z}}_t-\frac{{\beta }_t}{\sqrt{(1-{\beta }_t)(1-{\alpha }_t)}}{g}_t[{\mathbf{z}}_t,{\theta }_t]$

——如果$t>1$：

————从$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t^{\ast },{\theta }_t)$中采样出${\mathbf{z}}_{t-1}$：${\mathbf{z}}_{t-1}=\mathbf{\mu }+{\sigma }_t\mathbf{ϵ}$，其中$\mathbf{ϵ}\sim N(0,\mathbf{I})$

——否则：

————从$p(\mathbf{x}|{\mathbf{z}}_1,{\theta }_1)$中采样出$\mathbf{x}$：$\mathbf{x}=\mathbf{\mu }$

在实际训练时，并不会真的训练多个对应不同时刻的神经网络$g_t[{\mathbf{z}}_t,{\theta }_t],t=1,2,\cdots ,T$，而是用一个加上时间信息的神经网络或$g[{\mathbf{z}}_t,t,{\theta }_t]$来代替。