Denoising Score Matching with Langevin Dynamics
在自然图像等复杂数据集中,真实数据往往集中分布在一个低维流形上,概率密度函数的梯度(即得分函数)难以定义与估计。为缓解该问题,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》