概率的艺术：Diffusion Model理论推导

写在前面

大约在一年前的秋天，我听朋友们提到：一款名为StableDiffusion的AIGC生成图片工具正在席卷画师界。彼时的我也尝试着在本地部署了一下这个项目，并且拿它画了些奇妙的图片（doge），效果确实优秀得有些震撼。不过那时的我还没有想到去深究这款工具背后的知识，后来做科研时几次想过要彻底地学习一遍，却也都因为各种事情耽搁了。直到一年后的今天，赋闲在家，静下心来拿起纸笔推导一些数学，理清这强大工具背后坚实的理论基础，或许也是一种不错的消遣。

本文只是作者的一份学习笔记，主要内容大部分来自以下几位知乎大佬的文章：

一文解释 Diffusion Model (一) DDPM 理论推导( https://zhuanlan.zhihu.com/p/565901160 )

一文解释 Diffusion Model (二) Score-based SDE 理论推导( https://zhuanlan.zhihu.com/p/589106222 )

如何理解扩散模型中的SDE？( https://www.zhihu.com/question/616179189/answer/3230281054 )

以及宋飏博士关于Score-Based SDE的论文：

https//arxiv.org/abs/2011.13456

以及Google Research关于Diffusion Model框架的工作：

https//arxiv.org/abs/2208.11970

概念介绍

扩散模型（Diffusion）是一种基于马尔科夫过程，通过多步加噪/去噪过程实现内容生成的模型。它和我们之前学过的VAE十分类似，都需要经过原图片->高斯噪声->生成图片这三步过程。

注：从我的个人感受来看，这几种生成模型的big-picture其实都差不多：它们都经历了原图->白噪声->新图的流程，且它们在数学上都是一个从某个概率分布里采样出实例的采样器（这个后面会解释）。而Diffusion相比于先前模型的最大区别可能就是引入了 微分 的思想，使得生成过程更加细粒度（精细）和可控。

一位dalao( https//lilianweng.github.io/posts/2021-07-11-diffusion-models/ )制作了几种生成式模型的对比图如下：

在实践中，Diffusion Model有很多种实现方法，例如基于VAE压缩latent space（隐空间）思想的Stable Diffusion，或者不压缩latent space（也就是让隐空间大小和原图保持一致）的DDPM。我们这篇文章不去讨论具体的实现细节，只是从general的理论角度梳理一遍Diffusion的整个过程。不过如果对这个过程理解得足够透彻，再去看具体的实现方式想必也不会感到困难了。

流程概述

Diffusion Model的训练流程可以大致分为两个阶段：第一阶段是加噪阶段，给定一张原图 ，我们会对其进行 步加噪，最终将其变为高斯噪声 ；第二阶段是去噪阶段，对于刚加完噪的 ，我们会对其进行 步去噪，最终尝试将其复原成原图 。在训练结束进行推理时，我们不进行第一阶段，而是取一个白噪声然后进行第二阶段，最后 “复原” 得到的图片就是模型生成的图片。

注: 一张图片是怎么和随机变量产生联系的? 假定原图是一个 的张量，那么它在每一个位置 的值都是一个确定的数字，看起来好像和随机变量没什么关系。但是不要忘了，我们训练的对象并非一张图片，而是一整类图片。假如我们有 张人脸的图片（大小都是 )，那么现在的每一个位置 上就有 个数值，可以视作一个概率分布。这些概率分布事实上就代表了这类图片的所有特征，例如人脸在靠近屏幕中间的位置 通常是肉色的，那么这里的概率分布方差 就会较小，且RGB三通道的均值 拼起来会产生 “肉色”。

绝大多数的Diffusion Model都是基于上述流程实现的，它们的区别通常只在于“加噪”和“去噪”的方式不同。当然，仅凭抽象的框架很难深入推导模型的内容，在这里我们尝试具体一些，以比较流行的DDPM（Denoising Diffusion Probabilistic Models( https//arxiv.org/abs/2006.11239 )）为例进行一些推导。

加噪过程/前向过程

假设我们的加噪过程已经进行了 步，现在要进行 的加噪。一个自然的想法是把图片和噪声构造成一个线性关系:

其中

递归的形式不太好解，我们考虑能不能写成通式:

不妨假定 相互独立，根据正态分布的可加性，上式可以写成

其中 . 这个式子有点不好看，因为里面既有 又有 ，难以化简。我们是否可以尝试把这两个参数归一化，也就是令 ，以此来消掉其中的一个变量呢? 既然我们希望最终得到的图片 ，那么归一化显然不会影响我们达到这个目标，故我们尝试归一化，此时有

于是我们就可以令 ，从而 ，于是(3)式就可以改写成

同理，因为有了 ，我们就可以将(1)式也改写成

将上面两个式子表示成概率分布的形式，我们就有

整个前向过程可以表达为

去噪过程/反向过程

在去噪过程中，我们希望实现从纯噪声 到真实图片 的转换，为此我们同样考虑每一步 ，但与前向过程不同，此时我们不把它建模成一个简单的线性关系（但仍然建模成正态分布)，而是让模型通过神经网络的参数 去学习每一步去噪过程的期望和方差

事实上这看起来就是学了前向过程的反向分布，我们后面也会再提到这一点。

将所有步累积起来得到整个随机过程的概率分布

注：这是整个随机过程 的概率分布，不是 的概率分布。想要求 的分布，我们需要枚举所有可能的 并将它们所对应的随机过程 的概率求和（积分），这也就是所谓的全概率公式。

从概率论的视角看，我们模型的计算结果应当是 的分布 (就是上述积分之后的结果），而我们的优化目标就是希望 尽可能接近 的真实分布（比如人脸）。不过我们并不知道真实数据 的分布，但我们手里有若干个从 的真实分布里采样出来的样本 ，那么我们自然希望这些样本在我们建模出来的分布里能够具有尽量高的概率 (举个例子，假如我们有很多张人脸，但建模建成了熊猫脸的分布，那么我们的样本在模型里出现的概率就会很低)，也就是我们希望求出

(事实上这就是最大似然估计的思想，这里只是又用白话叙述了一遍)

现在我们考虑通过上面提到的积分方法求出 的表达式

我们知道，KL散度可以视作对两个概率分布度量距离。考察上面式子中的三项，第一项是在反向过程的最后一步最大化 ；第二项是在反向过程的第一步试图拉近 和 ，但这一步没有可训练的参数；第三步则是在反向过程的所有中间步都试图拉近 和 。

注: 可以理解为，我们希望模型在反向过程中能够尽可能准确地模拟出前向过程的每一步。我们前面提到过 ，这就是说前向过程的每一步都是一个正态分布，这个正态分布的均值包含了我们希望获得的原图特征的一个部分（微分）。所以事实上我们希望准确预测的是前向过程每一步的均值，这一点马上就会在后面提到。

现在我们已经可以通过梯度下降的办法去优化上面的式子了 (事实上我们优化的是 的一个下界，但这和对原目标函数的优化可以视作是等效的) ，但是有一个小问题：在第三项计算中我们每次需要同时求三项 才能算一个KL散度，这是否有点复杂? 能否进行一点优化?一个直接的想法是利用贝叶斯公式把 和 都改成同向的（比如都改成 ），根据

我们可以推得（仿照上面过程即可，这里就不推了)

现在我们就把第三项优化成了每次只需要算两项的形式，不过注意到 和前向过程是反着的，我们并不知道它的分布是什么，需要用贝叶斯公式算一下

现在我们再把前面建模好的 带进去，前面的 【注】提到我们希望估计的是均值，而方差并不携带信息，所以我们不妨就把 当做和 的方差完全一样 (因为优化它对我们学习图片特征没有任何帮助，所以不如不优化)。

现在我们再次考虑KL散度的计算，根据公式

可以计算得到

算了这么一堆，实际上我们只是从数学角度严格地证明了【注】中的观察是正确的。

现在考虑怎么学习这个 ，当然我们可以直接硬学，但是如果能做一点分解然后学其中的一个部分或许会更简单。

我们已经知道

在逆向过程中，我们的 是 的函数，所以上面的 那部分可以直接拿下来用，需要通过神经网络拟合的只有 那部分，也就是说

那么我们就可以算得

也就是说，我们最终学习的还是拟合原图 。

注：既然都是拟合原图 ，那Diffusion岂不是和之前的生成模型没有区别了? 但是值得注意的是: 我们是从 出发预测 的，而不是从啥也没有或者白噪声出发预测 的。从数学上来看，这里就是Diffusion和其他生成模型的最本质区别：它将整个生成过程进行了微分化，每次只预测一段过程而非整个过程。从直观上来看，这大大减小了预测的难度（提升了预测的质量)。后面我们还会从更加数学的角度分析这样做的优势。

但是这仍然没有达到最简形式，因为我们在前向过程中知道 ， 这意味着我们可以将 用 来表示，于是就有

那么我们又有了一种新的预测方法：预测噪声! 也就是建模

然后优化目标就变成了

预测噪声和预测 在数学上可以互相转化，但是实际操作中发现预测噪声比预测 表现要更好。

注1: 但是从数学直觉上来看预测噪声显然更好，因为 是一个相当复杂的分布，而噪声只是一个正态分布。

注2: 这里有的朋友可能会疑惑: 我们绕了一圈难道就是学了个正态分布生成器吗? 那我不用学不也可以生成正态分布? 请注意: 我们这篇文章里所有黑体的字母都代表一个具体的样本值，而不是一个随机变量。我们预测的不是随便一个噪声值，而是从具体值 变换到具体值