【生成模型新方向】score-based generative models

0、前言

本文主要翻译自此领域先驱Song Yang博士(斯坦福大学博士)的博客。并且对于重要知识点给出了表格形式的整理汇总,方便记忆和理解!

一言以蔽之:

我们可以在大量噪声扰动的数据分布上(on a large number of noise-perturbed data distributions)学习得分函数score functions(对数概率密度函数的梯度gradients of log probability density functions),然后使用 Langevin 型采样生成样本。

由此产生的生成模型,通常称为基于分数的生成模型,与现有模型系列相比具有几个重要优势:

  1. 无需对抗训练的 GAN 级样本质量、
  2. 灵活的模型架构、
  3. 精确的对数似然计算
  4. 唯一可识别表征学习(uniquely identifiable representation learning)
  5. 无需重新训练模型即可解决逆向问题。

在这篇博文中,我们将更详细地向您展示基于分数的生成模型的提出的动机、基本概念和潜在应用。

1、介绍

目前,生成模型(generative models)可以根据其表示概率分布的方式主要分为2个大类别:

  1. likelihood-based models基于似然的模型,它通过(近似)最大似然直接学习分布的概率密度(或质量)函数。典型的基于似然的模型包括自回归模型autoregressive models、归一化流模型normalizing flow models、基于能量的模型 (EBMs)energy-based models (EBMs) 和变分自动编码器 (VAEs)variational auto-encoders (VAEs)。
    在这里插入图片描述
  2. implicit generative models隐式生成模型,其中概率分布由其采样过程的模型隐式表示。最突出的例子是生成对抗网络 (GAN),其中通过使用神经网络转换随机高斯向量来合成来自数据分布的新样本。
    在这里插入图片描述

然而,这两大类生成模型,基于似然的模型和隐式生成模型都有很大的局限性。基于似然的模型要么需要对模型架构进行严格限制,以确保为似然计算提供易于处理的归一化常数,要么必须依赖替代目标来近似最大似然训练。另一方面,隐式生成模型通常需要对抗性训练,这是出了名的不稳定并可能导致模式崩溃。

本篇博客介绍的就是宋博士提出的score-based generative model, 用这种新的生成模型来解决/规避刚才提到的这些问题。score-based generative model的核心idea是:

对对数概率密度函数的梯度建模,这个量通常被称为 (Stein) 得分函数。model the gradient of the log probability density function, a quantity often known as the (Stein) score function

这种score-based generative models不需要处理类似likelihood-based models的规则化常数。而且,score-based generative models在噪声干扰的数据下训练的效果非常好。这类方法可以恢复被噪声干扰的图片本身,并且有着良好的sample quality(采样质量)。
Score function (the vector field) and density function (contours) of a mixture of two Gaussians.

两个高斯混合的得分函数(向量场)和密度函数(等高线)。

基于分数的模型在许多下游任务和应用程序上取得了最先进的性能。这些任务包括图像生成(是的,比 GAN 更好!)、音频合成、形状生成和音乐生成。

当噪声扰动的过程是由 随机可微分方程(stochastic differential equation (SDE)) 给出时, score-based generative models和FLOW这种模型在数学上联系起来了,因此可以进行精确的似然估计计算以及表征学习。

此外,建模和估计分数有助于解决逆向问题,应用程序包括

  1. 图像修复
  2. 图像着色
  3. 压缩传感
  4. 医学图像重建(例如 CT、MRI)。
    在这里插入图片描述

从基于分数的模型生成的分辨率为 1024 x 1024 的样本

2、The score function, score-based models, and score matching

一个表格给出总结,下面是正文。

名称 形式
score function ∇ x log ⁡ p ( x ) \nabla_\mathbf{x} \log p( \mathbf{x}) xlogp(x)
score-based model 使用这种score function的模型我们就统称为score-based model,用 s θ ( x ) \mathbf{s}_\theta(\mathbf{x}) sθ(x)来表示,这种模型的目标是在无需考虑规则化常数项的情况下,使得 s θ ( x ) ≈ ∇ x log ⁡ p ( x ) \mathbf{s}_\theta(\mathbf{x}) \approx \nabla_\mathbf{x} \log p( \mathbf{x}) sθ(x)xlogp(x)能被参数化。
score matching 因为 ∇ x log ⁡ p ( x ) \nabla_\mathbf{x} \log p( \mathbf{x}) xlogp(x)不可知,我们无法优化Fisher divergence,而score matching可以在不知道ground-truth data score的情况下最小化Fisher divergence。Score matching objectives可以直接在数据集上估计并使用随机梯度下降进行优化,类似于训练基于似然的模型(具有已知归一化常数)的log-likelihood objective。Commonly used score matching methods include denoising score matching and sliced score matching Sliced score matching,我们可以训练一个score-based模型来优化score-matching objective, 而不需要对抗学习!

假设我们有一个数据集{x1,x2,⋯,xN},其中的每个 x i , i ∈ 1 , . . . , N x_i, i \in {1, ..., N} xi,i1,...,N都是从一个潜在的数据分布 p θ ( x ) p_{\theta}(x) pθ(x)中独立取得的(i.i.d). 生成模型的目的是能够完美的建模这个数据分布 p θ ( x ) p_{\theta}(x) pθ(x),以便任意的采样生成符合这个分布的新数据。

为了构建这样的生成模型,我们首先需要一种可以表示这种概率分布的方法。与基于似然的模型一样,其中一种方法是直接对概率密度函数 (p.d.f.) 或概率质量函数 (p.m.f.) 建模。让 f θ ( x ) ∈ R f_\theta(x)∈\mathbb{R} fθ(x)R 是由可学习参数 θ 参数化的实值函数。我们可以定义一个 p.d.f.通过:

probability density function (p.d.f.) or probability mass function (p.m.f.)

在这里插入图片描述
在这里,Zθ>0 是一个依赖于 θ 的归一化常数,这样∫pθ(x)dx=1。这里的功能fθ(x) 通常称为非归一化概率模型或基于能量的模型。

这里, Z θ > 0 Z_{\theta} > 0 Zθ>0是一个依赖于θ的normalizing constant(规则化常数),其目的是让 ∫ p θ ( x ) d x = 1 \int p_{\theta}(x)dx = 1 pθ(x)dx=1. 函数 f θ ( x ) f_{\theta}(\bf{x}) fθ(x)是一个unnormalized 概率模型,或者叫做EBM能量模型。

我们可以训练 p θ ( x ) p_{\theta}(x) pθ(x)来最大化数据的对数似然.
在这里插入图片描述
然而,上面的这个公式要求 p θ ( x ) p_{\theta}(x) pθ(x)是一个归一化的概率密度函数PDF,然后这不是我们想要的:

因为为了计算 p θ ( x ) p_{\theta}(x) pθ(x)我们必须计算归一化常数 Z θ Z_{\theta} Zθ ,对于任何一般情况下的 f θ ( x ) f_{\theta}(\bf{x}) fθ(x),这是一个典型的难以处理的量

所以,为了使得maximum likelihood training的训练变得可行,likelihood-based models通过如下2种方式,而这2种方式,尤其是FLOW-based模型,会使得计算量极具增加:

  1. 限制模型结构(causal convolutions in autoregressive models, invertible networks in normalizing flow models)来使得 Z θ = 1 Z_{\theta}=1 Zθ=1(例如,自回归模型中的因果卷积,归一化流模型中的可逆网络)
  2. 近似归一化常数(variational inference in VAEs, or MCMC sampling used in contrastive divergence)(例如,VAEs中的变分推理,或用于对比发散的 MCMC 采样)

而score-based模型则是通过构造一个score function而非density function来避开处理这个归一化常数的问题。对一个分布P(x), 其score function定义为:
在这里插入图片描述
使用这种score function的模型我们就统称为score-based model,我们将其表示为 s θ ( x ) \mathbf{s}_\theta(\mathbf{x}) sθ(x)
这种模型的目标是在无需考虑规则化常数项的情况下,能够使得 s θ ( x ) ≈ ∇ x log ⁡ p ( x ) \mathbf{s}_\theta(\mathbf{x}) \approx \nabla_\mathbf{x} \log p( \mathbf{x}) sθ(x)xlogp(x)被参数化。

例如,我们可以使用等式 p θ ( x ) = e − f θ ( x ) Z θ p_{\theta}(x)=\frac{e^{-f_{\theta}(\bf{x})}}{Z_{\theta}} pθ(x)=Zθefθ(x)中定义的基于能量的模型轻松地参数化基于分数的模型,通过:(其实就对数函数的展开)
在这里插入图片描述
可以看出, 基于分数的模型 s θ ( x ) s_{\theta}(\bf{x}) sθ(x)和normalizing constant Z θ Z_{\theta} Zθ​相互独立。
这显着扩展了我们可以轻松使用的模型系列,因为我们不需要任何特殊的架构来使标准化常量 Z θ Z_{\theta} Zθ​变得易于处理(tractable)。例如上面提到的likelihood-based models
conv_ops

Parameterizing probability density functions(pdfs). No matter how you change the model family and parameters, it has to be normalized (area under the curve (AUC) must integrate to one).参数化概率密度函数。无论您如何更改模型族和参数,都必须对其进行归一化(曲线下面积必须积分为一)。

1

Parameterizing score functions. No need to worry about normalization.参数化分数函数。无需担心规范化。

与基于似然的模型类似,我们可以通过最小化模型和数据分布之间的Fisher 散度来训练基于分数的模型,定义为
在这里插入图片描述
其中Fisher 散度的定义为:在这里插入图片描述
直觉上来说,Fisher divergence (Fisher散度)比较的是ground-truth数据分数和score-based模型的lℓ2距离的平方。但是由于不知道数据得分 ∇ x l o g p ( x ) \nabla_{x} log p(x) xlogp(x), 我们没法直接计算出Fisher divergence.

幸运的是,现存了一系列称之为score matching的方法,这种方式可以在不知道ground-truth data score的情况下,minimize Fisher divergence.

得分匹配目标可以直接在数据集上估计并使用随机梯度下降进行优化,类似于训练基于似然的模型(具有已知归一化常数)的对数似然目标。

我们可以通过最小化分数匹配目标(score matching objective)来训练基于分数的模型,而不需要对抗性优化。

此外,使用score matching objective给了我们在模型结构设计的灵活性。Fisher Divergence不需要 s θ ( x ) s_{\theta}(\bf{x}) sθ(x)是任意的规则化分布的实际得分函数(actual score function). 即: 无需像之前的方法那样,对 s θ ( x ) s_{\theta}(\bf{x}) sθ(x)有一个较强的假设! 它只是比较ground-truth 数据分数和基于分数的模型之间的ℓ2 距离,没有对 s θ ( x ) s_{\theta}(\bf{x}) sθ(x)形式的额外假设。

在使用中,score-based model的唯一要求是:

score-based model should be a vector-valued function with the same input and output dimensionality, which is easy to satisfy in practice.它应该是具有相同输入和输出维度的向量值函数,这在实践中很容易满足。

作为一个简短的总结,我们可以通过对其评分函数(score function)建模来表示分布,score function可以通过训练具有score matching的并且是自由形式架构的基于分数的模型来估计。

3、Langevin dynamics朗之万动力学

一个表格给出总结,下面是正文。

名称 介绍
Langevin dynamics 在这里插入图片描述 将score-based model s θ ( x ) \mathbf{s}_\theta(\mathbf{x}) sθ(x)代入上式来进行采样。并且当 ϵ \epsilon ϵ 足够小且 K K K 足够大时,此公式采样得到的数据分布与真实数据分布的误差可以忽略不计。

一旦我们训练得到一个基于分数的模型 s θ ( x ) ≈ ∇ x log ⁡ p ( x ) \mathbf{s}_\theta(\mathbf{x}) \approx \nabla_\mathbf{x} \log p( \mathbf{x}) sθ(x)xlogp(x),我们可以使用称为 Langevin 动力学的迭代过程从中采样。

Langevin dynamics仅通过使用score function ∇ x l o g p ( x ) \nabla_{x} log p(x) xlogp(x)来对真实数据分布 p ( x ) p(x) p(x)进行MCMC(MCMC, 马尔科夫链蒙特卡洛(Markov Chain Monte Carlo)方法,是用于从复杂分布中获取随机样本的统计学算法)的采样。具体来说,它先从任意的先验的分布中 x 0 ∼ π ( x ) \bf{x}_{0} \sim \bf{\pi(x)} x0π(x), 初始化构造一个chain,然后按着下面公式所述的那样进行迭代:
在这里插入图片描述
其中 z i ∼ N ( 0 , I ) \mathbf{z}_i \sim \mathcal{N}(0, I) ziN(0,I)。当 ϵ → 0 \epsilon \to 0 ϵ0 K → ∞ K \to \infty K 时,在常规条件下,从上式中的过程获得的 x K \mathbf{x}_K xK 收敛到真实分布 p ( x ) p( \mathbf{x}) p(x)
实际上,当 ϵ \epsilon ϵ 足够小且 K K K 足够大时,两者的误差可以忽略不计。
在这里插入图片描述
因为Langevin dynamics仅通过 ∇ x l o g p ( x ) \nabla_xlogp(x) xlogp(x)来获得 p ( x ) p(x) p(x),而我们的score-based model为 s θ ( x ) ≈ ∇ x log ⁡ p ( x ) \mathbf{s}_\theta(\mathbf{x}) \approx \nabla_\mathbf{x} \log p( \mathbf{x}) sθ(x)xlogp(x),因此我们可以将score-based model s θ ( x ) \mathbf{s}_\theta(\mathbf{x}) sθ(x)代入上式来进行采样。

4、朴素的(Naive)score-based生成建模及其缺陷

一个表格给出总结,下面是正文。

问题 介绍
采样不准确 在这里插入图片描述 可以看到积分中乘了一个权重p(x),当在低密度区域,p(x)很小,那么该Fisher divergence 将被忽略,这种不准确可能使Langevin 动力学的采样过程从一开始就derail(出轨),从而无法产生高质量的样本。

到目前为止,我们已经讨论了如何使用分数匹配score matching训练基于分数score-based的模型,然后通过 Langevin 动力学生成样本。
然而,这种最朴素的方法在实践中取得的成功有限——我们将讨论一些在以前的工作中很少受到关注的分数匹配的陷阱pitfalls。
在这里插入图片描述
目前的版本,由于在score matching中的一些问题,导致出现了较为明显的失败情况,而这些问题,前人的文章并没有仔细的探究。

关键的挑战是估计的分数函数score functions在低密度区域不准确,其中很少有数据点可用于计算分数匹配目标score matching objective。

我们之前提到,score-based模型的分数匹配是通过最小化真实数据和模型输出的Fisher divergence进行的,其可以写成这种积分的形式。
在这里插入图片描述
可以发现ℓ2 真实数据评分函数 score function和score-based的模型之间的差异由 p(x) 加权,在 p(x) 较小的低密度区域中,它们在很大程度上被忽略了。这种行为可能导致结果不佳subpar ,如下图所示:
在这里插入图片描述
当使用 Langevin 动力学进行采样时,当真实数据位于高维空间时,我们的初始的采样的样本很可能位于低密度区域
因此,基于一个不准确的score-based模型,将从程序的一开始就会使 Langevin 动力学的采样过程derail(出轨),阻止它生成代表真实数据的高质量样本。

5、multiple noise perturbation后的score-based模型

一个表格给出总结,下面是正文。

名称 介绍
解决第四节中提出的问题 用多尺度的噪声扰乱数据点,然后在噪声数据点noisy data points上训练基于分数的模型 s θ ( x , i ) s_θ(x,i) sθ(x,i)。在得到我们训练的noise-conditional score-based mode s θ ( x , i ) s_θ(x,i) sθ(x,i)之后,我们就可以使用Langevin Dynamics来进行采样了.i=L,L−1,⋯,1 . (相当于一个加噪和去噪的过程,和扩散模型相同)

我们如何绕过在低数据密度区域进行准确分数估计的困难?
我们的解决方案是用噪声扰乱数据点,然后在噪声数据点noisy data points上训练基于分数的模型。
当噪声幅度足够大时,它可以填充低数据密度区域以提高估计分数的准确性。
例如,当我们扰动两个被额外高斯噪声扰动的高斯分布时,会发生这种情况。
在这里插入图片描述

Estimated scores are accurate everywhere for the noise-perturbed data distribution due to reduced low data density regions.由于低数据密度区域的减少,对于受噪声扰动的数据分布,估计分数在任何地方都是准确的。

然而这又引发了另一个问题:我们如何为扰动过程选择合适的噪声标度?
较大的噪声显然可以覆盖更多的低密度区域以获得更好的分数估计score estimation,但它会过度破坏数据并显着改变原始分布。
另一方面,较小的噪声会导致原始数据分布的损坏较少,但不会像我们希望的那样覆盖低密度区域。

为了实现两全其美,我们同时使用多尺度的噪声扰动。
假设我们总是用各向同性isotropic 高斯噪声扰动数据,并且共有 L 个递增的标准差 σ 1 < σ 2 < ⋯ < σ L σ_1<σ_2<⋯<σ_L σ1<σ2<<σL
首先我们用每个高斯噪声 N ( 0 , σ i 2 I ) , i = 1 , 2 , ⋯ , L N(0,σ_i^2I),i=1,2,⋯,L N(0,σi2I),i=1,2,,L 扰动数据分布 p(x)获得噪声扰动分布:
在这里插入图片描述
请注意,我们可以轻松地从中抽取样本 p σ i ( x ) p_{σi}(x) pσi(x) 通过采样 x∼p(x) 并计算 x + σ i z x+σ_iz x+σiz,其中 z∼N(0,I)。

接下来,我们估计每个噪声扰动分布的得分函数, ∇ x l o g p σ i ( x ) \nabla_xlogp_{\sigma i}(x) xlogpσi(x),通过评分匹配score matching训练Noise Conditional Score-Based Model s θ ( x , i ) s_θ(x,i) sθ(x,i)(当使用神经网络参数化时,也称为Noise Conditional Score Network或 NCSN),这样 s θ ( x , i ) ≈ ∇ x l o g p σ i ( x ) s_θ(x,i)≈∇xlogp_{σi}(x) sθ(x,i)xlogpσi(x) 对于所有 i=1,2,⋯,L。
在这里插入图片描述
在这里插入图片描述
训练目标为 s θ ( x , i ) s_θ(x,i) sθ(x,i) 是所有噪声尺度的 Fisher 散度的加权和。特别是,我们使用以下目标objective函数 :
在这里插入图片描述
其中 λ ( i ) ∈ R > 0 λ(i)∈\mathbb{R}>0 λ(i)R>0 是一个正的加权函数,通常被选为 λ ( i ) = σ i 2 λ(i)=σ_i^2 λ(i)=σi2
上式可以通过分数匹配score matching进行优化,就像优化朴素(无条件)基于分数的模型 s θ ( x ) s_θ(x) sθ(x)一样。

在得到我们训练的noise-conditional score-based mode s θ ( x , i ) s_θ(x,i) sθ(x,i)之后,我们就可以使用Langevin Dynamics来进行采样了.i=L,L−1,⋯,1 .
这种方法称为退火 Langevin 动力学annealed Langevin dynamics(由【3】中的算法 1 定义,并由后面的工作改进),之所以称之为退火,是因为噪声标度σi 随着时间的推移逐渐减小(退火)。
在这里插入图片描述
在这里插入图片描述在这里插入图片描述

在 CelebA(左)和 CIFAR-10(右)上训练的噪声条件评分网络 (NCSN) 模型(来自参考文献【3】)的退火 Langevin 动力学。我们可以从非结构化噪声开始,根据分数修改图像,生成好的样本。该方法在当时的 CIFAR-10 上取得了最先进的 Inception 分数。Annealed Langevin dynamics for the Noise Conditional Score Network (NCSN) model (from ref.[17]) trained on CelebA (left) and CIFAR-10 (right). We can start from unstructured noise, modify images according to the scores, and generate nice samples. The method achieved state-of-the-art Inception score on CIFAR-10 at its time.

以下是调整具有多个噪声尺度with multiple noise scales的score-based生成模型的一些实用建议:

  1. 选择σ1<σ2<⋯<σL 呈几何级数,有σ1 足够小并且σL 相当于所有训练数据点之间的最大成对距离。 L 通常约为数百或数千。
  2. 参数化基于分数的模型sθ(x,i) 用 U-Net skip connections 。
  3. 在测试使用时,对基于分数的模型的权重应用指数移动平均exponential moving average(EMA)。
    在这里插入图片描述

使用上面的说明,我们能够在各种数据集上生成与 GAN 质量相当的高质量图像样本,如下所示:
在这里插入图片描述

6、Score-based generative modeling with stochastic differential equations (SDEs)随机微分方程

名称 介绍
6.1在第五部分提到了用多尺度的噪声扰乱数据点,并且是人为定义了L个方差逐渐变大的高斯噪声,那么L扩展到无限大该如何处理? 使用SDE(随机微分方程)来扰动数据
6.2那么随之而来的问题是对应无限噪声尺度该如何进行采样? 我们可以通过使用反向 SDE 类似地反转样本生成的扰动过程。注:在有限数量的噪声尺度下,我们可以通过使用annealed Langevin dynamics退火朗之万动力学反转扰动过程来生成样本,即使用朗之万动力学从每个噪声扰动分布中顺序采样。
6.3如何训练一个Time-Dependent Score-Based Model基于时间依赖分数的模型 s θ ( x , t ) \mathbf{s}_{\theta}(\mathbf{x}, t) sθ(x,t),使得 s θ ( x , t ) ≈ ∇ x l o g p t ( x ) \mathbf{s}_{\theta}(\mathbf{x}, t) \approx \nabla_{\mathbf{x}}log p_t(\mathbf{x}) sθ(x,t)xlogpt(x) 和以前一样,我们的 Fisher divergences 的加权组合可以使用score matching methods进行有效优化,例如去噪分数匹配和切片分数匹配denoising score matching and sliced score matching。一旦我们的基于分数的模型 s θ ( x , t ) s_θ(x,t) sθ(x,t) 被训练到最优,我们可以将其代入6.2中的Reverse SDE的表达式,以获得估计的Reverse SDE。还提出了一个 λ ( t ) = g 2 ( t ) λ(t)=g^2(t) λ(t)=g2(t) 的似然加权函数likelihood weighting function,来获得比自回归模型还强的似然。
6.4最后是reverse SDE的求解 提出了最简单的Euler-Maruyama method,还有其他的 Milstein method和stochastic Runge-Kutta methods。宋博士也提出了一个类似于Euler-Maruyama method的求解方法叫Predictor-Corrector samplers,更适合与求解reverse-time SDEs。预测器predictor每次选择一个合适的步长Δt<0,然后根据当前样本x(t)预测x(t+Δt)。然后根据基于分数的模型 s θ ( x , t + Δ t ) s_θ(x,t+Δt) sθ(x,t+Δt)运行几个校正corrector步骤steps来改进样本 x(t+Δt),使 x ( t + Δ t ) x(t+Δt) x(t+Δt) 成为来自 p t + Δ t ( x ) p_{t+Δt}(x) pt+Δt(x)的高质量样本。借助于PC方法和score-based models,可以比现在最好的GAN模型还强。
6.5尽管能够生成高质量样本,但基于 Langevin MCMC 和 SDE 求解器的采样器并未提供一种方法来计算基于分数的生成模型的精确对数似然exact log-likelihood。 介绍了一个基于常微分方程 (ODE) 的采样器,它允许进行精确的似然计算。当使用我们之前讨论的似然加权likelihood weighting( λ ( t ) = g 2 ( t ) λ(t)=g^2(t) λ(t)=g2(t))训练基于分数的模型,并使用变分反量化variational dequantization来获得离散图像的似然时,我们可以获得与最先进的自回归模型相当甚至更好的似然(所有这些都没有任何数据增强) 。
用于反问题求解的可控生成 主要原理是贝叶斯公式,使用场景如CIFAR-10的类条件生成,LSUN的图像修复Image inpainting,LSUN和FFHQ的图像着色Image colorization

正如我们已经讨论过的,添加多个噪声尺度对于基于分数的生成模型的成功至关重要。
通过将噪声尺度的数量推广到无穷大infinity,我们可以基于score-based生成模型构造迄今为止最强大的框架。
我们不仅获得了更高质量的样本,而且还获得了精确的对数似然计算,并加快采样速度,使得学习的特征具有更好的,更加独立的表征,和用于逆向问题求解的可控生成(inverse problem solving)。

除了本介绍之外,宋博士还提供了 Google Colab 编写的教程,以提供在 MNIST 上训练玩具模型的分步指南。我们还有更高级的代码存储库,可为大型应用程序提供成熟的实现。
在这里插入图片描述

6.1使用SDE(随机微分方程)来扰动数据

当噪声尺度的数量接近无穷大时,我们基本上会随着噪声水平的不断增加而扰乱数据分布。在这种情况下,噪声扰动过程是一个连续时间随机过程(continuous-time stochastic process),如下所示:
在这里插入图片描述
GIF图具体看[1], 这里可以看到,随着随机过程的加深,原图的信息被大量的隐藏起来。

我们如何以简洁的方式表示随机过程?许多随机过程(尤其是扩散过程)是随机微分方程 (SDE) 的解。通常,SDE 具有以下形式:
在这里插入图片描述
其中 f ( ⋅ , t ) : R d → R d f(⋅,t):\mathbb{R}^d→\mathbb{R}^d f(,t):RdRd是向量值函数,称为漂移系数 g ( t ) ∈ R g(t)∈\mathbb{R} g(t)R是实值函数,称为扩散系数,w表示标准布朗运动,dw可以看作为无穷小的白噪声infinitesimal white noise。

这个随机微分方程的解是随机变量的连续集合 { x ( t ) } t ∈ [ 0 , T ] \left \{x(t)\right \}_{t∈[0,T]} { x(t)}t[0,T]。随着时间索引 t 从开始时间 0 增长到结束时间 T,这些随机变量跟踪随机轨迹。

p t ( x ) p_t(x) pt(x) 表示 x(t) 的边缘概率密度函数。这里 t∈[0,T] 类似于不同尺度下的噪声 i=1,2,⋯,L 当我们有有限数量的噪声尺度时,此时 p t ( x ) p_t(x) pt(x) 类似于 p σ i ( x ) p_{σi}(x) pσi(x)。显然, p 0 ( x ) = p ( x ) p_0(x)=p(x) p0(x)=p(x) 是原始数据分布,因为在 t=0 时没有对数据应用扰动。

在用随机过程扰动 p(x) 足够长的时间 T 后,pT(x) 变得接近易于处理的噪声分布 π(x),称为先验分布prior distribution。我们注意到 p T ( x ) p_T(x) pT(x) 类似于 p σ L ( x ) p_{σL}(x) pσL(x) 在有限噪声尺度的情况下,对应于应用最大的噪声扰动 σ L σ_L σL到数据。

上式中的 SDE 是手工设计的,类似于我们手工设计的方式在有限噪声尺度的情况下,σ1<σ2<⋯<σL。添加噪声扰动的方法有很多种,SDE 的选择也不是唯一的。比如下面的SDE:
在这里插入图片描述
用均值为零且方差呈指数增长的高斯噪声扰动数据,这类似于用 N ( 0 , σ 1 2 I ) , N ( 0 , σ 2 2 I ) , ⋯ , N ( 0 , σ L 2 I ) N(0,σ_1^2I),N(0,σ_2^2I),⋯,N(0,σ_L^2I) N(0,σ12I),N(0,σ22I),,N(0,σL2I)扰动数据,其中σ1<σ2<⋯<σL是等比级数。

因此,SDE 应该被视为模型的超参数,就像{σ1,σ2,⋯,σL}。
在【4】中,我们提供了三种通常适用于图像的 SDE:

  1. 方差爆炸 SDE (VE SDE)the Variance Exploding SDE (VE SDE)
  2. 方差保持 SDE (VP SDE) the Variance Preserving SDE (VP SDE)
  3. 子 VP SDE。and the sub-VP SDE

6.2Reverse SDE用于生成样本

回想一下,在有限数量的噪声尺度下,我们可以通过使用annealed Langevin dynamics退火朗之万动力学反转扰动过程来生成样本,即使用朗之万动力学从每个噪声扰动分布中顺序采样。
对于无限噪声尺度,我们可以通过使用反向 SDE 类似地反转样本生成的扰动过程。
在这里插入图片描述
在这里插入图片描述

在这里插入图片描述
重要的是,任何 SDE 都有对应的反向 SDE ,其封闭形式closed form由下式给出:
在这里插入图片描述
重要部分!!!:
这里 d t \mathrm{d}t dt 表示无穷小时间步长,因为 SDE 对于时间步而言是需要向后求解(从 t=T 到 t=0)。为了计算反向 SDE,我们需要估计 ∇ x l o g p t ( x ) \nabla_{\mathbf{x}}logp_t(\mathbf{x}) xlogpt(x),而这与 p t ( x ) p_t(\mathbf{x}) pt(x)的score function是一样的。

在这里插入图片描述

解决反向 SDE 会产生一个基于分数的生成模型。可以使用 SDE 将数据转换为简单的噪声分布。如果我们知道每个中间时间步的分布得分,则可以反过来从噪声中生成样本。

6.3Estimating the reverse SDE with score-based models and score matching

解决反向 SDE 需要我们知道终端分布 p T ( x ) p_T(x) pT(x)和得分函数 ∇ x l o g p t ( x ) \nabla_{\mathbf{x}}log p_t(\mathbf{x}) xlogpt(x)
通过设计,前者接近于完全易于处理的先验分布 π(x)。
为了估计 ∇ x l o g p t ( x ) \nabla_{\mathbf{x}}log p_t(\mathbf{x}) xlogpt(x),我们训练一个Time-Dependent Score-Based Model基于时间依赖分数的模型 s θ ( x , t ) \mathbf{s}_{\theta}(\mathbf{x}, t) sθ(x,t),使得 s θ ( x , t ) ≈ ∇ x l o g p t ( x ) \mathbf{s}_{\theta}(\mathbf{x}, t) \approx \nabla_{\mathbf{x}}log p_t(\mathbf{x}) sθ(x,t)xlogpt(x)
这类似于noise-conditional score-based model基于噪声条件分数的模型 s θ ( x , i ) \mathbf{s}_{\theta}(\mathbf{x}, i) sθ(x,i)用于有限噪声尺度,经过训练使得 s θ ( x , i ) ≈ ∇ x l o g p σ i ( x ) \mathbf{s}_{\theta}(\mathbf{x}, i) \approx \nabla_{\mathbf{x}}log p_{\sigma i}(\mathbf{x}) sθ(x,i)xlogpσi(x)

我们的训练目标是 s θ ( x , t ) s_θ(x,t) sθ(x,t) 是 Fisher 散度的连续加权组合,由下式给出:
在这里插入图片描述
其中 U ( 0 , T ) \mathcal{U}(0,T) U(0,T) 表示时间间隔 [0,T] 上的均匀分布,并且 λ : R → R > 0 \lambda:\mathbb{R}→\mathbb{R}>0 λ:RR>0 为正加权函数。
通常我们使用在这里插入图片描述
以平衡不同分数匹配损失在不同时间的量纲。

和以前一样,我们的 Fisher 散度的加权组合可以使用分数匹配方法进行有效优化,例如去噪分数匹配和切片分数匹配denoising score matching and sliced score matching。
一旦我们的基于分数的模型 s θ ( x , t ) s_θ(x,t) sθ(x,t) 被训练到最优,我们可以将其代入6.2中的Reverse SDE的表达式,以获得估计的Reverse SDE。
在这里插入图片描述
我们可以从 x ( T ) ∼ π x(T)∼π x(T)π开始,求解上面的反向SDE得到一个样本x(0)。
让我们将以这种方式获得的 x ( 0 ) x(0) x(0) 的分布表示为 p θ p_θ pθ
当基于分数的模型 s θ ( x , t ) sθ(x,t) sθ(x,t) 被训练好,我们有 p θ ≈ p 0 p_θ≈p_0 pθp0,在这种情况下 x(0) 可以看作近似是来自数据分布 p 0 p_0 p0的样本。

λ ( t ) = g 2 ( t ) λ(t)=g^2(t) λ(t)=g2(t)时,我们在 Fisher 散度的加权组合和 KL 散度之间发现了重要的联系对于 p 0 p_0 p0 p θ p_θ pθ 在某些规律性条件下:
在这里插入图片描述
由于与 KL 散度的这种特殊联系以及模型训练的最小化 KL 散度和最大化似然之间的等价性,我们称 λ ( t ) = g 2 ( t ) λ(t)=g^2(t) λ(t)=g2(t) 似然加权函数likelihood weighting function。
使用这种似然加权函数,我们可以训练score-based generative models来实现非常高的似然,与最先进的自回归模型相当甚至优于。

6.4How to solve the reverse SDE

通过用数值 SDE 求解器求解估计的反向 SDE,我们可以模拟样本生成的反向随机过程。或许最简单的数值 SDE 求解器是 Euler-Maruyama 方法。当应用于我们估计的反向 SDE 时,它使用有限的时间步长和小的高斯噪声对 SDE 进行离散化。具体来说,它选择一个小的负时间步长 Δt≈0,初始化 t←T,并迭代以下过程直到 t≈0:
在这里插入图片描述
这里 z t ∼ N ( 0 , I ) z_t∼\mathcal{N}(0,I) ztN(0,I)。 Euler-Maruyama 方法在性质上类似于 Langevin 动力学——两者都通过遵循受高斯噪声扰动的得分函数来更新 x。They both update by following score functions perturbed with Gaussian noise.

除了 Euler-Maruyama 方法外,其他数值 SDE 求解器也可以直接用于求解反向 SDE 以生成样本,例如包括 Milstein 方法和随机 Runge-Kutta 方法。
在【4】中,宋博士提供了一个类似于 Euler-Maruyama 的反向扩散求解器,但更适合求解reverse-time的 SDEs。最近,【5】作者引入了自适应步长 SDE 求解器,可以更快地生成质量更好的样本。

此外,我们的反向 SDE 有两个特殊属性,允许更灵活的采样方法:

  1. 我们估计 ∇ x l o g p t ( x ) \nabla_{\mathbf{x}}log p_t(\mathbf{x}) xlogpt(x)通过我们基于时间的基于分数的模型 s θ ( x , t ) s_θ(x,t) sθ(x,t)

  2. 我们只关心从每个边缘分布中抽样点 p t ( x ) p_t(x) pt(x)。在不同时间步长获得的样本可以具有任意相关性,并且不必形成从反向 SDE 采样的特定轨迹。
    在这里插入图片描述

由于这两个属性,我们可以应用 马尔科夫链蒙特卡洛方法MCMC 方法来微调从数值 SDE 求解器获得的轨迹trajectories 。
具体来说,宋博士提出预测-校正采样器Predictor-Corrector samplers
预测器predictor 可以是任何数值 SDE 求解器预测 x ( t + Δ t ) ∼ p t + Δ t ( x ) x(t+Δt)∼p_{t+Δt}(x) x(t+Δt)pt+Δt(x)来自现有样本 x ( t ) 〜 p t ( x ) x(t)〜p_t(x) x(t)pt(x)
校正器corrector可以是任何完全依赖于得分函数的 MCMC 程序,例如 Langevin 动力学和哈密顿蒙特卡洛。
在这里插入图片描述
在Predictor-Corrector采样器的每一步,我们首先使用预测器选择一个合适的步长Δt<0,然后根据当前样本x(t)预测x(t+Δt)。
接下来,我们根据基于分数的模型 s θ ( x , t + Δ t ) s_θ(x,t+Δt) sθ(x,t+Δt)运行几个校正步骤来改进样本 x(t+Δt),使 x ( t + Δ t ) x(t+Δt) x(t+Δt) 成为来自 p t + Δ t ( x ) p_{t+Δt}(x) pt+Δt(x)的高质量样本。

借助Predictor-Corrector方法和更好的基于分数的模型架构,我们可以在 CIFAR-10 上实现最先进的样本质量(以 FID 和 Inception 分数衡量),优于迄今为止最好的 GAN 模型(StyleGAN2 + ADA)
在这里插入图片描述
采样方法还可以针对极高维数据进行扩展。例如,它可以成功生成分辨率为 1024×1024 的高保真图像。例如一些常见的数据集:FFHQ,LSUN和Celeba-HQ
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述

6.5Probability flow ODE【ordinary differential equations (ODEs)】

尽管能够生成高质量样本,但基于 Langevin MCMC 和 SDE 求解器的采样器并未提供一种方法来计算基于分数的生成模型的精确对数似然exact log-likelihood。
下面,我们介绍了一个基于常微分方程 (ODE) 的采样器,它允许进行精确的似然计算。

在【4】中,我们证明 t 可以在不改变边缘分布 { p t ( x ) } t ∈ [ 0 , T ] \left \{p_t(x)\right \}_{t∈[0,T]} { pt(x)}t[0,T]的情况下将任何 SDE 转换为常微分方程 (ODE)。
因此,通过求解此 ODE,我们可以从与反向 SDE 相同的分布中采样。 SDE 的相应 ODE 被命名为概率流 ODEprobability flow ODE ,由下式给出
在这里插入图片描述
SDE公式是:
在这里插入图片描述
下图描绘了 SDEs 和概率流 ODEs 的轨迹trajectories 。尽管 ODE 轨迹明显比 SDE 轨迹更平滑,但它们将相同的数据分布转换为相同的先验分布,反之亦然,共享同一组边缘分布 { p t ( x ) } t ∈ [ 0 , T ] \left \{p_t(x)\right \}_{t∈[0,T]} { pt(x)}t[0,T]
换句话说,通过求解概率流 ODE 得到的轨迹与 SDE 轨迹具有相同的边缘分布。
在这里插入图片描述

我们可以使用 SDE 将数据映射到噪声分布(先验),并反转此 SDE 以进行生成建模。我们还可以反转关联的概率流 ODE,这会产生一个确定性过程,该过程从与 SDE 相同的分布中采样。逆时SDE和概率流ODE都可以通过估计得分函数得到。

这种概率流 ODE 公式具有几个独特的优势。

∇ x l o g p t ( x ) \nabla _xlogp_t(x) xlogpt(x)被替换为它的近似值 s θ ( x , t ) s_θ(x,t) sθ(x,t),概率流 ODE 成为神经 ODE(neural ODE)的特例。
特别地,它是连续归一化流continuous normalizing flows的示例,因为概率流 ODE 将数据分布 p 0 ( x ) p_0(x) p0(x) 转换到先验噪声分布 p T ( x ) p_T(x) pT(x)(因为它与 SDE 具有相同的边缘分布)并且是完全可逆的。

因此,概率流 ODE 继承了神经 ODE 或连续归一化流的所有属性,包括精确的对数似然计算。
具体来说,我们可以利用瞬时变量变化公式instantaneous change-of-variable formula( 【6】中的定理 1, 【7】中的等式(4))和数值 ODE 求解器来从自已知的先验密度 p T p_T pT 去计算未知数据密度 p 0 p_0 p0

事实上,我们的模型在均匀去量化uniformly dequantized的 CIFAR-10 图像上实现了最先进的对数似然,即使没有最大似然训练even without maximum likelihood training。
在这里插入图片描述
当使用我们之前讨论的似然加权likelihood weighting( λ ( t ) = g 2 ( t ) λ(t)=g^2(t) λ(t)=g2(t))训练基于分数的模型,并使用变分反量化variational dequantization来获得离散图像的似然时,我们可以获得与最先进的自回归模型相当甚至更好的似然(所有这些都没有任何数据增强) 。

在这里插入图片描述

6.6Controllable generation for inverse problem solving用于反问题求解的可控生成

基于分数的生成模型特别适合解决反问题。从本质上讲,逆问题与贝叶斯推理问题相同。令 x 和 y 为两个随机变量,假设我们知道从 x 生成 y 的正向过程,由转移概率分布 p(y|x) 表示。逆向问题是计算 p(x|y)。根据贝叶斯规则,我们有 p(x|y)=p(x)p(y|x)/∫p(x)p(y|x)dx。通过在两边取关于 x 的梯度,可以大大简化此表达式,从而得出以下得分函数的贝叶斯规则:
在这里插入图片描述

通过分数匹配,我们可以训练一个模型来估计无条件数据分布的分数函数,即, s θ ( x ) ≈ ∇ x l o g p ( x ) s_θ(x)≈∇_xlogp(x) sθ(x)xlogp(x)。这将使我们能够轻松计算后验得分函数 ∇ x l o g p ( x ∣ y ) ∇_xlogp(x|y) xlogp(xy) 从已知的前向过程 p ( y ∣ x ) p(y|x) p(yx) 通过上式,并使用朗之万型采样方法从中采样。

UT Austin【8】最近的一项工作表明,基于分数的生成模型可用于解决医学成像中的逆问题,例如加速磁共振成像 (MRI)。
同时在【9】中,我们展示了基于分数的生成模型的卓越性能,不仅在加速 MRI 上,而且在稀疏视图计算机断层扫描 (CT) 上。
我们能够实现与监督或展开的深度学习方法相当甚至更好的性能,同时在测试时对不同的测量过程更加稳健。

下面我们展示了一些解决计算机视觉逆问题的例子。
在这里插入图片描述

类条件生成使用的是无条件的基于时间的基于分数的模型,以及 CIFAR-10 上的预训练噪声条件图像分类器。

在这里插入图片描述

使用在 LSUN 卧室上训练的基于时间和分数的模型进行图像修复。最左边的列是基本事实。第二列显示蒙版图像(在我们的框架中为 y)。其余列显示不同的修复图像,通过求解条件逆时 SDE 生成。

在这里插入图片描述

使用在 LSUN church_outdoor 和 bedroom 上训练的基于时间的基于分数的模型进行图像着色。最左边的列是基本事实。第二列显示灰度图像(在我们的框架中为 y)。其余列显示不同的彩色图像,通过求解条件逆时 SDE 生成。

在这里插入图片描述

我们甚至可以使用在 FFHQ 上训练的基于时间的基于分数的模型为历史名人(亚伯拉罕林肯)的灰度肖像着色。图像分辨率为 1024 x 1024。

Reference

【1】https://yang-song.net/blog/2021/score/
【2】[生成模型新方向]: score-based generative models
【3】Generative Modeling by Estimating Gradients of the Data Distribution
【4】:Score-Based Generative Modeling through Stochastic Differential Equations
【5】:Gotta Go Fast When Generating Data with Score-Based Models
【6】:Neural Ordinary Differential Equations
【7】:Scalable Reversible Generative Models with Free-form Continuous Dynamics
【8】:Robust Compressed Sensing MRI with Deep Generative Priors
【9】:Solving Inverse Problems in Medical Imaging with Score-Based Generative Models

猜你喜欢

转载自blog.csdn.net/qq_45934285/article/details/129881669
今日推荐