在自然图像等复杂数据集中,真实数据往往集中分布在一个低维流形上,概率密度函数的梯度(即得分函数)难以定义与估计。为缓解该问题,SMLD 提出使用不同强度的高斯噪声对数据进行扰动,扰动后的数据不再集中于低维流形,从而提升学习鲁棒性。
多尺度噪声扰动
SMLD 假设存在一组递减的噪声标准差序列:
σ
1
>
σ
2
>
⋯
>
σ
L
,
\sigma_1 > \sigma_2 > \cdots > \sigma_L,
σ1>σ2>⋯>σL,噪声尺度之间具有等比关系:
σ
1
σ
2
=
σ
2
σ
3
=
⋯
=
σ
L
−
1
σ
L
>
1.
\frac{\sigma_1}{\sigma_2} = \frac{\sigma_2}{\sigma_3} = \cdots = \frac{\sigma_{L-1}}{\sigma_L} > 1.
σ2σ1=σ3σ2=⋯=σLσL−1>1.对于每个噪声级别
σ
i
\sigma_i
σi,对原始数据添加扰动后的分布为:
p
σ
i
(
x
~
)
=
∫
p
data
(
x
)
⋅
p
σ
i
(
x
~
∣
x
)
d
x
,
p_{\sigma_i}(\tilde{\mathbf{x}}) = \int p_{\text{data}}(\mathbf{x}) \cdot p_{\sigma_i}(\tilde{\mathbf{x}} \mid \mathbf{x}) \, d\mathbf{x},
pσi(x~)=∫pdata(x)⋅pσi(x~∣x)dx,
取条件概率分布
p
σ
i
(
x
~
∣
x
)
p_{\sigma_i}(\tilde{\mathbf{x}} \mid \mathbf{x})
pσi(x~∣x)为
N
(
x
~
∣
x
,
σ
i
2
I
)
\mathcal{N}(\tilde{\mathbf{x}} | {\mathbf{x}}, \sigma_{i}^2 I)
N(x~∣x,σi2I),当噪声强度最大时(
σ
1
\sigma_1
σ1),分布接近各向同性高斯分布:
p
σ
1
(
x
)
≈
N
(
0
,
σ
1
2
I
)
,
p_{\sigma_1}(\mathbf{x}) \approx \mathcal{N}(\mathbf{0}, \sigma_1^2 \mathbf{I}),
pσ1(x)≈N(0,σ12I),
当噪声强度最小时( σ L \sigma_L σL),扰动数据接近原始分布:
p σ L ( x ) ≈ p data ( x ) . p_{\sigma_L}(\mathbf{x}) \approx p_{\text{data}}(\mathbf{x}). pσL(x)≈pdata(x).
训练目标:得分匹配
SMLD 通过最小化每个噪声尺度下的 denoising score matching 目标函数来估计得分函数,即:
θ
∗
=
arg
min
θ
∑
i
=
1
N
σ
i
2
E
p
data
(
x
)
E
p
σ
i
(
x
~
∣
x
)
[
∥
s
θ
(
x
~
,
σ
i
)
−
∇
x
~
log
p
σ
i
(
x
~
∣
x
)
∥
2
2
]
.
\theta^* = \arg\min_{\theta} \sum_{i=1}^N \sigma_i^2 \mathbb{E}_{p_{\text{data}}(\mathbf{x})} \mathbb{E}_{p_{\sigma_i}(\tilde{\mathbf{x}}|\mathbf{x})} \big[ \| \mathbf{s}_{\theta}(\tilde{\mathbf{x}}, \sigma_i) - \nabla_{\tilde{\mathbf{x}}} \log p_{\sigma_i}(\tilde{\mathbf{x}} \mid \mathbf{x}) \|_2^2 \big].
θ∗=argθmini=1∑Nσi2Epdata(x)Epσi(x~∣x)[∥sθ(x~,σi)−∇x~logpσi(x~∣x)∥22].
条件概率分布
p
σ
i
(
x
~
∣
x
)
p_{\sigma_i}(\tilde{\mathbf{x}} \mid \mathbf{x})
pσi(x~∣x)为
N
(
x
~
∣
x
,
σ
i
2
I
)
\mathcal{N}(\tilde{\mathbf{x}} | {\mathbf{x}}, \sigma_{i}^2 I)
N(x~∣x,σi2I),那么
∇
x
~
log
p
σ
i
(
x
~
∣
x
)
=
−
x
~
−
x
σ
i
2
\nabla_{\tilde{\mathbf{x}}} \log p_{\sigma_i}(\tilde{\mathbf{x}} \mid \mathbf{x})=-\frac{\tilde{\mathbf{x}}-\mathbf{x}}{ \sigma_{i}^2}
∇x~logpσi(x~∣x)=−σi2x~−x此时优化目标为
θ
∗
=
arg
min
θ
∑
i
=
1
N
σ
i
2
E
p
data
(
x
)
E
x
~
∼
N
(
x
~
∣
x
,
σ
i
2
I
)
[
∥
s
θ
(
x
~
,
σ
i
)
+
x
~
−
x
σ
i
2
∥
2
2
]
.
\theta^* = \arg\min_{\theta} \sum_{i=1}^N \sigma_i^2 \mathbb{E}_{p_{\text{data}}(\mathbf{x})} \mathbb{E}_{\tilde{\mathbf{x}} \sim \mathcal{N}(\tilde{\mathbf{x}} | {\mathbf{x}}, \sigma_{i}^2 I)} \big[ \| \mathbf{s}_{\theta}(\tilde{\mathbf{x}}, \sigma_i) +\frac{\tilde{\mathbf{x}}-\mathbf{x}}{ \sigma_{i}^2} \|_2^2\big].
θ∗=argθmini=1∑Nσi2Epdata(x)Ex~∼N(x~∣x,σi2I)[∥sθ(x~,σi)+σi2x~−x∥22].
采样方法
基于朗之万动力学(Langevin Dynamics)
朗之万动力学(Langevin dynamics)能够仅利用概率密度函数
p
(
x
)
p ({ \mathbf { x } })
p(x)的得分函数
∇
x
log
p
(
x
)
\nabla _ { \mathbf { x } } \operatorname { l o g } p ({ \mathbf { x } })
∇xlogp(x)生成样本,给定固定的步长
ϵ
>
0
\epsilon > 0
ϵ>0和初始值
x
~
0
∼
π
(
x
)
\tilde{x}_0 \sim \pi(x)
x~0∼π(x)(其中
π
\pi
π为先验分布),朗之万方法通过以下递推公式进行计算:
x
~
t
=
x
~
t
−
1
+
ϵ
2
∇
x
log
p
(
x
~
t
−
1
)
+
ϵ
z
t
,
\tilde {\mathbf { x } } _ { t } = \tilde { \mathbf { x } } _ { t - 1 } + \frac { \epsilon } { 2 } \nabla _ { \mathbf { x } } \operatorname { l o g } p ( \tilde { \mathbf { x } } _ { t - 1 } ) + \sqrt { \epsilon } \, \mathbf { z } _ { t } ,
x~t=x~t−1+2ϵ∇xlogp(x~t−1)+ϵzt,
其中
z
t
∼
N
(
0
,
I
)
\mathbf{z}_t \sim \mathcal{N}(0, I)
zt∼N(0,I),当
ϵ
→
0
,
T
→
∞
\epsilon \rightarrow 0 ,T \rightarrow \infty
ϵ→0,T→∞时,
p
(
x
~
t
)
→
p
(
x
)
p(\tilde { \mathbf{x}} _ {t})\rightarrow p(x)
p(x~t)→p(x)
为了从高噪声逐步过渡到低噪声的精细采样,SMLD 采用“退火”策略:从最大噪声尺度开始,通过多步采样逐步降低噪声强度。
x
(
m
)
=
x
(
m
−
1
)
+
ϵ
i
s
θ
(
x
(
m
−
1
)
,
σ
i
)
+
2
ϵ
i
z
i
(
m
)
,
z
i
(
m
)
∼
N
(
0
,
I
)
,
\mathbf{x}^{(m)} = \mathbf{x}^{(m-1)} + \epsilon_i \, \mathbf{s}_\theta(\mathbf{x}^{(m-1)}, \sigma_i) + \sqrt{2 \epsilon_i} \, \mathbf{z}_i^{(m)}, \quad \mathbf{z}_i^{(m)} \sim \mathcal{N}(0, \mathbf{I}),
x(m)=x(m−1)+ϵisθ(x(m−1),σi)+2ϵizi(m),zi(m)∼N(0,I),
初始从先验分布中采样,通过 Langevin dynamics 得到噪声尺度为
σ
1
\sigma_1
σ1 的样本,再以此为初始点,迭代采样得到
σ
2
\sigma_2
σ2 的样本,如此逐步迭代,最终生成接近数据分布
p
data
(
x
)
p_{\text{data}}(\mathbf{x})
pdata(x) 的样本。
与 DDPM 类似的采样方法
SMLD 也可采用DDPM 的采样方式,将噪声扰动过程建模为马尔可夫链。
和《Score-Based Generative Modeling through Stochastic Differential Equations》的符号表示一致,噪声强度从小到大表示为
σ
1
<
σ
2
<
⋯
<
σ
N
\sigma_1 < \sigma_2 < \cdots < \sigma_N
σ1<σ2<⋯<σN,利用这些噪声尺度扰动
x
0
x_0
x0,得到马尔可夫链
x
0
→
x
1
→
.
.
.
→
x
n
x_0 \rightarrow x_1 \rightarrow...\rightarrow x_n
x0→x1→...→xn,由于
p
(
x
i
)
=
N
(
x
i
;
x
0
,
σ
i
2
I
)
p(x_i)=\mathcal{N}(\mathbf{x}_i;\mathbf{x}_{0},\sigma_i^2\mathbf{I})
p(xi)=N(xi;x0,σi2I),即
x
i
=
x
0
+
σ
i
z
,
x
i
−
1
=
x
0
+
σ
i
−
1
z
,
\mathbf{x}_i=\mathbf{x}_0+\sigma_i\mathbf{z}, \\ \mathbf{x}_{i-1}=\mathbf{x}_0+\sigma_{i-1}\mathbf{z},
xi=x0+σiz,xi−1=x0+σi−1z,其中
z
∼
N
(
0
,
I
)
\mathbf{z} \sim \mathcal{N}(0, I)
z∼N(0,I),那么
x
i
=
x
i
−
1
+
(
σ
i
−
σ
i
−
1
)
z
,
p
(
x
i
∣
x
i
−
1
)
=
N
(
x
i
;
x
i
−
1
,
(
σ
i
2
−
σ
i
−
1
2
)
I
)
,
i
=
1
,
2
,
⋯
,
N
.
\mathbf{x}_i=\mathbf{x}_{i-1}+(\sigma_i-\sigma_{i-1})\mathbf{z},\\ p(\mathbf{x}_i\mid\mathbf{x}_{i-1})=\mathcal{N}(\mathbf{x}_i;\mathbf{x}_{i-1},(\sigma_i^2-\sigma_{i-1}^2)\mathbf{I}),\quad i=1,2,\cdots,N.
xi=xi−1+(σi−σi−1)z,p(xi∣xi−1)=N(xi;xi−1,(σi2−σi−12)I),i=1,2,⋯,N.利用贝叶斯公式得到
q
(
x
i
−
1
∣
x
i
,
x
0
)
q(\mathbf{x}_{i-1}\mid\mathbf{x}_i,\mathbf{x}_0)
q(xi−1∣xi,x0),
q
(
x
i
−
1
∣
x
i
,
x
0
)
=
q
(
x
i
∣
x
i
−
1
,
x
0
)
q
(
x
i
−
1
∣
x
0
)
q
(
x
i
∣
x
0
)
=
N
(
x
i
−
1
;
σ
i
−
1
2
σ
i
2
x
i
+
(
1
−
σ
i
−
1
2
σ
i
2
)
x
0
,
σ
i
−
1
2
(
σ
i
2
−
σ
i
−
1
2
)
σ
i
2
I
)
.
\begin{align*} q(\mathbf{x}_{i-1}\mid\mathbf{x}_i,\mathbf{x}_0)&=\frac{q(\mathbf{x}_i|\mathbf{x}_{i-1},\mathbf{x}_0)q(\mathbf{x}_{i-1}|\mathbf{x}_0)}{q(\mathbf{x}_i|\mathbf{x}_0)} \\ &=\mathcal{N}\left(\mathbf{x}_{i-1}; \frac{\sigma_{i-1}^2}{\sigma_i^2}\mathbf{x}_i+\Big(1-\frac{\sigma_{i-1}^2}{\sigma_i^2}\Big)\mathbf{x}_0, \frac{\sigma_{i-1}^2(\sigma_i^2-\sigma_{i-1}^2)}{\sigma_i^2}\mathbf{I}\right). \end{align*}
q(xi−1∣xi,x0)=q(xi∣x0)q(xi∣xi−1,x0)q(xi−1∣x0)=N(xi−1;σi2σi−12xi+(1−σi2σi−12)x0,σi2σi−12(σi2−σi−12)I).
参数化
p
θ
(
x
i
−
1
∣
x
i
)
p _ { \theta } ( { \bf x } _ { i - 1 } \mid { \bf x } _ { i } )
pθ(xi−1∣xi):
p
θ
(
x
i
−
1
∣
x
i
)
=
N
(
x
i
−
1
;
μ
θ
(
x
i
,
i
)
,
τ
i
2
I
)
p _ { \theta } ( { \bf x } _ { i - 1 } \mid { \bf x } _ { i } ) = { \cal N } ( { \bf x } _ { i - 1 } ; \mu _ { \theta } ( { \bf x } _ { i } , i ) , \tau _ { i } ^ { 2 } { \bf I } )
pθ(xi−1∣xi)=N(xi−1;μθ(xi,i),τi2I)
损失函数项
L
t
−
1
L_{t-1}
Lt−1为:
L
t
−
1
=
E
q
[
D
KL
(
q
(
x
i
−
1
∣
x
i
,
x
0
)
)
∥
p
θ
(
x
i
−
1
∣
x
i
)
]
=
E
x
0
,
z
[
1
2
τ
i
2
∥
x
i
(
x
0
,
z
)
−
σ
i
2
−
σ
i
−
1
2
σ
i
z
−
μ
θ
(
x
i
(
x
0
,
z
)
,
i
)
∥
2
2
]
+
C
,
\begin{align*} L_{t-1} &= \mathbb{E}_q[D_{\text{KL}}(q(\mathbf{x}_{i-1} \mid \mathbf{x}_i, \mathbf{x}_0)) \parallel p_\theta(\mathbf{x}_{i-1} \mid \mathbf{x}_i)] \\ &= \mathbb{E}_{\mathbf{x}_0, \mathbf{z}} \left[ \frac{1}{2\tau_i^2} \left\| \mathbf{x}_i(\mathbf{x}_0, \mathbf{z}) - \frac{\sigma_i^2 - \sigma_{i-1}^2}{\sigma_i} \mathbf{z} - \boldsymbol{\mu}_\theta(\mathbf{x}_i(\mathbf{x}_0, \mathbf{z}), i) \right\|_2^2 \right] + C, \end{align*}
Lt−1=Eq[DKL(q(xi−1∣xi,x0))∥pθ(xi−1∣xi)]=Ex0,z[2τi21
xi(x0,z)−σiσi2−σi−12z−μθ(xi(x0,z),i)
22]+C,
根据
L
t
−
1
L_{t-1}
Lt−1的形式参数化
μ
θ
(
x
i
,
i
)
\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{x}_i,i)
μθ(xi,i):
μ
θ
(
x
i
,
i
)
=
x
i
+
(
σ
i
2
−
σ
i
−
1
2
)
s
θ
(
x
i
,
i
)
,
\begin{align*}\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{x}_i,i)=\mathbf{x}_i+(\sigma_i^2-\sigma_{i-1}^2)\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_i,i),\end{align*}
μθ(xi,i)=xi+(σi2−σi−12)sθ(xi,i),
与DDPM一致,取标准差为
σ
i
−
1
2
(
σ
i
2
−
σ
i
−
1
2
)
σ
i
2
\sqrt{\frac{\sigma_{i-1}^2(\sigma_i^2-\sigma_{i-1}^2)}{\sigma_i^2}}
σi2σi−12(σi2−σi−12),最终采样公式为:
x
i
−
1
=
x
i
+
(
σ
i
2
−
σ
i
−
1
2
)
s
θ
(
x
i
,
i
)
+
σ
i
−
1
2
(
σ
i
2
−
σ
i
−
1
2
)
σ
i
2
z
i
,
i
=
1
,
2
,
⋯
,
N
,
\begin{align*}\mathbf{x}_{i-1}=\mathbf{x}_i+(\sigma_i^2-\sigma_{i-1}^2)\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_i,i)+\sqrt{\frac{\sigma_{i-1}^2(\sigma_i^2-\sigma_{i-1}^2)}{\sigma_i^2}}\mathbf{z}_i,i=1,2,\cdots,N,\end{align*}
xi−1=xi+(σi2−σi−12)sθ(xi,i)+σi2σi−12(σi2−σi−12)zi,i=1,2,⋯,N,
参考
《Generative Modeling by Estimating Gradients of the
Data Distribution》
《Score-Based Generative Modeling through Stochastic Differential Equations》