论文链接:Score-Based Generative Modeling through Stochastic Differential Equations
文章目录
- 摘要
- 引文
- 背景
- 基于郎之动力学的去噪分数匹配/SMLD
- 去噪扩散概率模型/DDPM
- SDEs的基于分数的生成模型
- SDEs下的受扰动数据
- 逆转SDE生成样本
- 估计SDE的分数
- 例子:VE、VP SDE等
- 解决逆SDE
- 通用数值SDE求解器
- 预测-校正采样器
- Probability flow and connection to neural ODEs
- 精确似然计算
- 操纵潜在表征
- 唯一可识别编码
- 有效采样
- 架构改善
- 可控的生成
- 结论
摘要
从数据中制造噪音很容易;从噪声中创建数据是生成建模。本文提出了一个随机微分方程(SDE),通过缓慢注入噪声平滑地将复杂数据分布转换为已知的先验分布,以及一个相应的逆时序SDE,通过缓慢去除噪声将先验分布转换回数据分布。至关重要的是,逆时序SDE仅依赖于受扰数据分布的时变梯度场(即分数)。通过利用基于分数的生成建模的优势,可以用神经网络准确地估计这些分数,并使用数值SDE求解器生成样本。本文展示了这个框架在基于分数的生成建模和扩散概率建模中封装了以前的方法,允许新的抽样程序和新的建模能力。特别是,引入了一个预测器校正框架来纠正离散逆时序SDE演化中的错误。本文还推导出一个等效的神经ODE,它从与SDE相同的分布中采样,但也可以实现精确的似然计算并提高采样效率。此外,提供了一种基于分数的模型解决逆问题的新方法,通过类条件生成、图像嵌入和着色等实验证明了这一点。结合多个体系结构改进,在CIFAR-10上实现了无条件图像生成的破纪录性能,inception分数为9.89,FID为2.20,2.99 bits/dim的有竞争力的似然,并首次从基于分数的生成模型实现1024x1024图像的高保真生成。
引文
两类成功的概率生成模型涉及到用缓慢增加的噪声顺序破坏训练数据,然后学习逆转这种破坏以形成数据的生成模型。分数与朗格万动力学(SMLD)匹配估计每个噪声尺度上的分数(即对数概率密度相对于数据的梯度),然后使用朗格万动力学从生成过程中不断下降的噪声尺度序列中采样。去噪扩散概率建模(DDPM)训练一个概率模型序列来逆转噪声破坏的每一步,使用反向分布的函数形式的知识使训练变得可处理。对于连续状态空间,DDPM训练目标隐式计算每个噪声尺度上的分数。因此,本文将这两个模型类一起称为基于分数的生成模型。基于分数的生成模型和相关技术已被证明能有效生成图像、音频、和形状等。为了实现新的采样方法并进一步扩展基于分数的生成模型的功能,本文提出了一个统一的框架,通过随机微分方程(SDEs)的视角来概括以前的方法。
具体地说,不是用有限数量的噪声分布来扰动数据,而是考虑根据扩散过程随时间演变的分布的连续体。该过程将数据点逐步扩散为随机噪声,并由不依赖于数据且没有可训练参数的规定的SDE给出。通过反转这一过程,可以顺利地将随机噪声转换为生成样本的数据。至关重要的是,这个反向过程满足一个逆时序SDE,在给定边际概率密度分数作为时间函数的情况下,可以从正向SDE推导出该逆时SDE。因此,可以通过训练一个时间相关的神经网络来估计分数,从而近似逆时序SDE,然后使用数值SDE求解器生成样本。本文主要思想如图1所示。
图1 解决一个逆时序SDE生成一个基于分数的生成模型;利用连续时间SDE可以将数据转换为简单的噪声分布,如果知道在每个中间时间步上分布的分数
∇
x
log
p
t
(
x
)
\nabla_x\log p_t(x)
∇xlogpt(x),则可以反转这个SDE
本文提出的框架有以下几个理论和实践贡献:
- 灵活的采样和似然计算:可以使用任何通用SDE求解器来集成逆时序SDE进行采样。此外,提出了两种不适用于一般SDE的特殊方法:(1)将数值SDE求解器与基于分数的MCMC方法相结合的预测校正(PC)采样器,如Langevin MCMC和HMC;(2)基于概率流常微分方程(ODE)的确定性采样器。者对现有的基于分数模型的抽样方法进行了统一和改进。后者允许通过黑盒ODE求解器进行快速自适应采样,通过潜在代码进行灵活的数据操作,一个唯一可识别的编码,以及值得注意的精确似然计算。
- 可控的生成:可以通过对训练中不可用的信息进行条件反射来调节生成过程,因为条件逆时序SDE可以从无条件分数中有效地估计出来。这使得应用程序,如类别条件生成、图像修复、着色和其他逆向问题,都可以使用一个无条件的基于分数的模型,而无需重新训练。
- 统一的框架:本文框架提供了一种统一的方法来探索和调优各种SDEs,以改进基于分数的生成模型。SMLD和DDPM方法可以合并到本文框架中,作为两个独立的SDEs的离散化。尽管最近有报道称DDPM比SMLD获得了更高的样本质量,本文展示了本文框架允许的更好的架构和新的采样算法下,后者可以赶上——它在CIFAR-10上实现了最新的SOTA的Inception得分(9.89)和FID得分(2.20),以及第一次从基于分数的模型中高保真生成1024x1024图像。此外,在本文框架下提出了一个新的SDE,在统一去量化CIFAR-10图像上实现了2.99 bits/dim的似然值,创下了该任务的新记录。
背景
基于郎之动力学的去噪分数匹配/SMLD
设
p
σ
(
x
~
∣
x
)
:
=
N
(
x
~
;
x
,
σ
2
I
)
p_{\sigma}(\tilde{x}|x) := N(\tilde{x};x,\sigma^2I)
pσ(x~∣x):=N(x~;x,σ2I)是一个扰动核,并且有
p
σ
(
x
~
)
:
=
∫
p
d
a
t
a
(
x
)
p
σ
(
x
~
∣
x
)
d
x
p_{\sigma}(\tilde{x}) := \int p_{data}(x)p_{\sigma}(\tilde{x}|x)dx
pσ(x~):=∫pdata(x)pσ(x~∣x)dx,其中
p
d
a
t
a
(
x
)
p_{data}(x)
pdata(x)表示数据分布。考虑一系列正噪声尺度
σ
m
i
n
=
σ
1
<
σ
2
⋅
⋅
⋅
<
σ
N
=
σ
m
a
x
\sigma_{min}=\sigma_1<\sigma_2 \cdot \cdot \cdot <\sigma_N=\sigma_{max}
σmin=σ1<σ2⋅⋅⋅<σN=σmax。典型地,
σ
m
i
n
\sigma_{min}
σmin足够小,使得
p
σ
m
i
n
(
x
)
≈
p
d
a
t
a
(
x
)
p_{\sigma_{min}}(x) \approx p_{data}(x)
pσmin(x)≈pdata(x);并且
σ
m
a
x
\sigma_{max}
σmax足够大,使得
p
σ
m
a
x
(
x
)
≈
N
(
x
;
0
,
σ
m
a
x
2
I
)
p_{\sigma_{max}}(x) \approx N(x;0,\sigma_{max}^2I)
pσmax(x)≈N(x;0,σmax2I)。Generative Modeling by Estimating Gradients of the Data Distribution提出训练一个噪声条件分数网络(NCSN),用
S
θ
(
x
,
σ
)
S_θ(x,\sigma)
Sθ(x,σ) 表示,去噪分数匹配目标的加权和表示为:
θ
∗
=
a
r
g
m
i
n
θ
∑
i
=
1
N
σ
i
2
E
p
d
a
t
a
(
x
)
E
p
σ
i
(
x
~
∣
x
)
[
∣
∣
S
θ
(
x
~
,
σ
i
)
−
∇
x
~
log
p
σ
i
(
x
~
∣
x
)
∣
∣
2
2
]
.
(1)
\theta^*=\underset{\theta}{arg \,min}\sum_{i=1}^N\sigma_i^2E_{p_{data}(x)}E_{p_{\sigma_i}(\tilde{x}|x)}[||S_{\theta}(\tilde{x},\sigma_i)-\nabla_{\tilde{x}}\log p_{\sigma_i}(\tilde{x}|x)||^2_2].\tag{1}
θ∗=θargmini=1∑Nσi2Epdata(x)Epσi(x~∣x)[∣∣Sθ(x~,σi)−∇x~logpσi(x~∣x)∣∣22].(1)
给定足够的数据和模型容量,对于
σ
∈
{
σ
i
}
i
=
1
N
\sigma \in \lbrace\sigma_i\rbrace^N_{i=1}
σ∈{σi}i=1N,基于最佳分数的模型
S
θ
∗
(
x
,
σ
)
S_{θ^*}(x,σ)
Sθ∗(x,σ)几乎处与
∇
x
log
p
σ
(
x
)
\nabla_x\log p_{\sigma}(x)
∇xlogpσ(x)匹配。对于采样,运行M 步Langevin MCMC,按顺序为每个
p
σ
i
(
x
)
p_{σ_i}(x)
pσi(x)获得一个样本:
x
i
m
=
x
i
m
−
1
+
ϵ
i
S
θ
∗
(
x
i
m
−
1
,
σ
i
)
+
2
ϵ
i
z
i
m
,
m
=
1
,
2
,
⋅
⋅
⋅
,
M
,
(2)
x_i^m=x_i^{m-1}+\epsilon_iS_{\theta^*}(x_i^{m-1}, \sigma_i)+\sqrt{2\epsilon_i}z_i^m,\quad m=1,2,\cdot\cdot\cdot,M,\tag{2}
xim=xim−1+ϵiSθ∗(xim−1,σi)+2ϵizim,m=1,2,⋅⋅⋅,M,(2)
其中
ϵ
i
>
0
\epsilon_i>0
ϵi>0是步长,
z
i
m
z_i^m
zim符合标准正太分布。以
x
N
0
∼
N
(
x
∣
0
,
σ
m
a
x
2
I
)
x^0_N \sim N(x|0,\sigma^2_{max}I)
xN0∼N(x∣0,σmax2I)开始,上述计算步骤从
i
=
N
,
N
−
1
,
⋅
⋅
⋅
,
1
i=N,N-1,\cdot\cdot\cdot,1
i=N,N−1,⋅⋅⋅,1重复N次;当
i
<
N
i<N
i<N时,
x
i
0
=
x
i
+
1
m
x^0_i=x^m_{i+1}
xi0=xi+1m(本轮计算的初始值等于上一轮的最后值)。当对于所有
i
i
i满足
M
→
∞
,
ϵ
i
→
0
M \to \infty,\,\epsilon_i \to 0
M→∞,ϵi→0,
x
1
M
x_1^M
x1M在某些规律性条件下成为
p
σ
m
i
n
(
x
)
≈
p
d
a
t
a
(
x
)
p_{\sigma_{min}}(x) \approx p_{data}(x)
pσmin(x)≈pdata(x)分布的精确样本。
去噪扩散概率模型/DDPM
DDPM考虑一个正噪声尺度序列:
0
<
β
1
,
β
2
,
⋅
⋅
⋅
,
β
N
<
1
0<\beta_1,\beta_2,\cdot\cdot\cdot,\beta_N<1
0<β1,β2,⋅⋅⋅,βN<1。对于每个训练数据点
x
0
∼
p
d
a
t
a
(
x
)
x_0 \sim p_{data}(x)
x0∼pdata(x),基于
p
(
x
i
∣
x
i
−
1
)
=
N
(
x
i
;
1
−
β
i
x
i
−
1
,
β
i
I
)
p(x_i|x_{i-1})=N(x_i;\sqrt{1-\beta_i}x_{i-1},\beta_iI)
p(xi∣xi−1)=N(xi;1−βixi−1,βiI)可构造一个离散的马尔科夫链
{
x
0
,
x
1
,
⋅
⋅
⋅
,
x
N
}
\lbrace x_0,x_1,\cdot\cdot\cdot,x_N\rbrace
{x0,x1,⋅⋅⋅,xN},进而推导出
p
(
x
i
∣
x
0
)
=
N
(
x
i
;
α
i
x
0
,
(
1
−
α
i
)
I
)
p(x_i|x_0)=N(x_i;\sqrt{\alpha_i}x_0,(1-\alpha_i)I)
p(xi∣x0)=N(xi;αix0,(1−αi)I),其中
α
i
:
=
∏
j
=
1
i
(
1
−
β
j
)
\alpha_i := \prod^i_{j=1}(1-\beta_j)
αi:=∏j=1i(1−βj)。
与SMLD类似,可以将扰动数据分布表示为
p
α
i
(
x
~
)
:
=
∫
p
d
a
t
a
(
x
)
p
α
i
(
x
~
∣
x
)
d
x
p_{\alpha_i}(\tilde{x}) := \int p_{data}(x)p_{\alpha_i}(\tilde{x}|x)dx
pαi(x~):=∫pdata(x)pαi(x~∣x)dx。规定噪声尺度,使得
x
N
x_N
xN近似符合分布
N
(
0
,
I
)
N(0, I)
N(0,I)。逆向的变分马尔可夫链由
p
θ
(
x
i
−
1
∣
x
i
)
=
N
(
x
i
−
1
;
1
1
−
β
i
(
x
i
+
β
i
S
θ
(
x
i
,
i
)
)
,
β
i
I
)
p_{\theta}(x_{i-1}|x_i)=N(x_{i-1};\frac{1}{\sqrt{1-\beta_i}}(x_i+\beta_iS_{\theta}(x_i,i)),\beta_iI)
pθ(xi−1∣xi)=N(xi−1;1−βi1(xi+βiSθ(xi,i)),βiI)参数化,并用证据下限的重新加权变体(ELBO)进行训练:
θ
∗
=
a
r
g
m
i
n
θ
∑
i
=
1
N
(
1
−
α
i
)
E
p
d
a
t
a
(
x
)
E
p
α
i
(
x
~
∣
x
)
[
∣
∣
S
θ
(
x
~
,
i
)
−
∇
x
~
log
p
α
i
(
x
~
∣
x
)
∣
∣
2
2
]
.
(3)
\theta^*=\underset{\theta}{arg \, min}\sum^N_{i=1}(1-\alpha_i)E_{p_{data}(x)}E_{p_{\alpha_i}(\tilde{x}|x)}[||S_{\theta}(\tilde{x},i)-\nabla_{\tilde{x}}\log p_{\alpha_i}(\tilde{x}|x)||^2_2].\tag{3}
θ∗=θargmini=1∑N(1−αi)Epdata(x)Epαi(x~∣x)[∣∣Sθ(x~,i)−∇x~logpαi(x~∣x)∣∣22].(3)
在求解公式3得到最优模型
S
θ
∗
(
x
,
i
)
S_{\theta^*}(x,i)
Sθ∗(x,i)后,从
x
N
∼
N
(
0
,
I
)
x_N \sim N(0,I)
xN∼N(0,I)开始,按照估计的反向马尔可夫链生成样本如下:
x
i
−
1
=
1
1
−
β
i
(
x
i
+
β
i
S
θ
∗
(
x
i
,
i
)
)
+
β
i
z
i
,
i
=
N
,
N
−
1
,
⋅
⋅
⋅
,
1.
(4)
x_{i-1}=\frac{1}{\sqrt{1-\beta_i}}(x_i+\beta_iS_{\theta^*}(x_i,i))+\sqrt{\beta_i}z_i,\quad i=N,N-1,\cdot\cdot\cdot,1.\tag{4}
xi−1=1−βi1(xi+βiSθ∗(xi,i))+βizi,i=N,N−1,⋅⋅⋅,1.(4)
将此方法称为祖先采样,因为它相当于从图形模型
∏
i
=
1
N
p
θ
(
x
i
−
1
∣
x
)
\prod^N_{i=1}p_{\theta}(x_{i-1}|x)
∏i=1Npθ(xi−1∣x)中执行祖先采样。在DDPM论文中公式3记为
L
s
i
m
p
l
e
L_{simple}
Lsimple,整体构造与公式1很相似。与公式1相似,公式3也是去噪分数匹配目标的加权和,这意味着最优模型
S
θ
∗
(
x
,
i
)
S_{\theta^*}(x,i)
Sθ∗(x,i)与扰动数据分布
∇
x
log
p
α
i
(
x
)
\nabla_{x}\log p_{\alpha_i}(x)
∇xlogpαi(x)的得分相匹配。值得注意的是,公式1和公式3中第
i
i
i个和的权重
σ
i
2
σ^2_i
σi2和
(
1
−
α
i
)
(1-\alpha_i)
(1−αi)与相应的扰动核具有相同的泛函形式:
σ
i
2
∝
1
/
E
[
∣
∣
∇
x
log
p
σ
i
(
x
~
∣
x
)
∣
∣
2
2
]
;
(
1
−
α
i
)
∝
1
/
E
[
∣
∣
∇
x
log
p
α
i
(
x
~
∣
x
)
∣
∣
2
2
]
σ^2_i \propto 1/E[||\nabla_x \log p_{\sigma_i}(\tilde{x}|x)||^2_2]; \quad (1-\alpha_i) \propto 1/ E[||\nabla_x \log p_{\alpha_i}(\tilde{x}|x)||^2_2]
σi2∝1/E[∣∣∇xlogpσi(x~∣x)∣∣22];(1−αi)∝1/E[∣∣∇xlogpαi(x~∣x)∣∣22]
SDEs的基于分数的生成模型
对具有多个噪声尺度的数据进行扰动是以前方法成功的关键。将这一想法进一步推广到无限多个噪声尺度,以便随着噪声的加剧,扰动数据分布根据SDE而演变。图2中给出了本文框架的概述。
图2 通过SDEs的基于分数的生成建模概述。可以使用SDE将数据映射到噪声分布(先验),并将该SDE反向用于生成建模;还可以反向关联概率流ODE,它产生一个确定性过程,从与SDE相同的分布中取样;通过对分数
∇
x
log
p
t
(
x
)
\nabla_x \log p_t(x)
∇xlogpt(x)的估计,可以得到逆时SDE和概率流ODE
SDEs下的受扰动数据
目标是构建一个由连续时间变量
t
∈
[
0
,
T
]
t \in [0,T]
t∈[0,T]索引的扩散过程
{
x
(
t
)
}
t
=
0
T
\lbrace x(t) \rbrace^T_{t=0}
{x(t)}t=0T;例如
x
(
0
)
∼
p
0
x(0) \sim p_0
x(0)∼p0,一个样本独立同分布的数据集和
x
(
T
)
∼
p
T
x(T) \sim p_T
x(T)∼pT,一个易于处理的形式来有效地生成样本。换句话说,
p
0
p_0
p0是数据分布,
p
T
p_T
pT是先验分布。该扩散过程可以建模为一个SDE的解:
d
x
=
f
(
x
,
t
)
d
t
+
g
(
t
)
d
w
,
(5)
dx = f(x,t)dt+g(t)dw,\tag{5}
dx=f(x,t)dt+g(t)dw,(5)
其中
w
w
w是标准维纳过程(又称布朗运动),
f
(
⋅
,
t
)
:
R
D
→
R
D
f(\cdot,t):R^D \to R^D
f(⋅,t):RD→RD是一个向量值函数,称为
x
(
t
)
x(t)
x(t)的漂移系数,
g
(
⋅
)
:
R
D
→
R
D
g(\cdot):R^D \to R^D
g(⋅):RD→RD是一个标量函数,称为
x
(
t
)
x(t)
x(t)的扩散系数。为了便于表示,假设扩散系数是一个标量(而不是
d
∗
d
d * d
d∗d矩阵),并且不依赖于
x
x
x,但本文理论可以推广到这些情况(见附录A)。只要系数在状态和时间上都是全局Lipschitz的,SDE就有唯一的强解。在下文中用
p
t
(
x
)
p_t(x)
pt(x)表示
x
(
t
)
x(t)
x(t)的概率密度,并使用
p
s
t
(
x
(
t
)
∣
x
(
s
)
)
p_{st}(x(t)|x(s))
pst(x(t)∣x(s))表示从
x
(
s
)
x(s)
x(s)到
x
(
t
)
x(t)
x(t)的转换核,其中
0
≤
s
<
t
≤
T
0 ≤ s < t ≤ T
0≤s<t≤T。
通常, p T p_T pT是一个不包含 p 0 p_0 p0信息的非结构化先验分布,例如具有固定均值和方差的高斯分布。在公式5中有多种设计SDE的方法,使其将数据分布扩散为一个固定的先验分布。将在后文中提供来自SMLD和DDPM的连续泛化的几个示例。
逆转SDE生成样本
从
x
(
T
)
∼
p
T
x(T) \sim p_T
x(T)∼pT的样本开始,逆转SDE过程,能获得样本
x
(
0
)
∼
p
0
x(0) \sim p_0
x(0)∼p0。Reverse-time diffusion equation models中的一个显著结果表明,扩散过程的反向也是一个扩散过程,它在时间上反向运行,并由逆时间SDE给出:
d
x
=
[
f
(
x
,
t
)
−
g
(
t
)
2
∇
x
log
p
t
(
x
)
]
d
t
+
g
(
t
)
d
w
ˉ
,
(6)
dx=[f(x,t)-g(t)^2\nabla_x \log p_t(x)]dt+g(t)d\bar{w},\tag{6}
dx=[f(x,t)−g(t)2∇xlogpt(x)]dt+g(t)dwˉ,(6)
其中
w
ˉ
\bar{w}
wˉ为时间从
T
T
T向后流向
0
0
0时的标准维纳过程,
d
t
dt
dt为无穷小的负时间步长。一旦每个边际分布的分数
∇
x
log
p
t
(
x
)
∇_x \log p_t(x)
∇xlogpt(x)对于所有
t
t
t已知,可以公式6推导出逆向扩散过程,并模拟它从
p
0
p_0
p0中取样。
估计SDE的分数
分布的分数可以通过在分数匹配的样本上训练基于分数的模型来估计。为了估计
∇
x
log
p
t
(
x
)
∇_x \log p_t(x)
∇xlogpt(x),可以通过对公式1和3的连续泛化来训练基于时间的分数模型
S
θ
(
x
,
t
)
S_{\theta}(x,t)
Sθ(x,t):
θ
∗
=
a
r
g
m
i
n
θ
E
t
{
λ
(
t
)
E
x
(
0
)
E
x
(
t
)
∣
x
(
0
)
[
∣
∣
S
θ
(
x
(
t
)
,
t
)
−
∇
x
(
t
)
log
p
0
t
(
x
(
t
)
∣
x
(
0
)
)
∣
∣
2
2
]
}
.
(7)
\theta^*=\underset{\theta}{arg \, min}E_t\lbrace \lambda(t) E_{x(0)}E_{x(t)|x(0)}[||S_{\theta}(x(t),t)-\nabla_{x(t)}\log p_{0t}(x(t)|x(0))||^2_2]\rbrace.\tag{7}
θ∗=θargminEt{λ(t)Ex(0)Ex(t)∣x(0)[∣∣Sθ(x(t),t)−∇x(t)logp0t(x(t)∣x(0))∣∣22]}.(7)
此处
λ
:
[
0
,
T
]
→
R
>
0
\lambda:[0,T] \to R_{>0}
λ:[0,T]→R>0是一个正加权函数,
t
t
t是从
[
0
,
T
]
[0, T]
[0,T]中均匀采样,
x
(
0
)
∼
p
0
,
x
(
t
)
∼
p
0
t
(
(
x
(
t
)
∣
x
(
0
)
)
x(0) \sim p_0, \quad x(t) \sim p_{0t}((x(t)|x(0))
x(0)∼p0,x(t)∼p0t((x(t)∣x(0))。在有足够的数据和模型容量的情况下,分数匹配可以确保对于几乎所有的
x
x
x和
t
t
t,公式7的最优解
S
θ
∗
(
x
,
t
)
S_{\theta^*}(x,t)
Sθ∗(x,t)等于
∇
x
log
p
t
(
x
)
∇_x \log p_t(x)
∇xlogpt(x)。在SMLD和DDPM中,通常可以选择
λ
∝
1
/
E
[
∣
∣
∇
x
(
t
)
log
p
0
t
(
x
(
t
)
∣
x
(
0
)
)
∣
∣
2
2
]
λ \propto 1/E[||\nabla_{x(t)}\log p_{0t}(x(t)|x(0))||^2_2]
λ∝1/E[∣∣∇x(t)logp0t(x(t)∣x(0))∣∣22]。请注意,公式7使用去噪分数匹配,但其他分数匹配目标,如切片分数匹配和有限差分数匹配也适用于这里。
需要知道转换核 p 0 t ( x ( t ) ∣ x ( 0 ) ) p_{0t}(x(t)|x(0)) p0t(x(t)∣x(0))来有效地求解公式7。当 f ( ⋅ , t ) f(\cdot,t) f(⋅,t)是仿射时,转换核总是一个高斯分布,其中均值和方差通常以封闭形式已知,可以用标准技术获得。对于更一般的SDEs,可以解Kolmogorov的正向方程得到 p 0 t ( x ( t ) ∣ x ( 0 ) ) p_{0t}(x(t)|x(0)) p0t(x(t)∣x(0))。或者,可以模拟SDE从 p 0 t ( x ( t ) ∣ x ( 0 ) ) p_{0t}(x(t)|x(0)) p0t(x(t)∣x(0))中采样,将公式7中的去噪得分匹配替换为切片得分匹配进行模型训练,这样就绕过了 ∇ x ( t ) log p 0 t ( x ( t ) ∣ x ( 0 ) ) \nabla_{x(t)}\log p_{0t}(x(t)|x(0)) ∇x(t)logp0t(x(t)∣x(0))的计算(见附录A)。
例子:VE、VP SDE等
SMLD和DDPM中使用的噪声扰动可以看作是两个不同sde的离散化。下面为一个简短的讨论,更多细节见附录B。
当使用总共
N
N
N个噪声尺度时,SMLD的每个扰动核
p
σ
i
(
x
∣
x
0
)
p_{\sigma_i}(x|x_0)
pσi(x∣x0)对应于以下马尔可夫链中
x
i
x_i
xi的分布:
x
i
=
x
i
−
1
+
σ
i
2
−
σ
i
−
1
2
z
i
−
1
,
i
=
1
,
⋅
⋅
⋅
,
N
,
(8)
x_i=x_{i-1}+\sqrt{\sigma_i^2-\sigma_{i-1}^2}z_{i-1}, \quad i=1,\cdot\cdot\cdot,N,\tag{8}
xi=xi−1+σi2−σi−12zi−1,i=1,⋅⋅⋅,N,(8)
其中
z
i
−
1
∼
N
(
0
,
1
)
z_{i-1} \sim N(0,1)
zi−1∼N(0,1),设置
σ
0
=
0
\sigma_0=0
σ0=0简化符号。在
N
→
∞
N \to \infty
N→∞的限制的,
{
σ
i
}
i
=
1
N
\lbrace\sigma_i\rbrace^N_{i=1}
{σi}i=1N变为一个函数
σ
(
t
)
\sigma(t)
σ(t),
z
i
z_i
zi变为
z
(
t
)
z(t)
z(t),并且马尔科夫链
{
x
i
}
i
=
1
N
\lbrace x_i \rbrace^N_{i=1}
{xi}i=1N变为一个连续随机过程
{
x
(
t
)
}
t
=
0
1
\lbrace x(t)\rbrace^1_{t=0}
{x(t)}t=01,其中使用连续时间变量
t
∈
[
0
,
1
]
t \in [0, 1]
t∈[0,1]进行索引。而不是整数
i
i
i。
{
x
(
t
)
}
t
=
0
1
\lbrace x(t)\rbrace^1_{t=0}
{x(t)}t=01过程由以下SDE给出
d
x
=
d
[
σ
2
(
t
)
]
d
t
d
w
.
(9)
dx=\sqrt{\frac{d[\sigma^2(t)]}{dt}}dw.\tag{9}
dx=dtd[σ2(t)]dw.(9)
同样,对于DDPM的扰动核
{
p
α
i
(
x
∣
x
0
)
}
i
−
1
N
\lbrace p_{\alpha_i}(x|x_0) \rbrace^N_{i-1}
{pαi(x∣x0)}i−1N,离散马尔科夫链为
x
i
=
1
−
β
i
x
i
−
1
+
β
i
z
i
−
1
,
i
=
1
,
⋅
⋅
⋅
,
N
.
(10)
x_i=\sqrt{1-\beta_i}x_{i-1}+\sqrt{\beta_i}z_{i-1}, \quad i=1,\cdot\cdot\cdot,N.\tag{10}
xi=1−βixi−1+βizi−1,i=1,⋅⋅⋅,N.(10)
当
N
→
∞
N \to \infty
N→∞,公式10收敛于以下SDE,
d
x
=
−
1
2
β
(
t
)
x
d
t
+
β
(
t
)
d
w
.
(11)
dx=-\frac{1}{2}\beta(t)x\,dt+\sqrt{\beta(t)}dw.\tag{11}
dx=−21β(t)xdt+β(t)dw.(11)
因此,在SMLD和DDPM中使用的噪声扰动对应于SDE公式9和11的离散化。有趣的是,公式9的SDE在 t → ∞ t \to \infty t→∞时总是表现为一个方差爆炸的过程,而当初始分布具有单位方差时,公式11的SDE产生了一个具有固定方差的过程(附录B证明)。由于此差异,将公式9称为“方差爆炸”(Variance Exploding,VE)SDE,公式11称为“方差保持”(Variance preserving,VP)SDE。
受VP SDE的启发,本文提出了一种新的SDEs类型,它在可能性上的表现特别好,由下式给出
d
x
=
−
1
2
β
(
t
)
x
d
t
+
β
(
t
)
(
1
−
e
−
2
∫
0
t
β
(
s
)
d
s
)
d
w
.
(12)
dx=-\frac{1}{2}\beta(t)x\,dt+\sqrt{\beta(t)(1-e^{-2\int_0^t\beta(s)ds})}dw.\tag{12}
dx=−21β(t)xdt+β(t)(1−e−2∫0tβ(s)ds)dw.(12)
当使用相同的 β ( t ) \beta(t) β(t)并从相同的初始分布开始时,由公式12所诱导的随机过程的方差在每个中间时间步总是受VP SDE的限制(附录B证明)。因此,本文将公式12称为sub-VP SDE。
由于VE、VP和sub-VP SDE都具有仿射漂移系数,因此它们的摄动核 p 0 t ( x ( t ) ∣ x ( 0 ) ) p_{0t}(x(t)|x(0)) p0t(x(t)∣x(0))都符合高斯分布,可以用闭合形式计算。这使得用公式7进行训练特别有效。
解决逆SDE
训练一个基于时间的得分模型 S θ S_{\theta} Sθ后,可以用它构造逆时SDE,然后用数值方法模拟它,从 p 0 p_0 p0生成样本。
通用数值SDE求解器
数值求解器提供了SDE的近似轨迹。求解SDE存在许多通用数值方法,如Euler-Maruyama和随机Runge-Kutta方法,它们对应于随机动力学的不同离散化。可以将其中任何一个应用于反向时间 SDE 以生成样本。
祖先采样(ancestral sampling),DDPM(公式4)的采样方法实际上对应于逆向VP SDE的一种特殊离散化(公式11)(见附录E)。然而,推导出新SDE的祖先采样规则可能并非易事。为了解决这个问题,本文提出了逆向扩散采样器
r
e
v
e
r
s
e
d
i
f
f
u
s
i
o
n
s
a
m
p
l
e
r
s
reverse\,diffusion\,samplers
reversediffusionsamplers(详见附录 E),它以与前向 SDE 相同的方式离散化逆时 SDE,因此可以在给定前向离散化的情况下轻松导出。如表1所示,在CIFAR-10上,逆向扩散采样器在SMLD和DDPM模型上的性能都略好于祖先采样(DDPM-祖先采样也适用于SMLD模型,见附录F)。
表1 CIFAR-10上不同逆时SDE求解器的比较。阴影区域用同样的计算(分数函数评估的数量)得到。均值和标准偏差超过五次抽样运行。“P1000”或“P2000”:仅使用1000或2000步的预测器采样器。C2000”:仅使用2000步的校正采样器。“PC1000”:使用1000步预测器和1000步校正器步骤的预测校正器(PC)采样器。
预测-校正采样器
与一般的SDEs不同,拥有可用于改进解决方案的附加信息。由于有一个基于分数的模型 S θ ∗ ( x , t ) ≈ ∇ x log p t ( x ) S_{\theta^*}(x,t) \approx \nabla_x \log p_t(x) Sθ∗(x,t)≈∇xlogpt(x),可以使用基于分数的MCMC方法,如Langevin MCMC或HMC直接从 p t p_t pt中采样,并修正数值SDE求解器的解。
具体来说,在每个时间步中,数值SDE求解器首先给出下一个时间步样本的估计,起到“预测器”的作用。然后,基于分数的MCMC方法对估计样本的边际分布进行校正,起到“校正器”的作用。这个想法类似于Predictor-Corrector方法,这是一种用于解决方程组的数值延续技术,本文也将提出的混合采样算法命名为Predictor-Corrector (PC)采样器。在附录G中可找到伪代码和完整的描述。PC采样器推广了SMLD和DDPM的原始采样方法:前者使用恒等函数作为预测器,退火朗之万动力学作为校正器;后者使用祖先采样作为预测器,恒等函数作为校正器。
在SMLD和DDPM模型(参见附录 G 中的算法 2 和 3)上测试PC采样器,该模型使用方程式给出的原始离散目标公式1和3进行训练,展示了PC采样器与用固定数量的噪声尺度训练的基于分数的模型的兼容性。在表1中总结了不同采样器的性能,其中Probability flow是后文论述的一个预测器。详细的实验设置和其他结果在附录G中给出。观察到,本文的逆向扩散采样器总是优于祖先采样,并且只有校正器的方法(C2000)在相同的计算下比其他竞争对手(P2000, PC1000)表现更差(事实上,每个噪声尺度需要更多的校正器步骤,因此需要更多的计算,以匹配其他采样器的性能)。对于所有预测器,为每个预测器步骤(PC1000)添加一个校正步骤,可使计算量加倍,但始终提高样本质量(相对于P1000)。此外,它通常比在不添加校正器的情况下将预测器步骤数增加一倍(P2000)要好,在P2000中,必须以一种特别的方式在SMLD/DDPM模型的噪声尺度之间进行插值(详见附录G)。在附录G中,对连续目标公式7训练的模型进行了定性比较,该模型在尺寸为256x256的LSUN图像和VE SDE上进行了训练,当使用适当数量的校正步骤时,PC采样器在可比计算下明显优于仅预测器采样器。
Probability flow and connection to neural ODEs
基于分数的模型为求解逆时SDE提供了另一种数值方法。对于所有扩散过程,存在一个相应的确定性过程,其轨迹与SDE具有相同的边际概率密度
{
p
t
(
x
)
}
t
=
0
T
\lbrace p_t(x) \rbrace^T_{t=0}
{pt(x)}t=0T。这个确定性过程满足ODE(详见附录D.1):
d
x
=
[
f
(
x
,
t
)
−
1
2
g
(
t
)
2
∇
x
log
p
t
(
x
)
]
d
t
,
(13)
dx=[f(x,t)-\frac{1}{2}g(t)^2\nabla_x \log p_t(x)]dt,\tag{13}
dx=[f(x,t)−21g(t)2∇xlogpt(x)]dt,(13)
一旦分数已知,就可以从SDE中确定。将式(13)中的ODE命名为概率流ODE。当分数函数被时间依赖的基于分数的模型(通常是神经网络)近似时,这是一个神经ODE的例子。
精确似然计算
利用与神经ODEs的连接,可以通过瞬时变量变化公式计算公式13定义的密度。这允许计算任何输入数据的确切可能性(详见附录D.2)。作为一个例子,表2表示CIFAR-10数据集上以bits/dim为单位测量的负对数可能性(NLLs)。在统一去量化的数据上计算对数似然,并且只与以相同方式评估的模型进行比较(省略用变分去量化评估的模型离散数据),DDPM(
L
/
L
s
i
m
p
l
e
L/L_{simple}
L/Lsimple)除外,其ELBO值(带*注释)是在离散数据上计算的。主要结果:(i)对于与原始论文相同的DDPM模型,本文获得了比ELBO更好的bits/dim,因为可能性是准确的;(ii)使用相同的架构,公式7中的连续目标(即DDPM cont.)训练另一个DDPM模型,这进一步提高了似能;(iii) 使用sub-VP SDE,与VP SDE相比,总是获得更高的可能性; (iv) 通过改进的架构(即 DDPM++ cont.)和sub-VP SDE,即使在没有最大似然训练的情况下,也可以在均匀去量化的CIFAR-10上2.99 bits/dim的新记录。
操纵潜在表征
通过对公式13进行积分,可以将任何数据点 x ( 0 ) x(0) x(0)编码到潜在空间 x ( T ) x(T) x(T)。通过对逆时SDE的相应ODE进行积分,可以实现解码。与其他可逆模型(如神经ODE和归一化流)一样,可以操纵这种潜在表征进行图像编辑,如插值和温度缩放(见图3和附录D.4)。
唯一可识别编码
与大多数当前的可逆模型不同,本文的编码是唯一可识别的,这意味着在有足够的训练数据、模型容量和优化精度的情况下,输入的编码是由数据分布唯一确定的。这是因为前向SDE(公式5),没有可训练的参数,其相关的概率流ODE(公式13),在给定完美估计的分数的情况下提供相同的轨迹。在附录 D.5中提供了对此属性的额外实证验证。
有效采样
使用神经ODEs,可以通过从不同的最终条件
x
(
T
)
∼
p
T
x(T) \sim p_T
x(T)∼pT解公式13采样
x
(
0
)
∼
p
0
x(0) \sim p_0
x(0)∼p0。使用固定的离散化策略,可以生成有竞争力的样本,特别是在与校正器结合使用时(表1中“Probability flow”,详细信息见附录D.3)。使用黑盒ODE求解器不仅可以生成高质量的样本(表2,详见附录D.4),而且还允许显式地为效率权衡准确性。通过更大的容错,可以在不影响样品视觉质量的情况下减少90%以上的功能评估次数(图3)。
架构改善
本文探索了使用VE和VP SDEs的基于分数的模型的几个新的体系结构设计(详情见附录H),其中用SMLD/DDPM中相同的离散目标训练模型。由于VP SDEs的相似性,直接将VP SDEs的架构转移到sub-VP SDEs。本文对VE SDE的优化架构,命名为NCSN++,在CIFAR-10和PC采样器上实现了2.45的FID,而对VP SDE的优化架构,命名为DDPM++,实现了2.78的FID。
通过切换到公式7中的连续训练目标,并增加网络深度,可以进一步提高所有模型的样本质量。所得到的VE和VP/sub-VP SDEs的体系结构分别表示为表3中的NCSN++ cont.和DDPM++ cont.。表3中报告的结果是训练过程中FID最小的检查点,其中样本是用PC采样器生成的。相比之下,表2中报告了最后一个训练检查点的FID分数和NLL值,并且使用黑箱ODE求解器获得样本。如表3所示,VE SDE通常比VP/sub-VP SDE提供更好的样本质量,但经验地观察到,它们的似然比VP/sub-VP SDE更差。这表明从业者可能需要为不同的领域和体系结构试验不同的SDEs。
最好的样本质量模型,NCSN++ cont. (deep, VE),使网络深度翻了一番,并为CIFAR-10无条件生成的inception分数和FID创造了新记录。令人惊讶的是,可以实现比之前最好的条件生成模型更好的FID,而不需要标记数据。通过所有的改进,还从基于分数的模型中获得了CelebA-HQ 1024x1024数据集上的第一组高保真生成样本(见附录H.3)。最佳似然模型DDPM++ cont. (deep, sub-VP)同样将网络深度翻了一番,并在公式7中的连续目标下实现了2.99 bits/dim的对数似然。据本文所知,这是统一去量化CIFAR-10上的最高似然。
可控的生成
本文框架的连续结构允许不仅从
p
0
p_0
p0生成数据样本,而且在
p
t
(
y
∣
x
(
t
)
)
p_t(y|x(t))
pt(y∣x(t))已知的情况下,从
p
0
(
x
(
0
)
∣
y
)
p_0(x(0)|y)
p0(x(0)∣y)生成数据样本。给定前向SDE如公式5,可以从
p
T
(
x
(
T
)
∣
y
)
p_T(x(T)|y)
pT(x(T)∣y)开始,从
p
t
(
x
(
t
)
∣
y
)
p_t(x(t)|y)
pt(x(t)∣y)中取样,并求解一个条件逆时SDE:
d
x
=
{
f
(
x
,
t
)
−
g
(
t
)
2
[
∇
x
log
p
t
(
x
)
+
∇
x
log
p
t
(
y
∣
x
)
]
}
d
t
+
g
(
t
)
d
w
ˉ
.
(14)
dx=\lbrace f(x,t)-g(t)^2[\nabla_x \log p_t(x)+\nabla_x \log p_t(y|x)] \rbrace dt+g(t)d\bar{w}.\tag{14}
dx={f(x,t)−g(t)2[∇xlogpt(x)+∇xlogpt(y∣x)]}dt+g(t)dwˉ.(14)
一般来说,可以使用公式14用基于分数的生成模型来解决大量的逆问题,一旦给出正向过程的梯度估计,即 ∇ x log p t ( y ∣ x ( t ) ) \nabla_x \log p_t(y|x(t)) ∇xlogpt(y∣x(t))。在某些情况下,可以训练一个单独的模型来学习正向过程 log p t ( y ∣ x ( t ) ) \log p_t(y|x(t)) logpt(y∣x(t))并计算其梯度。否则,可以用启发式和领域知识估计梯度。在附录I.4中,提供了一个广泛适用的方法来获得这样的估计,而不需要训练辅助模型。
用这种方法考虑了可控生成的三个应用:类条件生成、图像修复和着色。当
y
y
y表示类标签时,可以训练一个依赖于时间的分类器
p
t
(
y
∣
x
(
t
)
)
p_t(y|x(t))
pt(y∣x(t))进行类条件采样。由于正向SDE易于处理,可以很容易地为时间依赖的分类器创建训练数据
(
x
(
t
)
,
y
)
(x(t),y)
(x(t),y),首先从数据集中采样
(
x
(
0
)
,
y
)
(x(0),y)
(x(0),y),然后采样
x
(
t
)
∼
p
0
t
(
x
(
t
)
∣
x
(
0
)
)
x(t) \sim p_{0t}(x(t)|x(0))
x(t)∼p0t(x(t)∣x(0))。然后,可以使用不同时间步长的交叉熵损失的混合,如公式7,来训练时间依赖分类器
p
t
(
y
∣
x
(
t
)
)
p_t(y|x(t))
pt(y∣x(t))。在图4(左)中提供了类条件CIFAR-10样本,更多细节和结果见附录I。
“修复”是条件抽样的一种特殊情况。假设有一个不完整的数据点 y y y,其中只有部分像素点即子集 Ω ( y ) Ω(y) Ω(y)是已知的。“修复”相当于从 p ( x ( 0 ) ∣ Ω ( y ) ) p(x(0)|Ω(y)) p(x(0)∣Ω(y))中采样,可以使用无条件模型来完成(见附录I.2)。除了已知的数据维度是耦合的之外,着色是一种特殊的“修复”。可以用正交线性变换将这些数据维解耦,并在变换后的空间中进行“修复”(详见附录I.3)。图4(右)显示了使用无条件时间依赖的基于分数的模型进行“修复”和着色的结果。
结论
本文提出了一个基于SDEs的基于分数的生成建模框架。本工作使读者能够更好地理解现有的方法、新的采样算法、精确的似然计算、唯一可识别的编码、潜在的代码操作,并为基于分数的生成模型家族带来新的条件生成能力。
虽然本文提出的采样方法改善了结果,并实现了更有效的采样,但在相同的数据集上,它们的采样速度仍然比GANs慢。确定将基于分数的生成模型的稳定学习与GANs等隐式模型的快速采样相结合的方法仍然是一个重要的研究方向。此外,当获得对分数函数的访问权时,可以使用的采样器的广度引入了许多超参数。未来的工作将受益于自动选择和调优这些超参数的改进方法,以及对各种采样器的优点和局限性进行更广泛的研究。