论文笔记（六十三）Understanding Diffusion Models: A Unified Perspective（五）

文章概括

引用：

@article{luo2022understanding,
  title={Understanding diffusion models: A unified perspective},
  author={Luo, Calvin},
  journal={arXiv preprint arXiv:2208.11970},
  year={2022}
}

Luo, C., 2022. Understanding diffusion models: A unified perspective. arXiv preprint arXiv:2208.11970.

原文： https://arxiv.org/abs/2208.11970
代码、数据和视频：https://arxiv.org/abs/2208.11970

系列文章：
请在 $《$文章$》$ 专栏中查找

基于得分的生成模型（Score-based Generative Models）

我们已经证明，可以通过优化神经网络$s_{\theta }(x_t,t)$来预测得分函数$\nabla \log p(x_t)$，从而学习一个变分扩散模型（VDM）。然而，在我们的推导中，得分项是通过应用Tweedie公式得出的；这并没有为我们提供关于得分函数的具体含义或为什么值得建模的深入直观理解。幸运的是，我们可以参考另一类生成模型，即基于得分的生成模型（Score-based Generative Models）[9, 10, 11]，以获得这些直觉。事实证明，我们可以证明之前推导的VDM公式与基于得分的生成模型公式是等价的，这使我们能够在这两种解释之间灵活切换。

为了开始理解为什么优化得分函数是有意义的，我们可以绕道回顾一下基于能量的模型（Energy-based Models）[12, 13]。任意灵活的概率分布可以写成以下形式：

$$p_{\theta }(x)=\frac{1}{Z_{\theta }}{e}^{-f_{\theta }(x)}\text{\hspace{1em}(152)}$$

其中，$f_{\theta }(x)$是一个可以任意灵活调整的参数化函数，称为能量函数，通常由神经网络建模，而$Z_{\theta }$是一个归一化常数，用于确保$\int p_{\theta }(x)dx=1$。一种学习这种分布的方法是最大似然估计（最大似然估计是什么？）；然而，这需要对归一化常数$Z_{\theta }=\int e^{-f_{\theta }(x)}dx$进行可计算的求解，但对于复杂的$f_{\theta }(x)$函数来说，这可能无法实现。

一种避免计算或建模归一化常数的方法是使用一个神经网络$s_{\theta }(x)$来学习分布$p(x)$的得分函数$\nabla \log p(x)$。这一方法的动机来自以下观察：对方程（152）两边的对数取导数，可得：

$$\begin{array}{rlr}{\nabla }_x\log p_{\theta }(x)& ={\nabla }_x\log \left(\frac{1}{Z_{\theta }}{e}^{-f_{\theta }(x)}\right)& \text{(153)}\\ & ={\nabla }_x\log \frac{1}{Z_{\theta }}+{\nabla }_x\log e^{-f_{\theta }(x)}& \text{(154)}\\ & =-{\nabla }_x{f}_{\theta }(x)& \text{(155)}\\ & \approx s_{\theta }(x)& \text{(156)}\end{array}$$

这可以自由地表示为一个神经网络，而无需涉及任何归一化常数。通过最小化与真实得分函数的Fisher散度，可以优化得分模型：

$${\mathbb{E}}_{p(x)}\left[\Vert s_{\theta }(x)-\nabla \log p(x){\Vert }_2^2\right]\text{\hspace{1em}(157)}$$

得分函数表示什么？对于每一个$x$，对其对数似然关于$x$的梯度，实质上描述了在数据空间中需要朝哪个方向移动以进一步增加其似然值。

直观地说，得分函数定义了数据$x$所在整个空间的一个向量场，指向数据分布的模式（modes）。在图6的右图中对此进行了直观展示。通过学习真实数据分布的得分函数，我们可以通过从同一空间中的任意点开始，反复跟随得分函数指引的方向，直到达到一个模式来生成样本。此采样过程被称为朗之万动力学（Langevin dynamics），其数学描述如下：

$$x_{i+1}\leftarrow x_i+c\nabla \log p(x_i)+\sqrt{2c}ϵ,i=0,1,\dots ,K\text{\hspace{1em}(158)}$$

其中，$x_0$是从一个先验分布（例如均匀分布）中随机采样得到的，而$ϵ\sim \mathcal{N}(ϵ;0,I)$是一个额外的噪声项，用于确保生成的样本不会总是收敛到某个模式上，而是围绕模式徘徊，以增加多样性。此外，由于学习到的评分函数是确定性的，加入噪声项进行采样为生成过程增加了随机性，这可以避免确定性轨迹的问题，特别是在采样从位于多个模式之间的位置开始时。图6中展示了Langevin动力学采样以及噪声项带来的好处的直观描述。

图6：使用Langevin动力学生成的三个随机采样轨迹的可视化，所有轨迹均从同一个初始化点开始，针对一个高斯混合分布进行采样。左图展示了这些采样轨迹在三维等高线图上的分布，而右图则将采样轨迹与真实评分函数进行对比。从相同的初始化点出发，由于Langevin动力学采样过程中引入的随机噪声项，我们能够从不同的模式生成样本；如果没有噪声项，从固定点采样的每次试验都会确定性地沿着评分收敛到相同的模式。

1. 什么是得分函数？

1.1 定义 得分函数定义为概率密度分布 $p(x)$ 的对数梯度：$${\nabla }_x\log p(x).$$

• 它描述了在数据空间中，如何朝着使概率密度 $p(x)$ 最大的方向移动。
• 得分函数实际上是一个向量场，在每个点 $x$ 指向分布 $p(x)$ 增大的方向。

1.2 直观意义

• 如果把 $\log p(x)$ 看作山脉的高度，那么得分函数是山脉的斜率，指向 “爬坡” 的方向。
• 在图 6 的右图中，得分函数形成了一个向量场，箭头表示每个点的梯度方向。

---

2. 归一化常数和得分函数

2.1 概率密度的形式 一个概率分布通常可以写成： $$p_{\theta }(x)=\frac{1}{Z_{\theta }}{e}^{-f_{\theta }(x)}.$$ 这里：

• $Z_{\theta }=\int e^{-f_{\theta }(x)}dx$ 是归一化常数，确保 $p_{\theta }(x)$ 满足积分为 1 的条件。
• $f_{\theta }(x)$ 是未归一化的负对数密度。

2.2 归一化常数的难点

• 在高维数据中，直接计算 $Z_{\theta }$ 是非常困难的，因为积分 $\int e^{-f_{\theta }(x)}dx$ 通常无法解析。
• 为了避免直接计算 $Z_{\theta }$，我们可以转而学习得分函数 ${\nabla }_x\log p_{\theta }(x)$，因为它不依赖于 $Z_{\theta }$。

2.3 得分函数公式的推导 从 $p_{\theta }(x)=\frac{1}{Z_{\theta }}{e}^{-f_{\theta }(x)}$ 出发，对其取对数： $$\log p_{\theta }(x)=-f_{\theta }(x)-\log Z_{\theta }.\text{\hspace{1em}(153)}$$

对 $x$ 求梯度： $${\nabla }_x\log p_{\theta }(x)={\nabla }_x\left(-f_{\theta }(x)-\log Z_{\theta }\right).$$

注意到：

• $\log Z_{\theta }$ 是一个常数，与 $x$ 无关，因此其梯度为 0。
• 只剩下 $-{\nabla }_x{f}_{\theta }(x)$。

于是： $${\nabla }_x\log p_{\theta }(x)=-{\nabla }_x{f}_{\theta }(x).\text{\hspace{1em}(155)}$$

通过神经网络 $s_{\theta }(x)$ 来近似 $-{\nabla }_x{f}_{\theta }(x)$，因此： $${\nabla }_x\log p_{\theta }(x)\approx s_{\theta }(x).\text{\hspace{1em}(156)}$$

---

3. 优化目标

3.1 如何优化得分函数？ 学习得分函数的目标是让模型 $s_{\theta }(x)$ 接近真实得分函数 ${\nabla }_x\log p(x)$。我们通过以下损失函数来实现： $${\mathbb{E}}_{p(x)}\left[\Vert s_{\theta }(x)-{\nabla }_x\log p(x){\Vert }_2^2\right].\text{\hspace{1em}(157)}$$

3.2 直观理解

• 我们希望神经网络 $s_{\theta }(x)$ 的输出与真实分布 $p(x)$ 的得分函数 ${\nabla }_x\log p(x)$ 尽可能一致。
• 这一损失函数被称为 Fisher 散度，是得分匹配方法的核心。

---

4. Langevin 动力学采样

4.1 什么是 Langevin 动力学？ Langevin 动力学是一种通过得分函数进行采样的算法。它的核心思想是：从一个随机点出发，沿着得分函数的方向逐步逼近数据分布的模式（modes）。其迭代公式为：
$$x_{i+1}\leftarrow x_i+c{\nabla }_x\log p(x_i)+\sqrt{2c}ϵ,ϵ\sim \mathcal{N}(0,\mathbf{I}).\text{\hspace{1em}(158)}$$

4.2 每一项的意义

• $x_{i+1}$：下一步的采样点。
• $x_i$：当前采样点。
• $c{\nabla }_x\log p(x_i)$：沿得分函数方向的梯度更新，使得样本更靠近数据分布的高密度区域。
• $\sqrt{2c}ϵ$：加入随机噪声，增加样本多样性，避免所有样本收敛到同一个模式。

4.3 Langevin 动力学的效果 在图 6 中：

• 左图展示了多个采样轨迹如何逐步逼近高斯混合分布的高密度区域。
• 右图展示了得分函数形成的向量场，表明 Langevin 动力学的梯度更新方向。

---

5. Langevin 动力学的意义

5.1 采样过程 通过 Langevin 动力学，我们可以从随机初始点（例如均匀分布或标准高斯分布）生成样本，这些样本最终将逼近数据分布 $p(x)$。

5.2 噪声的作用 噪声项 $\sqrt{2c}ϵ$ 的引入是为了：

1. 避免采样陷入某个模式，保证生成样本的多样性。
2. 在分布中探索不同的高密度区域。

---

6. 示例：从二维高斯混合分布采样

1. 假设数据分布是一个二维高斯混合分布，有两个主要模式。
2. 初始点 $x_0$ 从均匀分布中随机采样。
3. 使用 Langevin 动力学公式迭代：
   • 计算得分函数 ${\nabla }_x\log p(x_i)$，指示向高密度区域移动的方向。
   • 加入噪声项 $\sqrt{2c}ϵ$，避免收敛到固定点。
4. 最终，生成的样本分布将匹配数据的高斯混合分布。

---

总结

1. 得分函数：描述如何在数据空间中移动以增加概率密度。
2. 神经网络学习得分函数：避免直接计算归一化常数。
3. Langevin 动力学：通过得分函数采样，从随机点生成符合数据分布的样本。
4. 噪声的作用：增加多样性，避免收敛到单一模式。

请注意，公式(157)中的目标依赖于真实的评分函数，而对于建模自然图像等复杂分布，我们无法获得该函数。幸运的是，已经衍生出一些被称为评分匹配的替代技术（如文献[14, 15, 16, 17]），它们无需知道真实评分即可最小化Fisher散度，并能够通过随机梯度下降进行优化。

总体而言，将分布表示为评分函数并通过马尔可夫链蒙特卡洛技术（例如Langevin动力学）使用它来生成样本的过程被称为基于评分的生成建模（Score-based Generative Modeling）[9, 10, 11]。

1. 公式 (157) 的目标和问题
2. 评分匹配（Score Matching）的定义和动机
3. 评分匹配方法如何解决 Fisher 散度问题
4. 基于评分的生成建模（Score-based Generative Modeling）的关键概念

---

1. 公式 (157)：目标与问题

公式 (157) 是 Fisher 散度的定义，用于最小化模型得分函数与真实得分函数的差异： $${\mathbb{E}}_{p(x)}\left[\Vert s_{\theta }(x)-\nabla \log p(x){\Vert }_2^2\right].$$

• 目标：让模型得分函数 $s_{\theta }(x)$ 接近真实得分函数 $\nabla \log p(x)$。
• 真实得分函数的难点：在建模复杂分布（如自然图像）时，无法直接获得真实分布 $p(x)$ 或其得分函数 $\nabla \log p(x)$。
  • 原因：真实数据分布可能未知，或者处于复杂的低维流形上。

---

2. 评分匹配（Score Matching）逐步加噪评分匹配

2.1 定义
评分匹配（Score Matching）是一种无需直接访问 $\nabla \log p(x)$ 的方法，用于最小化 Fisher 散度。通过数学技巧，Fisher 散度可以被重新表达为一个可计算的目标函数，从而绕过对真实得分函数的依赖。

2.2 推导 Fisher 散度的可计算形式

目标函数： $${\mathbb{E}}_{p(x)}\left[\Vert s_{\theta }(x)-\nabla \log p(x){\Vert }_2^2\right].$$ 展开平方： $${\mathbb{E}}_{p(x)}\left[\Vert s_{\theta }(x){\Vert }_2^2\right]-2{\mathbb{E}}_{p(x)}\left[s_{\theta }(x{)}^T\nabla \log p(x)\right]+{\mathbb{E}}_{p(x)}\left[\Vert \nabla \log p(x){\Vert }_2^2\right].$$

关键点：

• 第三项 ${\mathbb{E}}_{p(x)}\left[\Vert \nabla \log p(x){\Vert }_2^2\right]$ 是一个常数，与 $\theta$ 无关。
• 第二项 ${\mathbb{E}}_{p(x)}\left[s_{\theta }(x{)}^T\nabla \log p(x)\right]$ 是核心，需要通过积分技巧消除对 $\nabla \log p(x)$ 的显式依赖。

2.3 无需真实评分的技术
在文献 [14, 15, 16, 17] 中，研究者提出了以下两种方法来优化 Fisher 散度：

1. 直接评分匹配 将目标函数通过积分分部法（Integration by Parts）重写，消除对 $\nabla \log p(x)$ 的依赖。

   • 假设数据分布 $p(x)$ 的支持域是全空间，且 $p(x)$ 在边界上迅速衰减为零。
   • 将期望项重写为：
     $${\mathbb{E}}_{p(x)}\left[s_{\theta }(x{)}^T\nabla \log p(x)\right]=-{\mathbb{E}}_{p(x)}\left[\text{div}{s}_{\theta }(x)\right],$$
     其中 $\text{div}{s}_{\theta }(x)$ 表示得分函数的散度。
   • 替换后：
     $${\mathbb{E}}_{p(x)}\left[\Vert s_{\theta }(x)-\nabla \log p(x){\Vert }_2^2\right]\propto {\mathbb{E}}_{p(x)}\left[\Vert s_{\theta }(x){\Vert }_2^2+2\text{div}{s}_{\theta }(x)\right].$$

2. 逐步加噪评分匹配 为数据添加逐步递增的高斯噪声，并训练得分函数 $s_{\theta }(x,t)$。

   • 噪声扰动后的分布：
     $$p_{{\sigma }_t}(x_t)=\int p(x)\mathcal{N}(x_t;x,{\sigma }_t^{2}I)dx.$$
   • 学习目标：
     $$\arg \underset{\theta }{\min }\sum _{t=1}^T\lambda (t){\mathbb{E}}_{p_{{\sigma }_t}(x_t)}\left[\Vert s_{\theta }(x,t)-\nabla \log p_{{\sigma }_t}(x_t){\Vert }_2^2\right].$$
   • 优点：通过逐步加噪扩大分布覆盖范围，使得低密度区域的训练信号变得更加丰富。

---

3. 基于评分的生成建模

3.1 定义 将数据分布 $p(x)$ 表示为得分函数 $\nabla \log p(x)$，并通过马尔可夫链蒙特卡洛（MCMC）技术（如 Langevin 动力学）采样，生成符合数据分布的样本。

3.2 Langevin 动力学 采样公式： $$x_{i+1}=x_i+c\nabla \log p(x_i)+\sqrt{2c}ϵ,$$

• $c$：步长。
• $ϵ\sim \mathcal{N}(0,I)$：随机噪声。

3.3 逐步退火 Langevin 动力学 为了解决原始评分匹配的低效问题，引入逐步加噪的退火 Langevin 动力学：

1. 初始化：从一个简单的先验分布（如高斯分布）采样初始点 $x_0$。
2. 逐步采样：依次降低噪声水平 ${\sigma }_t$，并在每一步运行 Langevin 动力学。
3. 最终样本收敛到真实分布 $p(x)$。

---

4. 与变分扩散模型的联系

4.1 训练目标的相似性 逐步加噪评分匹配的目标（公式 161）与变分扩散模型（VDM）的训练目标几乎一致：

• 都是同时学习多个噪声水平下的得分函数。
• 通过加噪扩展分布的支持域。

4.2 无限时间步的推广 当时间步数 $T\to \infty$ 时：

• 逐步加噪过程可以被视为一个连续的随机过程。
• 对应于一个随机微分方程（SDE）：
  • 正向 SDE：描述从数据分布逐步加噪的过程。
  • 逆向 SDE：描述从高斯噪声分布逐步采样生成数据的过程。

---

5. 示例：二维高斯分布的评分匹配

假设数据分布是一个二维高斯分布： $$p(x)=\mathcal{N}(x;\mu ,\Sigma ).$$

5.1 真实得分函数 得分函数为： $$\nabla \log p(x)=-{\Sigma }^{-1}(x-\mu ).$$

5.2 模型得分函数 假设模型的得分函数为 $s_{\theta }(x)=-x+{ϵ}_{\theta }(x)$，其中 ${ϵ}_{\theta }(x)$ 是神经网络的误差项。

5.3 训练目标 Fisher 散度： $${\mathbb{E}}_{p(x)}\left[\Vert s_{\theta }(x)-\nabla \log p(x){\Vert }_2^2\right]={\mathbb{E}}_{p(x)}\left[\Vert {ϵ}_{\theta }(x){\Vert }_2^2\right].$$

通过最小化 ${ϵ}_{\theta }(x)$ 的期望平方，可以优化得分函数。

---

总结

1. 公式 (157) 的问题：依赖于真实得分函数 $\nabla \log p(x)$，但真实分布未知。
2. 评分匹配方法：通过积分技巧和逐步加噪，绕过了对真实得分函数的显式依赖。
3. 基于评分的生成建模：利用学习到的得分函数进行采样，生成符合数据分布的样本。
4. 与扩散模型的联系：两者在训练目标和采样过程中密切相关。

然而，正如Song和Ermon在文献[9]中详细说明的，原始评分匹配存在三个主要问题。首先，当$x$位于高维空间中的低维流形上时，评分函数定义不明确。从数学上可以看出，这种情况下，低维流形以外的所有点的概率为零，其对数是未定义的。这在试图学习自然图像生成模型时尤为不便，因为自然图像已知位于整个环境空间的低维流形上。

其次，通过原始评分匹配训练的评分函数在低密度区域不够准确。从我们在公式(157)中最小化的目标中可以看出这一点。由于它是$p(x)$上的期望，并且明确地在其样本上训练，模型对于很少见或未见过的样本无法获得准确的学习信号。这是一个问题，因为我们的采样策略涉及从高维空间中的随机位置（最有可能是随机噪声）开始，并根据学习到的评分函数进行移动。由于我们依赖的是噪声或不准确的评分估计，最终生成的样本可能也会表现不佳，或者需要更多迭代才能收敛到准确的输出。

最后，即使Langevin动力学采样使用的是真实评分，其采样过程也可能无法混合。假设真实数据分布是两个不相交分布的混合体：

$$p(x)=c_1{p}_1(x)+c_2{p}_2(x)\text{\hspace{1em}(159)}$$

然后，当计算分数时，这些混合系数会丢失，因为对数操作将系数与分布分离，而梯度操作会将其清零。为了形象化这一点，请注意右图6中显示的真实分数函数无法区分三个分布之间的不同权重；从图示的初始化点进行Langevin动力学采样时，抵达每个模式的概率大致相等，尽管在实际的高斯混合分布中，右下角的模式权重更高。

事实证明，这三个缺点可以通过向数据添加多个层级的高斯噪声同时解决。首先，由于高斯噪声分布的支持范围是整个空间，因此被扰动的数据样本不再局限于低维流形。其次，添加大幅度的高斯噪声会增加数据分布中每个模式覆盖的区域，从而在低密度区域提供更多的训练信号。最后，添加具有递增方差的多层级高斯噪声将导致尊重真实混合系数的中间分布。

正式地，我们可以选择一个正的噪声水平序列 { σ t } t = 1 T \{\sigma_t\}_{t=1}^T {σt​}t=1T​，并定义一系列逐步扰动的数据分布：

$$p_{{\sigma }_t}(x_t)=\int p(x)\mathcal{N}(x_t;x,{\sigma }_t^{2}I)dx\text{\hspace{1em}(160)}$$

然后，通过使用得分匹配（score matching）训练一个神经网络$s_{\theta }(x,t)$，以同时学习所有噪声水平的得分函数：

$$\arg \underset{\theta }{\min }\sum _{t=1}^T\lambda (t){\mathbb{E}}_{p_{{\sigma }_t}(x_t)}\left[{\Vert s_{\theta }(x,t)-\nabla \log p_{{\sigma }_t}(x_t)\Vert }_2^2\right]\text{\hspace{1em}(161)}$$

其中，$\lambda (t)$是一个依赖于噪声水平$t$的正权重函数。需要注意的是，该目标函数与在公式148中为训练变分扩散模型推导出的目标几乎完全一致。此外，作者提出了退火Langevin动力学采样作为一种生成程序，其中通过依次对$t=T,T-1,\dots ,2,1$运行Langevin动力学来生成样本。初始化选择某个固定的先验分布（如均匀分布），而每个后续的采样步骤都从前一次模拟的最终样本开始。由于噪声水平随着时间步$t$逐渐减小，同时步长也随时间减少，样本最终会收敛到一个真实的模式。这与变分扩散模型的马尔科夫HVAE解释中所执行的采样过程直接类似，在该过程中，一个随机初始化的数据向量通过降低噪声水平被逐步优化。

因此，我们已经在训练目标和采样程序中明确建立了变分扩散模型与基于分数的生成模型之间的联系。

一个问题是如何将扩散模型自然地推广到无限数量的时间步。在马尔科夫HVAE视角下，这可以被解释为将层级数量扩展到无穷，即$T\to \infty$。从等价的基于分数的生成模型的视角来看，这一点更加清晰；在无限数量的噪声尺度下，图像在连续时间上的扰动可以表示为一个随机过程，因此可以用随机微分方程（SDE）(什么是随机微分方程（SDE）)来描述。采样则通过逆向SDE进行，这自然需要估计每个连续值噪声水平下的分数函数[10]。SDE的不同参数化本质上描述了不同的时间扰动方案，从而能够灵活地对加噪过程进行建模[6]。