前言
本文尝试总结score-based model的模型构建，训练目标的选择，以及采样生成方式。并且采用这一视角处理前向扩散的马尔可夫过程，给出了另一种看待扩散模型的视角。
模型构建
假设数据分布为\(p(x)\)，我们试图通过神经网络的输出\(q(x|\theta)\)拟合\(p(x)\)，即
\[\begin{align}
p(x) &\sim q(x|\theta)=\frac{1}{Z_{\theta}}\exp(-f(x, \theta)) \tag{1}\label{1}\\
&\Leftrightarrow ~~~\log{p(x)} \sim \log{q(x|\theta)}=-f(x, \theta) - \log{Z_{\theta}} \tag{2}\label{2}\\
&\Leftrightarrow ~~~\nabla_{x}\log{p(x)} \sim \nabla_{x}\log{q(x|\theta)}=s(x, \theta) \tag{3}\label{3}\\
\end{align}
\]
其中，\(f(x, \theta)\)为未归一化的概率密度函数；\(Z_{\theta}\)为归一化因子，满足\(Z_{\theta}=\int{\exp(-f(x, \theta))dx}\)；\(\nabla_{x}\log{p(x)}\)称作\(p(x)\)的得分函数(score function)。我们可以采用上述不同的公式，制定不同的优化目标。
对于式\eqref{1}，优化目标可以设置为最小化\(p(x)\)和\(q(x|\theta)\)的KL散度，即\(\min{KL(p(x)\|q(x|\theta))}\)；对于式\eqref{2}，优化目标可以设置为\(\max{\mathbb{E}_{x\sim p(x)}[\log{q(x|\theta)}]}\)，等价于最小化\(p(x)\)和\(q(x|\theta)\)的KL散度；对于式\eqref{3}，优化目标可以设置为最小化\(p(x)\)和\(q(x|\theta)\)的Fisher散度，即\(\min{\mathbb{E}_{x\sim p(x)}[\|\nabla_{x}\log{p(x)}-s(x, \theta)\|^{2}]}\)。对比上述三种优化目标，我们发现，最小化\(p(x)\)和\(q(x|\theta)\)的Fisher散度时，绕开了对于\(Z_{\theta}\)的计算，而计算\(Z_{\theta}\)时，由于存在积分问题往往比较困难。因此，基于得分的模型(score-based model)，采用神经网络\(s(x, \theta)\)近似数据分布的得分函数\(\nabla_{x}\log{p(x)}\)。
训练目标
\[\min{\mathbb{E}_{x \sim p(x)}[\|\nabla_{x}\log{p(x)} - \nabla_{x}\log{q(x|\theta)}\|^2]}=\min{\mathbb{E}_{x \sim p(x)}[\|\nabla_{x}\log{p(x)} - s(x, \theta)\|^{2}]} \tag{4}\label{4}
\]
正如上一节所说，基于得分的模型采用的训练目标为式\eqref{4}。然而，实际应用中，\(\nabla_{x}\log{p(x)}\)几乎无法得到。因此，直接利用式\eqref{4}不现实。考虑到全概率公式\(p(x)=\int{p(x|z)p(z)dz}=\mathbb{E}_{z\sim p(z)}[p(x|z)]\)，可得
\[\begin{align}
\nabla_{x}\log{p(x)}
&=\frac{\nabla_{x}p(x)}{p(x)}=\frac{\nabla_{x}\int{p(x|z)p(z)dz}}{p(x)}=\frac{\int{p(z)\nabla_{x}p(x|z)dz}}{p(x)} \notag\\
&=\frac{\int{p(z)p(x|z)\nabla_{x}\log{p(x|z)}dz}}{p(x)}=\int{p(z|x)\nabla_{x}\log{p(x|z)}dz} \notag\\
&=\mathbb{E}_{z \sim p(z|x)}[\nabla_{x}\log{p(x|z)}] \notag\\
\end{align}
\]
代入式\eqref{4}可得，
\[\begin{align}
\min{\mathbb{E}_{x \sim p(x)}[\|\nabla_{x}\log{p(x)} - s(x, \theta)\|^{2}]}
&= \min{\mathbb{E}_{x\sim p(x)}[\|\mathbb{E}_{z \sim p(z|x)}[\nabla_{x}\log{p(x|z)}]-s(x, \theta)\|^{2}]} \notag\\
&= \min{\mathbb{E}_{x\sim p(x)}[\|\mathbb{E}_{z \sim p(z|x)}[\nabla_{x}\log{p(x|z)}-s(x, \theta)]\|^{2}]} \notag\\
&\leq \min{\mathbb{E}_{x\sim p(x)}\mathbb{E}_{z \sim p(z|x)}[\|\nabla_{x}\log{p(x|z)}-s(x, \theta)\|^{2}]} \notag\\
&= \min{\mathbb{E}_{x,z \sim p(x,z)}[\|\nabla_{x}\log{p(x|z)}-s(x, \theta)\|^{2}]} \notag\\
&= \min{\mathbb{E}_{z \sim p(z), x \sim p(x|z)}[\|\nabla_{x}\log{p(x|z)}-s(x, \theta)\|^{2}]} \tag{5}\label{5}
\end{align}
\]
上式中的\(\leq\)部分，由Jesen不等式\(\|\mathbb{E}[f(x)]\|^{2} \leq \mathbb{E}[\|f(x)\|^{2}]\)得出。式\eqref{5}表明，条件得分匹配是得分匹配的上界。同时，式\(\eqref{5}\)与式\(\eqref{4}\)相比，\(\nabla_{x}\log{p(x|z)}\)更加容易计算，因此，在训练时应用更加方便。事实上
\[\mathbb{E}_{z \sim p(z), x \sim p(x|z)}[\|\nabla_{x}\log{p(x|z)}-s(x, \theta)\|^{2}]
= \mathbb{E}_{x \sim p(x)}[\mathbb{E}_{z\sim p(z|x)}[\|\nabla_{x}\log{p(x|z)}-s(x, \theta)\|^{2}]]
\]
\[\begin{align}
\mathbb{E}_{z\sim p(z|x)}[\|\nabla_{x}\log{p(x|z)}-s(x, \theta)\|^{2}]
&= \mathbb{E}_{z\sim p(z|x)}[\|\nabla_{x}\log{p(x|z)}\|^{2}+\|s(x, \theta)\|^{2}-2s(x, \theta) \cdot\nabla_{x}\log{p(x|z)}] \notag\\
&= \mathbb{E}_{z\sim p(z|x)}[\|\nabla_{x}\log{p(x|z)}\|^{2}]+\|s(x, \theta)\|^{2}-2s(x, \theta) \cdot\mathbb{E}_{z\sim p(z|x)}[\nabla_{x}\log{p(x|z)}] \notag\\
&=\mathbb{E}_{z\sim p(z|x)}[\|\nabla_{x}\log{p(x|z)}\|^{2}]+\|s(x, \theta)\|^{2}-2s(x, \theta) \cdot\nabla_{x}\log{p(x)} \notag\\
\end{align}
\]
而
\[\mathbb{E}_{x \sim p(x)}[\|\nabla_{x}\log{p(x)} - s(x, \theta)\|^{2}]=\mathbb{E}_{x \sim p(x)}[\|\nabla_{x}\log{p(x)}\|^{2} + \|s(x, \theta)\|^{2} - 2s(x, \theta) \cdot\nabla_{x}\log{p(x)}]
\]
因此，条件得分匹配与得分匹配公式只相差一个与参数\(\theta\)无关的常数\(\mathbb{E}_{x \sim p(x)}[\mathbb{E}_{z\sim p(z|x)}[\|\nabla_{x}\log{p(x|z)}\|^{2}] - \|\nabla_{x}\log{p(x)}\|^{2}]\)，式\eqref{5}不仅仅是式\eqref{4}的一个上界，并且二者在训练时完全等价。
得分匹配与扩散
我们将\eqref{5}中的隐变量\(z\)设置为\(x_{0}\)，将\(x\)设置为\(x_{t}\)，同时将时间\(t\)并入网络的输入中，上述目标改写为
\[\min{\mathbb{E}_{x_{0} \sim p(x_{0}), x_{t} \sim p(x_{t}|x_{0})}[\|\nabla_{x_{t}}\log{p(x_{t}|x_{0})}-s(x_{t}, t, \theta)\|^{2}]} \tag{6}\label{6}
\]
其中，\(s(x_{t}, t, \theta)\)用于拟合\(\nabla_{x_{t}}\log{p(x_{t})}\)。保持前向扩散的加噪公式不变，即
\[x_{t}=\bar{\alpha}_{t}x_{0}+\bar{\beta}_{t}\bar{\epsilon}_{t}
~~~~ \Leftrightarrow ~~~ p(x_{t}|x_{0}) \sim \mathcal{N}(x_{t}; \bar{\alpha}_{t}x_{0}, \bar{\beta}_{t}^{2}I)
\]
\[\nabla_{x_{t}}\log{p(x_{t}|x_{0})}=-\frac{x_{t}-\bar{\alpha}_{t}x_{0}}{\bar{\beta}_{t}^{2}}=-\frac{\bar{\epsilon}_{t}}{\bar{\beta}_{t}}
\]
带入式\eqref{6}可得最终的训练目标为
\[\min {\mathbb{E}_{x_{0} \sim p(x_{0}), \epsilon \sim \mathcal{N}(0, I)}[\|-\frac{\epsilon}{\bar{\beta}_{t}} - s(x_{t}, t, \theta)\|^{2}]}
\]
朗之万采样
当网络训练完毕后，即可通过\(s(x_{t}, t, \theta)\)近似\(\nabla_{x_{t}}\log{p(x_{t})}\)，此时，可以通过朗之万动力学进行采样。
\[x_{t-1} = x_{t} + \epsilon\nabla_{x_{t}}\log{p(x_{t})} + \sqrt{2\epsilon}z_{t}, ~~~ z_{t} \sim \mathcal{N}(0, I)
\]
后记
本文回顾了score-based model的模型构建，即采用神经网络\(s(x, \theta)\)拟合分布的得分函数\(\nabla_{x}\log{p(x)}\)。推导了条件得分匹配与得分匹配之间的关系，并给出了训练时采用的训练目标。给出了在得分匹配视角下处理前向扩散过程的方法，并采用朗之万采样公式进行反向采样生成。
参考文献
- https://yang-song.net/blog/2021/score/
- https://spaces.ac.cn/archives/9509
浙公网安备 33010602011771号