理解Difussion Model，以及遇到的一些问题

Diffusion Model

今天简单了解了一下扩散模型，主要是学习了最经典的那篇2020年发表的denoising diffusion probabilistic models（DDPM）。需要掌握一些高数基础、概率论以及KL散度相关的知识，数学公式推导起来才不至于费劲，而且里面的概率公式多是基于贝叶斯公式以及马尔科夫假设，运用的比较灵活，虽然数学原理推导不是很难，但是整个模型的理解会存在一些或大或小的问题，我将在这篇文章中记录一下我的一些疑惑，并尝试着去解答，欢迎大佬指出错误。

下面列出我学习过程中参考的一些文章：

• 原论文DDPM
• What are Diffusion Models? 大佬的博客，写的非常简练，值得阅读
• 源自知乎的深入讲解，基于上篇博客的翻译与理解
• 涉及一些基于tensorflow源码的讲解

问题1 模型学习到的是什么，它需要预测什么？

①在扩散过程中，输入的原始图片x0，在一步步加入高斯噪声后，会得到一个标准的高斯分布噪声。而且基于重参数（reparameterization trick）以及相关数学公式的推导，能够得到结论：任意时刻的x_t可以由x0和β表示。

因此在实际训练的过程中，在从数据集中采样的得到原始图片x0后，并不需要一步一步地添加高斯噪声，利用重参数的技巧，获取对应时间步的x_t （在原始图片的基础上添加了一系列的高斯噪声）。而是利用已经推导出的结论，快速利用公式从x0，α和β参数中得到任意时刻的x_t。
$$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)=\prod _{t=1}^{T}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})$$

$$\begin{array}{rlr}{\mathbf{x}}_t& =\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}+\sqrt{1-{\alpha }_t}{\mathbf{ϵ}}_{t-1}& \text{ ;where }{\mathbf{ϵ}}_{t-1},{\mathbf{ϵ}}_{t-2},\cdots \sim \mathcal{N}(0,\mathbf{I})\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{\mathbf{ϵ}}}_{t-2}& \text{ ;where }{\overline{\mathbf{ϵ}}}_{t-2}\text{ merges two Gaussians (*).}\\ & =\dots & \\ & =\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ}& \\ q({\mathbf{x}}_t|{\mathbf{x}}_0)& =\mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,(1-{\bar{\alpha }}_t)\mathbf{I})& \end{array}$$

②在Diffusion的逆扩散，逆向过程中，实际上就是一个逐渐去噪的过程。如果我们能够逐步得到逆转后的分布，即给定x_t推出x_t-1。（并且在论文中给出相关文献证明这样一个逆向分布其实也是高斯分布，只需要满足β参数足够小即可）
$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)(逆向分布，公式①)$$
，就可以从完全的标准高斯分布
$${\mathbf{x}}_T\sim \mathcal{N}(0,\mathbf{I})$$
还原出原图分布 x0 。但是这个逆向的概率分布求解起来比较麻烦，每步都需要遍历整个数据集，因此作者使用深度学习模型去预测这样一个逆向的分布p_θ。
$$p_{\theta }({\mathbf{x}}_{0:T})=p({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))(模型求解的逆向分布，公式②)$$
作者发现,虽然我们不能直接求逆向分布，但是引入x_0之后，我们可以直接求出一个给定x_0的逆向分布的解析式,也即后验条件概率。并且直接求出了对应高斯分布的均值与方差，相关的推导过程可以参考贴出的博客。
$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};\tilde{\mathbf{\mu }}({\mathbf{x}}_t,{\mathbf{x}}_0),{\tilde{\beta }}_t\mathbf{I})(给定x_0的逆向分布，公式③)$$

$$\begin{array}{rl}{\tilde{\beta }}_t& =1/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})=1/(\frac{{\alpha }_t-{\bar{\alpha }}_t+{\beta }_t}{{\beta }_t(1-{\bar{\alpha }}_{t-1})})=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t(方差，公式④)\\ {\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)& =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0)/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})\\ & =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0)\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t\\ & =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0\end{array}$$

结合前面扩散过程中，由x_0快速求出x_t的公式，我们可以将上面的均值μ改写化简成下列公式。
$${\mathbf{x}}_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}({\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_t)$$

$$\begin{array}{rl}{\tilde{\mathbf{\mu }}}_t& =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}\frac{1}{\sqrt{{\bar{\alpha }}_t}}({\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_t)\\ & =\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_t)\end{array}(均值，公式⑤)$$

重点来了：

经过上面的分析，我们的扩散模型主要的工作就是需要近似模拟一个真实的逆向分布（公式①），模型的预测得出的分布（公式②），给定x_0的逆向分布（公式③），这三个分布都是高斯分布，含有均值和方差。原论文中作者真正做的是，想办法让模型预测的分布来逼近 给定x_0的逆向分布（公式③），主要是通过MSE来拉近两个分布的均值，而方差并没有选择通过模型来预测，而是直接选择了 给定x_0的逆向分布的方差，即公式④，并且作者也给出了解释，方差通过网络预测，可能导致训练不稳定。

想要尽可能的近似真实分布，可以通过对真实数据分布下，最大化模型预测分布的对数似然。类似VAE的思想，使用变分下限(VLB)来优化负对数似然（涉及到KL散度的知识），经过一系列的公式推导，发现得到了一系列组合的KL散度，而最小化负对数似然就是在最小化这些组合的KL散度，拉近两个高斯分布，即模型预测的分布和给定 x_0的逆向分布。

进一步对多元高斯分布的KL散度求解，可以得到训练过程中损失函数的真实面貌，就是在让模型学习到指定分布的均值（给定x_0的逆向分布，可以通过公式推导求出解析式），更进一步，因为均值表达式中x_t和α这些参数在逆扩散过程中都是已知的、作为输入，因此唯一不同的就是高斯噪声ε。
$$\begin{array}{rl}{L}_t^{\text{simple}}& ={\mathbb{E}}_{t\sim [1,T],{\mathbf{x}}_0,{\mathbf{ϵ}}_t}[\Vert {\mathbf{ϵ}}_t-{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]\\ & ={\mathbb{E}}_{t\sim [1,T],{\mathbf{x}}_0,{\mathbf{ϵ}}_t}[\Vert {\mathbf{ϵ}}_t-{\mathbf{ϵ}}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_t,t){\Vert }^2]\end{array}$$

$$\begin{array}{rl}{L}_t& ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\frac{1}{2\Vert {\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t){\Vert }_2^2}\Vert {\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)-{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\frac{1}{2\Vert {\mathbf{\Sigma }}_{\theta }{\Vert }_2^2}\Vert \frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_t)-\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)\left){\Vert }^2\right]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\bar{\alpha }}_t)\Vert {\mathbf{\Sigma }}_{\theta }{\Vert }_2^2}\Vert {\mathbf{ϵ}}_t-{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\bar{\alpha }}_t)\Vert {\mathbf{\Sigma }}_{\theta }{\Vert }_2^2}\Vert {\mathbf{ϵ}}_t-{\mathbf{ϵ}}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_t,t){\Vert }^2]\end{array}$$

重点！！！

因此高斯噪声ε_θ就是最终模型需要学习预测的，并且训练的核心就是取高斯噪声ε和ε_θ之间的MSE。逆扩散过程中，模型预测出高斯噪声ε_θ后，带入公式⑥中得到均值，然后利用指定的方差（公式④），使用重参数的技巧，从x_t中推出x_t-1,一步步的迭代至求出原始图片x_0
$$\begin{array}{rl}{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t)& =\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t))(模型预测分布的均值，公式⑥)\\ \text{Thus }{\mathbf{x}}_{t-1}& =\mathcal{N}({\mathbf{x}}_{t-1};\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))\end{array}$$

总结：模型预测输出的就是一个高斯噪声ε_θ，输入是x_t(通过x_0和α相关参数带入公式快速求出的，可以理解为原始图片经过t步加噪后得到的一个含有高斯噪声的图片) 和t（一般都是通过embedding层后输入模型）。从直觉上来讲，模型对于一个带噪的图片（可能是一个标准的高斯分布，即X_T），预测出添加的噪声是什么，然后通过带入公式、重参数采样，逆向扩散过程一步步反向求出x_0。

问题2 为什么训练的过程中，对于从训练集中采样得到的一个输入x_0，只是随机产生一个t ，扩散添加噪声得到x_t，而不是从1到T?

一开始看论文以及一些博客的时候这个问题就让我有些迷惑，因为按作者的说法是可以添加噪声得到的x_1,x_2…x_T的啊，为什么不这样做，产生这些，而不只是x_t？仔细思考，其实可以简单理解为**一个输入多样性的问题？**但是在足够数量的epoch的训练迭代下这个问题似乎就可以理解了，在一轮epoch中，对于固定输入图片1，在改轮中随机产生t1，那么在下一轮中随机产生t2，只要足够epoch，产生1…T不是问题。在另一个方面，似乎也可以从计算效率上考虑，训练集中的所有样本，都生成1…T的带噪x_t，T也是一个比较大的数，那么计算量过大，而且这样训练也让模型被迫更加聪明一些。

问题3 在推理，逆扩散的过程中，为什么不根据扩散过程中得到的x_0推x_t的公式，直接模型预测输出高斯噪声后，直接从x_t推出x_0，而是一步步地来？

这个问题也困扰了我，相信有一部分人也存在类似的疑惑。下面给出相关公式，其实上文已经列出，这里重新给出。
$$\begin{array}{rlr}{\mathbf{x}}_t& =\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}+\sqrt{1-{\alpha }_t}{\mathbf{ϵ}}_{t-1}& \text{ ;where }{\mathbf{ϵ}}_{t-1},{\mathbf{ϵ}}_{t-2},\cdots \sim \mathcal{N}(0,\mathbf{I})\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{\mathbf{ϵ}}}_{t-2}& \text{ ;where }{\overline{\mathbf{ϵ}}}_{t-2}\text{ merges two Gaussians (*).}\\ & =\dots & \\ & =\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}\mathbf{ϵ}(公式⑦)& \\ q({\mathbf{x}}_t|{\mathbf{x}}_0)& =\mathcal{N}({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,(1-{\bar{\alpha }}_t)\mathbf{I})& \end{array}$$

$${\mathbf{x}}_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}({\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_t)(公式⑧)$$

个人觉得主要的原因在于生成的效果。这种直接预测x_0的做法，直观上理解，就是原始图片添加噪声后的x_t作为输入，模型直接预测添加的随机噪声，然后直接带入公式（公式⑧）求得x_0，不经过逆扩散一步步采样、重参数的过程。那么对于训练集中的固定的采样的样本x_0，相当于模型尝试去记住了训练过程中添加的噪声，这样一步到位效果不是很好，尤其对于大数据集而言，模型需要足够的复杂度来记忆大量随机噪声，最终训练的效果可能不够好？而让模型预测均值，方差这些参数，从直觉上来讲是记住了逆扩散过程中的高斯分布，然后一步步迭代，从x_T…x_t…x_0，更加平滑一些，能够得到高质量的样本。

总结：

算是扩散模型的入门，很多东西喜欢钻牛角尖，想办法弄清楚细枝末节，导致理解起来磕磕绊绊，但总归结果是好的，对Diffusion Model的原理有了一个大致的掌握。

注： 本文中的一系列公式以及插图主要参考 大佬的博客，其实上面已经贴出来了，还是强烈建议自习阅读一番。