Diffusion学习笔记

---

一些资料

主要部分是对该文章的翻译：
What are Diffusion Models?

提出扩散模型的文章：链接
DDPM：链接
score-based generative modeling：链接

一、Diffusion 模型概述

扩散模型受到非平衡热力学的启发。他们定义了扩散步骤的马尔可夫链，以缓慢地将随机噪声添加到数据中，然后学习反转扩散过程，从噪声中构建所需的数据样本。与VAE或流动模型不同，扩散模型是通过固定的过程学习的，并且潜在变量具有高维度（与原始数据相同）。
Diffusion Model的训练过程通常分为两个阶段：前向扩散过程和后向逆扩散过程。

1、前向扩散过程

这个阶段模拟了从真实数据到随机噪声的过程。给定一个原始数据样本，通过连续的T个时间步长，在每个时间步长上都添加一些随机噪声。这样，随着时间的推移，原始数据逐渐被随机噪声所覆盖，最终形成一个完全由噪声构成的样本。
设$x_0$为从实际数据分布$x_0\sim q\left(x\right)$采样得到的一个样本点。在前向扩散的$T$步过程中，我们通过逐渐向$x_0$中添加少量的高斯噪声，生成了一系列添加了噪声的样本：$x_1,\cdots ,x_T$。每一步的步长由 { β t ∈ ( 0 , 1 ) } t = 1 T \left\lbrace \beta _t \in \left( 0,1\right) \right\rbrace ^T_{t=1} {βt​∈(0,1)}t=1T​决定。
每一步的转移概率服从高斯分布：$$q\left(x_t|x_{t-1}\right)=\mathcal{N}\left(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I\right)$$经过T步有：$$q\left(x1:T|x_0\right)=\prod _{t=1}^{T}q\left(x_t|x_{t-1}\right)$$在这个过程中，随着t的增长，样本$x_0$逐渐失去了精确性。当$T\to \infty$时，$x_t$服从各向同性高斯分布（an Isotropic Gaussian distribution）。
各向同性高斯分布可参考：各向同性高斯分布
上述过程的一个很好的属性就是，我们可以使用重参数技巧在任意时间步$t$以一种“closed form”对$x_t$进行采样。重参数技巧使得高斯分布的随机性转移到一个参数上

重参数技巧（reparameterization trick）:
若$z\sim \mathcal{N}\left(z;\mu ,\sigma I\right)$，则$z=\mu +\sigma \odot ϵ$，其中$ϵ\sim \mathcal{N}\left(0,I\right)，\odot$表示element-wise 乘积。

设${\alpha }_t=1-{\beta }_t$并且$\overline{{\alpha }_t}={\prod }_{i=1}^t{\alpha }_i$，则
$$\begin{array}{rl}{x}_t=& \sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\\ =& \sqrt{{\alpha }_t}\left(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{ϵ}_{t-2}\right)+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\\ =& \sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t}\sqrt{1-{\alpha }_{t-1}}{ϵ}_{t-2}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\\ =& \sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t\left(1-{\alpha }_{t-1}\right)+1-{\alpha }_t}{\overline{ϵ}}_{t-2}\\ =& \sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{ϵ}}_{t-2}\\ =& \cdots \\ =& \sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ\\ q(x_t|x_0)=& \mathcal{N}\left(x_t;\sqrt{\overline{\alpha }}{x}_0,\left(1-{\overline{\alpha }}_t\right)I\right)\end{array}\text{\hspace{1em}（*）}$$其中，${\bar{\alpha }}_t={\prod }_{t=1}^T{\alpha }_t,且{ϵ}_{t-1},{ϵ}_{t-2},\cdots \sim \mathcal{N}\left(0,I\right)$。

解释上式第三行到第四行的转换：
两个独立高斯分布的和，仍然服从高斯分布：
设$x_1$与$x_2$为相互独立的两个随机变量。$x_1\sim \mathcal{N}\left(0,{\sigma }_1^{2}I\right),x_2\sim \mathcal{N}\left(0,{\sigma }_2^{2}I\right)$，则两者之和$x_3=x_1+x_2$满足$x_3\sim \mathcal{N}\left(0,\left({\sigma }_1^2+{\sigma }_2^2\right)I\right)$。
因此，由${ϵ}_{t-1},{ϵ}_{t-2},\cdots \sim \mathcal{N}\left(0,I\right)$，有
$$\sqrt{{\alpha }_t}\sqrt{1-{\alpha }_{t-1}}{ϵ}_{t-2}\sim \mathcal{N}\left(0,{\alpha }_t\left(1-{\alpha }_{t-1}\right)I\right),\sqrt{1-{\alpha }_t}{ϵ}_{t-1}\sim \mathcal{N}\left(0,\left(1-{\alpha }_t\right)I\right)$$因此有${\overline{ϵ}}_{t-2}\sim \mathcal{N}\left(0,\left(1-{\alpha }_t{\alpha }_{t-1}\right)I\right)$。

2、后向逆扩散过程

在后向过程中，模型学习如何通过去噪函数逐步地从随机噪声中恢复原始数据。这个过程可以看作是前向过程的反向操作，即逐步减少噪声并增加对原始数据的恢复。
如果我们能够反向进行上面的操作并且从$q\left(x_{t-1}|x_t\right)$中采样，我们就能够从高斯噪声$x_t\sim \mathcal{N}\left(0,I\right)$中还原原始数据。注意到${\beta }_t$足够小，这个逆过程仍为高斯分布。我们无法容易地估计$q\left(x_{t-1}|x_t\right)$，因为这需要得知整个数据集的分布，因此，我们采取学习模型$p_{\theta }$的方式来估计这个条件概率。
$$\begin{gathered} p_{\theta }\left(x_{0:T}\right)=p\left(x_T\right)\prod _{t=1}^T{p}_{\theta }\left(x_{t-1}|x_t\right) \\ {p}_{\theta }\left(x_{t-1}|x_t\right)=\mathcal{N}\left(x_{t-1};{\mu }_{\theta }\left(x_t,t\right),{\Sigma }_{\theta }\left(x_t,t\right)\right) \end{gathered}$$
值得注意：
当条件中包含$x_0$时，后向逆扩散过程的条件概率是可以求得的。
由贝叶斯公式，有：
$$\begin{array}{rl}q\left(x_{t-1}|x_t,x_0\right)=& q\left(x_t|x_{t-1},x_0\right)\frac{q\left(x_{t-1}|x_0\right)}{q\left(x_t|x_0\right)}\\ \propto & \exp \left(-\frac{1}{2}\left(\frac{{\left(x_t-\sqrt{{\alpha }_t}{x}_{t-1}\right)}^2}{{\beta }_t}+\frac{{\left(x_{t-1}-\sqrt{{\bar{\alpha }}_{t-1}}{x}_0\right)}^2}{1-{\bar{\alpha }}_{t-1}}-\frac{{\left(x_t-\sqrt{{\bar{\alpha }}_t}{x}_0\right)}^2}{1-{\bar{\alpha }}_t}\right)\right)\\ =& \exp \left(-\frac{1}{2}\left(\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{{\left(x_t-\sqrt{{\bar{\alpha }}_t}{x}_0\right)}^2}{1-{\bar{\alpha }}_t}\right)\right)\\ =& \exp \left(-\frac{1}{2}\left(\left(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}\right)x_{t-1}^2-\left(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{2\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0\right)x_{t-1}+C\left(x_t,x_0\right)\right)\right)\end{array}$$
其中，$C\left(x_t,x_0\right)$与$x_{t-1}$无关，因此略去细节。遵循高斯密度函数的形式，可以经过整理得出均值和方差：
$$\begin{array}{rl}{\tilde{\beta }}_t=& \frac{1}{\left(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}\right)}=\frac{1}{\left(\frac{{\alpha }_t-{\bar{\alpha }}_t+{\beta }_t}{{\beta }_t\left(1-{\bar{\alpha }}_{t-1}\right)}\right)}=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t\\ {\tilde{\mu }}_t\left(x_t,x_0\right)=& \frac{\left(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0\right)}{\left(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}\right)}\\ =& \left(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{x}_0\right)\cdot \frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t\\ =& \frac{\sqrt{{\alpha }_t}\left(1-{\bar{\alpha }}_{t-1}\right)}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0\end{array}$$
将（*）式即$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$代入上式消去$x_0$，有：
$${\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t\right)$$

二、训练过程

在训练时，需要使用神经网络来拟合从当前噪声状态一步步回到原始数据的条件概率分布。首先自行复习一下熵，交叉熵和KL散度。

1、变分下界

这个过程与VAE十分类似，因此我们可以使用变分下界（Variational lower bound,VLB）来优化负对数似然（negative log-likelihood）。
$$\begin{array}{rl}-\log p_{\theta }\left(x_0\right)\le & -\log p_{\theta }\left(x_0\right)+D_{KL}\left(q\left(x_{1:T}|x_0\right)\Vert p_{\theta }\left(x_{1:T}|x_0\right)\right)\\ =& -\log p_{\theta }\left(x_0\right)+{\mathbb{E}}_{x_{1:T}\sim q\left(x_{1:T}|x_0\right)}\left[\log \frac{q\left(x_{1:T}|x_0\right)}{p_{\theta }\left(x_{0:T}\right)/p_{\theta }\left(x_0\right)}\right]\\ =& -\log p_{\theta }\left(x_0\right)+{\mathbb{E}}_q\left[\log \frac{q\left(x_{1:T}|x_0\right)}{p_{\theta }\left(x_{0:T}\right)}+\log p_{\theta }\left(x_0\right)\right]\\ =& {\mathbb{E}}_q\left[\log \frac{q\left(x_{1:T}|x_0\right)}{p_{\theta }\left(x_{0:T}\right)}\right]\\ 令L_{VLB}=& {\mathbb{E}}_{q\left(x_{0:T}\right)}\left[\log \frac{q\left(x_{1:T}|x_0\right)}{p_{\theta }\left(x_{0:T}\right)}\right]\end{array}$$
也可以直接使用Jensen不等式：

Jenson不等式：若$f\left(x\right)$是区间$\left[a,b\right]$上的凸函数，则对任意的$x_1,x_2,\cdots ,x_n\in \left[a,b\right]$，则有：
$$f\left(\sum _{i=1}^n\frac{x_i}{n}\right)\ge \frac{{\sum }_{i=1}^{n}f\left(x_i\right)}{n}$$
即$f\left(\mathbb{E}\left(x\right)\right)\ge \mathbb{E}\left(f(x)\right)$

因此，交叉熵损失(Cross Entropy)可以像下面这样：
$$\begin{array}{rl}{L}_{CE}=& -{\mathbb{E}}_{q(x_0)}\log p_{\theta }(x_0)\\ =& -{\mathbb{E}}_{q(x_0)}\log (\int p_{\theta }(x_{0:T})d{x}_{1:T})\\ =& -{\mathbb{E}}_{q(x_0)}\log \left(\int q(x_{1:T}|x_0)\frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}d{x}_{1:T}\right)\\ =& -{\mathbb{E}}_{q(x_0)}\log \left({\mathbb{E}}_{q(x_{1:T}|x_0)}\frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}\right)\\ \le & -{\mathbb{E}}_{q(x_{0:T})}\log \left(\frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}\right)\\ =& {\mathbb{E}}_{q(x_{0:T})}\log \left(\frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right)\\ =& L_{VLB}\end{array}$$倒数第三行用到了Jensen不等式，对凸函数而言，期望的函数大于函数的期望，由于前面有负号，因此这里是小于号。于是，我们只需要优化$L_{VLB}$（很有意思，变分下界损失（$L_{VLB}$）实际上是交叉熵损失（$L_{CE}$）的上界）就可以间接地压缩$L_{CE}$了。
下面，为了能够更好地计算$L_{VLB}$，我们将其转化为多个KL散度和熵的和。
$$\begin{array}{rl}{L}_{VLB}=& {\mathbb{E}}_{q(x_{0:T})}\log \left(\frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right)\\ =& {\mathbb{E}}_q\log \left(\frac{{\prod }_{t=1}^{T}q\left(x_t|x_{t-1}\right)}{p_{\theta }\left(x_T\right){\prod }_{t=1}^T{p}_{\theta }\left(x_{t-1}|x_t\right)}\right)\\ =& {\mathbb{E}}_q\left[-\log p_{\theta }\left(x_T\right)+\sum _{t=1}^T\log \frac{q\left(x_t|x_{t-1}\right)}{p_{\theta }\left(x_{t-1}|x_t\right)}\right]\\ =& {\mathbb{E}}_q\left[-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \frac{q\left(x_t|x_{t-1}\right)}{p_{\theta }\left(x_{t-1}|x_t\right)}+\log \frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right]\\ =& {\mathbb{E}}_q\left[-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \frac{q\left(x_{t-1}|x_t,x_0\right)}{p_{\theta }\left(x_{t-1}|x_t\right)}\frac{q\left(x_t|x_0\right)}{q\left(x_{t-1}|x_0\right)}+\log \frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right]\\ =& {\mathbb{E}}_q\left[-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \frac{q\left(x_{t-1}|x_t,x_0\right)}{p_{\theta }\left(x_{t-1}|x_t\right)}+\sum _{t=2}^T\log \frac{q\left(x_t|x_0\right)}{q\left(x_{t-1}|x_0\right)}+\log \frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right]\\ =& {\mathbb{E}}_q\left[-\log p_{\theta }\left(x_T\right)+\sum _{t=2}^T\log \frac{q\left(x_{t-1}|x_t,x_0\right)}{p_{\theta }\left(x_{t-1}|x_t\right)}+\log \frac{q\left(x_T|x_0\right)}{q\left(x_1|x_0\right)}+\log \frac{q\left(x_1|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}\right]\\ =& {\mathbb{E}}_q\left[\log \frac{q\left(x_T|x_0\right)}{p_{\theta }\left(x_0|x_1\right)}+\sum _{t=2}^T\log \frac{q\left(x_{t-1}|x_t,x_0\right)}{p_{\theta }\left(x_{t-1}|x_t\right)}-\log p_{\theta }\left(x_T\right)\right]\\ =& {\mathbb{E}}_q\left[\underset{L_t}{\underbrace{D_{KL}\left(q\left(x_T|x_0\right)||p_{\theta }\left(x_0|x_1\right)\right)}}+\sum _{t=2}^T\underset{L_{t-1}}{\underbrace{D_{KL}\left(q\left(x_{t-1}|x_t,x_0\right)||p_{\theta }\left(x_{t-1}|x_t\right)\right)}}-\underset{L_0}{\underbrace{\log p_{\theta }\left(x_T\right)}}\right]\end{array}$$
使用上面几个符号来表达$L_{VLB}$更加简易：
$$\begin{array}{rl}{L}_{VLB}=& L_0+L_1+\cdots +L_t+\cdots +L_{T-1}+L_T\\ 其中，& \\ L_T=& D_{KL}\left(q\left(x_T|x_0\right)||p_{\theta }\left(x_0|x_1\right)\right)\\ L_0=& \log p_{\theta }\left(x_T\right)\\ L_t=& D_{KL}\left(q\left(x_{t-1}|x_t,x_0\right)||p_{\theta }\left(x_{t-1}|x_t\right)\right),t\in \left[1,T-1\right]\end{array}$$
上式中的每一个KL项都是两个高斯分布之间的比较，因此可以以“closed form”求解。$L_T$是一个常量，没有可学习的部分，因此在训练中直接忽略。(Ho et al. 2020)使用了一个从$\mathcal{N}\left(x_0;{\mu }_{\theta }\left(x_1,1\right),{\Sigma }_{\theta }\left(x_1,1\right)\right)$中产生的单独的离散decoder来给$L_0$建模。

2、$L_t$的重参数化

为了在后向逆过程中使用神经网络拟合条件概率分布：$p_{\theta }\left(x_{t-1}|x_t\right)=\mathcal{N}\left(x_{t-1};{\mu }_{\theta }\left(x_{t-1},t\right),{\Sigma }_{\theta }\left(x_t,t\right)\right)$，我们需要训练一个${\mu }_{\theta }$来拟合${\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t\right)$。$x_t$作为输入是已知的，因此我们可以对高斯噪声进行重参数化，在时间步$t$，从$x_t$中预测${ϵ}_t$。
$$\begin{array}{rl}{\mu }_{\theta }\left(x_t,t\right)=& \frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }\left(x_t,t\right)\right)\\ 因此，p_{\theta }\left(x_{t-1}|x_t\right)=& \mathcal{N}\left(x_{t-1};\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }\left(x_t,t\right)\right),{\Sigma }_{\theta }\left(x_t,t\right)\right)\end{array}$$
损失$L_t$被重参数化，用来缩小${\mu }_{\theta }$与${\tilde{\mu }}_t$之间的差距：
$$\begin{array}{rl}{L}_t=& {\mathbb{E}}_{x_0,ϵ}\left[\frac{1}{2\Vert {\Sigma }_{\theta }\left(x_t,t\right){\Vert }_2^2}\Vert {\mu }_{\theta }\left(x_t,x_0\right)-{\tilde{\mu }}_t\left(x_t,x_0\right){\Vert }^2\right]\\ =& {\mathbb{E}}_{x_0,ϵ}\left[\frac{1}{2\Vert {\Sigma }_{\theta }\left(x_t,t\right){\Vert }_2^2}\Vert \frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }\left(x_t,t\right)\right)-\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t\right){\Vert }^2\right]\\ =& {\mathbb{E}}_{x_0,ϵ}\left[\frac{{\left(1-{\alpha }_t\right)}^2}{2{\alpha }_t\left(1-{\bar{\alpha }}_t\right)\Vert {\Sigma }_{\theta }{\Vert }^2}\Vert {ϵ}_{\theta }\left(x_t,t\right)-{ϵ}_t{\Vert }^2\right]\\ =& {\mathbb{E}}_{x_0,ϵ}\left[\frac{{\left(1-{\alpha }_t\right)}^2}{2{\alpha }_t\left(1-{\bar{\alpha }}_t\right)\Vert {\Sigma }_{\theta }{\Vert }^2}\Vert {ϵ}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t,t\right)-{ϵ}_t{\Vert }^2\right]\end{array}$$

总结

推导过程不是很复杂，主要是找对什么部分是可学习的，训练就揪着这一部分可劲造。