大一统视角理解扩散模型（二）

4.『Variation Diffusion Model』

我们之所以在谈论扩散模型之前，要花如此大的篇幅介绍VAE，并引出MHVAE的证据下界推导是因为我们可以非常自然地将扩散模型视为一种特殊的MHVAE，该MHVAE满足以下三点限制（注意以下三点限制也是整个扩散模型推断的基础）：

1. 潜在向量Z的维度和输入X的维度保持一致。
2. 每一个时间步的潜在向量都被编码为一个仅依赖于上一个时间步的潜在向量的高斯分布。
3. 每一个时间步的潜在向量的高斯分布的参数，随时间步变化，且满足最终时间步的高斯分布满足标准高斯分布的限制。

因为第一点维度一致的原因，在不影响理解的基础上，我们将MHVAE里的Zt表示为Xt（其中x0为原始输入），则我们可以将MHVAE的层级潜在向量依赖图，重新画为以下形式（即将扩散模型的中间扩散过程当做潜在向量的层级建模过程）：

直至这里，我们终于见到了我们熟悉的扩散模型的形式。

而在将上面的公式25-28里的Zt与Xt替换后，我们可以得到VDM里证据下界的推导公式里的前四行，即公式34-37。并且在此基础上，我们可以继续往下推导。37至38行的变换是链式法则的等价替换（或上述公式23和24的变换），38至39行是连乘过程的重组，39至40行是对齐连乘符号的区间，40至41行应用了Log乘法的性质，41至42继续运用该性质进一步拆分，42至43行是因为和的期望等于期望的和，43至44是因为期望目标与部分时间步的概率无关可以直接省去，44至45步是应用了KL散度的定义进行了重组。

至此，我们又一次将原数据分布的对数似然，转化为了证据下界（公式37），并将其转化为了几项非常直观的损失函数的加和形式（公式45），他们分别为：

1. 重构项，即从潜在向量x1到原数据x0的变化。在VAE里该重构项写为logP(x|z),而在这里我们写做logP(x0|x1)
2. 先验匹配项。回忆我们上述提到的MHVAE里最终时间步的高斯分布应建立为标准高斯分布
3. 一致项。该项损失是为了使得前向加噪过程和后向去噪的过程中，Xt的分布保持一致。直观上讲，对一个更混乱图像的去噪应一致于对一个更清晰的图像的加噪。而因为一致项的损失是定义于所有时间步上的，这也是三项损失里最耗时计算的一项。

虽然以上的公式推导给了我们一个非常直观的证据下界，并且由于每一项都是以期望来计算，所以天然适用蒙特卡洛方法来近似，但如果优化该证据下界依然存在几个问题：

1. 我们的一致项损失是一项建立在两个随机变量(Xt-1, Xt+1)上的期望。他们的蒙特卡洛估计的方差大概率比建立在单个独立变量上的蒙特卡洛估计的方差大。
2. 我们的一致项是定义于所有时间步上的KL散度的期望和。对于T取值较高的情况（通常扩散模型T取2000左右），该期望的方差也会很大。

所以我们需要重新推导一个证据下界。而这个推导的关键将着眼于以下这个观察：我们可以将扩散过程的正向加噪过程q(xt|xt-1)重写为q(xt|xt-1, x0)。之所以这样重写的原因是基于马尔可夫假设，这两个式子完全等价。于是对这个式子使用贝叶斯法则，我们可以得到式46.

基于公式46，我们可以重写上面的证据下界（式37）为以下形式：其中式47，48和式37，38一致。式49开始，分母的连乘拆解由从T开始改为从1开始。式50基于上文提及的马尔可夫假设对分母添加了x0的依赖。式51用log的性质拆分了对数的目标。式52代入了式46做了替换。式53将划掉的分母部分连乘单独提取出来后发现各项可约剩下式54部分的log(q(x1|x0)/q(xT|x0))。式54用log的性质消去了q(x1|x0)得到了式55。式56用log的性质拆分重组了公式，式57如同前述式43-44的变换，省去了无关的时间步。式58则用了KL散度的性质。

至此，我们应用了马尔可夫假设得到了一个更优的证据下界推导。该证据下界同样包含几项直观的损失函数：

1. 重构项。该重构项与上面提及的重构项一致。
2. 先验匹配项。与上面提及的形式略有差别，但同样是基于最终时间步应为标准高斯的先验假设
3. 去噪匹配项。与上面提及的一致项的最大区别在于不再是对两个随机变量的期望。并且直观上理解p(xt-1|xt)代表的是后向的去噪过程，而q(xt-1|xt, x0)代表的是已知原始图像和目标噪声图像的前向加噪过程。该加噪过程作为目标信号，来监督后向的去噪过程。该项解决了期望建立于两个随机变量上的问题。

注意，以上的推导完全基于马尔可夫的性质所以适用于所有MHVAE，所以当T=1的时候，以上的证据下界和VAE所推导出的证据下界完全一致！并且本文之所以称为大一统视角，是因为对于该证据下界里的去噪匹配项，不同的论文有不同的优化方式。但归根结底，他们的本质互相等价，且皆由该式展开推导得到。下面我们会从扩散模型的角度做公式推导，来展开计算去噪匹配项。(注意第一版的推导里的一致项，也完全可以通过下一节的方式得到q和p的表达式，再通过KL来计算解析式)

5.『Diffusion Model recap』

在扩散模型里，有几个重要的假设。其中一个就是每一步扩散过程的变换，都是对前一步结果的高斯变换（上一节MHVAE的限制条件2）：

这一点和VAE有很大不同。VAE里编码器侧的潜在向量的分布是通过模型训练得到的。而扩散模型里，前向加噪过程里的每一步都是基于上一步结果的高斯变换。其中 alpha_t 一般当作超参设置得到。这点对于我们计算扩散模型的证据下界有很大帮助。因为我们可以基于输入x0确切地知道前向过程里的某一步的具体状态，从而监督我们的预测。

基于式31，我们可以递归式地对x0不断加噪变换，得到最终xt的表达式：

所以对于式58里噪音匹配项里的监督信号，我们可以重写成以下形式，其中根据式70，我们可以得到q(xt|x0)和q(xt-1|x0)的表达式，而q(xt|xt-1, x0)因为是前向扩散过程，可以应用马尔可夫性质看做q(xt|xt-1)使用式31得到具体表达式。

参考已经证明了前向加噪过程可以写为一个高斯分布了。在扩散模型的初始论文[2]里提到，对于一个连续的高斯扩散过程，其逆过程与前向过程的方程形式(functional form)一致。所以我们将对去噪匹配项里的p_theta(xt-1|xt)也采用高斯分布的形式（更加具体的一些推导放在了末尾的补充里）。注意式58里，对两个高斯分布求KL散度，其解析解的形式如下：

我们现在已知其中一个高斯分布（左侧）的参数，现在如果我们令右侧的高斯分布和左侧高斯分布的方差保持一致。那么优化该KL散度的解析式将简化为以下形式：

如此一来式58的噪音匹配项就被简化为最小化前后向均值的预测误差(式92)。读者请注意，以下的大一统的三个角度来看待Diffusion model，实质上都是对式92里mu_q的不同变形所推论出来的。 其中mu_q是关于xt, x0的函数，而mu_theta是关于xt和t的函数。其中通过式84，我们有mu_q的准确计算结果，而因为mu_theta是关于xt的函数。我们可以将其写为类似式84的形式（注意，有关为什么可以忽略方差并且让均值选取这个形式放在了最末尾的补充讨论里。但关于这个形式的选择的深层原因实质上开辟了一个全新的领域来研究，并且关于该领域的研究直接导向了扩散模型之后的一系列加速采样技术的出现）

比较式84与94可知，x_hat是我们通过噪音数据xt来预测原始数据x0的神经网络。那么我们可以将式58里证据下界的噪音匹配项，最终写为

那么，我们最后得到扩散模型的优化，最终表现为训练一个神经网络，以任意时间步的噪音图像为输入，来预测最初的原始图像！此时优化目标转化为了最小化预测误差。同时式58上的对所有时间步的噪音匹配项求和的优化，可以近似为对每一时间步上的预测误差的期望的最小值，而该优化目标可以通过随机采样近似：