DDPM交叉熵损失函数推导

$\mathrm{K}\mathrm{L}$散度

由于以下推导需要用到$\mathrm{K}\mathrm{L}$散度，这里先简单介绍一下。
$\mathrm{K}\mathrm{L}$散度一般用于度量两个概率分布函数之间的“距离”，其定义如下：
$KL[P(X)||Q(X)]={\sum }_{x\in X}[P(x)\log \frac{P(x)}{Q(x)}]=E_{x\sim P(x)}\left[\log \frac{P(x)}{Q(x)}\right]$
这里$P(X)$和$Q(X)$是两个概率分布函数，可以看到对于离散型随机变量，$\mathrm{K}\mathrm{L}$散度对$x$进行求和；对于连续型随机变量，$\mathrm{K}\mathrm{L}$散度对$x$进行积分(期望)。
高斯分布的$\mathrm{K}\mathrm{L}$散度
对于两个单一变量的高斯分布$p\sim \mathcal{N}({\mu }_1,{\sigma }_1^2)$和$q\sim \mathcal{N}({\mu }_2,{\sigma }_2^2)$而言，它们的KL散度为
$KL(p,q)=\log \frac{{\sigma }_2}{{\sigma }_1}+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}-\frac{1}{2}$

似然函数

下方是论文中给出的后向过程${\mathbf{x}}_{t-1}$的分布，其方差为常数。
$p_{\theta }({\mathbf{x}}_{0:T})=p({\mathbf{x}}_T){\prod }_{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t),p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mu }_{\theta }({\mathbf{x}}_t,t),{\sum }_{\theta }({\mathbf{x}}_t,t))$
推出扩散模型目标数据分布的似然函数，推出似然函数后才能优化模型。$p_{\theta }({\mathbf{x}}_0)$为目标数据分布，其对数似然下界越大，那么对数似然越大。为了方便推导，这里用其负对数似然$-\log p_{\theta }({\mathbf{x}}_0)$推导，其上界越小，负对数似然越小，相对应其对数似然越大。
$$\begin{array}{rl}-\log p_{\theta }({\mathbf{x}}_0)& \le -\log p_{\theta }({\mathbf{x}}_0)+D_{KL}(q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0))(1)\\ & =-\log p_{\theta }({\mathbf{x}}_0)+{\mathbb{E}}_{{\mathbf{x}}_{1:T}\sim q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}\left[\log \frac{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})/p_{\theta }({\mathbf{x}}_0)}\right](2)\\ & =-\log p_{\theta }({\mathbf{x}}_0)+{\mathbb{E}}_q[\log \frac{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}+\log p_{\theta }({\mathbf{x}}_0)](3)\\ & ={\mathbb{E}}_{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}\left[\log \frac{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}\right](4)\end{array}$$

公式推导

• $(1)$ : 不等式右边加上一个$\mathrm{K}\mathrm{L}$散度，由于$\mathrm{K}\mathrm{L}$散度始终大于等于0，所以不等号成立。也即不等式右边是左边的上界，我们只需要优化右边的式子使其达到最小，那么等式左边的对数似然就达到最小。
• $(1)\to (2)$ : 这一步是将$\mathrm{K}\mathrm{L}$散度展开，可以见上方$\mathrm{K}\mathrm{L}$散度的定义，定义中$P(x)$相当于$q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)$，$Q(x)$相当于$p_{\theta }({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)$。将$Q(x)$按照条件概率公式展开：$p_{\theta }({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)=p_{\theta }({\mathbf{x}}_{1:T},{\mathbf{x}}_0)/p_{\theta }({\mathbf{x}}_0)=p_{\theta }({\mathbf{x}}_{0:T})/p_{\theta }({\mathbf{x}}_0)$，这样就得到了第$(2)$步的式子。
• $(2)\to (3)$ : 将$\log$进行展开即可。
• $(3)\to (4)$ : 由于该期望是针对分布$q$的，则$\log p_{\theta }({\mathbf{x}}_0)$相对于$q$就是常数。所以${\mathbb{E}}_q[\log p_{\theta }({\mathbf{x}}_0)]=\log p_{\theta }({\mathbf{x}}_0)$，然后和前面的$-\log p_{\theta }({\mathbf{x}}_0)$约去，就得到了式子$(4)$。

推导结束

然后我们将不等式左边的$-\log p_{\theta }({\mathbf{x}}_0)$套上一个关于分布$q({\mathbf{x}}_0)$的期望，得到$-{\mathbb{E}}_{q({\mathbf{x}}_0)}\log p_{\theta }({\mathbf{x}}_0)$(交叉熵，也即loss)；相应的，不等式右边也要加上一个${\mathbf{x}}_0$，则由${\mathbb{E}}_{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}$变为${\mathbb{E}}_{q({\mathbf{x}}_{0:T})}$。如果我们想最小化loss，也就是最小化${\mathbb{E}}_{q({\mathbf{x}}_{0:T})}$。
$$\mathrm{L}\mathrm{e}\mathrm{t}{L}_{\mathrm{V}\mathrm{L}\mathrm{B}}={\mathbb{E}}_{q({\mathbf{x}}_{0:T})}\left[\log \frac{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}\right]\ge -{\mathbb{E}}_{q({\mathbf{x}}_0)}\log p_{\theta }({\mathbf{x}}_0)$$

化简loss上界

$$\begin{array}{rl}{L}_{\mathrm{V}\mathrm{L}\mathrm{B}}& ={\mathbb{E}}_{q({\mathbf{x}}_{0:T})}\left[\log \frac{q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}\right](1)\\ & =\mathbb{E}\left[\log \frac{{\prod }_{t=1}^{T}q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})}{p_{\theta }({\mathbf{x}}_T){\prod }_{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)}\right](2)\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=1}^T\log \frac{q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})}{p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)}](3)\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})}{p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)}+\log \frac{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}](4)\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log (\frac{q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)}\cdot \frac{q({\mathbf{x}}_t\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_0)})+\log \frac{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}](5)\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)}+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_t\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_0)}+\log \frac{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}](6)\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)}+\log \frac{q({\mathbf{x}}_T\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}+\log \frac{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}](7)\\ & ={\mathbb{E}}_q[\log \frac{q({\mathbf{x}}_T\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_T)}+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)}-\log p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)](8)\\ & ={\mathbb{E}}_q[\underset{L_T}{D_{\mathrm{K}\mathrm{L}}(q({\mathbf{x}}_T\mid {\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_T))}+\sum _{t=2}^T\underset{L_{t-1}}{D_{\mathrm{K}\mathrm{L}}(q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t))}-\underset{L_0}{\log p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}](9)\end{array}$$

公式推导

• $(1)\to (2)$ : 将条件概率展开。由于$q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)$是扩散过程，是从${\mathbf{x}}_0$逐步推导${\mathbf{x}}_T$得到过程，其符合马尔科夫假设，故$q({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)=q({\mathbf{x}}_1\mid {\mathbf{x}}_0)\cdot q({\mathbf{x}}_2\mid {\mathbf{x}}_1)\cdot ...\cdot q({\mathbf{x}}_T\mid {\mathbf{x}}_{T-1})={\prod }_{t=1}^{T}q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})$；对于$p_{\theta }({\mathbf{x}}_{0:T})$，我们先将其根据条件概率转换为$p_{\theta }({\mathbf{x}}_T)p_{\theta }({\mathbf{x}}_{0:T-1}\mid {\mathbf{x}}_T)$，然后将后面那一项和$q$一样，展开即可。
• $(2)\to (3)$ : 将$\log$进行展开，连乘展开后转换为求和。
• $(3)\to (4)$ : 将$\log \frac{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}$单独拿出来计算。
• $(4)\to (5)$ : 回忆一下，之前在讲逆扩散过程的时候我们得到了这样一个式子$q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)=q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})\frac{q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_t\mid {\mathbf{x}}_0)}$，通过这个式子，我们就能得到$q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})$的表达式，然后替换即可。
• $(5)\to (6)$ : 将$\log$进行展开。
• $(6)\to (7)$ : ${\sum }_{t=2}^T\log \frac{q({\mathbf{x}}_t\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_0)}=\log (\frac{q({\mathbf{x}}_2\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}\cdot \frac{q({\mathbf{x}}_3\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_2\mid {\mathbf{x}}_0)}\cdot ...\cdot \frac{q({\mathbf{x}}_T\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_T-1\mid {\mathbf{x}}_0)})=\log \frac{q({\mathbf{x}}_T\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}$
• $(7)\to (8)$ : $\log \frac{q({\mathbf{x}}_T\mid {\mathbf{x}}_0)}{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}+\log \frac{q({\mathbf{x}}_1\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}=\log q({\mathbf{x}}_T\mid {\mathbf{x}}_0)-\log p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)$，然后将$\log q({\mathbf{x}}_T\mid {\mathbf{x}}_0)$和$-\log p_{\theta }({\mathbf{x}}_T)$合并成$\log \frac{q({\mathbf{x}}_T\mid {\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_T)}$
• $(8)\to (9)$ : 对于$L_T$，$q({\mathbf{x}}_T\mid {\mathbf{x}}_0)$和$p_{\theta }({\mathbf{x}}_T)$都是不含参的，前者$q$分布是由${\beta }_t$求出的，不含有任何参数；后者是一个各向同性的高斯分布。故$L_T$是不含参的，在优化时可以将其舍弃。对于$L_{t-1}$，参见$\mathrm{K}\mathrm{L}$散度定义，可以将其表示为$\mathrm{K}\mathrm{L}$散度，如果这里我们将$t$取1，其转化为$\log \frac{q({\mathbf{x}}_0\mid {\mathbf{x}}_1,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}=\log \frac{1}{p_{\theta }({\mathbf{x}}_0\mid {\mathbf{x}}_1)}$。故当$t$为1时，得到的结果就是$L_{t-1}$后面那一项$L_0$，故我们可以将其合并。故我们只需要优化$L_{t-1}$即可。

推导结束

在论文中，作者将分布$p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$的方差看作与$\beta$相关的常数，那么可训练的参数就存在于其均值当中。在$L_{t-1}$中，$q({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)$是一个高斯分布，其方差和均值我们已经在之前后向过程推导中求出，均值为${\tilde{\mu }}_t({\mathbf{x}}_t)$，方差为和${\beta }_t$有关的常数。而$p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)$也是我们假设的高斯分布，它的方差也是常数，均值为${\mu }_{\theta }({\mathbf{x}}_t,t)$，所以参数只在${\mu }_{\theta }$当中。对于这两个高斯分布，我们可以运用高斯分布的$\mathrm{K}\mathrm{L}$散度公式，其中的方差我们可以不考虑。则我们可以得到如下的式子：
$$L_{t-1}={\mathbb{E}}_q[\frac{1}{2{\sigma }_t^2}\Vert {\tilde{\mu }}_t({\mathbf{x}}_t,{\mathbf{x}}_0)-{\mu }_{\theta }({\mathbf{x}}_t,t){\Vert }^2]+C$$

由这个式子，我们优化目标就很明确了，我们要优化${\mu }_{\theta }$，让其无线逼近于${\tilde{\mu }}_t$，这样才能使$L_{t-1}$最小。首先我们将${\tilde{\mu }}_t({\mathbf{x}}_t)$代入上述的式子中，原式中的${\tilde{z}}_t$用$ϵ$来表示，${\mathbf{x}}_t$用${\mathbf{x}}_t({\mathbf{x}}_0,ϵ)$替换，就能得到下方第二个等号的式子。
$$\begin{array}{rl}{L}_{t-1}-C& ={\mathbb{E}}_{{\mathbf{x}}_0,ϵ}[\frac{1}{2{\sigma }_t^2}\Vert {\tilde{\mu }}_t({\mathbf{x}}_t({\mathbf{x}}_0,ϵ),\frac{1}{\sqrt{{\bar{\alpha }}_t}}({\mathbf{x}}_t({\mathbf{x}}_0,ϵ)-\sqrt{1-{\bar{\alpha }}_t}ϵ))-{\mu }_{\theta }({\mathbf{x}}_t({\mathbf{x}}_0,ϵ),t){\Vert }^2]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,ϵ}[\frac{1}{2{\sigma }_t^2}\Vert \frac{1}{{\sqrt{\alpha }}_t}({\mathbf{x}}_t({\mathbf{x}}_0,ϵ)-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}ϵ)-{\mu }_{\theta }({\mathbf{x}}_t({\mathbf{x}}_0,ϵ),t){\Vert }^2]\end{array}$$
这里我们的${\mathbf{x}}_t$是已知的，那么为了使$L_{t-1}$最小，我们可以将${\mu }_{\theta }({\mathbf{x}}_t,t)$表示为${\tilde{\mu }}_t$的一个波动，其中的$ϵ$是未知的，则我们可以训练一个网络来预测$ϵ$。
${\mu }_{\theta }({\mathbf{x}}_t,t)={\tilde{\mu }}_t({\mathbf{x}}_t,\frac{1}{\sqrt{{\bar{\alpha }}_t}}({\mathbf{x}}_{\mathbf{t}}-\sqrt{1-{\bar{\alpha }}_{\mathbf{t}}}{ϵ}_{\theta }({\mathbf{x}}_{\mathbf{t}})))=\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }({\mathbf{x}}_t,t))$
于是$L_{t-1}$可以简化为如下形式
${\mathbb{E}}_{{\mathbf{x}}_0,ϵ}[\frac{{\beta }_t^2}{2{\sigma }_t^2{\alpha }_t(1-{\bar{\alpha }}_t)}\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t){\Vert }^2]$
作者又发现，将系数丢掉，训练更加稳定质量更好，于是就得到了下方的$L_{\mathrm{s}\mathrm{i}\mathrm{m}\mathrm{p}\mathrm{l}\mathrm{e}}$
$L_{\mathrm{s}\mathrm{i}\mathrm{m}\mathrm{p}\mathrm{l}\mathrm{e}}(\theta ):={\mathbb{E}}_{t,{\mathbf{x}}_0,ϵ}[\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t){\Vert }^2]$

---