DDPM扩散模型公式推理----损失函数

DDPM扩散模型公式推理----损失函数

2.3 损失函数推导

我们用极大似然估计的思想构建损失函数：
$$\mathcal{L}=-\log p_{\theta }\left(x_0\right)$$
即逆扩散网络的参数 $\theta$ 使刚开始采样的数据 $x_0$ 出现的概率最大。

接下来需要对上式进行变形，用到了一些ELBO和VAE的内容。

2.3.1 ELBO

已知极大似然估计 $p_{\theta }\left(x_0\right)$ 和观测值 $x_0$ ，以及由观测值扩散得到的 $x_{1:T}$ ,由边缘概率分布公式：
$$p_{\theta }({\mathbf{x}}_0)=\int p_{\theta }(x_{0:T})d{x}_{1:T}$$
因此
$$\begin{array}{rl}\log p_{\theta }({\mathbf{x}}_0)& =\log \int p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})d{\mathbf{x}}_{1:\mathbf{T}}\\ & =\log \int \frac{p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})q_{\varphi }({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}d{\mathbf{x}}_{1:\mathbf{T}}\\ & =\log {\mathbb{E}}_{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\frac{p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})}{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\right]\\ & \ge {\mathbb{E}}_{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})}{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\right]\end{array}$$
最后一步是由琴声不等式 $(Jense{n}^{\prime }sInequality)$ 得来的。
这样看来不是很直观，还有一种推导方式更简单一些：
$$\begin{array}{rl}\log p(\mathbf{x})& =\log p(\mathbf{x})\int q_{\mathbf{\varphi }}(\mathbf{z}\mid \mathbf{x})dz\\ & =\int q_{\mathbf{\varphi }}(\mathbf{z}\mid \mathbf{x})(\log p(\mathbf{x}))dz\\ & ={\mathbb{E}}_{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}[\log p(\mathbf{x})]\\ & ={\mathbb{E}}_{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{p(\mathbf{x},\mathbf{z})}{p(\mathbf{z}\mid \mathbf{x})}\right]\\ & ={\mathbb{E}}_{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{p(\mathbf{x},\mathbf{z})q_{\varphi }(\mathbf{z}\mid \mathbf{x})}{p(\mathbf{z}\mid \mathbf{x})q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\right]\\ & ={\mathbb{E}}_{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{p(\mathbf{x},\mathbf{z})}{q_{\mathbf{\varphi }}(\mathbf{z}\mid \mathbf{x})}\right]+{\mathbb{E}}_{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}{p(\mathbf{z}\mid \mathbf{x})}\right]\\ & ={\mathbb{E}}_{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{p(\mathbf{x},\mathbf{z})}{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\right]+D_{\mathrm{KL}}\left(q_{\mathbf{\varphi }}(\mathbf{z}\mid \mathbf{x})\Vert p(\mathbf{z}\mid \mathbf{x})\right)\\ & \ge {\mathbb{E}}_{q_{\varphi }(\mathbf{z}\mid \mathbf{x})}\left[\log \frac{p(\mathbf{x},\mathbf{z})}{q_{\mathbf{\varphi }}(\mathbf{z}\mid \mathbf{x})}\right]\\ & ={\mathbb{E}}_{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})}{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\right]\end{array}$$
这里的 $z$ 代表 $x_{1:T}$，$x$ 代表 $x_0$。中间用到了贝叶斯公式。

这里两种方法推导的结论都为:
$$\log p_{\theta }({\mathbf{x}}_0)\ge {\mathbb{E}}_{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})}{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\right]$$

其中 ${\mathbb{E}}_{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})}{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\right]$ 就是ELBO $(EvidenceLowerBound)$，即变分下界。

我们要使损失函数 $-\log p_{\theta }\left(x_0\right)$ 最小，就是另ELBO ${\mathbb{E}}_{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})}{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\right]$ 最大。

我们令 $L_{VLB}=-{\mathbb{E}}_{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})}{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\right]={\mathbb{E}}_{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}\left[\log \frac{q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:\mathbf{T}})}\right]$，接下来就转换为对该损失函数求解。再进一步分解：
$$\begin{array}{rl}{L}_{\mathrm{VLB}}& ={\mathbb{E}}_{q\left({\mathbf{x}}_{0:T}\right)}\left[\log \frac{q\left({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{0:T}\right)}\right]\\ & ={\mathbb{E}}_q\left[\log \frac{{\prod }_{t=1}^{T}q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1}\right)}{p_{\theta }\left({\mathbf{x}}_T\right){\prod }_{t=1}^T{p}_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+\sum _{t=1}^T\log \frac{q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1}\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+\sum _{t=2}^T\log \frac{q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1}\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)}+\log \frac{q\left({\mathbf{x}}_1\mid {\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+\sum _{t=2}^T\log \left(\frac{q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)}\cdot \frac{q\left({\mathbf{x}}_t\mid {\mathbf{x}}_0\right)}{q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_0\right)}\right)+\log \frac{q\left({\mathbf{x}}_1\mid {\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+\sum _{t=2}^T\log \frac{q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)}+\sum _{t=2}^T\log \frac{q\left({\mathbf{x}}_t\mid {\mathbf{x}}_0\right)}{q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_0\right)}+\log \frac{q\left({\mathbf{x}}_1\mid {\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)}\right]\\ & ={\mathbb{E}}_q\left[-\log p_{\theta }\left({\mathbf{x}}_T\right)+\sum _{t=2}^T\log \frac{q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)}+\log \frac{q\left({\mathbf{x}}_T\mid {\mathbf{x}}_0\right)}{q\left({\mathbf{x}}_1\mid {\mathbf{x}}_0\right)}+\log \frac{q\left({\mathbf{x}}_1\mid {\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)}\right]\\ & ={\mathbb{E}}_q\left[\log \frac{q\left({\mathbf{x}}_T\mid {\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_T\right)}+\sum _{t=2}^T\log \frac{q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)}-\log p_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)\right]\\ & ={\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_T\mid {\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_T\right)\right)}}+\sum _{t=2}^T\underset{L_{t-1}}{\underbrace{D_{\mathrm{KL}}\left(q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)\Vert p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)\right)}}-\underset{L_0}{\underbrace{\log p_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)}}]\end{array}$$

这里刚开始的下脚标从 $q_{\mathbf{\varphi }}({\mathbf{x}}_{1:\mathbf{T}}\mid {\mathbf{x}}_0)$ 换为了 $q\left({\mathbf{x}}_{0:T}\right)$，我认为 $x_0$ 为已知，所以这两个式子是等价的。

中间还有容易混淆的一点是，$q\left({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0\right)$代表前向传播的分布，$q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)$ 代表逆扩散过程的真实分布，这里的 $q$ 没有正向或反向的含义，只是代表概率和分布，可以理解为概率论中的概率 $P$。$p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)$ 代表我们要求解的逆扩散分布。

接下来我们对 $L_T,L_{t-1},L_0$ 这三种情况进行分类讨论：

2.3.2 $L_T$

$q\left({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0\right)$ 代表前向扩散过程，没有可学习参数；$p_{\theta }\left({\mathbf{x}}_T\right)$ 中的 $x_T$ 为服从标准高斯分布的噪声，$p_{\theta }$ 为逆扩散过程，对逆扩散过程而言，$x_T$ 为已知的，所以这一项 $L_T$ 可以当作常量。

2.3.3 $L_{t-1}$

而 $L_{t-1}$ 可以看出是真实的逆扩散分布 $q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)$ 和我们要求的逆扩散分布 $p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)$ 的KL散度。

1. 真实分布 $q\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0\right)$ 的均值和方差我们已经求出：
   $${\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\epsilon }_t\right),{\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t$$
2. 第二个分布 $p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)$ 是我们希望拟合的目标分布，也是一个高斯分布，均值用网络估计，方差被设置为了和 ${\beta }_t$ 有关：
   $$p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right),{\mathbf{\Sigma }}_{\theta }\left({\mathbf{x}}_t,t\right)\right)$$

所以为了使这两个分布接近，方差我们可以不用管，只需让两个分布的均值的距离最小，我们用二范数来表示：
$$\begin{array}{rl}{L}_t& ={\mathbb{E}}_q\left[{\Vert {\tilde{\mathbf{\mu }}}_t\left({\mathbf{x}}_t,{\mathbf{x}}_0\right)-{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right)\Vert }^2\right]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,ϵ}\left[{\Vert \frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t\left({\mathbf{x}}_0,ϵ\right)-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ\right)-{\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t\left({\mathbf{x}}_0,ϵ\right),t\right)\Vert }^2\right]ϵ\sim \mathcal{N}(0,1)\end{array}$$
在该公式中可以观察到，我们需要用 ${\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t\left({\mathbf{x}}_0,ϵ\right),t\right)$ 去拟合 $\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t\left({\mathbf{x}}_0,ϵ\right)-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ\right)$ ，所以干脆令：
$${\mathbf{\mu }}_{\theta }\left({\mathbf{x}}_t,t\right)=\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }\left({\mathbf{x}}_t,t\right)\right)$$
也就是直接用神经网络 ${ϵ}_{\theta }\left({\mathbf{x}}_t,t\right)$ 去预测噪声 $ϵ$。然后把预测出来的噪声带入到定义好的表达式去计算出预测的均值。

所以损失函数变为了：
$$\begin{array}{rl}{L}_t& ={\mathbb{E}}_{{\mathbf{x}}_0,ϵ}\left[{\Vert \frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ\right)-\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }\left({\mathbf{x}}_t,t\right)\right)\Vert }^2\right]ϵ\sim \mathcal{N}(0,1)\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,ϵ}\left[{\Vert ϵ-{ϵ}_{\theta }\left({\mathbf{x}}_t,t\right)\Vert }^2\right]ϵ\sim \mathcal{N}(0,1)\text{ 把常数项系数都扔了, 作者说这样更好训练 }\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,ϵ}\left[{\Vert ϵ-{ϵ}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t\right)\Vert }^2\right],ϵ\sim \mathcal{N}(0,1)\end{array}$$
网络的输入是一张和噪声线性组合的图片 $x_t$，与其组合的噪声真实值为 $ϵ$，我们需要用 ${ϵ}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t\right)$ 去拟合这个噪声。

2.3.4 $L_0$

最后的 $L_0=-\log p_{\theta }\left(x_0\mid x_1\right)$ 是将最后一步的加噪图像 $x_1$ 生成去燥图像 $x_0$ 的极大似然估计，为了生成更好的图像，我们需要对每个像素都运用极大似然估计，使得图像上每个像素值都满足离散的对数似然。

为了达到这个目的，将逆扩散过程中的最后从 $x_1$ 到 $x_0$ 的转换设置为独立的离散计算方式。 即在最后一个转换过程在给定 $x_1$ 下得到图像 $x_0$ 满足对数似然，假设像素与像素之间是相互独立的：
$$p_{\theta }\left(x_0\mid x_1\right)=\prod _{i=1}^D{p}_{\theta }\left(x_0^i\mid x_1^i\right)$$
$D$ 代表 $x$ 的维度，上标 $i$ 表示图像中的一个坐标位置。现在的目标是确定给定像素的值可能性有多大，也就是想要知道对应时间步 $t=1$ 下噪声图像 $x$ 中相应像素值的分布：
$$\mathcal{N}\left(x;{\mu }_{\theta }^i\left(x_1,1\right),{\sigma }_1^2\right)$$
其中 $t=1$ 的像素分布来自多元高斯分布，其对角协方差矩阵允许我们将分布拆分为单变量高斯的乘积：
$$\mathcal{N}\left(x;{\mu }_{\theta }\left(x_1,1\right),{\sigma }_1^2\mathbb{I}\right)=\prod _{i=1}^D\mathcal{N}\left(x;{\mu }_{\theta }^i\left(x_1,1\right),{\sigma }_1^2\right)$$
现在假设图像已经从0-255的数值之间，经过归一化在[-1,1]的范围内。在 t=0 时给定每个像素的像素值，最后一个时间步 t=1 的转换概率分布 $p_{\theta }\left(x_0\mid x_1\right)$ 的值就是每个像素值的乘积。所以：
$$\begin{array}{rl}{p}_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)& =\prod _{i=1}^D{\int }_{{\delta }_{-}\left(x_0^i\right)}^{{\delta }_{+}\left(x_0^i\right)}\mathcal{N}\left(x;{\mu }_{\theta }^i\left({\mathbf{x}}_1,1\right),{\sigma }_1^2\right)dx\\ {\delta }_{+}(x)& =\{\begin{array}{ll}\infty & \text{ if }x=1\\ x+\frac{1}{255}& \text{ if }x<1\end{array}{\delta }_{-}(x)=\{\begin{array}{ll}-\infty & \text{ if }x=-1\\ x-\frac{1}{255}& \text{ if }x>-1\end{array}\end{array}$$
这个公式来自原论文，这里解析一下它的含义。就是我们要将最后一步的加噪图像 $x_1$ 拟合去燥图像 $x_0$ ，把图像的每一个像素点都设为一个高斯分布，一共 $D$ 个像素点。而 $x_0$ 每个像素点原本的取值范围为 { 0 , 1 , … , 255 } \{0,1, \ldots, 255\} { 0,1,…,255} ，经过归一化映射到了 $[-1,1]$ 范围内。

现在我们把单独取出一个 $x_1$ 上的像素点 $x_1^i$，它服从分布 $\mathcal{N}\left(x;{\mu }_{\theta }^i\left(x_1,1\right),{\sigma }_1^2\right)$，需要拟合的目标为 $x_0$ 上的对应位置像素点 $x_0^i$ ，而 $x_0^i$ 的取值范围由原本的离散空间 { 0 , 1 , … , 255 } \{0,1, \ldots, 255\} { 0,1,…,255} 映射到了连续空间 $[-1,1]$，所以每个原本的离散值在连续空间中对应一个区间，而区间映射的公式就是：
$$\begin{array}{rl}{\delta }_{+}(x)& =\{\begin{array}{ll}\infty & \text{ if }x=1\\ x+\frac{1}{255}& \text{ if }x<1\end{array}{\delta }_{-}(x)=\{\begin{array}{ll}-\infty & \text{ if }x=-1\\ x-\frac{1}{255}& \text{ if }x>-1\end{array}\end{array}$$

以上就是对DDPM得扩散和逆扩散过程中涉及到的所有公式的解析和推导，包括损失函数构建部分。

参考文献和博客

Understanding Diffusion Models: A Unified Perspective
Denoising Diffusion Probabilistic Models
https://yinglinzheng.netlify.app/diffusion-model-tutorial
https://zhuanlan.zhihu.com/p/549623622