扩散模型：DDPM的两种数学推导-轻识

在上一篇文章最后结尾时，笔者提到扩散模型(Diffusion Models)实际上就是一种在限定条件下的马尔可夫层次变分自编码器(Markovian Hierarchical Variational Autoencoder, MHVAE)，是融入了马尔可夫过程的VAE模型。其本质上和VAE的流程有相似之处，都是经过原始图像->高斯噪声图->生成图像的三个步骤。本文将接续上一篇关于ELBO的推导逻辑，分别从两个角度阐述扩散模型的详细数学推导过程。

(注：全文公式用markdown编排后到公众号会出现格式错乱，所以本文公式编号会在正文中预先描述。)

扩散模型概述

扩散模型是一种深度学习生成模型，能够通过一系列逐渐加入噪声的步骤，将数据从原始分布转换为随机噪声，然后再通过一个逐渐去除噪声的过程来生成新的数据样本。这种模型的关键思想是模拟数据的分布与纯噪声之间的转换过程，以此来学习数据的底层结构。扩散模型主要包括两个过程，前向过程和反向过程。

前向过程(forward process)。前向过程也叫加噪过程，目的是将原始图像转化为随机的高斯噪声。这一过程是逐步进行的，并且是预定义好的，直到原始图像转化为随机噪声为止。如图1从左往右的箭头显示，通过对原始图像逐步经过个时间步加噪后得到随机噪声图像，每一步的转换分布描述为，即给定上一步噪声图像条件下，当前加噪图像的概率分布。
反向过程(reverse process)。反向过程也叫去噪过程，目的是将完全随机的噪声图像还原为原始图像。这一过程同样也是逐步进行的，是通过神经网络来直接预测原图或者间接的预测每一步的噪声来实现，直到噪声图像还原为原始图像为止。如图1中从右往左的箭头所示，通过对噪声图像逐步经过个时间步去噪后得到原始图像，每一步的转换分布描述为，即给定上一步噪声图像条件下，当前去噪图像的概率分布。

图1 扩散模型的前向与反向流程

扩散模型的前向与反向过程如图1所示。不同的扩散模型可能在加噪与去噪方式上有所差异，但上述前向-反向过程是大多数扩散模型的核心思想，所以深度理解前向/反向过程的详细推导过程非常有必要。

扩散模型的ELBO

变分下界(ELBO)的推导永远是生成模型的核心步骤，扩散模型自然也不例外。在推导扩散模型的ELBO前，我们先从扩散模型的三个限定条件出发，给出每个限定条件的数学表达。

限定条件1：扩散模型的潜变量维度恰好与输入数据维度相等。假设输入图像和中间过程图像都表达为，当时表示为原始图像，表示为潜变量图像。根据MHVAE的后验分布，扩散模型的后验可以表示为(公式1)：

限定条件2：潜变量编码器(前向过程)在每个时间步并非训练得到，而是预定义为一个线性高斯模型。换句话说，就是当前时间步以上一个时间步的输出为基础的高斯分布。我们将高斯分布的均值参数化表示为，方差表示为，那么前向过程的状态转移分布可以表示为(公式2)：

限定条件3：潜变量编码器的高斯分布参数随时间变化，直到最后一个时间步时的分布为标准正态分布。根据该条件我们可以将MHVAE的联合分布改写为(公式3~4)：

我们可以上述限定条件的数学表达和上一篇关于对数似然的ELBO推导，可以初步将扩散模型的ELBO推导为(公式5~16)：

根据上述推导，扩散模型的ELBO可以划分为三个部分：

第一项为重构项，用于预测在给定第一个时间步时原图的对数概率。
第二项为先验匹配项，用于最小化最终的潜变量分布与高斯先验之间的差异，用一个KL散度来表示。
第三项为一致性项，也是一个KL散度项，用于衡量前向过程和反向过程中的的分布一致性，这一项是我们需要重点关注的。

图2 初步的扩散模型ELBO优化目标

第一版的扩散模型ELBO如图2所示。在这一版我们的优化目标是最小化图中粉色箭头与绿色箭头两个分布之间的差异。

上述ELBO表达式还是过于复杂，特别是第三项，涉及了和两个随机变量的期望，我们接下来尝试看看能否对每一项化简到只计算一个随机变量期望的情形。我们先根据马尔可夫链的无记忆性将前向过程中编码器转换分布改写为(公式17)：

所谓无记忆性，即在一个马尔可夫过程中，下一状态的概率分布只能由当前状态决定，与它前面的历史状态均无关。

再根据贝叶斯定理，我们进一步地将改写为(公式18)：

然后我们重新对前述ELBO的推导过程进行梳理(公式19~30)：

上述推导给出扩散模型最终的ELBO表达式，前两项与之前推导的两项没有太大差异，我们同样是重点关注第三项。该项也叫去噪匹配项，目的是将需要学习的去噪转换步与作为groundtruth的去噪转换步进行无限近似，所以最终扩散模型是要最小化这两个分布的KL散度。

图3 最终的扩散模型ELBO推导的优化目标

图3给出了扩散模型的最终优化目标，最终我们要优化的目标仍然是最小化绿色箭头和粉色箭头之间的分布差异，不过这次粉色箭头也表示了反向过程。

以上就是扩散模型的ELBO推导过程，看起来有点复杂，但是实际推导起来并没有太多难以理解的地方。

扩散模型的推导：直接预测原图

基于扩散模型ELBO的推导，我们实际感兴趣的是这个分布，那么这个分布我们怎么知道它的表达式呢？还是基于贝叶斯定理(公式31)：

其中(公式32~33):

所以只需要重点推导一下和的表达式。我们先来看一下，与之类似，或许可以在推导出来之后直接使用相似表达。表示的是在任意时间步时潜变量图像与原图之间的关系。直接看肯定看不出二者之间有什么关联，我们已知之间的关系，所以尝试从该式入手或许是个出路。

已知(公式34~35)：

根据式(35)进行递推，有(公式36~44)：

同理可以表示为(公式45)：

然后将和的表达式代入到式(31)中，可得(公式46)：

可以看到，最终也服从一个正态分布，为简化公式，我们可以将该正态分布的均值设为，方差设为。类似地，我们也可以将同样也作为一个服从均值为，方差同样为的正态分布。

那么现在问题变成了如何最小化两个正态分布的KL散度问题。这里正好有公式可以套用，根据两个正态分布之间的KL散度定义(公式47~48)：

将前述两个正态分布套用上述公式，有(公式49~55)：

其中即，即，根据式(46)我们可以分别将和表示为(公式56~57)：

其中可以通过一个参数化的神经网络(比如UNet)从噪声图像和时间步来直接预测。所以，我们从噪声图像中预测原图的最终优化问题可以推导为(公式58~63)：

最终，我们经过一系列推导，我们可以认为扩散模型通过一个参数化的神经网络来从任意噪声图像和时间步中预测出原始图像。

扩散模型的推导：预测噪声

虽然我们可以通过一系列的推导证明了扩散模型可以直接拟合原图来建模，但实际训练时拟合原图表现并不好。一种自然的想法就是能否先拟合图像中的噪声，然后将图像减去预测的噪声后得到去噪后的图像。所以本小节我们从拟合噪声的角度来推导扩散模型。

我们将式(34)表示为的表达式(公式64)：

将式(64)代入式(55)的的表达式中，可推导得(公式65)：

进一步地，我们也可以将近似表示为(公式66)：

所以相应的优化问题就变成了(公式67~72)：

最后式(72)中的就是一个用于预测噪声的神经网络模型。通过一顿操作，预测原图的问题就可以转换为预测噪声图像了。

DDPM的训练与采样DDPM论文给出了扩散模型的训练与采样算法流程。如图4所示，可以看到这里的训练与采样算法流程都是基于预测噪声的视角阐述的。

图4 DDPM的训练与采样算法

这里笔者借用李宏毅老师的扩散模型PPT插图，让大家对DDPM的训练与采样流程有一个更直观的理解。具体如图5-6所示。

图5 DDPM训练图解

基于预测噪声视角的DDPM训练阶段是使用噪声预测模型来预测图像中的噪声，然后最小化目标噪声与预测噪声之间的损失。

图6 DDPM采样图解

在采样阶段，噪声预测模型以噪声图像和当前时间步作为输入得到一个预测噪声图像，通过噪声图像减去预测噪声(注意不要遗漏系数项)即可得到去噪后的图像。对该图像加上一个从标准正态分布中采样一个噪声即可得到最终的去噪图像。

那么这里为什么还要再加上一个采样出来的噪声呢？我的理解是在逆向过程中，噪声的加入使得每次采样都略有不同，从而引入随机性，允许模型生成多样化的样本。

总结

本文重点阐述了DDPM扩散模型的详细推导过程，在推导出扩散模型ELBO的基础上，分别从预测原图和预测噪声图像两个角度推导出了扩散模型的优化表达式，并对DDPM的训练与采样流程进行了详细解释。本文可作为深入理解扩散模型的入门材料。

参考文献

[1] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020.

[2] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696–21707, 2021.

[3] Calvin Luo. Understanding diffusion models: A unified perspective. arXiv preprint arXiv:2208.11970, 2022.

[4] https://speech.ee.ntu.edu.tw/~hylee/ml/2023-spring.php

[5] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022.