生成扩散模型漫谈：W距离 ≤ 得分匹配

2023-02-24 04:02

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 追一科技

研究方向 | NLP、神经网络

Wasserstein 距离（下面简称“W距离”），是基于最优传输思想来度量两个概率分布差异程度的距离函数，笔者之前在《从Wasserstein距离、对偶理论到WGAN》等文章中也做过介绍。

对于很多读者来说，第一次听说 W 距离，是因为 2017 年出世的 WGAN [1]，它开创了从最优传输视角来理解 GAN 的新分支，也提高了最优传输理论在机器学习中的地位。很长一段时间以来，GAN [2] 都是生成模型领域的“主力军”，直到最近这两年扩散模型异军突起，GAN 的风头才有所下降，但其本身仍不失为一个强大的生成模型。

从形式上来看，扩散模型和 GAN 差异很明显，所以其研究一直都相对独立。不过，去年底的一篇论文《Score-based Generative Modeling Secretly Minimizes the Wasserstein Distance》[3] 打破了这个隔阂：它证明了扩散模型的得分匹配损失可以写成 W 距离的上界形式。这意味着在某种程度上，最小化扩散模型的损失函数，实则跟 WGAN 一样，都是在最小化两个分布的 W 距离。

结论分析

具体来说，原论文的结果，是针对《生成扩散模型漫谈：一般框架之SDE篇》中介绍的 SDE 式扩散模型的，其核心结论是不等式（其中是的非负函数，具体含义我们后来再详细介绍）

那么怎样理解这个不等式呢？首先，扩散模型可以理解为 SDE 从到的一个运动过程，最右边的是时刻的随机采样分布，通常就是标准正态分布，而实际应用中一般都有，所以，原论文之所以显式写出它，只是为了从理论上给出最一般的结果。

接着，左边的，是从采样的随机点出发，经反向 SDE

求解得到的时刻的值的分布，它实际上就是要生成的数据分布；而，则是从采样的随机点出发，经过 SDE

求解得到的时刻的值的分布，其中是的神经网络近似，所以实际就是扩散模型生成的数据分布。因此，的含义就是数据分布与生成分布的 W 距离。

最后，剩下的积分项，其关键部分是

这也正好是扩散模型的“得分匹配”损失。所以，当我们用得分匹配损失去训练扩散模型的时候，其实也间接地最小化了数据分布与生成分布的距离。跟 WGAN 不同的是，WGAN 优化的距离是而这里是。

注：准确来说，式（4）还不是扩散模型的损失函数，扩散模型的损失函数应该是“条件得分匹配”，它跟得分匹配的关系是：

最后的结果才是扩散模型的损失函数“条件得分匹配”。第一个等号是因为恒等式

，第二个不等号则是因为平方平均不等式的推广或者詹森不等式，第三个等号则是贝叶斯公式了。也就是说，条件得分匹配是得分匹配的上界，所以也是距离的上界。

从式（1）中我们也可以简单理解为什么扩散模型的目标函数要舍去模长前面的系数了，因为距离是概率分布的良好度量，而式（1）右端的是关于的单调递增函数，这意味着我们要适当加大当较小时的得分匹配损失。而在《生成扩散模型漫谈：一般框架之SDE篇》我们推导过得到匹配的最终形式为：

舍去系数等价于乘以，而也是的单调递增函数。也就是说，可以简单地认为舍去系数是让训练目标更加接近两个分布的距离。

准备工作

尽管原论文给出了不等式（1）的证明过程，但涉及到较多的最优传输相关知识，如连续性方程、梯度流等，特别是它不加证明引用的一个定理，还是放在一本梯度流专著的第 8 章或另一本最优传输专著的第 5 章，这对笔者来说阅读难度实在太大。

经过一段时间的尝试，笔者终于在上周笔者完成了自己关于不等式（1）的（一部分）证明，其中只需要用到 W 距离的定义、微分方程基础以及柯西不等式，相比原论文的证明理解难度应该是明显降低了。经过几天的修改完善，给出如下的证明过程。

在开始证明之前，我们先做一下准备，先整理一下接下来会用到的一些基本概念和结论。首先是 W 距离，它定义为

其中是指所有以为边缘分布的联合概率密度函数，它描述了具体的传输方案。本文只考虑，因为只有这种情形方便后续推导。注意到距离的定义包含了下确界运算，这就意味着对于任意我们能写出的，都有

这是笔者所给证明的核心思想。证明过程的放缩，主要用到柯西不等式：

证明过程中我们会假设函数满足“单侧 Lipschitz 约束”，其定义为

可以证明它比常见的 Lipschitz 约束（参考深度学习中的 Lipschitz 约束：泛化与生成模型）更弱，即如果函数满足 Lipschitz 约束，那么它一定满足单侧 Lipschitz 约束。

牛刀小试

不等式（1）过于一般了，一上来就试图分析一般化的结果并不利于我们的思考和理解。所以，我们先将问题简化一下，看能不能先证明一个稍弱一些的结果。怎么简化呢？首先，不等式（1）考虑了初始分布（提示，扩散模型是到的演化过程，所以是初始时刻，是终止时刻）的差异，而这里我们先考虑相同初始分布；此外，原本的反向方程（2）是一个 SDE，这里先考虑确定性的 ODE。

具体来说，我们考虑从同一个分布出发采样作为时刻的初始值，然后分别沿着两个不同的 ODE

进行演化，设时刻的分布为、的分布为，我们尝试去估计的一个上界。

我们知道，都是以为初始值通过各自的 ODE 演化而来，所以它们其实都是的确定性函数，更准确的记号应该是，简单起见我们才略去了。这就意味着对应于同一个的构成了的样本之间的一个对应关系（传输方案），如下图（这个图不大好画，就随便手画了一下）：

▲ 近似最优传输方案示意图

于是根据式（8），我们可以写出

下面我们对进行放缩。为了将它跟联系起来，我们对它求导：

其中第一个不等号用到了柯西不等式的向量版，以及单侧 Lipschitz 约束假设（10），第二个不等号则用到了柯西不等式的期望版，的意思是最终得到的不等关系，不管取+还是-都是成立的，下面的推导只用到了-这一侧。结合，我们得到

用常数变易法，设

，代入上式得到

两边在积分，并结合（初始时刻两个分布相等，距离为 0），得到

于是

其中。根据式（12），这也是的上界。最后，由于求期望的式子只是的函数，又是的确定性函数，对于它关于的期望等价于直接关于的期望，于是：

一鼓作气

实际上，简化版的不等式（18）已经和更一般的（1）没有本质区别了，它的推导过程已经包含了导出完整结果的一般思路，下面我们来完成剩余的推导过程。

首先，我们将式（18）推广到不同初始分布的场景，假设两个初始分布为，从采样初始值演化，从采样初始值演化，所以此时分别是的函数，而不是像之前那样是同一个的函数，所以无法直接构造一个传输方案。

所以，我们还需要之间的一个对应关系（传输方案），我们将它选择为之间的一个最优传输方案。于是，我们可以写出类似式（12）的结果：

由于定义的一致性，那么放缩过程（13）同样是成立的，只不过期望换成了，所以不等式（14）、（15）也是成立的。不同的是在对（15）两端在积分时，不再有，而是根据定义有

。所以，最终的结果是

最后，我们回到扩散模型。在《生成扩散模型漫谈：一般框架之ODE篇》我们已经推导过，同一个前向扩散过程，实际上对应一簇反向过程：

其中是可以自由选择的标准差函数，当时，那么就是方程（2）。由于我们上面分析的是 ODE，所以我们先考虑的情形，此时结果（20）依然可用，只不过将换成

、将换成，代入式（20）后就得到文章开头的结论（1）了。当然别忘了我们推导过程中对所做的单侧 Lipschitz 约束假设（10），现在可以分别对、做出假设，这些细节就不展开了。

艰难收尾

按照流程，接下来我们应该再接再厉，完成的收尾证明。不过很遗憾，本文的思路不能完全证明 SDE 的情形，下面给出笔者的分析过程。事实上，对于大部分读者来说，了解到上一节的 ODE 例子就可以窥见式（20）的精髓了，完整的细节也不是太重要。

简单起见，下面我们以（2）为例，更一般的（21）也可以类似地分析。我们需要估算的是如下两个 SDE 的演化轨迹分布差异：

也就是将准确的换成近似的，对最终分布的影响有多大。笔者的证明思路同样是将它转化为 ODE，继而用回前面的证明过程。首先，根据式（21），我们知道第一个 SDE 对应的 ODE 为：

至于第二个 SDE 对应的 ODE 的推导有些技巧，需要先变为的形式，然后再利用式（21）：

对这两个 ODE 重复放缩过程（13）（取负号），那么主要的区别是多出来一项

如果这一项小于等于 0，那么放缩过程（13）依然成立，后面的所有结果同样也成立，最终结论的形式跟式（20）一致。

所以，现在剩下的问题就是能否证明

很遗憾，可以举出反例表明它一般是不成立的。原论文的证明过程也出现了类似的一项，不过求期望的分布不是，而是的最优传输分布，在此前提之下，原论文直接抛出两篇文献的结论作为引理，寥寥几行便完成了证明。不得不说原论文作者们真的很熟悉最优传输相关内容，各种文献结论“信手拈来”，就是苦了笔者这样的新手读者，想要彻底理解却难以下手，只能到此为止了。