Diffusion model(二): 训练推导详解

接上文

Diffusion model(一): 公式推导详解
Diffusion model(二): 训练推导详解
Diffusion model(三): 公式结论

Diffusion的训练推导

1. 最小化负对数似然与变分下界

在弄懂diffusion model前向和反向过程之后，最后我们需要了解其训练推导过程，即用什么loss以及为什么。在diffusion的反向过程中，根据$(3)$式我们需要预测${\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t)$，如何得到一个合理的均值和方差？类似于VAE，在对真实数据分布情况下，最大化模型预测分布的对数似然，即优化$x_0\sim q(x_0)$下的$p_{\theta }(x_0)$的交叉熵
$$\mathcal{L}={\mathbb{E}}_{q(x_0)}[-\log p_{\theta }(x_0)]\text{\hspace{1em}(12)}$$

注意，我们的真实目标是最大化$p_{\theta }(x)$，$(3)$对$p_{\theta }(x)$取了负$\log$，所以这里要最小化$\mathcal{L}$。与VAE类似，我们通过变分下限VLB来优化$(12)$的负对数似然，如下
$$\begin{array}{rl}-\log p_{\theta }(x_0)& \le -\log p_{\theta }(x_0)+D_{KL}(q(x_{1:T}|x_0)||p_{\theta }(x_{1:T}|x_0))\\ & =-\log p_{\theta }(x_0)+{\mathbb{E}}_{q(x_{1:T}|x_0)}[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{1:T}|x_0)}]\\ & =-\log p_{\theta }(x_0)+{\mathbb{E}}_{q(x_{1:T}|x_0)}[\log \frac{q(x_{1:T}|x_0)}{\frac{p_{\theta }(x_{0:T})}{P_{\theta }(x_0)}}]\\ & =-\log p_{\theta }(x_0)+{\mathbb{E}}_{q(x_{1:T}|x_0)}[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}+\underset{此项与q(x_{1:T}|x_0)无关，可直接提出}{\underbrace{\log p_{\theta }(x_0)}}]\\ & ={\mathbb{E}}_{q(x_{1:T}|x_0)}\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\end{array}\text{\hspace{1em}(13)}$$

现在我们得到了最大似然的变分下界，即式$(13)$，但是根据式$(12)$，我们要优化${\mathbb{E}}_{q(x_0)}[-\log p_{\theta }(x_0)]$，比$(13)$多了个期望，所以这里我们加上期望，同时根据重积分中的Fubini定理，得到
$$\begin{array}{rl}{\mathbb{E}}_{q(x_0)}[-\log p_{\theta }(x_0)]& \le {\mathcal{L}}_{VLB}\\ & ={\mathbb{E}}_{q(x_0)}\left({\mathbb{E}}_{q(x_{1:T}|x_0)}\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right)\\ & ={\mathbb{E}}_{q(x_{0:T})}\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\end{array}\text{\hspace{1em}(14)}$$

2. 变分下界的详细推导优化

再次回到$(12)$式，我们的目标是最小化交叉熵，现在我们得到了他的变分下界，因此我们将最小化$\mathcal{L}$转移为最小化变分下界，即${\mathcal{L}}_{VLB}$，下面我们对变分下界进行化简，化简过程比较长，对一些具有疑惑的化简最后附带详解
$$\begin{array}{rl}{\mathcal{L}}_{VLB}& ={\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 \frac{{\prod }_{t=1}^{T}q(x_t|x_{t-1})}{p_{\theta }(x_T){\prod }_{t=1}^T{p}_{\theta }(x_{t-1}|x_t)}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[-\log p_{\theta }(x_T)+\log \frac{{\prod }_{t=1}^{T}q(x_t|x_{t-1})}{{\prod }_{t=1}^T{p}_{\theta }(x_{t-1}|x_t)}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[-\log p_{\theta }(x_T)+\sum _{i=1}^T\log \frac{q(x_t|x_{t-1})}{p_{\theta }(x_{t-1}|x_t)}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[-\log p_{\theta }(x_T)+\sum _{i=2}^T\log \frac{q(x_t|x_{t-1})}{p_{\theta }(x_{t-1}|x_t)}+\underset{1.提出此项是因为无意义}{\underbrace{\log \frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}}}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[-\log p_{\theta }(x_T)+\sum _{i=2}^T\underset{2.此项与上一项的变化是额外考虑x_0}{\underbrace{\log \frac{q(x_t|x_{t-1},x_0)q(x_t|x_0)}{p_{\theta }(x_{t-1}|x_t)q(x_{t-1}|x_0)}}}+\log \frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[-\log p_{\theta }(x_T)+\sum _{i=2}^T\log \left(\frac{q(x_t|x_{t-1},x_0)}{p_{\theta }(x_{t-1}|x_t)}\cdot \frac{q(x_t|x_0)}{q(x_{t-1}|x_0)}\right)+\log \frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[-\log p_{\theta }(x_T)+\sum _{i=2}^T\left(\log \frac{q(x_t|x_{t-1},x_0)}{p_{\theta }(x_{t-1}|x_t)}+\log \frac{q(x_t|x_0)}{q(x_{t-1}|x_0)}\right)+\log \frac{q(x_1|x_0)}{p_{\theta }(x_0|x_1)}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[-\log p_{\theta }(x_T)+\sum _{i=2}^T\log \frac{q(x_t|x_{t-1},x_0)}{p_{\theta }(x_{t-1}|x_t)}+\underset{3.上一项展开即可}{\underbrace{\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)}\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[-\log p_{\theta }(x_T)+\sum _{i=2}^T\log \frac{q(x_t|x_{t-1},x_0)}{p_{\theta }(x_{t-1}|x_t)}+\log q(x_T|x_0)-\log q(x_1|x_0)+\log q(x_1|x_0)-\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)}+\sum _{i=2}^T\log \frac{q(x_t|x_{t-1},x_0)}{p_{\theta }(x_{t-1}|x_t)}-\log p_{\theta }(x_0|x_1)\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\underset{L_T}{\underbrace{D_{KL}(q(x_T|x_0)||p_{\theta }(x_T))}}+\sum _{i=2}^T\underset{L_t}{\underbrace{D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))}}-\underset{L_0}{\underbrace{\log p_{\theta }(x_0|x_1)}}\right]\end{array}$$

注意：这里化简有一个疑问是从倒数第2行到倒数第1行的过程中，如果要得到KL散度，是需要加上期望符号的，即$\mathbb{E}$，但是推导中没有加，这里不知道为啥，但是就这么用好了

公式解释部分，上述公式懂的话可以不看
1.公式中的红笔部分$q(x_t|x_{t-1})$化简
由$x_{t-1}$得$x_t$比较困难，但是当提供额外的$x_0$时，$x_{t-1}$和$x_t$的候选会减少，选择更加确定，因此这里我们加上$x_0$，这样推导如下
$$q(x_t|x_{t-1})=\frac{q(x_{t-1}|x_t)q(x_t)}{q(x_{t-1})}\stackrel{add{x}_0}{⟶}\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}$$
2. 公式中蓝色部分为什么无意义？
在红笔公式化简之后，给定$x_0$，$q(x_t|x_{t-1})$可以得到$\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}$，如果我们不将蓝色部分提出，那么$q(x_1|x_0)$会变为$\frac{q(x_0|x_1,x_0)q(x_1|x_0)}{q(x_0|x_0)}$，此表达式是无意义的

3. 由变分下界得到优化loss

去掉繁琐的推导过程以及最后的期望符号，我们再回头看一下VLB的化简结果
$${\mathcal{L}}_{VLB}\propto \underset{L_T}{\underbrace{D_{KL}(q(x_T|x_0)||p_{\theta }(x_T))}}+\sum _{i=2}^T\underset{L_t}{\underbrace{D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))}}-\underset{L_0}{\underbrace{\log p_{\theta }(x_0|x_1)}}$$

• $L_T=D_{KL}(q(x_T|x_0)||p_{\theta }(x_T))$: 这一项中的前向过程$q(x_T|x_0)$是一个前向加噪的过程，而$p_{\theta }(x_T)$是一个纯粹的高斯噪声，在我们前向过程中，$q(x_T|x_0)$很容易变为高斯噪声，因此$L_T$可以看为常量，计算时忽略即可
• $L_t=D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))$：这一项可以看做拉近2个高斯分布$q(x_{t-1}|x_t,x_0)=(x_{t-1};\tilde{\mu }(x_t,x_0),{\tilde{\beta }}_{t}I)$和$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$之间的距离，根据多元高斯分布的KL散度求解
  $$L_t={\mathbb{E}}_q\left[\frac{1}{2||{\sum }_{\theta }(x_t,t)|{|}_2^2}||{\tilde{\mu }}_t(x_t,x_0)-{\mu }_{\theta }(x_t,t)|{|}^2\right]+C$$
  把式$(10)$和$(11)$求得的${\tilde{\mu }}_t(x_t,x_0)$，${\mu }_{\theta }(x_t,t)$带入得
  $$\begin{array}{rl}{L}_t& ={\mathbb{E}}_{x_0,{\overline{z}}_t}\left[\frac{1}{2||{\Sigma }_{\theta }(x_t,t)|{|}_2^2}||{\tilde{\mu }}_t(x_t,x_0)-{\mu }_{\theta }(x_t,t)|{|}^2\right]\\ & ={\mathbb{E}}_{x_0,{\overline{z}}_t}\left[\frac{1}{2||{\Sigma }_{\theta }(x_t,t)|{|}_2^2}||\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\overline{z}}_t)-\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{z}_{\theta }(x_t,t))|{|}^2\right]\\ & ={\mathbb{E}}_{x_0,{\overline{z}}_t}\left[\frac{{\beta }_t^2}{2{\alpha }_t(1-{\overline{\alpha }}_t||{\Sigma }_{\theta }|{|}_2^2)}||{\overline{z}}_t-z_{\theta }(x_t,t)|{|}^2\right]\\ & ={\mathbb{E}}_{x_0,{\overline{z}}_t}\left[\frac{{\beta }_t^2}{2{\alpha }_t(1-{\overline{\alpha }}_t||{\Sigma }_{\theta }|{|}_2^2)}||{\overline{z}}_t-z_{\theta }(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}{\overline{z}}_t,t)|{|}^2\right]\end{array}$$
• $L_0=-\log p_{\theta }(x_0|x_1)$：其他博客解释这个相当于最后一步的熵，但是这个我不是很理解，如果想详细了解可以参考DDPM论文

所有的$L$就是最终的变分下界，DDPM将其进一步化简如下
$$L_t^{simple}={\mathbb{E}}_{x_0,{\overline{z}}_t}\left[||{\overline{z}}_t-z_{\theta }(\sqrt{{\overline{\alpha }}_t}{x}_0+\sqrt{1-{\overline{\alpha }}_t}{\overline{z}}_t,t)|{|}^2\right]$$

论文没有将方差${\Sigma }_{\theta }$考虑在训练和推断中，而是将untrained的${\beta }_t$代替$\widetilde{{\beta }_t}$，因为${\Sigma }_{\theta }$可能会导致训练不稳定

到这里我们知道了Diffusion model的训练其实也是去预测每一步的噪声，就像反向过程中对均值推导的那样，${\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\overline{z}}_t)$，这里的均值不依赖$x_0$，而他的求解本质上就是$x_t$减去随机噪声

4. 训练与推断

最后我们附上论文中训练和推断的过程

在看论文的时候有一点困惑：解释一下为什么训练的时候reverse一步到位，sample的时候得一步一步的来。因为训练的时候我们对于均值u的预测，是建立在$x_0$已知的基础上的，所以能够通过公式直接预测噪声进行训练，测试的时候我们只有采样的高斯随机噪声$x_t$，并不知道他的$x_0$是什么，所以需要一步一步的预测噪声

知识点补充

1. 重参数化技巧

重参数化技巧在VAE中被应用过，此技巧主要用来使采样可以进行反向传播，假设我们随机采样时从任意一个高斯分布$\mathcal{N}(\mu ,{\sigma }^2)$中采样，然后预测结果，最终结果是无法反向传播的（不可导），通常做法是使用标准高斯分布$\mathcal{N}(0,I)$作为引导

具体做法是首先从标准高斯分布中采样一个变量，然后根据高斯分布的均值$\mu$和方差${\sigma }^2$来对采样变量进行线性变换，如下
$$z=\mu +\sigma \odot ϵ,ϵ\sim \mathcal{N}(0,I)$$

重参数化之后得到的变量$z$具有随机性的，满足均值为$\mu$，方差为${\sigma }^2$的高斯分布，这样采样过程就可导了，随机性加到了$ϵ\sim \mathcal{N}(0,I)$上，而不是$\mathcal{N}(\mu ,{\sigma }^2)$

再通俗解释一下，就是如果我们直接对原来高斯分布采样的话，采样之后的所有计算是可以反向传播的，但是是传不到<在采样这个步骤之前的过程>的，因为数据具有随机性了，反向传播不能传播随机性的梯度，当我们将随机性转移到$ϵ\sim \mathcal{N}(0,I)$上时，因为标准高斯分布之前就没有数据了，所以不用继续传播了，传播的是对其采样之后的$\mu$和$\sigma$

到这里我们知道了Diffusion model的训练其实也是去预测每一步的噪声，就像反向过程中对均值推导的那样，${\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\overline{z}}_t)$，这里的均值不依赖$x_0$，而他的求解本质上就是$x_t$减去随机噪声