扩散模型（三）：Score-based Generative Models
Score function
为建模数据分布，我呢吧将数据分布写作能量模型EBM形式：
\[p_{\theta}(x)=\frac{\exp(-E_{\theta}(x))}{Z_{\theta}}
\]
其中， \(Z_\theta\) 为使其满足概率密度分布的归一化项。我们使用最大似然估计对参数 \(\theta\) 进行训练：
\[max_\theta\sum_i^N\log{p_\theta(x_i)}=max_\theta\sum_i^N-E_\theta(x)-\log{Z_\theta}
\]
此时会出现归一化项 \(Z_\theta\) 无法计算的问题，为解决该问题，我们引入score function \(s_\theta(x)=\nabla_x\log{p(x)}\) ，score function的计算与\(Z_\theta\)无关：
\[\nabla_x\log{p(x)}=-\nabla_xf_\theta(x)-s_\theta(x)
\]
那么，假设我们已经获得了数据分布 \(p(x)\) 的score function \(\nabla_x\log{p(x)}\) ，该如何获得该数据分布的样本呢？
此时，我们可以采用朗之万动力学采样法，从某个先验分布中随机采样的初始样本开始，利用如下迭代公式，逐渐得到服从数据分布 \(p(x)\) 的样本：
\[x_{i+1}=x_i+\epsilon\nabla_x\log{p(x)}+\sqrt{2\epsilon}z_i,\ z_i\sim\mathcal{N}(0,\mathbf{I}),\ i=0,1,...K
\]
当迭代次数 \(K\) 足够多时，样本 \(x_K\) 收敛为从 \(p(x)\) 采样的样本。
Score matching
下面我们讨论如何训练网络估计出score function，我们设计如下损失函数：
\[\mathcal{L}=\mathbb{E}_{x\sim p(x)}\left[||\nabla_x\log{p(x)}-s_{\theta}(x)||^2\right]
\]
然而，由于我们不知道真实的数据分布 \(p(x)\) ，该损失函数无法计算。此时，我们可以使用分数匹配score matching的方法在不知道真实数据分布的基础上计算该损失。
分数匹配的相关方法有很多，在此我们讨论使用条件得分匹配替代得分匹配的方法。
令 \(x_0\sim p_0(x_0)\) 表示原始数据分布，我们定义一个已知解析式的条件分布（例如给定方差系数的正态分布） \(x_0,x\sim p_0(x_0)p(x|x_0)\) ，那么该条件分布的score function \(\nabla_x\log{p(x|x_0)}\) 便可以通过采样一对 \((x_0,x)\) 计算。但我们想要网络学会估计的是分布 \(x\sim p(x)\) 的score function。
我们先给出结论，得分匹配 \(\mathbb{E}_{x\sim p(x)}\left[||\nabla_x\log{p(x)}-s_\theta(x)||^2\right]\) 、条件得分匹配 \(\mathbb{E}_{x_0,x\sim p_0(x_0)p(x|x_0)}\left[||\nabla_x\log{p(x|x_0)}-s_\theta(x)||^2\right]\) 二者作为优化目标是等价的。因此，我们可以将后者作为优化目标。下面给出该结论的证明：
首先，对得分匹配，我们有：
\[\begin{align*}
&\mathbb{E}_{x\sim p(x)}\left[||\nabla_x\log{p(x)}-s_\theta(x)||^2\right] \\
= &\mathbb{E}_{x\sim p(x)}\left[||\nabla_x\log{p(x)}||^2+||s_\theta(x)||^2-2s_{\theta}(x)\nabla_x\log{p(x)}\right]
\end{align*}
\]
对条件得分匹配，我们有：
\[\begin{align*}
&\mathbb{E}_{x_0,x\sim p_0(x_0)p(x|x_0)}\left[||\nabla_x\log{p(x|x_0)}-s_\theta(x)||^2\right] \\
= &\mathbb{E}_{x_0,x\sim p(x)p(x_0|x)}\left[||\nabla_x\log{p(x|x_0)}||^2+||s_\theta(x)||^2-2s_{\theta}(x)\nabla_x\log{p(x|x_0)}\right] \\
= &\mathbb{E}_{x\sim p(x),x_0\sim p(x_0|x)}\left[||\nabla_x\log{p(x|x_0)}||^2+||s_\theta(x)||^2-2s_{\theta}(x)\nabla_x\log{p(x|x_0)}\right] \\
= &\mathbb{E}_{x\sim p(x)}\left[\mathbb{E}_{x_0\sim p(x_0|x)}\left[||\nabla_x\log{p(x|x_0)}||^2\right]+||s_\theta(x)||^2-2s_\theta(x)\mathbb{E}_{x_0\sim p(x_0|x)}\left[\nabla_x\log{p(x|x_0)}\right]\right]
\end{align*}
\]
针对最后一项 \(\mathbb{E}_{x_0\sim p(x_0|x)}\left[\nabla_x\log{p(x|x_0)}\right]\) ，我们进行如下化简：
\[\begin{align*}
\mathbb{E}_{x_0\sim p(x_0|x)}\left[\nabla_x\log{p(x|x_0)}\right]
&= \int{p(x_0|x)\nabla_x\log{p(x|x_0)}dx_0} \\
&= \int{\frac{p(x|x_0)p_x(x_0)\nabla_x\log{p(x|x_0)}dx_0}{p(x)}} \\
&= \int{\frac{p(x|x_0)\nabla_xp(x|x_0)dx_0}{p(x)}} \\
&= \frac{\nabla_x\int{p(x|x_0)}p_0(x_0)dx}{p(x)} \\
&= \frac{\nabla_x p(x)}{p(x)} \\
&= \nabla_x\log{p(x)}
\end{align*}
\]
因此可以得到如下化简结果：
\[\begin{align*}
&\mathbb{E}_{x_0,x\sim p_0(x_0)p(x|x_0)}\left[||\nabla_x\log{p(x|x_0)}-s_\theta(x)||^2\right] \\
= &\mathbb{E}_{x\sim p(x)}\left[\mathbb{E}_{x_0\sim p(x_0|x)}\left[||\nabla_x\log{p(x|x_0)}||^2\right]+||s_\theta(x)||^2-2s_\theta(x)\nabla_x\log{p(x)}\right]
\end{align*}
\]
我们将得分匹配和条件得分匹配二者的化简结果相减，发现得到一个与网络参数 \(\theta\) 无关的结果：
\[\mathbb{E}_{x\sim p(x)}\left[\mathbb{E}_{x_0\sim p(x_0|x)}\left[||\nabla_x\log{p(x|x_0)}||^2\right]-||\nabla_x\log{p(x)}||^2\right]
\]
自此，我们证明了得分匹配和条件得分匹配的优化目标等价。
存在问题
上述给定的理论基础可以确定设计score-based generative model的基本方案，但实际存在中还存在一些问题。
首先，我们给出流形假设（manifold hypothesis）的定义，现实世界的数据分布倾向于聚集在内嵌在一个高维空间（ambinet space）的低维流形（manifold）上。
基于该假设，直接训练score-based model会遇到如下问题。对于数据密度较低的区域，采样到的训练样本过少，导致该区域的score无法被准确估计，对于分布之外的区域，score的估计更加不准确，而我们使用的朗之万采样法的起点为从某一分布中采样的随机样本，大概率会落在低概率区域或是分布之外，从而无法准确地生成样本。
解决方案：Noise Conditional Score Networks (NCSN)
为了解决上述问题，NSCN采样了如下方法，利用不同尺度的高斯噪声扰动训练数据，使其覆盖整个概率空间，并训练网络估计不同尺度噪声扰动下数据分布的score。
为了权衡噪声过小无法有效填充低密度区域和噪声过大使其严重偏离原始数据分布，NCSN定义了一系列从大到小的噪声级别 \(\{\sigma_i\}^L_{i=1}\) 的高斯分布作为不同采样步数的条件分布：
\[q_{\sigma}(\tilde{x}|x)=\mathcal{N}(\tilde{x}|x,\sigma^2\mathbf{I})
\]
NCSN训练一个网络 \(s_\theta(\tilde{x},\sigma)\) 用于估计不同尺度噪声扰动下的conditional score \(\nabla_{\tilde{x}}\log{q_\sigma(\tilde{x})}\) 。
由上文证明的条件得分匹配等价于得分匹配，我们可以把训练最小化目标转化为：
\[\begin{align*}
&\mathbb{E}_{x,\tilde{x}\sim p(x)q_\sigma(\tilde{x}|x)}\left[||s_\theta(\tilde{x},\sigma)-\nabla_{\tilde{x}}\log{q_\sigma(\tilde{x}|x)}||^2\right] \\
= &\mathbb{E}_{x,\tilde{x}\sim p(x)q_\sigma(\tilde{x}|x)}\left[||s_\theta(\tilde{x},\sigma)+\frac{\tilde{x}-x}{\sigma^2}||^2\right]
\end{align*}
\]
为不同噪声级别的目标函数增加权重项 \(\lambda(\sigma_i)\) （一般是 \(\sigma_i^2\) ），我们得到最终的目标函数：
\[\mathcal{L}=\sum_{i=1}^L{\lambda(\sigma_i)}\mathbb{E}_{x,\tilde{x}\sim p(x)q_\sigma(\tilde{x}|x)}\left[||s_\theta(\tilde{x},\sigma)+\frac{\tilde{x}-x}{\sigma^2}||^2\right]
\]
NCSN的采样方法为退火朗之万采样法。首先，从某个固定的先验分布（原作者采用均匀分布）采样一个初始样本 \(\tilde{x}_0\) ，之后，从 \(\sigma_L\) 开始到 \(\sigma_1\) ，在每种噪声强度 \(\sigma_i\) 下，设定步长为：
\[\alpha_i = \epsilon \cdot \sigma_i^2 / \sigma^2_L
\]
之后按如下公式进行 \(T\) 次采样：
\[\tilde{x}_t = \tilde{x}_{t-1}+\frac{\alpha_i}{2}s_\theta(\tilde{x}_{t-1},\sigma_i)+\sqrt{\alpha_i}z_t,z_t\sim\mathcal{N}(0,\mathbf{I})
\]
其中 \(\{\sigma_i\}^L_{i=1},\epsilon,T\) 为预先设定的超参数。
