概率扩散模型（Denoising Diffusion Probabilistic Models）基本原理

引入：大自然中的扩散现象都遵循一个熵增的过程，比如一滴墨水滴到清水里，比如罐子里的气体一步一步扩散到空气中，都是一个不可逆的过程，但当前状态是包含着上一个状态的信息且扩散规律是可遵循的。Diffusion Models 就是基于一些假设条件，通过找到并掌握扩散过程的每一个步骤与当前的状态，从一个初始的各向同性的高斯噪声分布来一步一步逆推上一个的状态，经过不断的逆推得到了粒子原来的状态。

  不妨设 ${\mathrm{x}}^0\sim q_{\mathcal{X}}\left({\mathrm{x}}^0\right)$ 是从输入空间 $\mathcal{X}={\mathbb{R}}^D$ 采样而来的样本，且用概率密度分布函数 $p_{\theta }({\mathrm{x}}^0)$ 来逼近真实分布 $q_{\mathcal{X}}({\mathrm{x}}^0)$。Diffusion Model 是一个包含隐变量的生成式模型 $p_{\theta }\left({\mathbf{x}}^0\right):=\int p_{\theta }\left({\mathbf{x}}^{0:N}\right)\mathrm{d}{\mathbf{x}}^{1:N}$ 其中${\mathbf{x}}^1,\dots ,{\mathbf{x}}^N$ 是 D 维的隐变量。
扩散过程：
扩散模型的后验分布 $q\left({\mathbf{x}}^{1:N}\mid {\mathbf{x}}^0\right)$表示为：$$q\left({\mathbf{x}}^{1:N}\mid {\mathbf{x}}^0\right)={\Pi }_{n=1}^{N}q\left({\mathbf{x}}^n\mid {\mathbf{x}}^{n-1}\right)\text{\hspace{1em}(1)}$$
  扩散模型的后验分布 $q\left({\mathbf{x}}^{1:N}\mid {\mathbf{x}}^0\right)$ 是一个不包含训练参数的确定的马尔科夫过程（当前状态只与前一个时刻的状态有关）。扩散模型的扩散过程是一个逐步加高斯噪声的过程：
$$q\left({\mathbf{x}}^n\mid {\mathbf{x}}^{n-1}\right):=\mathcal{N}\left({\mathbf{x}}^n;\sqrt{1-{\beta }_n}{\mathbf{x}}^{n-1},{\beta }_n\mathbf{I}\right)\text{\hspace{1em}(2)}$$
  扩散过程中使用的方差 ${\beta }_1,{\beta }_2,...,{\beta }_N\in [0,1]$ 是一个单调递增的序列。在此过程中使用重参数化技巧（Re parameter Trick）先从一个标准的高斯分布中采样得到噪声 $ϵ$ 然后将 $ϵ$ 乘以方差 ${\beta }_n\mathbf{I}$ 再加上均值 $\sqrt{1-{\beta }_n}{\mathbf{x}}^{n-1}$ 来获得样本 ${\mathbf{x}}^n$。给定输入样本 ${\mathbf{x}}^0$，经 N 步的扩散过程得到加噪后的样本 ${\mathbf{x}}^N$。
  扩散过程中要经过多步的加高斯噪声的过程得到最终的 ${\mathbf{x}}^N$， 经过证明，其中每一步加噪后的采样值 ${\mathbf{x}}^n$ 可以由 ${\mathbf{x}}^0$ 直接计算（将中间的加高斯噪声的步骤合并）如式（3）所示：
$$q\left({\mathbf{x}}^n\mid {\mathbf{x}}^0\right)=\mathcal{N}\left({\mathbf{x}}^n;\sqrt{{\bar{\alpha }}_n}{\mathbf{x}}^0,\left(1-{\bar{\alpha }}_n\right)\mathbf{I}\right)\text{\hspace{1em}(3)}$$
其中， ${\alpha }_n:=1-{\beta }_n$ 且 ${\bar{\alpha }}_n:={\Pi }_{i=1}^N{\alpha }_i$。

逆扩散过程：
  假设逆扩散过程的联合概率分布 $p_{\theta }({\mathbf{x}}^{0:N})$ 是一个马尔科夫过程， 且每一步的转移概率服从一个带有学习参数的高斯分布（通过神经网络拟合）。此外，逆扩散过程的起始点 $p\left({\mathbf{x}}^N\right)=\mathcal{N}\left({\mathbf{x}}^N;0,\mathbf{I}\right)$ 是一个标准的高斯分布（扩散过程不断加噪声，最后得到一个标准的噪声分布），则逆扩散过程的联合分布如式（4）所示：
$$p_{\theta }\left({\mathbf{x}}^{0:N}\right):=p\left({\mathbf{x}}^N\right){\Pi }_{n=N}^1{p}_{\theta }\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n\right)\text{\hspace{1em}(4)}$$
由于扩散过程的每一步都加了很小的高斯噪声，不妨假设每一步的逆扩散过程也服从高斯分布：
$$p_{\theta }\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n\right):=\mathcal{N}\left({\mathbf{x}}^{n-1};{\mu }_{\theta }\left({\mathbf{x}}^n,n\right),{\Sigma }_{\theta }\left({\mathbf{x}}^n,n\right)\mathbf{I}\right)\text{\hspace{1em}(5)}$$
其中均值 ${\mu }_{\theta }\left({\mathbf{x}}^n,n\right)$ 与方差 ${\Sigma }_{\theta }\left({\mathbf{x}}^n,n\right)\mathbf{I}$ 是输入为当前时刻的 ${\mathbf{x}}^n$ 与 时刻$n$ 的函数（用神经网络来逼近这个函数，参数为 $\theta$），$p_{\theta }\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n\right)$ 是为了消除在扩散过程中增加的高斯噪声。
  模型训练使得似然函数达到最大，也就是负对数似然函数最小，使用Jensen’s 不等式可以得到负对数似然函数的上界：
$$\underset{\theta }{\min }{\mathbb{E}}_{q\left({\mathbf{x}}^0\right)}\left[-\log p_{\theta }\left({\mathbf{x}}^0\right)\right]\le \underset{\theta }{\min }{\mathbb{E}}_{q\left({\mathbf{x}}^{0:N}\right)}\left[-\log p_{\theta }\left({\mathbf{x}}^{0:N}\right)+\log q\left({\mathbf{x}}^{1:N}\mid {\mathbf{x}}^0\right)\right]\text{\hspace{1em}(6)}$$
等价于优化：
$$\underset{\theta }{\min }{\mathbb{E}}_{q\left({\mathbf{x}}^{0:N}\right)}\left[-\log p\left({\mathbf{x}}^N\right)-\sum _{n=1}^N\log \frac{p_{\theta }\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n\right)}{q\left({\mathbf{x}}^n\mid {\mathbf{x}}^{n-1}\right)}\right]\text{\hspace{1em}(7)}$$
由于这些过程都假设为马尔科夫过程，式（7） 可以简化为：
$$-\log p_{\theta }\left({\mathbf{x}}^0\mid {\mathbf{x}}^1\right)+D_{\mathrm{KL}}\left(q\left({\mathbf{x}}^N\mid {\mathbf{x}}^0\right)\Vert p\left({\mathbf{x}}^N\right)\right)+\sum _{n=2}^N{D}_{\mathrm{KL}}\left(q\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n,{\mathbf{x}}^0\right)\Vert p_{\theta }\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n\right)\right)\text{\hspace{1em}(8)}$$
其中，$\log p_{\theta }\left({\mathbf{x}}^0\mid {\mathbf{x}}^1\right)$ 为逆扩散过程最后一步的条件概率，$D_{\mathrm{KL}}\left(q\left({\mathbf{x}}^N\mid {\mathbf{x}}^0\right)\Vert p\left({\mathbf{x}}^N\right)\right)$ 不包含训练参数，对于模型优化目标而言是个常量，可以省略，$D_{\mathrm{KL}}\left(q\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n,{\mathbf{x}}^0\right)\Vert p_{\theta }\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n\right)\right)$ 为分布 $q\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n,{\mathbf{x}}^0\right)$ 与高斯分布 $p_{\theta }\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n\right)$ 之间的KL散度。

对于前向扩散过程中的后验概率 $q\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n,{\mathbf{x}}^0\right)$ 服从高斯分布：
$$q\left({\mathrm{x}}^{n-1}\mid {\mathrm{x}}^n,{\mathrm{x}}^0\right)=\mathcal{N}\left({\mathrm{x}}^{n-1};{\tilde{\mu }}_n\left({\mathbf{x}}^n,{\mathbf{x}}^0\right),{\tilde{\beta }}_n\mathbf{I}\right)\text{\hspace{1em}(9)}$$
其中均值与方差：
$$\begin{array}{c}{\tilde{\mu }}_n\left({\mathbf{x}}^n,{\mathbf{x}}^0\right):=\frac{\sqrt{{\bar{\alpha }}_{n-1}}{\beta }_n}{1-{\bar{\alpha }}_n}{\mathbf{x}}^0+\frac{\sqrt{{\alpha }_n}\left(1-{\bar{\alpha }}_{n-1}\right)}{1-{\bar{\alpha }}_n}{\mathbf{x}}^n\\ \\ {\tilde{\beta }}_n:=\frac{1-{\bar{\alpha }}_{n-1}}{1-{\bar{\alpha }}_n}{\beta }_n.\end{array}\text{\hspace{1em}(10)}$$
则式（8）中表示两高斯分布的KL散度可以表示为：

$$D_{\mathrm{KL}}\left(q\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n,{\mathbf{x}}^0\right)||p_{\theta }\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n\right)\right)={\mathbb{E}}_q\left[\frac{1}{2{\Sigma }_{\theta }}{\Vert {\tilde{\mu }}_n\left({\mathbf{x}}^n,{\mathbf{x}}^0\right)-{\mu }_{\theta }\left({\mathbf{x}}^n,n\right)\Vert }^2\right]+C\text{\hspace{1em}(11)}$$

我们的目标是优化模型的权重$\theta$ 使得${\mu }_{\theta }\left({\mathbf{x}}^n,n\right)={\tilde{\mu }}_n\left({\mathbf{x}}^n,{\mathbf{x}}^0\right)$，由于直接预测均值得到的效果不好，由公式（3）有${\mathbf{x}}^n({\mathbf{x}}^0,ϵ)=\sqrt{{\bar{\alpha }}_n}{\mathbf{x}}^0+\sqrt{\left(1-{\bar{\alpha }}_n\right)}ϵ$， 其中$ϵ\sim \mathcal{N}(0,\mathbf{I})$，则用${\mathbf{x}}^n$来表示${\mathbf{x}}^0$有$${\mathbf{x}}^0=\left({\mathbf{x}}^n-{\beta }_nϵ/\sqrt{1-{\bar{\alpha }}_n}\right)/\sqrt{{\alpha }_n}\text{\hspace{1em}(12)}$$
我们可以训练模型，使得：
$${\mu }_{\theta }\left({\mathbf{x}}^n,n\right)=\frac{1}{\sqrt{{\alpha }_n}}\left({\mathbf{x}}^n-\frac{{\beta }_n}{\sqrt{1-{\bar{\alpha }}_n}}{ϵ}_{\theta }\left({\mathbf{x}}^n,n\right)\right)\text{\hspace{1em}(13)}$$
其中，${ϵ}_{\theta }\left({\mathbf{x}}^n,n\right)$ 是一个神经网络，这个网络用来预测$ϵ\sim \mathcal{N}(0,\mathbf{I})$，最后模型的优化目标可以表示为：
$${\mathbb{E}}_{{\mathbf{x}}^0,ϵ}\left[\frac{{\beta }_n^2}{2{\Sigma }_{\theta }{\alpha }_n\left(1-{\bar{\alpha }}_n\right)}{\Vert ϵ-{ϵ}_{\theta }\left(\sqrt{{\bar{\alpha }}_n}{\mathbf{x}}^0+\sqrt{1-{\bar{\alpha }}_n}ϵ,n\right)\Vert }^2\right]\text{\hspace{1em}(14)}$$
其中 $ϵ$ 为给定 ${\mathbf{x}}^0$ 到 ${\mathbf{x}}^n$ 所增加的高斯噪声，并且构建神经网络${ϵ}_{\theta }\left(\sqrt{{\bar{\alpha }}_n}{\mathbf{x}}^0+\sqrt{1-{\bar{\alpha }}_n}ϵ,n\right)$，使得网络的输出与 $ϵ$ 相等，等价于 ${\mu }_{\theta }\left({\mathbf{x}}^n,n\right)={\tilde{\mu }}_n\left({\mathbf{x}}^n,{\mathbf{x}}^0\right)$。有了 ${\tilde{\mu }}_n\left({\mathbf{x}}^n,{\mathbf{x}}^0\right)$便可以由公式(9)得到$q\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n,{\mathbf{x}}^0\right)$，由于训练过程中，目标是损失最小化，等价于（8）中第三项的两个分布的KL散度为0，也就是$p_{\theta }\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n\right)=q\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n,{\mathbf{x}}^0\right)$, 便可以由 ${\mathbf{x}}^n$不断计算得到 ${\mathbf{x}}^{n-1}$，最后得到 ${\mathbf{x}}^0$,也就是训练完成，则有 ${\mathbf{x}}^{n-1}\sim p_{\theta }({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n)$:

$${\mathbf{x}}^{n-1}=\frac{1}{\sqrt{{\alpha }_n}}\left({\mathbf{x}}^n-\frac{{\beta }_n}{\sqrt{1-{\bar{\alpha }}_n}}{ϵ}_{\theta }\left({\mathbf{x}}^n,n\right)\right)+\sqrt{{\Sigma }_{\theta }}\mathbf{Z}\text{\hspace{1em}(15)}$$
其中 $\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})\text{ for }n=N,\dots ,2\text{ and }\mathbf{z}=0\text{ if }n=1$

Diffusion Model 的核心思想是使用以 $({\mathbf{x}}^n,n)$ 为输入的神经网络来逼近后验分布 $q\left({\mathbf{x}}^{n-1}\mid {\mathbf{x}}^n,{\mathbf{x}}^0\right)$ , 最后达到使用 ${\mathbf{x}}^n与n$ 来预测 ${\mathbf{x}}^{n-1}$ 的效果。