扩散模型的数学原理

快速理解

[说明]:此节使用通俗语言进行扩散模型的总体概述，用词较为冗杂，某些概念为简易性做了些许牺牲，可能略有偏颇。深入理解请继续阅读后续章节，概念定义也以后续章节为准。

扩散模型是一类生成模型，在训练此类模型时，它可以被分为两个阶段，分别称为扩散阶段和逆扩散阶段。扩散阶段主要对训练数据样本$x_0$不断添加高斯噪声，在重复足够的次数（一般表示为$T$步）后，此数据样本变为各项独立的高斯样本$x_T$，这与VAE中的编码过程神似。逆扩散过程是主要是对一个高斯样本$x_T$不断的进行减高斯噪声的操作，在$T$步后将获得一个符合训练数据分布的样本$x_0^{\prime }$，这与VAE中解码神似。如果此过程被确定，则在推理阶段我们能使用任意的高斯样本逆扩散为一个符合训练数据分布的样本（这达到了生成模型的目的）。此训练过程可以用下图表示：

在训练过程，我们将确定一个加入噪声的策略，这被称作nois_schedule，由于高斯噪声的主要影响参数为均值($\mu$)和标准差($\sigma$)，此策略主要决定的是当前步数($t$)时，加入高斯噪声的均值和标准差。在此策略和训练样本$x_0$被确定的情况下，我们的扩散过程能被逐步推出。所以更重要的是逆扩散过程。
在逆扩散过程中，需要做的是逐步减高斯噪声，如果我们将其定义为直接减去扩散过程中加入的那个高斯噪声则在经过$T$步后就能还原数据样本$x_0$。但这样的定义将毫无意义，因为在推理阶段（从随机高斯样本逆扩散到样本数据），我们并没有加噪声策略的记录，也就无法获得减噪声的策略，逆扩散过程将无法进行。
所以，在训练阶段的逆扩散过程，我们需要使用神经网络来预测这个减噪声的策略，当此策略满足训练集样本的极大似然估计的时侯，推理阶段就可使用此神经网络来生成减噪声策略，从而用随机高斯样本生成符合训练集分布的生成样本。
神经网络损失函数的定义至关重要。在训练过程，由于已知样本数据$x_0$和加噪声策略，我们能轻易求出$x_T$和$x_{T-1}$。所以神经网络可以设定为接收已知的信息$x_T$和$T$为输入，预测其减噪声后的结果$x_{T-1}^{\prime }$。已知预测值$x_{T-1}^{\prime }$和真实值$x_{T-1}$，损失函数则呼之欲出。（注意，此处不具有数学严谨性，只为描述神经网络的建模想法）

数学基础

[说明]:下章节开始介绍扩散模型数学原理，对此数学基础的掌握必不可少。为简单起见，读者可不深究此节相关内容，但需要在阅读下节时实时翻阅，以作为推到的理论基础。

高斯分布表示 （数1）

符号表示：$N(\mu ,{\sigma }^2)$。 其中$\mu$为均值，$\sigma$为标准差。
概率密度函数：
$$f(x)=\frac{1}{\sqrt{2\pi }\sigma }exp(-\frac{1}{2}\frac{(x-\mu {)}^2}{{\sigma }^2})=\frac{1}{\sqrt{2\pi }\sigma }exp(-\frac{1}{2}(\frac{1}{{\sigma }^2}{x}^2-\frac{2\mu }{{\sigma }^2}x+\frac{{\mu }^2}{{\sigma }^2}))$$

高斯分布线性叠加 (数2)

两个高斯分布$X\sim N({\mu }_1,{\sigma }_1^2)$，$Y\sim N({\mu }_2,{\sigma }_2^2)$进行线性组合仍然为高斯分布，表示为：
$$aX+bY\sim N(a{\mu }_1+b{\mu }_2,a^2{\sigma }_1^2+b^2{\sigma }_2^2)$$

重参数化技巧 (数3)

设标准高斯分布$N(0,1)$随机采样的样本为$z$,高斯分布$N(\mu ,{\sigma }^2)$随机采样的样本表示为$z^{\prime }$。使用重参数化技巧，可以将$z^{\prime }$表示为：
$$z^{\prime }=\sigma z+\mu$$

贝叶斯公式 (数4)

$$p(A|B)=p(B|A)\frac{p(A)}{p(B)}$$

指数乘除法 （数5）

$$exp(a)\ast exp(b)/exp(c)=exp(a+b-c)$$

扩散过程

给定初始分布$x_0\sim q(x)$（其中$x_0$代表某各训练样本，$q(x)$代表训练数据集分布），可以不断向其中加入高斯噪声。规定此策略为：在$t$时刻，对$x_{t-1}$进行$\sqrt{1-{\beta }_t}$倍的数值缩放，并向其中加入${\beta }_t$倍标准高斯噪声$z_{t-1}$得到数据$x_t$。此过程表示为：
$$x_t=\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}{z}_{t-1}(1)$$
随着t不断增大，最终达到T时，样本数据($x_T$)成了一个各项独立的高斯分布。且由于策略已知，$x_T$能通过迭代计算求得。为了节约计算资源并为一些后续工作做准备，可以通过公式推导求得$x_t$和$x_0$之间的关系，这使得任意时刻的$x_t$能通过公式一步求得。具体公式推导如下：
为方便表示，令${\alpha }_t=1-{\beta }_t$，${\bar{\alpha }}_t={\prod }_{i=1}^T{\alpha }_i$。要求给定$x_{t-1}$情况下$x_t$的分布函数$q(x_t|x_{t-1})$。由(1)和(数3)进行比对可得：
$$\begin{array}{rl}q(x_t|x_{t-1})& =N(\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_t)(2)\\ & =N(\sqrt{{\alpha }_t}{x}_{t-1},1-{\alpha }_t)(3)\end{array}$$
对$x_t$进行递归拆解：
$$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_{t-1}(5)where{z}_{t-1},z_{t-2},...,z\sim N(0,1)\\ & =\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{z}_{t-2})+\sqrt{1-{\alpha }_t}{z}_{t-1}(6)\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{z}_{t-2}+\sqrt{1-{\alpha }_t}{z}_{t-1}(7)\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{z}_{t-2}(8)\\ & =...\\ & =\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}z(9)\end{array}$$
式(7)到式(8)的解释：
由(数3)及(数2)可知，$\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{z}_{t-2}+\sqrt{1-{\alpha }_t}{z}_{t-1}$可重参数化为一个只含标准高斯$z$的$\sqrt{1-{\alpha }_t{\alpha }_{t-1}}z$,又因为$z_{t-2},z$都是标准高斯，进行等效替换后为$\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{z}_{t-2}$。

逆扩散过程

如果你已经能理解上一节的内容，那祝贺你拥有非凡的语文阅读能力及数学理解能力。这一节的内容可能略难，我对自己的表达能力与数学水平并不是很自信，但拥有如此理解能力的你想必也能轻松阅读这一节，让我们开始吧。
故名思意，逆扩散过程必定是扩散过程的反过程。在上一节中，我们成功求出了$q(x_t|x_{t-1})$的表达式，这使得我们可以轻松求得任何扩散时刻下的结果。同样，如果我们在逆过程中能求得$q(x_{t-1}|x_t)$的表达式，这个逆过程也将成为我们的囊中之物。下面让我们开始吧：
$$\begin{array}{rl}q(x_{t-1}|x_t)& =q(x_t|x_{t-1})\ast \frac{q(x_{t-1})}{q(x_t)}(10)...由(数4)可知\\ & =(\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_{t-1})\ast \frac{\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}}z}{\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}z}(11)...由(5),(9)可知\\ & =(\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{{\beta }_t}{z}_{t-1})\ast \frac{\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}}z}{\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}z}(12)\\ & \propto exp(-\frac{(x_t-\sqrt{{\alpha }_t}{x}_{t-1}{)}^2}{2{\beta }_t})\ast exp(-\frac{(x_{t-1}-\sqrt{{\bar{\alpha }}_{t-1}}{x}_0{)}^2}{2(1-{\bar{\alpha }}_{t-1})})/exp(-\frac{(x_t-\sqrt{{\bar{\alpha }}_t}{x}_0{)}^2}{2(1-{\bar{\alpha }}_t)})(13)...由(数3),(数1)可知\\ & =exp[-\frac{1}{2}(\frac{(x_t-{\sqrt{\alpha }}_t{x}_{t-1}{)}^2}{{\beta }_t}+\frac{(x_{t-1}-\sqrt{{\bar{\alpha }}_{t-1}}{x}_0{)}^2}{1-{\bar{\alpha }}_{t-1}}-\frac{(x_t-\sqrt{{\bar{\alpha }}_t}{x}_0{)}^2}{1-{\bar{\alpha }}_t})](14)...由(数5)可知\\ & =exp[-\frac{1}{2}(\frac{x_t^2-2{\sqrt{\alpha }}_t{x}_t{x}_{t-1}+{\alpha }_t{x}_{t-1}^2}{{\beta }_t}+\frac{x_{t-1}^2-2\sqrt{{\bar{\alpha }}_{t-1}}{x}_0{x}_{t-1}+{\bar{\alpha }}_{t-1}{x}_0^2}{1-{\bar{\alpha }}_{t-1}}-\frac{(x_t-\sqrt{{\bar{\alpha }}_t}{x}_0{)}^2}{1-{\bar{\alpha }}_t})](15)\\ & =exp[-\frac{1}{2}((\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})x_{t-1}^2-(\frac{2{\sqrt{\alpha }}_t}{{\beta }_t}{x}_t+\frac{2\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0)x_{t-1}+C)](16)...由于x_0,x_t不含x_{t-1}项被整合为常数C\end{array}$$
与(数1)公式进行对比，可知：
$$\begin{array}{rl}\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}& \iff \frac{1}{{\sigma }^2}(17)\\ & \\ \frac{2{\sqrt{\alpha }}_t}{{\beta }_t}{x}_t+\frac{2\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0& \iff \frac{2\mu }{{\sigma }^2}(18)\\ & \\ q(x_{t-1}|x_t)& =N(\mu ,\sigma )\end{array}$$
将(17)带入(18)得：
$$\frac{{\sqrt{\alpha }}_t}{{\beta }_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0=(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})\mu (19)$$
进行整理得：
$$\begin{array}{rl}\mu & =\frac{\frac{{\sqrt{\alpha }}_t}{{\beta }_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0}{\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}}(20)\\ & =\frac{(1-{\bar{\alpha }}_{t-1})\sqrt{{\alpha }_t}{x}_t+{\beta }_t\sqrt{{\bar{\alpha }}_{t-1}}{x}_0}{(1-{\bar{\alpha }}_{t-1}){\alpha }_t+{\beta }_t}\\ & =\frac{(1-{\bar{\alpha }}_{t-1})\sqrt{{\alpha }_t}{x}_t+{\beta }_t\sqrt{{\bar{\alpha }}_{t-1}}{x}_0}{{\alpha }_t-{\bar{\alpha }}_{t-1}{\alpha }_t+1-{\alpha }_t}(21)\\ & =\frac{(1-{\bar{\alpha }}_{t-1})\sqrt{{\alpha }_t}{x}_t+{\beta }_t\sqrt{{\bar{\alpha }}_{t-1}}{x}_0}{1-{\bar{\alpha }}_{t-1}{\alpha }_t}(22)\\ & =\frac{(1-{\bar{\alpha }}_{t-1})\sqrt{{\alpha }_t}{x}_t+{\beta }_t\sqrt{{\bar{\alpha }}_{t-1}}{x}_0}{1-{\bar{\alpha }}_t}(23)\end{array}$$
将(9)变形得$x_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}z)$，(23)中可用$x_t$代替$x_0$变为：
$$\begin{array}{rl}\mu & =\frac{(1-{\bar{\alpha }}_{t-1})\sqrt{{\alpha }_t}{x}_t+{\beta }_t\sqrt{{\bar{\alpha }}_{t-1}}\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}z)}{1-{\bar{\alpha }}_t}(24)\\ & =\frac{(1-{\bar{\alpha }}_{t-1})\sqrt{{\alpha }_t}{x}_t+{\beta }_t\frac{1}{\sqrt{{\alpha }_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}z)}{1-{\bar{\alpha }}_t}(25)\\ & =\frac{(1-{\bar{\alpha }}_{t-1}){\alpha }_t{x}_t+{\beta }_t(x_t-\sqrt{1-{\bar{\alpha }}_t}z)}{(1-{\bar{\alpha }}_t)\sqrt{{\alpha }_t}}(26)\\ & =\frac{(1-{\bar{\alpha }}_{t-1}){\alpha }_t+{\beta }_t}{(1-{\bar{\alpha }}_t)\sqrt{{\alpha }_t}}{x}_t-\frac{{\beta }_t\sqrt{1-{\bar{\alpha }}_t}}{(1-{\bar{\alpha }}_t)\sqrt{{\alpha }_t}}z(27)\\ & =\frac{1}{\sqrt{a_t}}(\frac{(1-{\bar{\alpha }}_{t-1}){\alpha }_t+{\beta }_t}{1-{\bar{\alpha }}_t}{x}_t-\frac{{\beta }_t\sqrt{1-{\bar{\alpha }}_t}}{1-{\bar{\alpha }}_t}z)(28)\\ & =\frac{1}{\sqrt{a_t}}(\frac{{\alpha }_t-{\bar{\alpha }}_t+1-{\alpha }_t}{1-{\bar{\alpha }}_t}{x}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}z)(29)\\ & =\frac{1}{\sqrt{a_t}}(\frac{1-{\bar{\alpha }}_t}{1-{\bar{\alpha }}_t}{x}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}z)(30)\\ & =\frac{1}{\sqrt{a_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}z)(31)\end{array}$$
$\sigma$的求解更为简单，直接对(17)进行变形即可：
$$\begin{array}{rl}\frac{1}{{\sigma }^2}& =\frac{{\alpha }_t(1-{\bar{\alpha }}_{t-1})+{\beta }_t}{{\beta }_t(1-{\bar{\alpha }}_{t-1})}(32)\\ {\sigma }^2& =\frac{{\beta }_t(1-{\bar{\alpha }}_{t-1})}{{\alpha }_t(1-{\bar{\alpha }}_{t-1})+{\beta }_t}(33)\\ & \approx \frac{{\beta }_t}{{\alpha }_t+{\beta }_t}={\beta }_t(34)...由{\alpha }_t=1-{\beta }_t可知\\ \sigma & =\sqrt{{\beta }_t}(35)\end{array}$$
由于$0<{\alpha }_t<1$，当t足够大有${\bar{\alpha }}_t\approx 0$，所以(33)到(34)公式成立。
如果$\mu ,\sigma$表达式已知，则条件概率分布$q(x_{t-1}|x_t)=N(\mu ,\sigma )$已知，使得我们从一个随机高斯样本$x_t$通过循环迭代能求得符合真实数据分布的样本$x_0$。
事实上，我们离上诉描述已经很接近了，但在$\mu$的表达式中，我们存在唯一一个未知量，这就是$z$。在训练时由于策略已知，$z$可以人为是已知量。但在推理阶段，减噪声策略未知，$z$被视为未知量。此时，神经网络出场了。在训练阶段，我们只需训练一个模型用来预测$z^{\prime }$，通过与已知的$z$做损失对比，完成训练后，使用此网络来求得一系列的$z^{\prime }$就能完成整个逆扩散过程。根据论文，模型的输入被定义为$x_t和t$，损失函数则可表示为下式：
$$\begin{array}{rl}{\mathcal{L}}_{\text{simple}}& ={\mathbb{E}}_{t,{\mathbf{x}}_t,z}\left[\Vert z-{ϵ}_{\theta }(x_t,t){\Vert }^2\right]\\ & ={\mathbb{E}}_{t,{\mathbf{x}}_0,z}\left[\Vert z-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}z,t){\Vert }^2\right]\end{array}$$
至此，我们的扩散过程与逆扩散过程可被公式完整表达（其中$z$表示随机高斯采样）：
正向：
$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}z$$
逆向:
$$x_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-\overline{{\alpha }_t}}}{ϵ}_{\theta }(x_t,t))+\sqrt{{\beta }_t}z$$

代码

https://colab.research.google.com/drive/179Se6GXKB30RQnVt7uqvS-vKsBQnRAee

参考链接

https://t.bilibili.com/700526762586538024?spm_id_from=333.999.0.0
https://zhuanlan.zhihu.com/p/586362713