生成扩散模型漫谈：DDPM = 自回归式VAE

2022-07-10 10:07

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

在文章《生成扩散模型漫谈：DDPM = 拆楼 + 建楼》中，我们为生成扩散模型 DDPM 构建了“拆楼-建楼”的通俗类比，并且借助该类比完整地推导了生成扩散模型 DDPM 的理论形式。在该文章中，我们还指出 DDPM 本质上已经不是传统的扩散模型了，它更多的是一个变分自编码器 VAE，实际上 DDPM 的原论文中也是将它按照 VAE 的思路进行推导的。

所以，本文就从 VAE 的角度来重新介绍一版 DDPM，同时分享一下自己的 Keras 实现代码和实践经验。

多步突破

在传统的 VAE 中，编码过程和生成过程都是一步到位的：

这样做就只涉及到三个分布：编码分布、生成分布以及先验分布，它的好处是形式比较简单，与之间的映射关系也比较确定，因此可以同时得到编码模型和生成模型，实现隐变量编辑等需求；但是它的缺点也很明显，因为我们建模概率分布的能力有限，这三个分布都只能建模为正态分布，这限制了模型的表达能力，最终通常得到偏模糊的生成结果。

为了突破这个限制，DDPM 将编码过程和生成过程分解为步：

这样一来，每一个和仅仅负责建模一个微小变化，它们依然建模为正态分布。可能读着就想问了：那既然同样是正态分布，为什么分解为多步会比单步要好？这是因为对于微小变化来说，可以用正态分布足够近似地建模，类似于曲线在小范围内可以用直线近似，多步分解就有点像用分段线性函数拟合复杂曲线，因此理论上可以突破传统单步 VAE 的拟合能力限制。

联合散度

所以，现在的计划就是通过递归式分解（2）来增强传统 VAE 的能力，每一步编码过程被建模成，每一步生成过程则被建模成，相应的联合分布就是：

别忘了代表真实样本，所以就是数据分布；而代表着最终的编码，所以就是先验分布；剩下的、就代表着编码、生成的一小步。（提示：经过考虑，这里还是沿用本网站介绍 VAE 一直用的记号习惯，即“编码分布用 、生成分布用 ”，所以这里的 、 含义跟 DDPM 论文是刚好相反的，望读者知悉。）

在《变分自编码器（二）：从贝叶斯观点出发》中笔者就提出，理解 VAE 的最简洁的理论途径，就是将其理解为在最小化联合分布的 KL 散度，对于 DDPM 也是如此，上面我们已经写出了两个联合分布，所以 DDPM 的目的就是最小化

这就是 DDPM 的优化目标了。到目前为止的结果，都跟 DDPM 原论文的结果一样的（只是记号略有不同），也跟更原始的论文《Deep Unsupervised Learning using Nonequilibrium Thermodynamics》[1] 一致。接下来，我们就要将、具体形式定下来，然后简化 DDPM 的优化目标（4）。

分而治之

首先我们要知道，DDPM 只是想做一个生成模型，所以它只是将每一步的编码建立为极简单的正态分布：

，其主要的特点是均值向量仅仅由输入乘以一个标量得到，相比之下传统 VAE 的均值方差都是用神经网络学习出来的，因此 DDPM 是放弃了模型的编码能力，最终只得到一个纯粹的生成模型；至于，则被建模成均值向量可学习的正态分布。其中都不是可训练参数，而是事先设定好的值（怎么设置我们稍后讨论），所以整个模型拥有可训练参数的就只有。（提示：本文 的定义跟原论文不一样。）

由于目前分布不含任何的可训练参数，因此目标（4）中关于的积分就只是贡献一个可以忽略的常数，所以目标（4）等价于

由于先验分布一般都取标准正态分布，也是没有参数的，所以这一项也只是贡献一个常数。因此需要计算的就是每一项

其中第一个等号是因为至多依赖到，因此到的分布可以直接积分为 1；第二个等号则是因为也不依赖于，所以关于它们的积分我们也可以事先算出，结果为

，该结果可以参考下一节的式（9）。

场景再现

接下来的过程就跟上一篇文章的“又如何建”一节基本上是一样的了：

1、除去优化无关的常数，这一项所贡献的就是；

2、意味着

又意味着，其中；

3、由则启发我们将参数化为

。

这一系列变换下来，优化目标等价于

随后按照“降低方差”一节做换元，结果就是

这就得到了 DDPM 的训练目标了（原论文通过实验发现，去掉上式前面的系数后实际效果更好些）。它是我们从 VAE 的优化目标出发，逐步简化积分结果得到的，虽然有点长，但每一步都是有章可循的，有计算难度，但没有思路上的难度。

相比之下，DDPM 的原论文中，很突兀引入了一个（原论文记号）来进行裂项相消，然后还要假设（原论文记号）才能得到最终结果，可是最后又说可以不等于。整个过程之莫名其妙，对笔者来说相当难以接受。

超参设置

这一节我们来讨论一下的选择问题。

对于来说，习惯上约定，这样就减少了一半的参数了，并且有助于简化形式，这其实在上一篇文章我们已经推导过了，由于正态分布的叠加性，在此约束之下我们有

其中，而，这样一来就具有比较简约的形式。可能读者又想问事前是怎么想到这个约束呢？我们知道意味着

，如果也是的话，我们就希望也是，所以就确定了了。

前面说了，一般都取标准正态分布。而我们的学习目标是最小化两个联合分布的 KL 散度，即希望，那么它们的边缘分布自然也相等，所以我们也希望

由于数据分布是任意的，所以要使上式恒成立，只能让，即退化为与无关的标准正态分布，这意味着我们要设计适当的，使得。同时这再次告诉我们，DDPM 是没有编码能力了，最终的可以说跟输入无关的。用上一篇文章的“拆楼-建楼”类比就是说，原来的楼已经被完全拆成原材料了，如果用这堆材料重新建楼的话，可以建成任意样子的楼，而不一定是拆之前的样子。DDPM取了，关于该选择的性质，我们在上一篇文章的“超参设置”一节也分析过了。

至于，理论上不同的数据分布来说对应不同的最优，但我们又不想将设为可训练参数，所以只好选一些特殊的来推导相应的最优，并认为由特例推导出来的可以泛化到一般的数据分布。我们可以考虑两个简单的例子：

1、假设训练集只有一个样本是狄拉克分布，可以推出最优的；

2、假设数据分布服从标准正态分布，这时候可以推出最优的。

实验结果显示两个选择的表现是相似的，因此可以选择任意一个进行采样。两个结果的推导过程有点长，我们后面再择机讨论。

参考实现

这么精彩的模型怎么可以少得了 Keras 实现？下面提供笔者的参考实现：

Github地址：

https://github.com/bojone/Keras-DDPM

注意，笔者的实现并非严格按照 DDPM 原始开源代码来进行，而是根据自己的设计简化了 U-Net 的架构（比如特征拼接改为相加、去掉了 Attention 等），使得可以快速出效果。经测试，在单张 24G 显存的 3090 下，以 blocks=1,batch_size=64 训练 128*128 大小的 CelebA HQ 人脸数据集，半天就能初见成效。训练 3 天后的采样效果如下：

▲ 笔者训练的 DDPM 采样结果演示

在调试过程中，笔者总结出了如下的实践经验：

1. 损失函数不能用 mse，而必须用欧氏距离，两者的差别是 mse 在欧氏距离基础上除以图片的，这会导致损失值过小，部分参数的梯度可能会被忽略为 0，从而导致训练过程先收敛后发散，该现象也经常出现于低精度训练中，可以参考《在 bert4keras 中使用混合精度和 XLA 加速训练》[2]；

2. 归一化方式可以用 Instance Norm、Layer Norm、Group Norm 等，但不要用 Batch Norm，因为 Batch Norm 存在训练和推理不一致的问题，可能出现训练效果特别好，预测效果特别差的问题；

3. 网络结构没有必要照搬原论文，原论文是为了刷 SOTA 发论文，照搬的话肯定是又大又慢的，只需要按照 U-Net 的思路设计自编码器，就基本上可以训练出个大概效果了，因为就相当于是个纯粹的回归问题，还是很好训练的；

4. 关于参数的传入，原论文用了 Sinusoidal 位置编码，笔者发现直接换为可训练的 Embedding，效果也差不多；

5. 按照以往搞语言模型预训练的习惯，笔者用了 LAMB 优化器，它更方便调学习率，基本上的学习率可以适用于任意初始化方式的模型训练。

综合评价

结合《生成扩散模型漫谈：DDPM = 拆楼 + 建楼》和本文的介绍，想必读者都已经对 DDPM 有自己的看法了，能基本看出 DDPM 优点、缺点以及相应的改进方向在哪了。

DDPM 的优点很明显，就是容易训练，并且生成的图片也清晰。这个容易训练是相对 GAN 而言的，GAN 是一个过程，训练中的不确定性很大，容易崩溃，而 DDPM 就纯粹是一个回归的损失函数，只需要纯粹的最小化，因此训练过程非常平稳。同时，经过“拆楼-建楼”的类比，我们也可以发现 DDPM 在通俗理解方面其实也不逊色于 GAN。

不过，DDPM 的缺点也很明显。首先最突出的就是采样速度太慢，需要执行模型步（原论文才能完成采样），可以说这比 GAN 的一步到位的采样要慢上倍，后面有很多工作对这一点进行改进；其次，在 GAN 中，从随机噪声到生成样本的训练是一个确定性的变换，随机噪声是生成结果的一个解耦的隐变量，我们可以进行插值生成，或者对之编辑以实现控制生成等，但是 DDPM 中生成过程是一个完全随机的过程，两者没有确定性的关系，这种编辑生成就不存在了。DDPM 原论文虽然也演示了插值生成效果，但那只是在原始图片上进行插值的，然后通过噪声来模糊图片，让模型重新“脑补”出新的图片，这种插值很难做到语义上的融合。

除了针对上述缺点来做改进外，DDPM 还有其他一些可做的方向，比如目前演示的 DDPM 都是无条件的生成，那么很自然就想到有条件的 DDPM 的，就好比从 VAE 到 C-VAE、从 GAN 到 C-GAN 一样，这也是当前扩散模型的一个主流应用，比如用 Google 的 Imagen 就同时包含了用扩散模型做文本生成图片以及做超分辨率，这两者本质上就是条件式扩散模型了；再比如，目前的 DDPM 是为连续型变量设计的，但从其思想来说应该也是适用于离散型数据的，那么离散型数据的 DDPM 怎么设计呢？

相关工作

说到 DDPM 的相关工作，多数人会想到传统扩散模型、能量模型等工作，又或者是去噪自编码器等工作，但笔者接下来想说的不是这些，而是本博客之前介绍过的、甚至可以认为 DDPM 就是它的特例的《强大的 NVAE：以后再也不能说 VAE 生成的图像模糊了》。

站在 VAE 的视角来看，传统 VAE 生成的图片都偏模糊，而 DDPM 只能算是（笔者所了解到的）第二个能生成清晰图像的 VAE，第一个正是 NVAE。翻看 NVAE 的形式，我们可以发现它跟 DDPM 有非常多的相似之处，比如 NVAE 也是引入了一大堆隐变量，这些隐变量也呈递归关系，所以 NVAE 的采样过程跟 DDPM 也是很相似的。

从理论形式来说，DDPM 可以看成是一个极度简化的 NVAE，即隐变量的递归关系仅仅建模为马尔可夫式的条件正态分布，而不是像 NVAE 的非马尔科夫式，生成模型也只是同一个模型的反复迭代，而不是 NVAE 那样用一个庞大的模型同时用上了，但 NVAE 在利用众多之时，也加入了参数共享机制，这跟同一个模型反复迭代也异曲同工了。