Denoising Score Matching with Langevin Dynamics

在自然图像等复杂数据集中，真实数据往往集中分布在一个低维流形上，概率密度函数的梯度（即得分函数）难以定义与估计。为缓解该问题，SMLD 提出使用不同强度的高斯噪声对数据进行扰动，扰动后的数据不再集中于低维流形，从而提升学习鲁棒性。

多尺度噪声扰动

SMLD 假设存在一组递减的噪声标准差序列：
$${\sigma }_1>{\sigma }_2>\cdots >{\sigma }_L,$$噪声尺度之间具有等比关系：
$$\frac{{\sigma }_1}{{\sigma }_2}=\frac{{\sigma }_2}{{\sigma }_3}=\cdots =\frac{{\sigma }_{L-1}}{{\sigma }_L}>1.$$对于每个噪声级别${\sigma }_i$，对原始数据添加扰动后的分布为：
$$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_{{\sigma }_i}(\tilde{\mathbf{x}}\mid \mathbf{x})$为$\mathcal{N}(\tilde{\mathbf{x}}|\mathbf{x},{\sigma }_i^{2}I)$，当噪声强度最大时（${\sigma }_1$），分布接近各向同性高斯分布：
$$p_{{\sigma }_1}(\mathbf{x})\approx \mathcal{N}(0,{\sigma }_1^2\mathbf{I}),$$

当噪声强度最小时（${\sigma }_L$），扰动数据接近原始分布：

$$p_{{\sigma }_L}(\mathbf{x})\approx p_{\text{data}}(\mathbf{x}).$$

训练目标：得分匹配

SMLD 通过最小化每个噪声尺度下的 denoising score matching 目标函数来估计得分函数，即：
$${\theta }^{\ast }=\arg \underset{\theta }{\min }\sum _{i=1}^N{\sigma }_i^2{\mathbb{E}}_{p_{\text{data}}(\mathbf{x})}{\mathbb{E}}_{p_{{\sigma }_i}(\tilde{\mathbf{x}}|\mathbf{x})}[\Vert {\mathbf{s}}_{\theta }(\tilde{\mathbf{x}},{\sigma }_i)-{\nabla }_{\tilde{\mathbf{x}}}\log p_{{\sigma }_i}(\tilde{\mathbf{x}}\mid \mathbf{x}){\Vert }_2^2].$$
条件概率分布$p_{{\sigma }_i}(\tilde{\mathbf{x}}\mid \mathbf{x})$为$\mathcal{N}(\tilde{\mathbf{x}}|\mathbf{x},{\sigma }_i^{2}I)$，那么
$${\nabla }_{\tilde{\mathbf{x}}}\log p_{{\sigma }_i}(\tilde{\mathbf{x}}\mid \mathbf{x})=-\frac{\tilde{\mathbf{x}}-\mathbf{x}}{{\sigma }_i^2}$$此时优化目标为
$${\theta }^{\ast }=\arg \underset{\theta }{\min }\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)}[\Vert {\mathbf{s}}_{\theta }(\tilde{\mathbf{x}},{\sigma }_i)+\frac{\tilde{\mathbf{x}}-\mathbf{x}}{{\sigma }_i^2}{\Vert }_2^2].$$

采样方法

基于朗之万动力学（Langevin Dynamics）

朗之万动力学（Langevin dynamics）能够仅利用概率密度函数$p(\mathbf{x})$的得分函数${\nabla }_{\mathbf{x}}\log p(\mathbf{x})$生成样本，给定固定的步长$ϵ>0$和初始值${\tilde{x}}_0\sim \pi (x)$（其中$\pi$为先验分布），朗之万方法通过以下递推公式进行计算：
$${\tilde{\mathbf{x}}}_t={\tilde{\mathbf{x}}}_{t-1}+\frac{ϵ}{2}{\nabla }_{\mathbf{x}}\log p({\tilde{\mathbf{x}}}_{t-1})+\sqrt{ϵ}{\mathbf{z}}_t,$$
其中${\mathbf{z}}_t\sim \mathcal{N}(0,I)$，当$ϵ\to 0,T\to \infty$时，$p({\tilde{\mathbf{x}}}_t)\to p(x)$

为了从高噪声逐步过渡到低噪声的精细采样，SMLD 采用“退火”策略：从最大噪声尺度开始，通过多步采样逐步降低噪声强度。
$${\mathbf{x}}^{(m)}={\mathbf{x}}^{(m-1)}+{ϵ}_i{\mathbf{s}}_{\theta }({\mathbf{x}}^{(m-1)},{\sigma }_i)+\sqrt{2{ϵ}_i}{\mathbf{z}}_i^{(m)},{\mathbf{z}}_i^{(m)}\sim \mathcal{N}(0,\mathbf{I}),$$
初始从先验分布中采样，通过 Langevin dynamics 得到噪声尺度为 ${\sigma }_1$ 的样本，再以此为初始点，迭代采样得到 ${\sigma }_2$ 的样本，如此逐步迭代，最终生成接近数据分布 $p_{\text{data}}(\mathbf{x})$ 的样本。

与 DDPM 类似的采样方法

SMLD 也可采用DDPM 的采样方式，将噪声扰动过程建模为马尔可夫链。
和《Score-Based Generative Modeling through Stochastic Differential Equations》的符号表示一致，噪声强度从小到大表示为${\sigma }_1<{\sigma }_2<\cdots <{\sigma }_N$，利用这些噪声尺度扰动$x_0$，得到马尔可夫链$x_0\to x_1\to ...\to x_n$，由于$p(x_i)=\mathcal{N}({\mathbf{x}}_i;{\mathbf{x}}_0,{\sigma }_i^2\mathbf{I})$，即
$$\begin{gathered} {\mathbf{x}}_i={\mathbf{x}}_0+{\sigma }_i\mathbf{z}, \\ {\mathbf{x}}_{i-1}={\mathbf{x}}_0+{\sigma }_{i-1}\mathbf{z}, \end{gathered}$$其中$\mathbf{z}\sim \mathcal{N}(0,I)$，那么$$\begin{gathered} {\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}),i=1,2,\cdots ,N. \end{gathered}$$利用贝叶斯公式得到$q({\mathbf{x}}_{i-1}\mid {\mathbf{x}}_i,{\mathbf{x}}_0)$,
$$\begin{array}{rl}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+(1-\frac{{\sigma }_{i-1}^2}{{\sigma }_i^2}){\mathbf{x}}_0,\frac{{\sigma }_{i-1}^2({\sigma }_i^2-{\sigma }_{i-1}^2)}{{\sigma }_i^2}\mathbf{I}\right).\end{array}$$
参数化$p_{\theta }({\mathbf{x}}_{i-1}\mid {\mathbf{x}}_i)$:
$$p_{\theta }({\mathbf{x}}_{i-1}\mid {\mathbf{x}}_i)=\mathcal{N}({\mathbf{x}}_{i-1};{\mu }_{\theta }({\mathbf{x}}_i,i),{\tau }_i^2\mathbf{I})$$
损失函数项$L_{t-1}$为：
$$\begin{array}{rl}{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}{\Vert {\mathbf{x}}_i({\mathbf{x}}_0,\mathbf{z})-\frac{{\sigma }_i^2-{\sigma }_{i-1}^2}{{\sigma }_i}\mathbf{z}-{\mathbf{\mu }}_{\theta }({\mathbf{x}}_i({\mathbf{x}}_0,\mathbf{z}),i)\Vert }_2^2\right]+C,\end{array}$$
根据$L_{t-1}$的形式参数化${\mathbf{\mu }}_{\mathbf{\theta }}({\mathbf{x}}_i,i)$:
$$\begin{array}{r}{\mathbf{\mu }}_{\mathbf{\theta }}({\mathbf{x}}_i,i)={\mathbf{x}}_i+({\sigma }_i^2-{\sigma }_{i-1}^2){\mathbf{s}}_{\mathbf{\theta }}({\mathbf{x}}_i,i),\end{array}$$
与DDPM一致，取标准差为$\sqrt{\frac{{\sigma }_{i-1}^2({\sigma }_i^2-{\sigma }_{i-1}^2)}{{\sigma }_i^2}}$，最终采样公式为：
$$\begin{array}{r}{\mathbf{x}}_{i-1}={\mathbf{x}}_i+({\sigma }_i^2-{\sigma }_{i-1}^2){\mathbf{s}}_{\mathbf{\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{array}$$

参考

《Generative Modeling by Estimating Gradients of the
Data Distribution》
《Score-Based Generative Modeling through Stochastic Differential Equations》