生成模型 | 扩散模型损失函数公式推导

接上文：生成模型 | 扩散模型公式推导，原始正文补充内容后过长了，网站需要切分。

损失函数

考虑到扩散模型生成图片的过程是一步一步的序列，所以生成真实样本$x_0$的概率可以写成所有中间转筒序列的联合概率的积分：
$$p_{\theta }(x_0)=\int p_{\theta }(x_0,x_1,\dots ,x_T)d{x}_1\dots d{x}_T$$

而训练模型的目标是让它生成的样本分布$p_{\theta }(x)$更符合训练集分布（$q(x_0)$）。所以模型整体损失可以描述为：
$$\begin{array}{rl}L& ={\mathbb{E}}_{q(x_0)}\left[-\log p_{\theta }(x_0)\right]\\ & =-\int q(x_0)\log p_{\theta }(x_0)d{x}_0\\ & =-\int q(x_0)\log \left(\int p_{\theta }(x_{0:T})d{x}_{1:T}\right)d{x}_0(本身需要考虑所有噪声轨迹难以计算)\\ & =-\int 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)d{x}_0(引入容易采样的前向分布来简化计算)\\ & \le -\int q(x_0)\left(\int q(x_{1:T}|x_0)\log \frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}d{x}_{1:T}\right)d{x}_0(\int q(x_{1:T}|x_0)d{x}_{1:T}=1)\\ & =-\int q(x_0)q(x_{1:T}|x_0)\log \frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}d{x}_{0:T}\\ & =-\int q(x_{0:T})\log \frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}d{x}_{0:T}\\ & =-{\mathbb{E}}_{q(x_{0:T})}\left[\log \frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\log q(x_{1:T}|x_0)-\log p_{\theta }(x_{0:T})\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\log \prod _{t=1}^{T}q(x_t|x_{t-1})-\log [p_{\theta }(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t)]\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\sum _{t=1}^T\log q(x_t|x_{t-1})-\log p_{\theta }(x_T)-\sum _{t=1}^T\log p_{\theta }(x_{t-1}|x_t)\right](q(x_t|x_{t-1})=q(x_t|x_{t-1},x_0)=\frac{q(x_t,x_{t-1},x_0)}{q(x_{t-1},x_0)}=\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)q(x_0)}{q(x_{t-1}{x}_0)}=\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)})\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\sum _{t=2}^T\log \frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\frac{q(x_t|x_0)}{q(x_{t-1}|x_0)}+\log \frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}-\log p_{\theta }(x_T)\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\sum _{t=2}^T\log \frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}+\sum _{t=2}^T\log \frac{q(x_t|x_0)}{q(x_{t-1}|x_0)}+\log \frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}-\log p_{\theta }(x_T)\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\sum _{t=2}^T\log \frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}+\log \frac{q(x_T|x_0)}{q(x_1|x_0)}+\log \frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}-\log p_{\theta }(x_T)\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\sum _{t=2}^T\log \frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}+\log \frac{q(x_T|x_0)}{p_{\theta }(x_T)p_{\theta }(x_0|x_1)}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\log \frac{q(x_T|x_0)}{p_{\theta }(x_T)}+\sum _{t=2}^T\log \frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}-\log p_{\theta }(x_0|x_1)\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\log \frac{q(x_T|x_0)}{p_{\theta }(x_T)}\right]+\sum _{t=2}^T{\mathbb{E}}_{q(x_{0:T})}\left[\log \frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\right]+{\mathbb{E}}_{q(x_{0:T})}[-\log p_{\theta }(x_0|x_1)]\\ & ={\mathbb{E}}_{q(x_0)}{\mathbb{E}}_{q(x_T|x_0)}\left[\log \frac{q(x_T|x_0)}{p_{\theta }(x_T)}\right]+\sum _{t=2}^T{\mathbb{E}}_{q(x_t,x_0)}{\mathbb{E}}_{q(x_{t-1}|x_t,x_0)}\left[\log \frac{q(x_{t-1}|x_t,x_0)}{p_{\theta }(x_{t-1}|x_t)}\right]+{\mathbb{E}}_{q(x_0,x_1)}[-\log p_{\theta }(x_0|x_1)]\\ & ={\mathbb{E}}_{q(x_0)}[D_{\text{KL}}(q(x_T|x_0)|p_{\theta }(x_T))]+\sum _{t=2}^T{\mathbb{E}}_{q(x_t,x_0)}[D_{\text{KL}}(q(x_{t-1}|x_t,x_0)|p_{\theta }(x_{t-1}|x_t))]+{\mathbb{E}}_{q(x_0,x_1)}[-\log p_{\theta }(x_0|x_1)]\end{array}$$

注意，这里用到Jensen 不等式：对于凹函数 (向上凸的) $f(x)$（例如这里的$\log$函数），有：
$$f(\sum _{i=1}^n{\lambda }_i{x}_i)\ge \sum _{i=1}^n{\lambda }_{i}f(x_i),\sum _{i=1}^n{\lambda }_i=1$$
即自变量线性组合对应的函数值大于自变量对应函数值的线性组合。

损失中的第一项当作$L_T$：由于$q$无学习参数且$x_T$为固定高斯噪声，该项当做常量可以忽略。

损失中的最后一项当作$L_0$，考虑到$p_{\theta }(x_{t-1}|x_t)\approx q(x_{t-1}|x_t)=\mathcal{N}({\tilde{\mu }}_{\theta }(x_t,t),{\tilde{\beta }}_t\mathbf{I})$，可以近似写成：
$$\begin{array}{rl}{p}_{\theta }(x_0|x_1)& =\mathcal{N}({\tilde{\mu }}_{\theta }(x_1,1),{\tilde{\beta }}_1\mathbf{I})\\ & =\frac{1}{(2\pi {\beta }_1{)}^{d/2}}\exp (-\frac{1}{2{\beta }_1}\Vert x_0-{\tilde{\mu }}_{\theta }(x_1,1){\Vert }^2)\\ -\log p_{\theta }(x_0|x_1)& =\frac{d}{2}\log (2\pi {\beta }_1)+\frac{1}{2{\beta }_1}\Vert x_0-{\tilde{\mu }}_{\theta }(x_1,1){\Vert }^2\\ & \propto \Vert x_0-{\tilde{\mu }}_{\theta }(x_1,1){\Vert }^2\end{array}$$
很多实验表明即使不单独强调该项，生成样本质量仍然很好，尤其当$T$很大且中间项训练充分时。

损失中的其余项当作$L_{t-1}$，这可以看做在拉近两个高斯分布$q(x_{t-1}|x_t,x_0)=\mathcal{N}({\tilde{\mu }}_t(x_t,x_0),{\tilde{\beta }}_t\mathbf{I})$和$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}({\tilde{\mu }}_{\theta }(x_t,t),{\tilde{\beta }}_t\mathbf{I})$之间的距离。

考虑到两个$d$维高斯分布，同时协方差矩阵各向同性，二者之间的KL散度基于如下通用公式计算https://en.wikipedia.org/wiki/Kullback–Leibler_divergence#Multivariate_normal_distributions：
$$D_{\text{KL}}(p\Vert q)=\frac{1}{2}\left(d\frac{{\sigma }_p^2}{{\sigma }_q^2}+\frac{\Vert {\mu }_p-{\mu }_q{\Vert }^2}{{\sigma }_q^2}-d-d\log \frac{{\sigma }_p^2}{{\sigma }_q^2}\right)$$

所以具体形式表示为（${\tilde{\beta }}_{\theta }={\tilde{\beta }}_t$）：
$$\begin{array}{rl}{D}_{\text{KL}}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))& =\frac{1}{2}\left(d\frac{{\tilde{\beta }}_t}{{\tilde{\beta }}_{\theta }}+\frac{\Vert {\tilde{\mu }}_t-{\tilde{\mu }}_{\theta }(x_t,t){\Vert }^2}{{\tilde{\beta }}_{\theta }}-d-d\log \frac{{\tilde{\beta }}_t}{{\tilde{\beta }}_{\theta }}\right)\\ & =\frac{\Vert {\tilde{\mu }}_t-{\tilde{\mu }}_{\theta }(x_t,t){\Vert }^2}{2{\tilde{\beta }}_t}\\ & =\frac{1}{2{\tilde{\beta }}_t}\Vert \frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\overline{ϵ}}_t)-\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\overline{ϵ}}_{\theta }(x_t,t)){\Vert }^2\\ & =\frac{{\beta }_t^2}{2{\alpha }_t(1-{\bar{\alpha }}_t){\tilde{\beta }}_t}\Vert {\overline{ϵ}}_t-{\overline{ϵ}}_{\theta }(x_t,t){\Vert }^2\\ & =\frac{{\beta }_t^2}{2{\alpha }_t(1-{\bar{\alpha }}_t){\tilde{\beta }}_t}\Vert {\overline{ϵ}}_t-{\overline{ϵ}}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{\overline{ϵ}}_t,t){\Vert }^2\\ & \propto \Vert {\overline{ϵ}}_t-{\overline{ϵ}}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{\overline{ϵ}}_t,t){\Vert }^2\end{array}$$

不同于前向过程那样直接可以一步获得目标值，这里需要从对应于时间步$T$的随机噪声$x_T$逐步迭代去噪来获得对应于时间步$0$的目标数据$x_0$。

值得注意的是，如果在前面关于均值的推导中，不将$x_0$替换，那么这里的损失形式实际上可以写成如下对真值图像的逼近（与这里提到的另一种损失推导思路存在相关之处：https://spaces.ac.cn/archives/9164#去噪过程）：
$$\begin{array}{rl}{\tilde{\mu }}_t& =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0\\ {\tilde{\mu }}_{\theta }& =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{f}_{\theta }(x_t,t)\\ D_{\text{KL}}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))& =\frac{\Vert {\tilde{\mu }}_t-{\tilde{\mu }}_{\theta }(x_t,t){\Vert }^2}{2{\tilde{\beta }}_t}\\ & =\frac{1}{2{\tilde{\beta }}_t}\Vert \frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0-\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t-\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{f}_{\theta }(x_t,t){\Vert }^2\\ & =\frac{{\bar{\alpha }}_{t-1}{\beta }_t^2}{2{\tilde{\beta }}_t(1-{\bar{\alpha }}_t{)}^2}\Vert x_0-f_{\theta }(x_t,t){\Vert }^2\\ & \propto \Vert x_0-f_{\theta }(x_t,t){\Vert }^2\end{array}$$

两个一个是估计噪声，一个是估计图像。DDPM中的实验发现，估计噪声的效果是更好的。同时原论文提到省略掉前面的系数，以及把时间t在训练中也随机采样，能带来更好的生成结果

参考资料

• 由浅入深了解Diffusion Model - ewrfcas的文章 - 知乎：https://zhuanlan.zhihu.com/p/525106459
• 保姆级教程，基于pytorch从零实现生成扩散模型DDPM https://zhuanlan.zhihu.com/p/617895786
• 生成扩散模型漫谈（一）：DDPM = 拆楼 + 建楼 https://spaces.ac.cn/archives/9119