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)$推导，其上界越小，负对数似然越小，相对应其对数似然越大。
− log ⁡ p θ ( x 0 ) ≤ − log ⁡ p θ ( x 0 ) + D K L ( q ( x 1 : T ∣ x 0 ) ∥ p θ ( x 1 : T ∣ x 0 ) ) ( 1 ) = − log ⁡ p θ ( x 0 ) + E x 1 : T ∼ q ( x 1 : T ∣ x 0 ) [ log ⁡ q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) / p θ ( x 0 ) ] ( 2 ) = − log ⁡ p θ ( x 0 ) + E q [ log ⁡ q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) + log ⁡ p θ ( x 0 ) ] ( 3 ) = E q ( x 1 : T ∣ x 0 ) [ log ⁡ q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) ] ( 4 )

−logpθ​(x0​)​≤−logpθ​(x0​)+DKL​(q(x1:T​∣x0​)∥pθ​(x1:T​∣x0​))(1)=−logpθ​(x0​)+Ex1:T​∼q(x1:T​∣x0​)​[logpθ​(x0:T​)/pθ​(x0​)q(x1:T​∣x0​)​](2)=−logpθ​(x0​)+Eq​[logpθ​(x0:T​)q(x1:T​∣x0​)​+logpθ​(x0​)](3)=Eq(x1:T​∣x0​)​[logpθ​(x0:T​)q(x1:T​∣x0​)​](4)​

公式推导

• $(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上界

L V L B = E q ( x 0 : T ) [ log ⁡ q ( x 1 : T ∣ x 0 ) p θ ( x 0 : T ) ] ( 1 ) = E [ log ⁡ ∏ t = 1 T q ( x t ∣ x t − 1 ) p θ ( x T ) ∏ t = 1 T p θ ( x t − 1 ∣ x t ) ] ( 2 ) = E q [ − log ⁡ p θ ( x T ) + ∑ t = 1 T log ⁡ q ( x t ∣ x t − 1 ) p θ ( x t − 1 ∣ x t ) ] ( 3 ) = E q [ − log ⁡ p θ ( x T ) + ∑ t = 2 T log ⁡ q ( x t ∣ x t − 1 ) p θ ( x t − 1 ∣ x t ) + log ⁡ q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ] ( 4 ) = E q [ − log ⁡ p θ ( x T ) + ∑ t = 2 T log ⁡ ( q ( x t − 1 ∣ x t , x 0 ) p θ ( x t − 1 ∣ x t ) ⋅ q ( x t ∣ x 0 ) q ( x t − 1 ∣ x 0 ) ) + log ⁡ q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ] ( 5 ) = E q [ − log ⁡ p θ ( x T ) + ∑ t = 2 T log ⁡ q ( x t − 1 ∣ x t , x 0 ) p θ ( x t − 1 ∣ x t ) + ∑ t = 2 T log ⁡ q ( x t ∣ x 0 ) q ( x t − 1 ∣ x 0 ) + log ⁡ q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ] ( 6 ) = E q [ − log ⁡ p θ ( x T ) + ∑ t = 2 T log ⁡ q ( x t − 1 ∣ x t , x 0 ) p θ ( x t − 1 ∣ x t ) + log ⁡ q ( x T ∣ x 0 ) q ( x 1 ∣ x 0 ) + log ⁡ q ( x 1 ∣ x 0 ) p θ ( x 0 ∣ x 1 ) ] ( 7 ) = E q [ log ⁡ q ( x T ∣ x 0 ) p θ ( x T ) + ∑ t = 2 T log ⁡ q ( x t − 1 ∣ x t , x 0 ) p θ ( x t − 1 ∣ x t ) − log ⁡ p θ ( x 0 ∣ x 1 ) ] ( 8 ) = E q [ D K L ( q ( x T ∣ x 0 ) ∥ p θ ( x T ) ) ⏟ L T + ∑ t = 2 T D K L ( q ( x t − 1 ∣ x t , x 0 ) ∥ p θ ( x t − 1 ∣ x t ) ) ⏟ L t − 1 − log ⁡ p θ ( x 0 ∣ x 1 ) ⏟ L 0 ] ( 9 )

LVLB​​=Eq(x0:T​)​[logpθ​(x0:T​)q(x1:T​∣x0​)​](1)=E[logpθ​(xT​)∏t=1T​pθ​(xt−1​∣xt​)∏t=1T​q(xt​∣xt−1​)​](2)=Eq​[−logpθ​(xT​)+t=1∑T​logpθ​(xt−1​∣xt​)q(xt​∣xt−1​)​](3)=Eq​[−logpθ​(xT​)+t=2∑T​logpθ​(xt−1​∣xt​)q(xt​∣xt−1​)​+logpθ​(x0​∣x1​)q(x1​∣x0​)​](4)=Eq​[−logpθ​(xT​)+t=2∑T​log(pθ​(xt−1​∣xt​)q(xt−1​∣xt​,x0​)​⋅q(xt−1​∣x0​)q(xt​∣x0​)​)+logpθ​(x0​∣x1​)q(x1​∣x0​)​](5)=Eq​[−logpθ​(xT​)+t=2∑T​logpθ​(xt−1​∣xt​)q(xt−1​∣xt​,x0​)​+t=2∑T​logq(xt−1​∣x0​)q(xt​∣x0​)​+logpθ​(x0​∣x1​)q(x1​∣x0​)​](6)=Eq​[−logpθ​(xT​)+t=2∑T​logpθ​(xt−1​∣xt​)q(xt−1​∣xt​,x0​)​+logq(x1​∣x0​)q(xT​∣x0​)​+logpθ​(x0​∣x1​)q(x1​∣x0​)​](7)=Eq​[logpθ​(xT​)q(xT​∣x0​)​+t=2∑T​logpθ​(xt−1​∣xt​)q(xt−1​∣xt​,x0​)​−logpθ​(x0​∣x1​)](8)=Eq​[LT​ DKL​(q(xT​∣x0​)∥pθ​(xT​))​​+t=2∑T​Lt−1​ DKL​(q(xt−1​∣xt​,x0​)∥pθ​(xt−1​∣xt​))​​−L0​ logpθ​(x0​∣x1​)​​](9)​

公式推导

• $(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,ϵ)$替换，就能得到下方第二个等号的式子。
L t − 1 − C = E x 0 , ϵ [ 1 2 σ t 2 ∥ μ ~ t ( x t ( x 0 , ϵ ) , 1 α ˉ t ( x t ( x 0 , ϵ ) − 1 − α ˉ t ϵ ) ) − μ θ ( x t ( x 0 , ϵ ) , t ) ∥ 2 ] = E x 0 , ϵ [ 1 2 σ t 2 ∥ 1 α t ( x t ( x 0 , ϵ ) − β t 1 − α ˉ t ϵ ) − μ θ ( x t ( x 0 , ϵ ) , t ) ∥ 2 ]

Lt−1​−C​=Ex0​,ϵ​[2σt2​1​∥∥∥​μ~​t​(xt​(x0​,ϵ),αˉt​ ​1​(xt​(x0​,ϵ)−1−αˉt​ ​ϵ))−μθ​(xt​(x0​,ϵ),t)∥∥∥​2]=Ex0​,ϵ​[2σt2​1​∥∥∥​α ​t​1​(xt​(x0​,ϵ)−1−αˉt​ ​βt​​ϵ)−μθ​(xt​(x0​,ϵ),t)∥∥∥​2]​
这里我们的${\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]$

---