score_based

基于分数的生成模型（Score-based Generative Model ）

Overview

​ 首先简单讲讲分数模型(score-based models) 是怎么回事，以一句话总结来说就是：

它估计 数据分布相关的梯度 并基于 朗之万动力学(Langevin Dynamics) 的思想来生成样本。

基于分数的生成模型（Score-based Generative Model ）

Score Matching

​ 如果有一组数据 {x1,x1,…,xn}\{\boldsymbol x_1, \boldsymbol x_1, \dots, \boldsymbol x_n\}，我们的目的是最大化对数似然：

max⁡θ∑i=1nlog⁡pθ(x)\max_\theta \sum_{i=1}^n \log p_\theta(\boldsymbol x)

​ 由于概率的性质的硬约束（和为 1，非负性），我们让

pθ(x)=e−fθ(x)∫Xe−fθ(x)dx⇒log⁡pθ(x)=−fθ(x)−log⁡Zθp_\theta(\boldsymbol x) = \frac{e^{-f_\theta(\boldsymbol x)}}{\int_\mathcal{X} e^{-f_\theta(\boldsymbol x)}d\boldsymbol x} \Rightarrow \log p_\theta(\boldsymbol x)= -f_\theta(\boldsymbol x) - \log Z_\theta

其中 Zθ=∫Xe−fθ(x)dxZ_\theta = \int_\mathcal{X} e^{-f_\theta(\boldsymbol x)}d\boldsymbol x，优化问题变为了：

max⁡θ∑i=1nlog⁡Zθ−∑i=1nfθ(x)\max_\theta \sum_{i=1}^n \log Z_\theta - \sum_{i=1}^n f_\theta(\boldsymbol x)

​ 在这里 zθ=∑i=1nlog⁡∫Xe−fθ(x)dxz_\theta = \sum_{i=1}^n \log \int_\mathcal{X} e^{- f_\theta(\boldsymbol x)}d\boldsymbol x 是 intractable 的，像 VAE 中就是通过 ELBO 规避了这个问题，在这篇文章中我们寻求另一种解决方法，即 score based 的方法来最大化对数似然。

score 的定义：

​ 我们定义 score 为：

sθ(x)=∂log⁡pθ(x)∂x=−∇xfxs_\theta(\boldsymbol x) = \frac{\partial \log p_\theta(\boldsymbol x)}{\partial \boldsymbol x} = -\nabla_\boldsymbol x f_{\boldsymbol x}

优化目标为：

θ∗=12arg⁡min⁡θEpdata∥s(x)−sθ(x)∥2\theta^* = \frac{1}{2}\arg\min_\theta \mathbb{E}_{p_\text{data}}\Vert s(\boldsymbol x) -s_\theta(\boldsymbol x) \Vert^2

其中 s(x)s(\boldsymbol x) 为真实数据的分数函数，关于为什么能转化优化目标，我们在下一节的 score 的物理意义讨论。再回到如何解这个优化问题上，实际上真实数据的分数函数是不可知的，这对接下来的推导造成了困难，下面用一种巧妙的方法规避这个问题：

设概率路径pt(x)p_t(\boldsymbol x)满足p1(x)=pdata(x)p_1(\boldsymbol x)=p_{\text{data}}(\boldsymbol x)，定义速度场vt(x)=Ept(x∣x0)[v(x0)]\boldsymbol v_t(\boldsymbol x)=\mathbb E_{p_t(\boldsymbol x|\boldsymbol x_0)}[\boldsymbol v(\boldsymbol x_0)]，目标函数推导如下：

步骤 1：初始目标函数

J(θ)=12∥Ep(x)[S(x)−S0(x)]∥2J(\theta)=\frac{1}{2}\left\|\mathbb E_{p(\boldsymbol x)}[\boldsymbol S(\boldsymbol x)-\boldsymbol S_0(\boldsymbol x)]\right\|^2

步骤 2：展开平方项

J(θ)=12Ep(x)∥sθ(x)∥2−Ep(x)[sθ(x)⊤s(x)]+const\begin{aligned} J(\theta)&=\frac{1}{2}\mathbb E_{p(\boldsymbol x)}\|s_\theta(\boldsymbol x)\|^2 - \mathbb E_{p(\boldsymbol x)}[s_\theta(\boldsymbol x)^\top s(\boldsymbol x)] + \text{const} \end{aligned}

步骤 3：代入得分函数定义
设s(x)=∇xlog⁡p(x)s(\boldsymbol x)=\nabla_{\boldsymbol x}\log p(\boldsymbol x)，则交叉项可改写为：

Ep(x)[S0(x)⊤S(x)]=∫p(x)S0(x)⊤∇xlog⁡p(x)dx\mathbb E_{p(\boldsymbol x)}[\boldsymbol S_0(\boldsymbol x)^\top\boldsymbol S(\boldsymbol x)]=\int p(\boldsymbol x)\boldsymbol S_0(\boldsymbol x)^\top\nabla_{\boldsymbol x}\log p(\boldsymbol x)d\boldsymbol x

步骤 4：应用分部积分法

∫p(x)S0⊤∇xlog⁡pdx=∫S0⊤∇xpdx=p(x)S0(x)∣−∞+∞⏟=0−∫p(x)∇x⋅S0(x)dx\begin{aligned} \int p(\boldsymbol x)\boldsymbol S_0^\top\nabla_{\boldsymbol x}\log p d\boldsymbol x &= \int \boldsymbol S_0^\top\nabla_{\boldsymbol x}p d\boldsymbol x \\ &= \underbrace{p(\boldsymbol x)\boldsymbol S_0(\boldsymbol x)\big|_{-\infty}^{+\infty}}_{=0} - \int p(\boldsymbol x)\nabla_{\boldsymbol x}\cdot\boldsymbol S_0(\boldsymbol x)d\boldsymbol x \\ \end{aligned}

⇒J(θ)=−12Ep(x)∥S0(x)∥2−Ep(x)[∇x⋅S0(x)]+const\Rightarrow J(\theta)=-\frac{1}{2}\mathbb E_{p(\boldsymbol x)}\|\boldsymbol S_0(\boldsymbol x)\|^2 - \mathbb E_{p(\boldsymbol x)}[\nabla_{\boldsymbol x}\cdot\boldsymbol S_0(\boldsymbol x)] + \text{const}

​ 其实，上面说的“数据分布相关的梯度”实质上是对数概率密度函数对于输入数据的梯度，第一次提出是 Stein Score ，这就是分数。score based 的模型就是训练它让它学会估计(预测)分数。

物理意义

​ score 的具体物理意义是：对于输入数据(样本)来说，其对数概率密度增长最快的方向。如果我们已经训练好了一个分数预测模型，那么我们可以类似 energy  based model 的推理方法，使用梯度下降法来最大化一个样本的对数似然，让这个样本的概率最大化

img

​ 但是我们训练网络的目的就是具有泛化性，我们需要避免网络的输出和原来的数据一模一样的情况，因此我们需要在采样的过程中加入噪声。最终的效果应该是在保证多样性的同时也要符合(靠近)原数据分布。于是就采用了 Langevin Dynamics 采样方法

Langevin Dynamics

​ Langevin dynamics 指朗之万动力学方程，是描述物理学中布朗运动（悬浮在液体/气体中的微小颗粒所做的无规则运动）的微分方程，借鉴到这里作为一种生成样本的方法。概括地来说，该方法首先从先验分布随机采样一个初始样本，然后利用模型估计出来的分数逐渐将样本向数据分布的高概率密度区域靠近。为保证生成结果的多样性，我们需要采样过程带有随机性。朗之万动力学采样刚好满足这一点。

dX(t)=−∇xU(X(t))dt+σdtztzt∼N(0,I)d X(t) = -\nabla_x U(X(t)) dt + \sigma \sqrt{dt} \boldsymbol z_t \\ \boldsymbol z_t \sim \mathcal{N}(0, I)

其中 UU 代表能量函数，这个式子其实很好理解，我们把它写为离散型，令 dt=ϵdt=\epsilon：

xϵ+t−xt=−∇xU(xt)ϵ+σϵzt⇒xt+ϵ=xt−∇xU(xt)ϵ+σϵzt\boldsymbol x_{\epsilon+t} - \boldsymbol x_t = -\nabla_\boldsymbol x U(\boldsymbol x_t)\epsilon + \sigma\sqrt{\epsilon}z_t \\ \Rightarrow \boldsymbol x_{t+\epsilon} = \boldsymbol x_t - \nabla_\boldsymbol x U(\boldsymbol x_t)\epsilon + \sigma\sqrt{\epsilon}z_t

​ 也就是梯度下降加上一个随机项，随机项服从一个均值为 0 方差为 ϵσ2\epsilon \sigma^2 的正态分布。

​ 套入我们的分数函数，让 x0∼N(0,I)\boldsymbol x_0 \sim \mathcal{N}(\boldsymbol 0, \boldsymbol I)，目标为最终的 xN∼pdata(x)\boldsymbol x_N \sim p_\text{data}(\boldsymbol x)，我们用一个 p(x,t)p(\boldsymbol x, t) 表述，则：

P(x,0)∼N(0,I)P(x,T)∼pθ(x)≈pdata(x)∂p(x,T)∂t=0\begin{aligned} P(\boldsymbol x, 0) \sim& \mathcal{N}(\boldsymbol 0, \boldsymbol I) \\ P(\boldsymbol x, T) \sim& p_\theta(\boldsymbol x) \approx p_\text{data}(\boldsymbol x) \\ \frac{\partial p(\boldsymbol x, T)}{\partial t} =& 0 \end{aligned}

根据 Fokker-Planck 方程：

∂P(x,t)∂t=∇xP(x,t)T∇U(x)+P(x,t)div⁡∇U(x)+12σ2ΔP(x,t)\frac{\partial P(\boldsymbol x,t)}{\partial t}=\nabla_{x}P(\boldsymbol x,t)^T\nabla U(\boldsymbol x)+P(\boldsymbol x,t)\operatorname{div}\nabla U(\boldsymbol x)+\frac{1}{2}\sigma^2 \Delta P(\boldsymbol x,t)

之前的工作结论是 Fokker-Planck 方程为零时，可以求得，T=σ2/2T = \sigma^2/2，当当在我们定义能量函数的时候，其实已经默认设置了 T=1T=1（Gibbs 分布的一般形式为 pθ(x)=exp⁡(−fθ(x)T)/Zθp_\theta(x) = \exp(-\frac{f_\theta(x)}{T})/Z_\theta）。因此代入 σ=2\sigma=\sqrt{2}：

xt+ϵ=xt−∇xfθ(x)ϵ+2ϵzt=xt+sθ(x)ϵ+2ϵzt\begin{aligned} \boldsymbol x_{t+\epsilon} =& \boldsymbol x_t - \nabla_\boldsymbol x f_\theta(\boldsymbol x) \epsilon + \sqrt{2\epsilon} \boldsymbol z_t \\ =& \boldsymbol x_t + s_\theta(\boldsymbol x)\epsilon + \sqrt{2\epsilon} \boldsymbol z_t \end{aligned}

根据这个采样公式，就可以满足最终的 xT∼pθ(x)\boldsymbol x_T \sim p_\theta(\boldsymbol x) 了。

Denoising Score Matching

Noise Conditional Score Network(NCSN)

​ 根据流形假设，我们假设真实图像概率分布空间为高维空间内的一个流形，这样会导致一个问题，就是在训练的时候一些空间内的分布 p(x)=0p(\boldsymbol x)=0，就会导致无法训练对应区域的 sθ(x)s_\theta(\boldsymbol x)。但是我们采样的过程中是有可能经过这些点的，因此这些点的 score 也需要正确估计，即需要训练到（还是经典的一句话，数据分布要尽可能覆盖大的空间，有没有对应方向的梯度和梯度小不小是两码事）

​ 为了解决这个问题，作者引入了增噪的手段，假设增加噪声满足概率分布 N(0,σ2)\mathcal{N}(\boldsymbol 0, \boldsymbol \sigma^2)，增噪后真是数据概率分布为

img

优化目标推导

​ 根据加噪条件分布 qσ(x∣x′)=N(x′,σ2I)q_\sigma(x|x') = \mathcal{N}(x', \sigma^2 I)，其对数梯度为：

log⁡qσ(x∣x′)=−∥x−x′∥22σ2+const∇xlog⁡qσ(x∣x′)=−x−x′σ2\begin{aligned} \log q_\sigma(x|x') &= -\frac{\|x - x'\|^2}{2\sigma^2} + \text{const} \\ \nabla_x \log q_\sigma(x|x') &= -\frac{x - x'}{\sigma^2} \end{aligned}

优化目标是最小化估计分数 sθ(x,σ)s_\theta(x,\sigma) 与真实分数 ∇xlog⁡qσ(x∣x′)\nabla_x \log q_\sigma(x|x') 的均方误差：

J(θ)=12Eqσ(x∣x′)Ep(x′)∥sθ(x,σ)−∇xlog⁡qσ(x∣x′)∥2=12Ex′∼p(x′)Ex∼N(x′,σ2I)∥sθ(x,σ)+x−x′σ2∥2\begin{aligned} J(\theta) &= \frac{1}{2} \mathbb{E}_{q_\sigma(x|x')}\mathbb{E}_{p(x')} \left\| s_\theta(x,\sigma) - \nabla_x \log q_\sigma(x|x') \right\|^2 \\ &= \frac{1}{2} \mathbb{E}_{x'\sim p(x')} \mathbb{E}_{x\sim\mathcal{N}(x',\sigma^2 I)} \left\| s_\theta(x,\sigma) + \frac{x - x'}{\sigma^2} \right\|^2 \end{aligned}

此时，qσ(x)=∫qσ(x∣x′)p(x′)dx′q_\sigma(x) = \int q_\sigma(x|x') p(x') dx' 是加噪后的真实数据分布

多噪声尺度训练：

​ 在前面的分析中提及到，为了让扰动后的分布能够更多地“填充”原来的低概率密度区域，但是强度过大的噪声反之会干扰到原始数据的分布进而造成估计分数的误差增大，从而基于加噪后的分数使用朗之万动力学采样生成的结果也就不太符合原数据分布；相反地，噪声强度小能够获得与原数据分布较为近似的效果，但是却不能够很好地“填充”低概率密度区域。因此训练的时候使用了多尺度的噪声等级序列，即最开始的噪声水平足够大以能够充分“填充"低概率密度区域；而最后的噪声水平足够小以获得对原数据分布良好的近似，避免过度扰动对

​ LL 个噪声级别 {σi}i=1L\{\sigma_i\}_{i=1}^L，引入权重 λ(σi)\lambda(\sigma_i) 平衡不同噪声的贡献，联合损失为：

L(θ;{σi})=1L∑i=1Lλ(σi)J(θ;σi)\mathcal{L}(\theta; \{\sigma_i\}) = \frac{1}{L} \sum_{i=1}^L \lambda(\sigma_i) J(\theta; \sigma_i)

其中 λ(σi)\lambda(\sigma_i) 的常见选择是 λ(σi)=σi2\lambda(\sigma_i) = \sigma_i^2，以抵消 1σ2\frac{1}{\sigma^2} 的尺度效应

核心代码分析

Loss损失函数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def anneal_dsm_score_estimation(scorenet, samples, labels, sigmas, anneal_power=2.):
    # 取出每个样本对应噪声级别下的噪声分布的标准差，即公式中的sigma_i，
    # 这里的 labels 是用于标识每个样本的噪声级别的，就是 i，实际是一种索引标识
    # (bs,)->(bs,1,1,1) 扩展至与图像一致的维度数
    used_sigmas = sigmas[labels].view(samples.shape[0], *([1] * len(samples.shape[1:])))
    # 加噪：x' = x + sigma * z (z ~ N(0,1))
    perturbed_samples = samples + torch.randn_like(samples) * used_sigmas
    
    # 目标score，本质是对数条件概率密度 log(p(x'|x)) 对噪声数据 x' 的梯度
    # 由于这里建模为高斯分布，因此可计算出结果最终如下，见前文公式(vii)
    target = - 1 / (used_sigmas ** 2) * (perturbed_samples - samples)
    # 模型预测的 score
    scores = scorenet(perturbed_samples, labels)
    target = target.view(target.shape[0], -1)
    scores = scores.view(scores.shape[0], -1)

    # 先计算每个样本在所有维度下分数估计的误差总和，再对所有样本求平均
    # 见前文公式(vii)
    loss = 1 / 2. * ((scores - target) ** 2).sum(dim=-1) * used_sigmas.squeeze() ** anneal_power

    return loss.mean(dim=0)

​ annealed Langevin dynamics：这部分是根据上面的伪代码实现的，具体算法流程参见上面的伪代码公式。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def anneal_Langevin_dynamics(self, x_mod, scorenet, sigmas, n_steps_each=100, step_lr=0.00002):
    images = []

    with torch.no_grad():
        # 依次在每个噪声级别下进行朗之万动力学采样生成，噪声强度递减
        for c, sigma in tqdm.tqdm(enumerate(sigmas), total=len(sigmas), desc='annealed Langevin dynamics sampling'):
            # 噪声级别
            labels = torch.ones(x_mod.shape[0], device=x_mod.device) * c
            labels = labels.long()

            # 这个步长并非 Algorithm 1 中的 alpha，而是其中第6步的 alpha/2
            # 对应朗之万动力学采样公式(见公式(vi))的 epsilon/2
            step_size = step_lr * (sigma / sigmas[-1]) ** 2
            
            # 每个噪声级别下进行一定步数的朗之万动力学采样生成
            for s in range(n_steps_each):
                images.append(torch.clamp(x_mod, 0.0, 1.0).to('cpu'))
                # 对应公式(vi)最后一项
                noise = torch.randn_like(x_mod) * np.sqrt(step_size * 2)
                # 网络估计的分数
                grad = scorenet(x_mod, labels)
                # 朗之万动力方程
                x_mod = x_mod + step_size * grad + noise

        return images

img