Score-Based Diffusion 学习笔记

$$\gdef\vx{\mathbf{x}} \gdef\vy{\mathbf{y}} \gdef\vw{\mathbf{w}} \gdef\vL{\mathbf{L}} \gdef\vP{\mathbf{P}} \gdef\vm{\mathbf{m}} \gdef\vbeta{\boldsymbol{\beta}} \gdef\vf{\mathbf{f}} \gdef\md{\mathrm{d}} \gdef\norm2#1{\left\Vert { #1 } \right\Vert_2} \gdef\grad{\nabla} \gdef\tr{\mathrm{tr}}$$

Score Function

Score function：对于$\mathbb{R}^n$上的概率密度函数（probability density function）为$p: \mathbb{R}^n \to \mathbb{R}$，则$p$之 score function 定义为其对数函数的梯度：

$$s(\vx) = \frac{\partial \log p(\vx)}{\partial \vx} = \grad_{\vx} \log p(\vx) = \grad \log p(\vx)$$

（省略的梯度下标一般可以通过上下文直接推断）

这是 Score-Based 生成模型的核心概念。类比一个实函数的导数可以相当程度地刻画原函数的特征，那么 p. d. f. 之 score func. 也可以刻画原本的概率分布。

Langevin Algorithm¹

该算法的名称有很多： Metropolis-adjusted Langevin algorithm (MALA) or Langevin Monte Carlo (LMC)。Langevin 算法可以利用 score function 来生成服从该分布的随机变量。用到的公式是（$\dot{\vx}$表示$\md\vx$）

$$\dot{\vx} = \grad_{\vx}\log p(\vx) + \sqrt{2}\dot{\vw}$$

其中$\vw$表示一个标准布朗运动。使用离散化的算法实现它的方法有很多，比较简单的是 Euler–Maruyama method：

$$\vx(k + 1) = \vx(k) + \tau \grad p(\vx(k)) + \sqrt{2\tau} \xi(k)$$

其中$\xi(k)$服从$n$维高斯分布$\mathcal{N}_n(0, \mathbf{I})$²。

Generative Modeling

生成模型可以理解为：输入的数据$\vx$服从某个先验分布$\pi$（例如均匀分布或者正态分布），输出的数据$\vy$服从目标分布$p_{\pi}$（例如人脸的分布）。因此生成模型的问题主要分为两个阶段：

1. 训练一个模型$f(\vx)$能够很好地表达分布$p_{\vy | \vx}(\vx)$的特征，例如去拟合$p$的 p. d. f.
2. 根据模型$f$生成服从目标分布的值。

而事实上，对于实际问题来说，目标分布常常由大量的样本点来代表，因此目标分布本身其实是一组样本点的最大似然。基于此，又可以对于生成模型的训练思路进一步分类：

1. likelihood-based models：以最大似然估计为目标来拟合。
2. implicit generative models：使用一些别的最优化目标来拟合（比如 GAN）。

Score-based 则代表另一种流派：拟合 p. d. f. 的梯度。同样的，我们面临两个阶段的问题：如何训练模型，以及如何生成服从目标分布的值。后者在上文中已经给出了一种思路，放在前面的原因是不涉及太多的背景知识。下文主要关注模型的训练。

Trainning Objective

任何模型的学习过程都是以最优化一个目标为导向。假设我们训练的模型是$s(\vx; \theta) = s_{\theta}(\vx)$，表示带有参数$\theta$的一个函数（分号用来区分自变量和参数）。那么一个最优的参数$\theta^{\ast}$即为

$$\arg\min_{\theta} E_{p(\vx)}\left[\norm2{s_{\theta}(\vx) - \grad \log p(\vx)}^2\right]$$

更清楚地可以写成

$$\arg\min_{\theta} \int_{\vx \in \mathbb{R}^n} p(\vx)\norm2{s_{\theta}(\vx) - \grad_{\vx} \log p(\vx)}^2\md\vx$$

如果$p$满足一些条件 ³，那么我们可以把范数的部分写成离散求和的形式：

$$\arg\min_{\theta} \int_{\vx \in \mathbb{R}^n} p(\vx) \sum_{i = 1}^n \left[ [\grad_{\vx} (s_{\theta}(\vx)_i)]_i + \frac{1}{2} s_{\theta}(\vx)_i^2 \right] \md\vx$$

可以发现$\grad \log p(\vx)$被消了，但是出现了一个$s_{\theta}(\vx)$的梯度（也就是 p. d. f. 的二阶梯度），这东西是一个矩阵，非常不好求。

考虑现实情况的$T$个样本点$\vx^{(1)}, \ldots, \vx^{(T)}$，那么上式的积分部分就变成了离散形式

$$\arg\min_{\theta} \frac{1}{T} \sum_{t = 1}^T \sum_{i = 1}^n \left[ [\grad_{\vx} (s_{\theta}(\vx^{(t)})_i)]_i + \frac{1}{2} s_{\theta}(\vx^{(t)})_i^2 \right]$$

简写一下可以变成

$$\arg\min_{\theta} \frac{1}{T} \sum_{t = 1}^T \tr\left( \grad_{\vx} s_{\theta}(\vx^{(t)}) \right) + \frac{1}{2} \norm2{s_{\theta}(\vx^{(t)})}^2$$

这个式子显然还是没法直接求，对此有人开发了一套 score matching 理论来求解。

Denoise Score Matching⁴

这个方法的 Trainning Objective 是

$$\frac{1}{2} E _{p_{\mathrm{data}}(\vx)} E_{p_{\sigma}(\tilde\vx | \vx)} \left[\norm2{ s_{\theta}(\tilde\vx) - \grad_{\tilde\vx}\log p_{\sigma}(\tilde\vx | \vx) }^2\right]$$

其中$p_{\sigma}(\tilde\vx | \vx) \triangleq \mathcal{N}_n(\tilde\vx; \vx, \sigma^2\mathbf{I})$是一个扰动核（perturbation kernel），是$\vx$的一个微小高斯扰动。$p_{\mathrm{data}}$表示数据分布。扰动后的数据分布可以表示为$p_{\sigma}(\tilde\vx) = \int_{\vx} p_{\sigma}(\tilde\vx | \vx)p_{\mathrm{data}}(\vx)\md\vx$。

对于这里的高斯分布，性质比较好，带入$\mathbf{\Sigma} = \sigma^2\mathbf{I}$可以写出它的数学形式⁵

$$p_{\sigma}(\tilde\vx | \vx) = \frac{\exp (-\frac{1}{2\sigma^2} \norm2{ \tilde\vx - \vx }^2 )}{\sqrt{(2\pi)^k\sigma^2}}$$

这样就可以自然推出

$$\grad_{\tilde\vx} \log p_{\sigma} (\tilde\vx | \vx) =\frac{1}{\sigma^2}(\vx - \tilde\vx)$$

因此上面的 Objective 可以进一步写为

$$\frac{1}{2} E _{p_{\mathrm{data}}(\vx)} E_{p_{\sigma}(\tilde\vx | \vx)} \left[\norm2{ s_{\theta}(\tilde\vx) - \frac{1}{\sigma^2}(\vx - \tilde\vx) }^2\right]$$

DSM 的 Objective 被证明，在$\sigma$充分小的时候，与原始的 Objective 在最优化的角度是等价的。

Annealed Langevin Dynamics

顾名思义，这是一个在 Langevin Dynamics 采样方法上加入了退火思想的一个变种。它想要解决的问题是：即使我们训练出一个优秀的$s_{\theta} \approx\grad_{\vx}\log p(\vx)$，也不一定能够在采样时得到目标分布。尤其是空间维度很大时，目标分布有时候会很稀疏，而如果迭代时的初始值位置不够好，就会很难在可接受的时间内收敛到对应的目标分布。

但注意到我们不一定非得一步到位。取$\sigma_1, \ldots, \sigma_L$为一个递减的序列（例如等比数列）。假设初始的随机变量$\vx$服从先验分布$\pi(\vx)$。我们首先使用 Langevin 算法将$\vx$转化为$p_{\sigma_1}(\vx)$的一个采样，再以此为基础去生成$p_{\sigma_2}(\vx)$的一个采样，以此类推，一步一步得到最终的目标分布$p_{\sigma_L}(\vx)$的一个采样。这可以理解为，从前一个分布到达下一个分布是相对简单的，因为初始值距离下一个目标分布足够“近”，有较大可能收敛。

这个采样方法需要我们求出$\grad_{\vx} \log p_{\sigma_i}(\vx)$（$i = 1, \ldots, L$），我们可以直接把$\sigma_i$作为网络的输入，训练出一个$s_{\theta}(\vx, \sigma) \approx \grad_{\vx} \log p_{\sigma}(\vx)$的模型。

Stochastic Differential Equation

从 Annealed Langevin Dynamics 进一步出发。当$L\to \infty$时，实际上$\sigma$就变成了一个连续变化的值。为了刻画$\sigma$的变化过程，我们可以把离散的迭代指标$i$替换为时间$t$，并不妨设$t \in[0, 1]$。那么$\sigma = \lambda(t)$就是一个随时间变化的扰动系数。引入了时间，很容易想到使用物理的方法去描述采样的过程。于是我们引入以下形式的随机微分方程：

$$\md\vx = \vf(\vx, t)\md t + g(t)\md \vw$$

其中$\vw$是一个标准布朗运动（Wiener process，类似标准高斯分布），$\vf(\cdot, t)$称作平移系数（drift coef.），而$g(t)$称作扩散系数（diffusion coef.）。在 SBGM 的情境下，SDE 描述的是从一个数据分布$p_{\mathrm{data}}$出发，逐渐随机运动到先验分布的过程。不妨$p_t(\vx)$表示在$t$时刻$\vx$服从的分布，我们有$p_0 = p_{\mathrm{data}}$，$p_1 = \pi$。

要描述采样的过程，我们需要求解上述 SDE 的时间倒流过程⁶：

$$\md\vx = [\vf(\vx, t) - g^2(t)\grad_{\vx} \log p_{t}(\vx)] \md t + g(t)\md \bar\vw$$

要训练$s_{\theta}(\vx, t) \approx \grad_{\vx} p_t(\vx)$，记$\vx(t)$表示$t$时刻的随机变量$\vx$，$p_{0t}$表示从$p_0$到$p_t$的变换核（即上文的扰动核），$\lambda(t)$是一个正的系数，那么最优的参数可以写成：

$$\arg\min_{\theta}E_t E _{\vx(0)} E_{\vx(t) | \vx(0)} \left[\lambda(t)\norm2{ s_{\theta}(\vx(t), t) - \grad_{\vx(t)}\log p_{0t}(\vx(t) | \vx (0)) }^2\right]$$

当$\vf(\cdot, t) = 0$时，$p_{0t}$始终是一个高斯扰动（和$g$有关），此时 Trainning Objective 可以写成

$$\min E_t E_{\vx(0)} E_{\vx(t) | \vx(0)}\left[ \lambda(t) \norm2{ s_{\theta}(\vx(t), t) - \frac{\vx(0) - \vx(t)}{\sigma_{0t}^2} }^2\right]$$

其中$\sigma_{0t}^2 = 1/ E\left[\norm2{\grad_{\vx(t)} \log p_{0t}(\vx(t) | \vx(0))}^2\right]$表示$p_{0t}(\vx(t) | \vx(0))$的方差。另外论文中$E\left[\norm2{\grad_{\vx(t)} \log p_{0t}(\vx(t) | \vx(0))}^2\right] = \frac{1}{\sigma_{0t}^2}$。

取$\lambda(t) \propto \sigma^2_{0t}$就会得到

$$\min E_t E_{\vx(0)} E_{\vx(t) | \vx(0)} \norm2{ s_{\theta}(\vx(t), t)\sigma_{0t} + \frac{\vx(t) - \vx(0)}{\sigma_{0t}} }^2$$

并且由经验可以得出$\norm2{\sigma_{0t} s_{\theta}(\vx, t)} \propto 1$. 因此在训练神经网络时可以手动 normalize.

此时可以得到原本的 SDE 为

$$\md\vx = \sqrt{\lambda(t)}\md\vw$$

对应的 Reverse SDE：

$$\md \vx = -\lambda(t)\grad_{\vx}\log p_t(\vx)\md t + \sqrt{\lambda(t)}\md\bar\vw$$

More SDE

接下来我们探索一些 SDE 的具体内容，可能会作为上文的引用。主要参考 Applied Stochastic Differential Equation⁷。

Brownian Motion

首先我们给出布朗运动（Brownian motion）的定义：$\vbeta(t) \in \mathbb{R}^n$是一个具有以下性质的随机过程：

1. $\Delta\vbeta_k = \vbeta(t_{k + 1}) - \vbeta(t_k)$服从一个均值为$0$，协方差⁸矩阵为$\mathbf{Q}(t_{k + 1} - t_k)$的高斯分布，其中$\mathbf{Q}$是布朗运动的扩散矩阵（diffusion matrix）。
2. 时间段不相交的$\Delta \vbeta_k$是独立的。
3. $\vbeta(0) = 0$。

关于它有两点需要注意：$\vbeta(t)$处处不可微，不过白噪声可以作为布朗运动的形式上的导数：$\vw(t) = \frac{\md \vbeta(t)}{\md t}$。

The Stochastic Integral of Itô

一个随机过程，在任何时刻运动到的位置和速度等等都没有上下界，不是一个确定性的收敛的过程，因此没法用黎曼积分去刻画。

SDE 所对应的积分过程是 The Stochastic Integral of Itô，定义为$L_2$意义下的极限：

$$\int_{t_0}^t \vL(\vx(t), t) \md\vbeta(t) = \lim_{n\to \infty} \sum_k \vL(\vx(t), t)[\vbeta(t_{k + 1}) - \vbeta(t_k)]$$

其中$t_0 < t_1 < \cdots < t_n = t$，$\vL$是一个矩阵函数。注意到这个积分不像黎曼积分，在分划的小区间里面随便取一个值。如果这样做其实会导致积分不收敛。

这个积分也解释了上文$\vbeta(t)$的形式导数的具体含义。因此随机微分方程

$$\md\vx = \vf(\vx, t)\md t + \vL(\vx, t)\md \vbeta$$

的具体定义也就清楚了，也可以写成

$$\frac{\md\vx}{\md t} = \vf(\vx, t) +\vL(\vx, t) \frac{\md \vbeta}{\md t}$$

¹. https://en.wikipedia.org/wiki/Metropolis-adjusted_Langevin_algorithm# ↩

². https://en.wikipedia.org/wiki/Multivariate_normal_distribution ↩

³. Estimation of non-normalized statistical models by score matching. A. Hyvarinen. https://www.jmlr.org/papers/volume6/hyvarinen05a/hyvarinen05a.pdf ↩

⁴. A Connection Between Score Matching and Denoising Autoencoders. P. Vincent. https://www.iro.umontreal.ca/~vincentp/Publications/smdae_techreport.pdf ↩

⁵. https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Non-degenerate_case ↩

⁶. Reverse-Time Diffusion Equation Models. Brian D.O. Anderson. https://core.ac.uk/download/pdf/82826666.pdf ↩

⁷. Simo Särkkä and Arno Solin (2019). Applied Stochastic Differential Equations. Cambridge University Press. https://users.aalto.fi/~asolin/sde-book/sde-book.pdf ↩

⁸. https://en.wikipedia.org/wiki/Covariance ↩

⁹. https://en.wikipedia.org/wiki/Partial_derivative#Definition ↩