【学习笔记】扩散模型的基本原理与训练方法

1. Diffusion前向过程

给定采样自数据集的真实图片$x_0\sim q(x)$，并对其添加高斯噪声，共进行$T$步，称该过程为$q$过程，能够得到添加噪声后的图片分布$x_1,x_2,...,x_T$。

将加噪的过程看作一个马尔可夫过程，即$t$时刻的状态只与$t-1$时刻有关，设置超参数${\beta }_t\in (0,1),t\in (1,T)$ ，本质上该超参数即为每一时刻下添加的高斯分布的方差。至此可以将该前向$q$过程写成以下形式，表示在$x_{t-1}$满足的分布前提下，$x_t$的方差为${\beta }_t$，均值则受到了前一个状态的影响。
$$q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_t\mathbf{I})$$
实际场景中，$\beta$是随着$t$的增大而递增的，那么当$t$很大时，${\beta }_t$ 趋近于1，则$x_t$满足的分布接近于标准正态分布。【${\beta }_t$这一超参数的存在本质是描述方差，但却在均值上乘上了$\sqrt{1-{\beta }_t}$，这能够使得均值在$t$很大的时候趋向于0，整个分布则为标准正态分布。】

由于以上过程看作一个马尔可夫过程，因此根据乘法公式可以写出以下表达式：
$$q(x_{1:T}|x_0)=\prod _{t=1}^{T}q(x_t|x_{t-1})$$

重参数技巧（reparameterization trick）

以上前向的$q$过程的表达形式是采样形式的，表示在某个分布中进行随机采样，这回导致过程有随机性，无法反向传播梯度，为了使该采样过程变得可导，使用冲参数技巧，引入一个固定的随机变量$ϵ$实现。

例如需要进行以下采样：$z\sim \mathcal{N}(z;{\mu }_{\theta },{\sigma }_{\theta }^2\mathbf{I})$，可以写成：$z={\mu }_{\theta }+{\sigma }_{\theta }\odot ϵ$，这样的话$z$依旧是一个随机变量，但对于${\mu }_{\theta }$和${\sigma }_{\theta }$等含有网络参数的参数，能够通过$z$进行梯度的求取，随机性完全来自于固定不变服从标注正态分布的$\epsilon $中了。

那么对于原先的$q$过程，可以使用重参数技巧重写$x_t$，即：
$$x_t=\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}z，z\sim \mathcal{N}(0,\mathbf{I})$$

使用统一的方法表示$x_t$

由于上述扩散过程是一步步进行的，为了能够快速得到$x_t$，考虑使用$x_0$和$\beta$进行统一表示。

首先假设${\alpha }_t=1-{\beta }_t$，并记$\overline{{\alpha }_t}={\prod }_{i=1}^T{\alpha }_i$。

则对$x_t$有以下推导过程：

$$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_1\\ & =\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{z}_2)+\sqrt{1-{\alpha }_t}{z}_1\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{z}_2+\sqrt{1-{\alpha }_t}{z}_1\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}\overline{z_2}\\ & ...\\ & =\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-\overline{{\alpha }_t}}\overline{z_t}\end{array}$$

其中对于$\sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{z}_2+\sqrt{1-{\alpha }_t}{z}_1=\sqrt{1-{\alpha }_t{\alpha }_{t-1}}\overline{z_2}$的推导过程，有以下分析：

$$\begin{gathered} \sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{z}_2\sim \mathcal{N}(0,{\alpha }_t(1-{\alpha }_{t-1})\mathbf{I}) \\ \sqrt{1-{\alpha }_t}{z}_1\sim \mathcal{N}(0,(1-{\alpha }_t)\mathbf{I}) \\ \sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{z}_2+\sqrt{1-{\alpha }_t}{z}_1\sim \mathcal{N}(0,(1-{\alpha }_t{\alpha }_{t-1})\mathbf{I}) \end{gathered}$$

对于得到的混合高斯分布$\mathcal{N}(0,(1-{\alpha }_t{\alpha }_{t-1})\mathbf{I})$，可以表达为$\sqrt{1-{\alpha }_t{\alpha }_{t-1}}\overline{z_2}$，对于$\overline{z_2}$依然服从的是标准高斯分布。

综上，对于$x_t$的表达式形式为$x_t=\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-\overline{{\alpha }_t}}z$，采样形式为$q(x_t|x_0)\sim \mathcal{N}(x_t;\sqrt{\overline{{\alpha }_t}}{x}_0,(1-\overline{{\alpha }_t})\mathbf{I})$。可以看成原始图片与高斯噪声的加权求和结果。

2. Diffusion逆向推断

该过程可以看成前向$q$过程的逆过程，即去噪过程。目前在这里我们已知所有的前向过程$q(x_t|x_{t-1})$，想要从标准正态分布中逐步去噪得到最终的原图分布，即从$x_T$得到$x_0$，那么就需要知道$q(x_{t-1}|x_t)$。但这是十分困难的，因此考虑使用深度神经网络来对该分布进行预测，即训练一个模型使其能够做到$p_{\theta }(x_{t-1}|x_t)$。【即让神经网络去学习去噪的过程】

对于$p_{\theta }(x_{t-1}|x_t)$，写成正态分布的表达形式，即为下式，可以看到我们需要使网络根据$x_t$和$t$学习得到分布的均值与方差。

$$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$$

虽然$q(x_{t-1}|x_t)$难以直接得到，但是可以引入已知的原图像分布$x_0$的先验知识，尝试得到$q(x_{t-1}|x_t,x_0)$，可以进行以下过程推导：

$$\begin{array}{rl}q(x_{t-1}|x_t,x_0)& =\frac{q(x_{t-1},x_t,x_0)}{q(x_t,x_0)}\\ & =\frac{q(x_{t-1},x_t,x_0)}{q(x_{t-1},x_0)}\frac{q(x_{t-1},x_0)}{q(x_t,x_0)}\\ & =q(x_t|x_{t-1},x_0)\frac{\frac{q(x_{t-1},x_0)}{q_{x_0}}}{\frac{q(x_t,x_0)}{q_{x_0}}}\\ & =q(x_t|x_{t-1},x_0)\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}\end{array}$$

由此将加入先验知识的逆向分布转换为正向分布的表达式。

对于几个前向过程进行表达式说明：

$$\begin{gathered} q(x_{t-1}|x_0)=\sqrt{\overline{{\alpha }_{t-1}}}{x}_0+\sqrt{1-\overline{{\alpha }_{t-1}}}z\sim \mathcal{N}(\sqrt{\overline{{\alpha }_{t-1}}}{x}_0,1-\overline{{\alpha }_{t-1}}) \\ q(x_t|x_0)=\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-\overline{{\alpha }_t}}z\sim \mathcal{N}(\sqrt{\overline{{\alpha }_t}}{x}_0,1-\overline{{\alpha }_t}) \end{gathered}$$

由于$q$过程是马尔可夫过程，则有：

$$q(x_t|x_{t-1},x_0)=q(x_t|x_{t-1})=\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}z\sim \mathcal{N}(\sqrt{{\alpha }_t}{x}_{t-1},1-{\alpha }_t)$$

而对于高斯分布，我们能够使用概率密度函数的形式进行表达，即：$\mathcal{N}(\mu ,{\sigma }^2)\propto exp(-\frac{(x-\mu {)}^2}{2{\sigma }^2})$，则考虑使用概率密度函数描述$q(x_{t-1}|x_0)、q(x_t|x_0)、q(x_t|x_{t-1},x_0)$，如下表达式所示：

$$\begin{gathered} q(x_{t-1}|x_0)\propto exp(-\frac{(x_{t-1}-\sqrt{\overline{{\alpha }_{t-1}}}{x}_0{)}^2}{2(1-\overline{{\alpha }_{t-1}})}) \\ q(x_t|x_0)\propto exp(-\frac{(x_t-\sqrt{\overline{{\alpha }_t}}{x}_0{)}^2}{2(1-\overline{{\alpha }_t})}) \\ q(x_t|x_{t-1},x_0)\propto exp(-\frac{(x_t-\sqrt{{\alpha }_t}{x}_{t-1}{)}^2}{2(1-{\alpha }_t)}) \end{gathered}$$

则最终可以表达$q(x_{t-1}|x_t,x_0)$，如下推导所示。由于$q(x_{t-1}|x_t,x_0)$是关于$x_{t-1}$的表达式，则需要进行同类相合并

$$\begin{array}{rl}q(x_{t-1}|x_t,x_0)& =q(x_t|x_{t-1},x_0)\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}\\ & \propto exp(-\frac{(x_{t-1}-\sqrt{\overline{{\alpha }_{t-1}}}{x}_0{)}^2}{2(1-\overline{{\alpha }_{t-1}})})+exp(-\frac{(x_t-\sqrt{{\alpha }_t}{x}_{t-1}{)}^2}{2(1-{\alpha }_t)})-exp(-\frac{(x_t-\sqrt{\overline{{\alpha }_t}}{x}_0{)}^2}{2(1-\overline{{\alpha }_t})})\\ & =exp(-\frac{1}{2}(\frac{(x_{t-1}-\sqrt{\overline{{\alpha }_{t-1}}}{x}_0{)}^2}{1-\overline{{\alpha }_{t-1}}}+\frac{(x_t-\sqrt{{\alpha }_t}{x}_{t-1}{)}^2}{1-{\alpha }_t}-\frac{(x_t-\sqrt{\overline{{\alpha }_t}}{x}_0{)}^2}{1-\overline{{\alpha }_t}}))\\ & =exp(-\frac{1}{2}((\frac{{\alpha }_t}{1-{\alpha }_t}+\frac{1}{1-\overline{{\alpha }_{t-1}}})x_{t-1}^2-(\frac{2\sqrt{{\alpha }_t}}{1-{\alpha }_t}{x}_t+\frac{2\sqrt{\overline{{\alpha }_{t-1}}}}{1-\overline{{\alpha }_{t-1}}}{x}_0)x_{t-1}+C(x_t,x_0)))\end{array}$$

其中$C(x_t,x_0)$不含$x_{t-1}$。

考虑到对于概率密度的表达式，可以进行展开操作，即$\mathcal{N}(\mu ,{\sigma }^2)\propto exp(-\frac{(x-\mu {)}^2}{2{\sigma }^2})=exp(-\frac{1}{2}(\frac{1}{{\sigma }^2}{x}^2-\frac{2\mu }{{\sigma }^2}x+\frac{{\mu }^2}{{\sigma }^2}))$，与$q(x_{t-1}|x_t,x_0)$的表达式一一对应可知。

$$\begin{gathered} \frac{1}{{\sigma }^2}=\frac{{\alpha }_t}{1-{\alpha }_t}+\frac{1}{1-\overline{{\alpha }_{t-1}}} \\ {\sigma }^2=\frac{(1-{\alpha }_t)(1-\overline{{\alpha }_{t-1}})}{{\alpha }_t(1-\overline{{\alpha }_{t-1}})+(1-{\alpha }_t)}=\frac{{\beta }_t(1-\overline{{\alpha }_{t-1}})}{{\alpha }_t(1-\overline{{\alpha }_{t-1}})+{\beta }_t} \\ =\frac{1-\overline{{\alpha }_{t-1}}}{1-\overline{{\alpha }_t}}{\beta }_t \end{gathered}$$

对于均值，有求得表达式：

$$\begin{gathered} \frac{2\mu }{{\sigma }^2}=\frac{2\sqrt{{\alpha }_t}}{1-{\alpha }_t}{x}_t+\frac{2\sqrt{\overline{{\alpha }_{t-1}}}}{-\overline{{\alpha }_{t-1}}}{x}_0 \\ \mu =\frac{\sqrt{{\alpha }_t}(1-\overline{{\alpha }_{t-1}})}{1-\overline{{\alpha }_t}}{x}_t+\frac{\sqrt{\overline{{\alpha }_{t-1}}}{\beta }_t}{1-\overline{{\alpha }_t}}{x}_0 \end{gathered}$$

至此，我们在加入$x_0$的先验知识后，能够描述出分布$q(x_{t-1}|x_t,x_0)$的均值和方差的表达式，即

$$\begin{gathered} \mu =\frac{\sqrt{{\alpha }_t}(1-\overline{{\alpha }_{t-1}})}{1-\overline{{\alpha }_t}}{x}_t+\frac{\sqrt{\overline{{\alpha }_{t-1}}}{\beta }_t}{1-\overline{{\alpha }_t}}{x}_0 \\ {\sigma }^2=\frac{1-\overline{{\alpha }_{t-1}}}{1-\overline{{\alpha }_t}}{\beta }_t \end{gathered}$$

但实际上我们需要通过神经网络来训练出一个分布$p_{\theta }(x_{t-1}|x_t)$，该分布的均值和方差均是含参数的，即上述的${\mu }_{\theta }(x_t,t)$与${\Sigma }_{\theta }(x_t,t)$。我们考虑使用$q(x_{t-1}|x_t,x_0)$来近似估计$p_{\theta }(x_{t-1}|x_t)$，这显然是可行的，因为$x_t$可以由$x_0$来表示，同理也可以使用$x_t$来表示$x_0$。综上可得以下关于均值和方差的表达式：

$$\begin{gathered} {\mu }_{\theta }(x_t,t)=\frac{\sqrt{{\alpha }_t}(1-\overline{{\alpha }_{t-1}})}{1-\overline{{\alpha }_t}}{x}_t+\frac{\sqrt{\overline{{\alpha }_{t-1}}}{\beta }_t}{1-\overline{{\alpha }_t}}{x}_0 \\ {\Sigma }_{\theta }(x_t,t)=\frac{1-\overline{{\alpha }_{t-1}}}{1-\overline{{\alpha }_t}}{\beta }_t \end{gathered}$$

对于其中出现的先验知识$x_0$，考虑使用$x_t$来表达，即有：
$$x_0=\frac{x_t-\sqrt{1-\overline{{\alpha }_t}}z}{\sqrt{\overline{{\alpha }_t}}}$$
代入均值${\mu }_{\theta }(x_t,t)$的表达式，有以下推导过程：

$$\begin{array}{rl}{\mu }_{\theta }(x_t,t)& =\frac{\sqrt{{\alpha }_t}(1-\overline{{\alpha }_{t-1}})}{1-\overline{{\alpha }_t}}{x}_t+\frac{\sqrt{\overline{{\alpha }_{t-1}}}{\beta }_t}{1-\overline{{\alpha }_t}}{x}_0\\ & =\frac{\sqrt{{\alpha }_t}(1-\overline{{\alpha }_{t-1}})}{1-\overline{{\alpha }_t}}{x}_t+\frac{\sqrt{\overline{{\alpha }_{t-1}}}{\beta }_t}{1-\overline{{\alpha }_t}}\frac{x_t-\sqrt{1-\overline{{\alpha }_t}}z}{\sqrt{\overline{{\alpha }_t}}}\\ & =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-\overline{{\alpha }_t}}}z)\end{array}$$

观察表达式可以发现，${\alpha }_t$、$x_t$、$\overline{{\alpha }_t}$均为已知量，模型需要确定${\mu }_{\theta }(x_t,t)$，本质上是个噪声预测的过程，即对于上式中的$z$，需要交给神经网络去预测，记为$z_{\theta }(x_t,t)$。

综上，反向去噪过程可以概括为：

1. 根据$x_t$和$t$预测高斯噪声$z_{\theta }(x_t,t)$，则能够得到模型预测的去噪后的分布均值${\mu }_{\theta }(x_t,t)$，其实就是去噪后的图像。
2. 得到方差${\Sigma }_{\theta }(x_t,t)$，在DDPM中方差是untrained的，即${\Sigma }_{\theta }(x_t,t)=\frac{1-\overline{{\alpha }_{t-1}}}{1-\overline{{\alpha }_t}}{\beta }_t$，但对于方差也可以进行训练估计。
3. $p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$进行去噪，获得$x_{t-1}$，通过重参数技巧。

3. 训练方式与策略

常规的训练流程

1. 从数据集中采样得到$x_0$，$x_0\sim q(x)$，并在$\mathrm{1...}T$中随机采样一个 $t$。
2. 从标准高斯分布中采样一个噪声$z\sim \mathcal{N}(0,\mathbf{I})$。
3. 根据重参数技巧得到$x_t=\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-\overline{{\alpha }_t}}z$。
4. 训练UNet，输入$x_t$和$t$，模型输出$z_{\theta }(x_t,t)$，将$z_{\theta }(x_t,t)$和$z$做Loss。

推理流程

1. 从标准高斯分布中采样得到一个噪声$x_t\sim \mathcal{N}(0,\mathbf{I})$。
2. 从$T$到1遍历变量$t$，如果$t==1$则$z=0$，否则采样噪声$z\sim \mathcal{N}(0,\mathbf{I})$。
3. UNet推理得到$z_{\theta }(x_t,t)$，进行去噪$x_{t-1}=\frac{1}{\sqrt{\overline{{\alpha }_t}}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-\overline{{\alpha }_t}}}{z}_{\theta }(x_t,t))+{\Sigma }_{\theta }(x_t,t)z$。
4. 得到$x_0$。

4. Diffusion的应用拓展

Diffusion做分割（SegDiff）

要点：将待分割图像作为condition image进行特征的提取，在进入UNet之前进行特征的融合，扩散生成的为mask。

训练${\mu }_{\theta }(x_t,t)$时，需要加入图像的特征，即${\mu }_{\theta }(x_t,t,I)$，推理时同理。

Diffusion做分割（SegDiff）

[外链图片转存中…(img-fszHBj2E-1715260213416)]

要点：将待分割图像作为condition image进行特征的提取，在进入UNet之前进行特征的融合，扩散生成的为mask。

训练${\mu }_{\theta }(x_t,t)$时，需要加入图像的特征，即${\mu }_{\theta }(x_t,t,I)$，推理时同理。