Lecture 02: Score Matching
在Lecture01里，我们了解了什么是DDPM(Ho, Jain, and Abbeel 2020), 以及如何通过DDIM(J. Song, Meng, and Ermon 2022)来加速Inference的过程，在这节课中，我们将了解什么是Score Matching(Y. Song and Ermon 2020).
1 Problem Set Up
我们知道，对于一个数据分布 \(p_{data}\)，我们想要的是，对于一个样本 \(x\)，我们能够计算出它的概率密度 \(p_{data}(x)\)，并且希望\(x\) 在\(p_{data}\)中具有较高的概率密度。这很容易让我们联想到利用Gradient Ascent来最大化\(p_{data}(x)\)，通过将 \(x\) 沿着 \(\nabla_x \log p_{data}(x)\) 的方向进行更新，我们可以使得 \(p_{data}(x)\) 变得更大。
TIP: Gradient Descent
这个思想和我们优化模型参数 \(\theta\) 的过程非常相似，我们通过Gradient Descent来最小化Loss Function \(L(\theta)\)，通过将 \(\theta\) 沿着 \(-\nabla_\theta L(\theta)\) 的方向进行更新，我们可以使得 \(L(\theta)\) 变得更小。
我们先来看一下 \(p_{data}(x)\) 的是如何定义的，假如我们有一个函数 \(f(x)\)，它可以将任意的输入 \(x\) 映射到一个实数值，那么我们可以通过以下的方式来定义 \(p_{data}(x)\)：
\[ p_{data}(x) = \frac{e^{f(x)}}{\int e^{f(x')} dx'} \tag{1}\]
其中，分母 \(\int e^{f(x')} dx'\) 被称为Partition Function，它是一个常数，可以看作是一个归一化的常数，使得 \(p_{data}(x)\) 成为一个合法的概率密度函数。它的定义为：
\[ Z = \int e^{f(x')} dx' \tag{2}\]
但现在有一个显然的问题是：\(Z\) 是一个非常复杂的积分，是intractable的，我们无法直接计算出 \(Z\) 的值，因此我们无法直接计算出 \(p_{data}(x)\) 的值。除了这个问题之外，另一个问题就是：Numerical Instability，当 \(f(x)\) 的值非常大或者非常小时，\(e^{f(x)}\) 的值可能会变得非常大或者非常小，这会导致数值不稳定的问题。
解决Numerical Instability的常见方法就是使用Logarithm来进行计算，我们可以将 \(p_{data}(x)\) 的定义改写为：
\[ \log p_{data}(x) = f(x) - \log Z \tag{3}\]
那么它的Gradient就可以表示为： \[ \nabla_x \log p_{data}(x) = \nabla_x f(x) - \nabla_x \log Z \tag{4}\]
其中，\(\nabla_x \log Z\) 是一个常数项，因为 \(Z\) 是一个常数，所以它的Gradient为0，因此我们可以将上面的式子简化为：
\[ \nabla_x \log p_{data}(x) = \nabla_x f(x) \tag{5}\]
通过logarithm的变换，我们有以下的几个好处：
- Tractability: 通过使用logarithm，我们可以避免直接计算Partition Function \(Z\)，因为在计算Gradient的时候，\(Z\) 的Gradient为0，因此我们不需要计算 \(Z\) 的值。
- Points in same direction: \(\nabla_x \log p_{data}(x)\) 和 \(\nabla_x f(x)\) 是在同一个方向上的，因此我们可以通过计算 \(\nabla_x f(x)\) 来近似计算 \(\nabla_x \log p_{data}(x)\)。
- Numerical Stability: 通过使用logarithm，我们可以避免数值不稳定的问题
因此，我们给 \(\nabla_x \log p_{data}(x)\) 起了一个名字，叫做Score Function，记为 \(s(x)\)，它的定义为：
\[ s(x) = \nabla_x \log p_{data}(x) \tag{6}\]
不过显然，这也有一个问题，因为\(s(x)\) 指向的是highest的概率密度方向，因此在某个特定的region内，我们最终都会走到相同的点，这不是我们想要的，我们希望这个Model 有 diversity。因此，我们需要用特定的Sampling的方法来从 \(p_{data}\) 中采样，这个Sampling的方法就是 Langevin Dynamics:
\[ x_{t} = x_{t-1} + \frac{\alpha}{2} s(x_{t-1}) + \textcolor{red}{\sqrt{\alpha} \epsilon_t} \tag{7}\]
其中\(\textcolor{red}{\sqrt{\alpha} \epsilon_t}\) 是一个Gaussian Noise项，它的作用是为了增加采样的多样性，避免所有的采样点都集中在同一个位置。
现在，我们已经知道了什么是Score Function，以及如何通过Langevin Dynamics来从 \(p_{data}\) 中采样。但是最重要的问题是：Score Function \(s(x)\) 是未知的，我们需要通过某种方法来估计它，这就是Score Matching的核心问题。
2 Score Estimation
根据Score FunctionEquation 6的定义，我们可以知道，Score Function \(s(x)\) 是 \(\nabla_x \log p_{data}(x)\)，因此我们可以通过Minimize the distance between \(s(x)\) and \(\nabla_x \log p_{data}(x)\) 来估计 \(s(x)\)，我们可以使用以下的Loss Function来进行优化：
\[ \mathcal{L}_{\text{SM}}(\theta) = \mathbb{E}_{x \sim p_{data}} \left[ \| s_\theta(x) - \nabla_x \log p_{data}(x) \|^2 \right] \tag{8}\]
其中，\(s_\theta(x)\) 是我们用来估计 \(s(x)\) 的模型，它是一个参数化的函数，我们通过优化参数 \(\theta\) 来使得 \(s_\theta(x)\) 尽可能地接近 \(\nabla_x \log p_{data}(x)\)。
但是上面的Loss Function有一个问题，就是我们无法直接计算出 \(\nabla_x \log p_{data}(x)\)，因此我们需要通过一些技巧来改写这个Loss Function，使得我们不需要直接计算 \(\nabla_x \log p_{data}(x)\)。有以下几个方法可以实现这个目标：
- Implicit Score Matching: Integrate by parts to eliminate the need for \(\nabla_x \log p_{data}(x)\).
\[ \mathcal{L}_{\text{ISM}}(\theta) = \mathbb{E}_{x \sim p_{data}} \left[ \| s_\theta(x) \|^2 + 2 \nabla_x \cdot s_\theta(x) \right] \tag{9}\]
- Sliced Score Matching: Project the score function onto random directions to avoid computing the full gradient.
\[ \mathcal{L}_{\text{SSM}}(\theta) = \mathbb{E}_{x \sim p_{data}, v \sim \mathcal{N}(0, I)} \left[ 2 v^T s_\theta(x) + \| v^T s_\theta(x) \|^2 \right] \tag{10}\]
在这里就不展开这两种方法的具体细节了，感兴趣的同学可以参考相关的论文来了解更多。在这里我们主要关注一下Denoising Score Matching(DenoisingScoreMatching2011vincent?)，它是Score Matching的一种变体方法，它的核心思想是通过添加噪声来使得Score Matching的训练过程更加稳定。
3 Denoising Score Matching
Denoising Score Matching的核心思想是通过添加噪声来使得Score Matching的训练过程更加稳定。我们先来看一下简单的例子： Score of Gaussian Distribution
3.1 Score of Gaussian Distribution
对于一个Gaussian Distribution \(p(x) = \mathcal{N}(x; \mu, \sigma^2 I)\)，我们可以通过以下的方式来计算它的Score Function：
\[ s(x) = \nabla_x \log p(x) = -\frac{1}{\sigma^2} (x - \mu) \tag{11}\]
现在，假如我们在数据上添加一个Gaussian Noise \(\epsilon \sim \mathcal{N}(0, \sigma^2 I)\)，那么我们得到的Noisy Data \(\tilde{x}\) 可以表示为：
\[ \tilde{x} = x + \sigma \epsilon \tag{12}\]
那么这个Noisy Data \(\tilde{x}\) 的Probability Density Function \(p(\tilde{x})\) 可以表示为：
\[ p_{\sigma}(\tilde{x}) = \mathcal{N}(\tilde{x}, \sigma^2 I) \tag{13}\]
它的Score Function \(s_{\sigma}(\tilde{x})\) 可以表示为：
\[ \begin{split} s_{\sigma}(\tilde{x}) &= \nabla_{\tilde{x}} \log p_{\sigma}(\tilde{x}) \\ &= -\frac{1}{\sigma^2} (\tilde{x} - \mu) \\ &= -\frac{1}{\sigma^2} (x + \sigma \epsilon - \mu) \\ &= -\frac{1}{\sigma^2} (x - \mu) - \frac{1}{\sigma} \epsilon \\ &= s(x) - \frac{1}{\sigma} \epsilon \end{split} \tag{14}\]
对于这个Noisy Data \(\tilde{x}\)，我们可以通过以下的方式来计算它的Score Function \(s_{\sigma}(\tilde{x})\)，我们可以看到，\(s_{\sigma}(\tilde{x})\) 可以看作是 \(s(x)\) 的一个Noisy Version，因此我们可以通过优化以下的Loss Function来估计 \(s(x)\)：
\[ \mathcal{L}_{\text{DSM}}(\theta) = \mathbb{E}_{x \sim p_{data}, \epsilon \sim \mathcal{N}(0, I)} \left[ \| s_\theta(x + \sigma \epsilon) + \frac{1}{\sigma} \epsilon \|^2 \right] \tag{15}\]
现在有个问题就是，我们要如何选择噪声的标准差 \(\sigma\) 呢？
- 如果\(\sigma \ll 1\)，那么我们添加的噪声非常小，这样我们得到的Noisy Data \(\tilde{x}\) 就非常接近于原始数据 \(x\)，因此我们得到的Score Function \(s_{\sigma}(\tilde{x})\) 就非常接近于 \(s(x)\)，但是这些在Low-Density Regions的样本点的Score Function \(s(x)\) 的值可能会非常大，这会导致训练过程中的数值不稳定的问题。
- 如果\(\sigma \gg 1\)，那么我们添加的噪声非常大，这样我们得到的Noisy Data \(\tilde{x}\) 就非常远离原始数据 \(x\)，因此我们得到的Score Function \(s_{\sigma}(\tilde{x})\) 就非常远离 \(s(x)\)，这会导致训练过程中的Bias问题。
将这两个问题结合起来，我们可以得出一个结论：我们需要选择一个合适的 \(\sigma\) 来平衡Bias和Variance之间的关系，这样我们才能够得到一个稳定的训练过程，这个就是Annealed Langevin Dynamics的核心思想，通过逐渐减小 \(\sigma\) 的值来平衡Bias和Variance之间的关系，从而得到一个稳定的训练过程。
我们可以通过以下的方式来实现Annealed Langevin Dynamics的训练： 1. 选择一个初始的 \(\sigma_0\)，并且设置一个衰减率 \(\alpha\)。 2. 在每个训练步骤中，按照以下的方式来更新 \(\sigma\) 的值： 3. \[ \sigma_t = \sigma_0 \cdot \alpha^t \] 4. 在每个训练步骤中，按照以下的方式来更新模型参数 \(\theta\) 的值： 5. \[ \theta_{t} = \theta_{t-1} - \eta \nabla_\theta \mathcal{L}_{\text{DSM}}(\theta_{t-1}) \]
用Python代码来实现Annealed Langevin Dynamics的训练过程：
# Assume we have a dataset of samples from p_data
# and a model s_theta that estimates the score function
# step 1: Samping Sigma
u = torch.rand(x.size(0), device=device)
sigma = sigma_min * (sigma_max / sigma_min) ** u # log-uniform
sigma_view = sigma.view(-1, 1, 1, 1)
# 2. generate random noise image
noise = torch.randn_like(x)
x_noisy = x + sigma_view * noise
# 3. compute the loss
target_score = - (x_noisy - x) / (sigma_view ** 2)
pred_score = model(x_noisy)
loss = F.mse_loss(pred_score, target_score)
# 4. update model parameters
optimizer.zero_grad()
loss.backward()
optimizer.step()我们看一下如何通过Annealed Langevin Dynamics来进行Sampling的过程：
4 Differential Formulation of Score Matching
目前为止，Score Matching和DDPM都在一个离散的时间步长上进行训练和采样的，但是我们也可以将它们扩展到一个连续的时间步长上，这就是Score-Based Generative Modeling(Y. Song et al. 2021)，它的核心思想是通过一个Stochastic Differential Equation (SDE) 来描述数据的生成过程。不过在此之前，我们先来介绍一下Wiener Process
4.1 Wiener Process
Wiener Process，也被称为Brownian Motion，是一个连续时间的随机过程，它的定义如下：
\[ W(t) = \int_0^t \epsilon(s) ds \tag{16}\]
其中，\(\epsilon(s)\) 是一个标准的Gaussian Noise，它的均值为0，方差为1。Wiener Process具有以下的性质：
- \(W(0) = 0\)，即在时间0的时候，Wiener Process的值为0。
- \(W(t)\) 是一个连续的函数，即在任意的时间点上，Wiener Process的值都是连续的。
- \(W(t)\) 是一个独立增量的过程，即对于任意的时间点 \(t_1 < t_2 < ... < t_n\)，\(W(t_2) - W(t_1)\) 和 \(W(t_3) - W(t_2)\) 是独立的随机变量。
- \(W(t)\) 的增量是Gaussian分布的，即对于任意的时间点 \(t_1 < t_2\)，\(W(t_2) - W(t_1)\) 服从一个Gaussian分布，均值为0，方差为 \(t_2 - t_1\)。
接下来，我们看一下如何通过Wiener Process将DDPM变成一个连续时间的模型。
4.2 Continuous-Time Formulation of DDPM
在DDPM中，我们有一个离散的时间步长 \(t = 0, 1, ..., T\)，在每个时间步长上，我们都有一个对应的数据分布 \(p_t(x)\)，我们通过一个Markov Chain来描述数据的生成过程：
\[ x_t = \sqrt{1- \beta_t} x_{t-1} + \sqrt{\beta_t} \epsilon_t \tag{17}\]
其中，\(\beta_t\) 是一个时间步长相关的噪声强度，\(\epsilon_t\) 是一个标准的Gaussian Noise。我们可以将上面的式子改写为： \[ x_t - x_{t-1} = \left[\sqrt{1 - \beta_t} - 1\right] x_{t-1} + \sqrt{\beta_t} \epsilon_t \tag{18}\]
接下来我们定义一下\(\beta(t)\)，它是一个连续的函数，表示在时间 \(t\) 上的噪声强度，\(\beta_t=\beta(t)\)，我们可以将上面的式子改写为：
\[ x(t + \Delta t) - x(t) = \left[\sqrt{1 - \beta(t)} - 1\right] x(t) + \sqrt{\beta(t)} \epsilon(t) \tag{19}\]