扩散模型详细推导过程——编码与解码

符号表

符号	含义
${\mathbf{x}}^{(i)}={\mathbf{z}}_0^{(i)}$	第$i$个训练数据，其为长度为$d$的向量
${\mathbf{z}}_t^{(i)}$	第$i$个训练数据在第$t$时刻的加噪版本
${\mathbf{ϵ}}_t^{(i)}$	第$i$个训练数据在第$t$时刻所添加的高斯噪声
${\beta }_t$	噪声计划（noise schedule），范围为[0,1]
${\alpha }_t$	${\alpha }_t={\prod }_{s=1}^t(1-{\beta }_s)$
$N(\mathbf{\mu },{\sigma }^2\mathbf{I})$	均值为$\mathbf{\mu }$，标准差为$\sigma$的高斯分布
$q\left(\cdot \right)$	正向过程的转移核
$p(\cdot )$	反向过程的转移核
$p\left(\cdot |\theta \right)$	受参数$\theta$影响，用于拟合反向过程的真实概率密度函数
$f_t$	反向过程中$t$时刻对应的神经网络
$g_t$	反向过程中$t$时刻对应的神经网络
${\theta }_t$	神经网络$f_t$或$g_t$的参数
${\theta }_{1:T}$	${\theta }_1,{\theta }_2,\cdots ,{\theta }_T$
$d{\mathbf{z}}_{1:T}$	$d{\mathbf{z}}_{1}d{\mathbf{z}}_2\cdots d{\mathbf{z}}_T$
注：	如没有上标${}^{(i)}$，则表明在此语境下不特别指明对应某个样本

扩散模型的扩散过程（编码器）

扩散模型的编码器所做的工作如下:

设有原数据$\mathbf{x}$，经过如下的逐步编码（添加噪声）过程可以得到一个符合标准高斯分布的噪声

$$\begin{array}{c}{\mathbf{z}}_t=\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1}+\sqrt{{\beta }_t}{\mathbf{ϵ}}_t,t=1,2,\cdots ,T\end{array}$$

其中${\mathbf{z}}_0=\mathbf{x}$，${\mathbf{ϵ}}_1,{\mathbf{ϵ}}_2,\cdots ,{\mathbf{ϵ}}_t\sim N(0,\mathbf{I})$，${\beta }_1,{\beta }_2,\cdots ,{\beta }_t\in [0,1]$为噪声计划（noise schedule），一般逐级递增。当$T\to \infty$，${\mathbf{z}}_T$将服从高斯分布，该推导在下面会涉及。

由于每一步的扩散结果${\mathbf{z}}_t$仅依赖于上一个扩散结果${\mathbf{z}}_{t-1}$，也即只要已知${\mathbf{z}}_{t-1}$（不需要再知道${\mathbf{z}}_1,{\mathbf{z}}_2,\cdots ,{\mathbf{z}}_{t-2}$），再经过计算便可以得到${\mathbf{z}}_t$。该扩散特点符合马尔科夫链的性质，即每一时刻的状态仅依赖于上一时刻的状态，而与之前的状态无关。

现在用一个马尔科夫链表达该扩散过程。在${\mathbf{z}}_{t-1}$是已知的情况下，${\mathbf{z}}_t$的均值

$$\begin{array}{rl}E[{\mathbf{z}}_t]& =E[\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1}]+E[\sqrt{{\beta }_t}{\mathbf{ϵ}}_t]\\ & =\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1}+0\\ & =\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1}\end{array}$$

${\mathbf{z}}_t$的方差

$$\begin{array}{rl}Cov[{\mathbf{z}}_t]& =Cov[\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1}]+Cov[\sqrt{{\beta }_t}{\mathbf{ϵ}}_t]\\ & =0+(\sqrt{{\beta }_t}{)}^2\mathbf{I}\\ & ={\beta }_t\mathbf{I}\end{array}$$

以上推导源自于：①${\mathbf{z}}_{t-1}$是已知的，它不是分布，而是常量②${\mathbf{ϵ}}_t$是标准的高斯分布③若$\mathbf{x}\sim N({\mathbf{m}}_x,{\mathbf{\Sigma }}_x)$，$\mathbf{y}\sim N({\mathbf{m}}_y,{\mathbf{\Sigma }}_y)$，则$\mathbf{Ax}+\mathbf{By}+\mathbf{c}\sim N({\mathbf{Am}}_x+{\mathbf{Bm}}_y+\mathbf{c},{\mathbf{A}{\mathbf{\Sigma }}_{\mathbf{x}}\mathbf{A}}^T+{\mathbf{B}{\mathbf{\Sigma }}_{\mathbf{y}}\mathbf{B}}^T)$

根据前面的分析，在已知${\mathbf{z}}_{t-1}$的情况下，${\mathbf{z}}_t$的概率分布，即转移核的表达式如下：

$$\begin{array}{c}q({\mathbf{z}}_t|{\mathbf{z}}_{t-1})=N(\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1},{\beta }_t\mathbf{I})=\frac{1}{(2\pi {)}^{\frac{d}{2}}\sqrt{{\beta }_t}}\exp \left(-\frac{({\mathbf{z}}_t-\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1}{)}^2}{2{\beta }_t}\right)\end{array}$$

该表达式使用了多元高斯分布的定义，即若随机变量 $X={\left[\begin{array}{c}{X}_1\cdots X_n\end{array}\right]}^T$ 服从均值为 $\mathbf{\mu }\in {\mathbb{R}}^n$ ，协方差为 $\mathbf{\Sigma }\in {\mathbb{S}}_{++}^n$ 的多元高斯分布，则其概率密度函数为：
$$\begin{array}{r}\frac{1}{(2\pi {)}^{n/2}|\mathbf{\Sigma }{|}^{1/2}}\exp \left(-\frac{1}{2}(\mathbf{x}-\mathbf{\mu }{)}^T{\mathbf{\Sigma }}^{-1}(\mathbf{x}-\mathbf{\mu })\right).\end{array}$$

因此，在已知$\mathbf{x}$的情况下，将通过$q({\mathbf{z}}_1|\mathbf{x})$采样得到${\mathbf{z}}_1$；则${\mathbf{z}}_1$变为已知，再通过$q({\mathbf{z}}_2|{\mathbf{z}}_1)$采样得到${\mathbf{z}}_2$，类似地递推，最后得到${\mathbf{x}}_T$。当$T$非常大的时候，该过程十分耗时，但可以将${\mathbf{z}}_t$中的${\mathbf{z}}_{t-1}$逐层次替换为$\mathbf{x}$的表达式，得到
$$\begin{array}{rl}{\mathbf{z}}_t& =\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1}+\sqrt{{\beta }_t}{\mathbf{ϵ}}_t\\ & =\sqrt{1-{\beta }_t}\left(\sqrt{1-{\beta }_{t-1}}{\mathbf{z}}_{t-2}+\sqrt{{\beta }_{t-1}}{\mathbf{ϵ}}_{t-1}\right)+\sqrt{{\beta }_t}{\mathbf{ϵ}}_t\\ & =\sqrt{(1-{\beta }_t)(1-{\beta }_{t-1})}{\mathbf{z}}_{t-2}+\sqrt{1-{\beta }_t-(1-{\beta }_t)(1-{\beta }_{t-1})}{\mathbf{ϵ}}_{t-1}+\sqrt{{\beta }_t}{\mathbf{ϵ}}_t\end{array}$$
再根据高斯分布的混合公式，将${\mathbf{ϵ}}_{\mathbf{t}-1}$和${\mathbf{ϵ}}_t$的项混合为$\mathbf{ϵ}$的分布，得到
$$\begin{array}{rl}{\mathbf{z}}_t& =\sqrt{(1-{\beta }_t)(1-{\beta }_{t-1})}{\mathbf{z}}_{t-2}+\sqrt{{\sqrt{(1-{\beta }_t)-(1-{\beta }_t)(1-{\beta }_{t-1})}}^2+{\sqrt{{\beta }_t}}^2}\mathbf{ϵ}\\ & =\sqrt{(1-{\beta }_t)(1-{\beta }_{t-1})}{\mathbf{z}}_{t-2}+\sqrt{1-{\beta }_t-(1-{\beta }_t)(1-{\beta }_{t-1})+{\beta }_t}\mathbf{ϵ}\\ & =\sqrt{(1-{\beta }_t)(1-{\beta }_{t-1})}{\mathbf{z}}_{t-2}+\sqrt{1-(1-{\beta }_t)(1-{\beta }_{t-1})}\mathbf{ϵ}\\ & =\dots \\ & =\sqrt{\prod _{i=1}^t(1-{\beta }_i)}\mathbf{x}+\sqrt{1-\prod _{i=1}^t(1-{\beta }_i)}\mathbf{ϵ}\\ & =\sqrt{{\alpha }_t}\mathbf{x}+\sqrt{1-{\alpha }_t}\mathbf{ϵ},t=1,2,\cdots ,T\end{array}$$
为了区分不同时刻所对应的噪声，对$\mathbf{ϵ}$添加下标$t$，可得
$$\begin{array}{c}{\mathbf{z}}_t=\sqrt{{\alpha }_t}\mathbf{x}+\sqrt{1-{\alpha }_t}{\mathbf{ϵ}}_t,t=1,2,\cdots ,T\end{array}$$

其中，${\alpha }_t={\prod }_{s=1}^t(1-{\beta }_s)$，$\mathbf{ϵ}\sim N(0,\mathbf{I})$。

所以，一旦已知$\mathbf{x}$，便可以得到${\mathbf{z}}_t$的分布，故：

$$\begin{array}{c}q({\mathbf{z}}_t|\mathbf{x})=N(\sqrt{{\alpha }_t}\mathbf{x},(1-{\alpha }_t)\mathbf{I})=\frac{1}{(2\pi {)}^{\frac{d}{2}}\sqrt{(1-{\alpha }_t)}}\exp \left(-\frac{({\mathbf{z}}_t-\sqrt{{\alpha }_t}\mathbf{x}{)}^2}{1-{\alpha }_t}\right)\end{array}$$

因此，${\mathbf{z}}_t$可以通过先从标准的高斯分布中采样$\mathbf{ϵ}$，然后和${\mathbf{z}}_0$进行混合得到。另外可以观察到，因为${\beta }_t$在$t$很大的时候近似为$1$，那么${\alpha }_t$在$t$很大的时候近似等于0，此时$q({\mathbf{z}}_t|\mathbf{x})$近似为一个标准的高斯分布。

扩散模型的去噪过程（解码器）

扩散模型的解码器是为了反转编码过程。如果知道逆向转移核$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t)$，那么就可以先从$p({\mathbf{z}}_T)=N(0,\mathbf{I})$采样出${\mathbf{z}}_T$，再通过$p({\mathbf{z}}_{T-1}|{\mathbf{z}}_T)$采样出${\mathbf{z}}_{T-1}$，依次类推，直到采样出${\mathbf{z}}_0$，即$\mathbf{x}$。

贝叶斯公式给出了根据$q({\mathbf{z}}_t|{\mathbf{z}}_{t-1})$求出$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t)$的方法，即

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

观察该式可知，由于$q({\mathbf{z}}_{t-1})/q({\mathbf{z}}_t)$是未知的，所以求不出任何结果，而且实际上该逆向转移核不一定是高斯分布。

但是，如果给定额外条件$\mathbf{x}$，由（15），可以得到

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

根据马尔科夫链的性质$q({\mathbf{z}}_t|{\mathbf{z}}_{t-1},\mathbf{x})=q({\mathbf{z}}_t|{\mathbf{z}}_{t-1})$，结合公式（8）和（10），经过很复杂的一段化简（省略过程）得到：

$$\begin{array}{rl}p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})& =\frac{q({\mathbf{z}}_t|{\mathbf{z}}_{t-1})q({\mathbf{z}}_{t-1}|\mathbf{x})}{q({\mathbf{z}}_t|\mathbf{x})}\\ & \propto q({\mathbf{z}}_t|{\mathbf{z}}_{t-1})q({\mathbf{z}}_{t-1}|\mathbf{x})\\ & =N_{{\mathbf{z}}_t}\left(\sqrt{1-{\beta }_t}\cdot {\mathbf{z}}_{t-1},{\beta }_t\mathbf{I}\right)N_{{\mathbf{z}}_{t-1}}\left(\sqrt{{\alpha }_{t-1}}\cdot \mathbf{x},(1-{\alpha }_{t-1})\mathbf{I}\right)\end{array}$$

根据高斯随机变量的变量替换定理，即

$$N_{\mathbf{v}}\left[\mathbf{A}\mathbf{w},\mathbf{B}\right]\propto N_{\mathbf{w}}\left[{\left({\mathbf{A}}^T{\mathbf{B}}^{-1}\mathbf{A}\right)}^{-1}{\mathbf{A}}^T{\mathbf{B}}^{-1}\mathbf{v},{\left({\mathbf{A}}^T{\mathbf{B}}^{-1}\mathbf{A}\right)}^{-1}\right]$$
可得，
$$N_{{\mathbf{z}}_t}\left(\sqrt{1-{\beta }_t}\cdot {\mathbf{z}}_{t-1},{\beta }_t\mathbf{I}\right)N_{{\mathbf{z}}_{t-1}}\left(\sqrt{{\alpha }_{t-1}}\cdot \mathbf{x},(1-{\alpha }_{t-1})\mathbf{I}\right)\propto N_{{\mathbf{z}}_{t-1}}\left(\frac{1}{\sqrt{1-{\beta }_t}}{\mathbf{z}}_t,\frac{{\beta }_t}{1-{\beta }_t}\mathbf{I}\right)N_{{\mathbf{z}}_{t-1}}\left(\sqrt{{\alpha }_{t-1}}\cdot \mathbf{x},(1-{\alpha }_{t-1})\mathbf{I}\right)$$

再根据

$$\begin{array}{rl}{N}_{\mathbf{w}}[\mathbf{a},\mathbf{A}]\cdot N_{\mathbf{w}}[\mathbf{b},\mathbf{B}]\propto N_{\mathbf{w}}& \left[{\left({\mathbf{A}}^{-1}+{\mathbf{B}}^{-1}\right)}^{-1}({\mathbf{A}}^{-1}\mathbf{a}+{\mathbf{B}}^{-1}\mathbf{b}),{\left({\mathbf{A}}^{-1}+{\mathbf{B}}^{-1}\right)}^{-1}\right]\end{array}$$

最终得到

$$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})=N_{{\mathbf{z}}_{t-1}}\left[\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},\frac{{\beta }_t(1-{\alpha }_{t-1})}{1-{\alpha }_t}\mathbf{I}\right]$$

由此可知$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})$是一个高斯分布。

因此，尽管$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t)$不是高斯分布，但给定条件$\mathbf{x}$后得到的$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})$是高斯分布。另外，如果$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})$的均值和方差被确定，那么进一步可以写出从中采样的公式，得到${\mathbf{z}}_{t-1}$。因此，可以考虑用神经网络来近似$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})$（在后续训练目标的推导中可以看出网络的目标实际上是近似$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,\mathbf{x})$），记作$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)$。为了简化该分布，将其方差设为固定值，神经网络仅仅估计其均值。

$$\begin{array}{c}p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)=N(f_t({\mathbf{z}}_t,{\theta }_t),{\sigma }_t^2\mathbf{I})\end{array}$$

其中$f_t$为神经网络，其接受输入${\mathbf{z}}_t$并输出一个估计的均值，${\theta }_t$为该网络的参数，${\sigma }_t$为人为设定的标准差。

如果能训练出使得原数据${\mathbf{z}}_0$总体出现概率最大的神经网络$f_t({\mathbf{z}}_t,{\theta }_t)$，进而得到$p({\mathbf{z}}_{t-1}|{\mathbf{z}}_t,{\theta }_t)$，那么就可以先从$N(0,\mathbf{I})$采样出${\mathbf{z}}_T$，再通过$p({\mathbf{z}}_{T-1}|{\mathbf{z}}_T,{\theta }_t)$采样出${\mathbf{z}}_{T-1}$，依次类推，直到采样出$\mathbf{x}$，即${\mathbf{z}}_0$。