Generative Modeling by Estimating Gradients of the Data Distribution(paper reading)

Generative Modeling by Estimating Gradients of the Data Distribution (Paper reading)

Yang Song, Stanford University, NeurIPS2019, Cited:498, Code, Paper

简简介

Score-based model的开篇之作？扩散模型的理解点这里, 从VAE到Diffusion Moedl的理解。

1. 什么是分数Score和朴素分数模型。

1.1 Score

假设我们的数据集包含了来自未知分布$p_{data}(x)$的独立同分布的样本$x_i$，那么分数Score的定义就是概率密度函数$p(x)$的对数似然的梯度${▽}_{x}logp(x)$。分数网络Score network就是训练一个神经网络用来估计真实分布$p_{data}(x)$。基于分数的生成模型通常有两个成分：分数匹配和Langevin动力学。

1.2 Langevin动力学

Langevin动力学是一个采样算法，这个算法仅仅依赖于分数函数${▽}_{x}logp(x)$就可以从$p(x)$中生成样本。具体算法如下，给定一个特点的步长$ϵ>0$和一个初始值$\tilde{x}\sim \pi (x)$，这里$\pi$指任意的先验分布(高斯分布等)。Langevin的方法可以写成迭代的形式：
$$\widetilde{x_t}={\tilde{x}}_{t-1}+\frac{ϵ}{2}{▽}_{x}logp({\tilde{x}}_{t-1})+\sqrt{ϵ}{z}_t\text{\hspace{1em}(1)}$$
其中$z_t\sim N(0,I)$，当$ϵ⟶0$时且$T⟶\infty$时，我们可以认为${\tilde{x}}_T$的分布等于$p(x)$，也就是说当$ϵ$很小，$T$很大，${\tilde{x}}_T$就是从$p(x)$中生成的样本。我们可以从公式（1）中看到，整个Langevin算法仅仅依赖于上一步概率密度函数的分数${▽}_{x}logp({\tilde{x}}_{t-1})$。真实分布$p_{data}(x)$我们无法获得，因此我们需要训练分数网络来近似$s_{\theta }(x)\approx {▽}_{x}log{p}_{data}(\tilde{x})$并替换公式（1）中的分数。这就是基于分数生成模型的主要思想。

1.3 去噪分数匹配 Denoising Score Matching

首先对数据样本$x$加上一个预先设定好的噪声分布$q_{\sigma }(\tilde{x}|x)$（在$x$上加上噪声强度为$\sigma$的噪声后再加到$\tilde{x}$的分布上），然后用分数匹配的方法估计带噪声数据$q_{\sigma }(\tilde{x})=\int q_{\sigma }(\tilde{x}|x)p_{data}(x)dx$的分数。目标函数可以被证明等同于：
$$\frac{1}{2}{\mathbb{E}}_{q_{\sigma }(\tilde{x}|x)p_{data}(x)}[{\Vert s_{\theta }(\tilde{x})-{▽}_{\tilde{x}}log{q}_{\sigma }(\tilde{x}|x)\Vert }_2^2]\text{\hspace{1em}(2)}$$
此时，训练的最优的分数网络在加的噪声非常小的情况下，也就是说$q_{\sigma }(x)\approx p_{data}(x)$，近似等同于$q_{\sigma }(x)$的分数：$s_{{\theta }^{\ast }}(x)={▽}_{x}log{q}_{\sigma }(x)\approx {▽}_{x}log{p}_{data}(x)$。

1.4 挑战

来自B站UP主VictorYuki

可以看到socre整体指向数据密度更好的位置，假设我们已经获得了数据空间的score，也就是说训练好了对应的模型，那么每次我们从初始值根据score的方向就可以一步一步的走到数据密度高的位置，也就是完成了一次生成。当然，这一切的前提是我们可以准确的求score。但是在低密度区域，score的估计往往不够准确。

左一深色代表数据密度较高的位置

在模型推理的早期，数据分布密度较低的地方的梯度估计通常不准确，因此早期的推理会沿着错误的梯度进行，导致模型效果不好。解决的方法是对数据加噪声，扩大数据范围，让原本低密度的数据区域“膨胀”，如图：

但是，加了噪声后，你最后得到的是噪声分布，这里就会有一个权衡，加的噪声越强，你的估计范围越大，估计越准确，但更多破坏了原有分布；如果加的噪声很小就会导致score的估计不准确，因此解决的方法是，在推理的不同阶段加不同强度的噪声，噪声由大到小。在推理早期时，加强度大的噪声，因为早期的估计距离目标分布较远，所以大噪声更容易估计；而在后期，距离目标分布很近时，就可以减少噪声强度估计出准确的score。具体做法是：1）利用不同的噪声强度扰动数据分布；2）利用一个分数网络去预测不同噪声强度下的分数；3）训练好分数网络后就可以通过退火Langevin方法生成样本，也就是说在从大噪声强度下开始生成样本并逐渐减少噪声生成样本。

2. Noise Conditional Score Networks (NCSN)

2.1 NCSN定义

取 { σ i } i = 1 L \left \{\sigma_{i} \right \}_{i=1}^{L} {σi​}i=1L​满足$\frac{{\sigma }_1}{{\sigma }_2}=\cdot \cdot \cdot =\frac{{\sigma }_{L-1}}{{\sigma }_L}>1$的正数等比数列。令$q_{\sigma }(x)\triangleq \int p_{data}(t)N(x|t,{\sigma }^{2}I)dt$（指在原始数据上加入了噪声）为扰动的数据分布。${\sigma }_1$取一个很大的值，${\sigma }_L$取一个很小的值减少后期对数据的影响。我们的目标就是训练一个条件分数网络估计任一噪声下的分数，即 ∀ σ ∈ { σ i } i = 1 L : s θ ( x , σ ) ≈ ∇ x l o g q σ ( x ) \forall \sigma \in \left \{\sigma_{i} \right \}_{i=1}^{L}:s_{\theta}(x,\sigma) \approx \nabla_{x}logq_{\sigma}(x) ∀σ∈{σi​}i=1L​:sθ​(x,σ)≈∇x​logqσ​(x)。这里的$s_{\theta }(x,\sigma )$就叫做NCSN。

与去噪分数网络的区别就是，1.3讲的内容是仅仅加了噪声的扰动，很小很小的噪声，给分数网络得到分数，并近似的认为是$p_{data}(x)$的分数，而NCSN额外的以噪声$\sigma$作为条件输入，能够预测不同噪声强度下的分数。这里的分数网络是基于U-Net的神经网络。

2.2 训练NCSN

这里我们选择噪声分布为$q_{\sigma }(\tilde{x}|x)=N(\tilde{x}|x,{\sigma }^{2}I)$，即在给定$x$样本的条件下，$\tilde{x}$的分布符合均值为$x$，方差为${\sigma }^2$的高斯分布：
$$\begin{array}{rl}{q}_{\sigma }(\tilde{x}|x)& =\frac{1}{\sqrt{2\pi {\sigma }^2}}{e}^{-\frac{(\tilde{x}-x{)}^2}{2{\sigma }^2}}\\ log{q}_{\sigma }(\tilde{x}|x)& =-\frac{(\tilde{x}-x{)}^2}{2{\sigma }^2}-log\sqrt{2\pi {\sigma }^2}\\ {\nabla }_{\tilde{x}}log{q}_{\sigma }(\tilde{x}|x)& =--\frac{(\tilde{x}-x)}{{\sigma }^2}\end{array}$$
因此给定$\sigma$，去噪分数匹配的目标函数可以写作：
$$l(\theta ;\sigma )\triangleq \frac{1}{2}{\mathbb{E}}_{p_{data}(x)}{\mathbb{E}}_{\tilde{x}\sim N(x,{\sigma }^{2}I)}\left[{\Vert s_{\theta }(\tilde{x},\sigma )+\frac{\tilde{x}-x}{{\sigma }^2}\Vert }_2^2\right]\text{\hspace{1em}(3)}$$
对于不同的噪声强度，得到最终的目标函数：
L ( θ ; { σ i } i = 1 L ) ≜ 1 L ∑ i = 1 L λ ( σ i ) l ( θ ; σ i ) (4) \mathcal{L}(\theta; \left \{\sigma_{i} \right \}_{i=1}^{L})\triangleq\frac{1}{L}\sum_{i=1}^{L}\lambda(\sigma _{i})\mathscr{l}(\theta ;\sigma_{i}) \tag{4} L(θ;{σi​}i=1L​)≜L1​i=1∑L​λ(σi​)l(θ;σi​)(4)
这里的$\lambda ({\sigma }_i)>0$，表示关于${\sigma }_i$的权重，目的是在不同噪声强度下做出平衡，使得$l(\theta ;{\sigma }_i)$与特点噪声无关。那么我们就希望$\lambda ({\sigma }_i)l(\theta ;{\sigma }_i)$在每个噪声强度下都具有同样的量级，即无论噪声多大，两者相乘都是差不多的。作者经验性的发现网络收敛后${\Vert s_{\theta }(x,\sigma )\Vert }_2\propto \frac{1}{\sigma }$，因此$\lambda (\sigma )={\sigma }^2$，则${\sigma }^{2}l(\theta ;\sigma )=\frac{1}{2}\mathbb{E}{\Vert \sigma s_{\theta }(\tilde{x},\sigma )+\frac{\tilde{x}-x}{\sigma }\Vert }_2^2$。因为$\frac{\tilde{x}-x}{\sigma }\sim N(0,I)$且${\Vert \sigma s_{\theta }(\tilde{x},\sigma )\Vert }_2\propto 1$，因此$\lambda ({\sigma }_i)l(\theta ;{\sigma }_i)$是不依赖于噪声$\sigma$的。当训练好分数网络，最佳的情况是分数网络可以精准的预测不同噪声强度下的分数$s_{{\theta }^{\ast }}(x,{\sigma }_i)\approx {\nabla }_{x}log{q}_{{\sigma }_i}(x)$ 。

2.3 NCSN的推理和退火Langevin动力学

当训练好NCSN后，采用退火Langenvin动力学采样方法生成样本。首先设置一组由大到小的噪声${\sigma }_i$，和步长$ϵ$以及$T$。在先验分布中初始化一个$x_0$。遍历每一个噪声强度，并根据每一个噪声计算出步长${\alpha }_i=ϵ{\sigma }_i^2/{\sigma }_L^2$，可以理解为当噪声大的时候用步数较大去采样，而噪声小的时候就需要步数小一点，这跟学习率很像。在每个噪声下都需要通过Langevin动力学采样出当前加噪数据的分布，做T次的迭代后的数据${\tilde{x}}_T$做为下一次，也就是下一个噪声强度的初始值${\tilde{x}}_0$，当噪声强度这个循环结束后，得到的$q_{{\sigma }_L}(x)$可以近似的认为是$p_{data}(x)$。