Score based model学习笔记

内容来源于：

人民大学 李崇轩老师 《深度生成模型：原理与应用》

CW的博客

《深度生成模型》 Jakub M.T

Energy based model

一种最灵活的概率建模方式，其通过神经网络定义了样本的能量。其中，能量越低状态越稳定，应当拥有更高的概率。

但是，对于一个概率分布来说，其要满足归一化和非负性的要求。

也就是

$$p(x)>0$$

和

$$\int p(x)dx = 1$$

可以通过在神经网络外套非负函数$g$的方法构造出非负的$p_{\theta}(x)$，也就是$g(f_{\theta}(x))$。

到这里，虽然能构造出满足非负性的函数，但是很难保证其归一化，也就是相加为1。

因此，为了探究每个样本取值的概率密度，应当构造配分函数$Z(\theta) = \int g_{\theta}(x)dx$。

$$p_\theta(\mathbf{x})=\frac1{Volume(g_\theta)}g_\theta(\mathbf{x})=\frac1{\int g_\theta(\mathbf{x})d\mathbf{x}}g_\theta(\mathbf{x})$$

这样就能满足$\int p_{\theta}(x)dx = 1$的归一化要求。

最终，对于EBM来说，选取的$g$函数为$exp$指数函数，即：

$$Z(\theta)=\int\exp(f_\theta(\mathbf{x}))d\mathbf{x}$$

其中$-f_{\theta}(x)$为能量函数，能量越低概率越高。

优点：
高度的灵活性，可以使用任意的神经网络去定义$f$

缺点：

• 1. 采样困难，根本没有定义如何采样
• 2. 计算似然$p_{\theta}(x)$困难，因为计算配分函数$Z(\theta)$很困难，高纬度积分。
• 3. 没有表征学习的过程

但是，这种概率的表示方法同时很擅长对概率密度做比较，即：

$$\frac{p_\theta(\mathbf{x})}{p_\theta(\mathbf{x}^\prime)}=\exp(f_\theta(\mathbf{x})-f_\theta(\mathbf{x}^\prime))$$

External：乘积专家模型

如果，已经训练好了多个模型$p,q,r$,则可以用以下表示其$g_{\theta}(f(x))$：

$$p_{\theta_1}(\mathbf{x})q_{\theta_2}(\mathbf{x})r_{\theta_3}(\mathbf{x})$$

最终其概率即为：

$$p_{\theta_1,\theta_2,\theta_3}(\mathbf{x})=\frac1{Z(\theta_1,\theta_2,\theta_3)}q_{\theta_1}(\mathbf{x})r_{\theta_2}(\mathbf{x})t_{\theta_3}(\mathbf{x})$$

通过这样的方法，可以实现类似于AND与的操作。（$p,q,r$三个模型都认可的样本才会有较高的概率，混合专家模型MOE更像是或的操作）

应用：利用强化学习和乘积专家模型对已经训练好的模型，如GPT，Diffusion等做价值对齐。

目标，最大化：

$$\frac{\exp\{f_\theta(\mathbf{x}_{train})\}}{Z(\theta)}$$

这意味着，训练过程需要相对于不正确的数据点相对地抬高正确的样本的概率。

这里，EBM会使用当前模型的采样作为负样本：随机的噪声对于模型不具有代表性。因为大部分时候从随机点走的过程，都不会经过这些随机噪声。

而当前模型的采样，代表了其采样时可能会到达的位置，更有代表性（我猜的）。并且当前模型的采样（假设认为可以训练到采样数据在真实数据周围的情况），这样采样出的样本与真实数据类似，但是作为负样本可以有很高的效率（相对于随机噪声的负样本）。

而当模型训练的较好时，即采样与真实数据相似，真实数据拉高概率，与真实值类似的数据拉低概率，逐渐形成模型的收敛。

训练优化目标

最大化似然

$$\max_\theta f_\theta(x_{train})-\log Z(\theta)$$

这里，直接对优化目标计算对参数$\theta$的梯度

$$\begin{aligned} &\nabla_\theta f_\theta(x_{train})-\nabla_\theta\log Z(\theta) \\ &= \nabla_\theta f_\theta(x_{train})-\frac{\nabla_\theta Z(\theta)}{Z(\theta)} && \mathrm{logZ(\theta)看作复合函数}\\ &= \nabla_\theta f_\theta(x_{train})-\frac1{Z(\theta)}\int\nabla_\theta\exp\{f_\theta(x)\}\mathrm{d}x && \mathrm{对\theta求导与x无关，积分求导换序}\\ &= \nabla_\theta f_\theta(x_{train})-\frac1{Z(\theta)}\int\exp\{f_\theta(x)\}\nabla_\theta f_\theta(x)\mathrm{d}x && \mathrm{把exp看作复合函数}\\ &= \nabla_\theta f_\theta(x_{train})-\int\frac{\exp\{f_\theta(x)\}}{Z(\theta)}\nabla_\theta f_\theta(x)\mathrm{d}x && \mathrm{Z(\theta)对于积分是常数}\\ &= \nabla_\theta f_\theta(x_{train})-E_{x_{sample}}[\nabla_\theta f_\theta(x_{sample})] && \mathrm{用采样代替全局样本}\\ &\approx \nabla_\theta f_\theta(x_{train})-\nabla_\theta f_\theta(x_{sample})&& \mathrm{蒙特卡洛估计} \end{aligned}$$

Note：

优化目标函数的时候，不能随便加减项，可能会导致求梯度时，梯度方向产生严重错误。而确定了目标函数下界，优化下界，就能间接最大化目标函数。

而如果先求梯度再近似梯度，只要期望意义下，梯度无偏即可。

那么，对于EBM，如何进行采样？

马尔科夫链蒙特卡洛方法，利用能量函数或者概率比值去指导从随机点出发的采样过程。

这里引入了Score function，得分函数中对$logZ(\theta)$求导的过程因为与$x$无关，所以直接等于0。

$$s_\theta(\mathbf{x}):=\nabla_x\log p_\theta(\mathbf{x})=\nabla_xf_\theta(\mathbf{x})-\underbrace{\nabla_x\log Z(\theta)}_{=0}=\nabla_xf_\theta(\mathbf{x})$$

Score function的几何含义为，描述了一个场，代表当前样本点能最快提升概率密度的方向。样本点沿着这个场的方向走，最终会到达一个概率的极大值点。并且，在随时间转移的每一步都加上一定程度的高斯噪声，可以带来一定程度的随机性。

$$x_{i+1}\leftarrow x_{i}+\epsilon\nabla_{x}\log p_{\theta}(x)+\sqrt{2\epsilon}z_{i},z_{i}\sim N(0,I)$$

Likelihood free learning

最大似然的方法直接催生的就是KL散度最小化的优化任务。这里尝试，不基于似然，但是仍然基于散度的方法去定义优化目标。

这里，由于score function和概率密度是一体两面的，尝试基于score function定义散度：

Fisher 散度定义如下：

$$D_{F}(p,q):=\frac{1}{2}E_{x\sim p}[\|\nabla_{x}\log p(x)-\nabla_{x}\log q(x)\|_{2}^{2}]$$

数据分布$p_{data}(x)$和EBM模型的分布$p_{\theta}(x)$之间的散度定义如下：

$$\frac12E_{\mathbf{x}\sim p_{\mathrm{data}}}[\|\nabla_{\mathbf{x}}\log p_{\mathrm{data}}(\mathbf{x})-s_{\theta}(\mathbf{x})\|_{2}^{2}]\\=\frac12E_{\mathbf{x}\sim p_{\mathrm{data}}}[\|\nabla_{\mathbf{x}}\log p_{\mathrm{data}}(\mathbf{x})-\nabla_{\mathbf{x}}f_{\theta}(\mathbf{x})\|_{2}^{2}]$$

Score matching：

可以将fisher 散度进行如下变换，证明在后面（利用分步积分将交叉项转为二阶梯度，避免了计算真实的似然，不想计算一个东西的梯度就用分步积分转换）

$$\begin{aligned}&\frac{1}{2}E_{\mathbf{x}\sim p_{\mathrm{data}}}[\|\nabla_{\mathbf{x}}\log p_{\mathrm{data}}(\mathbf{x})-\nabla_{\mathbf{x}}\log p_{\theta}(\mathbf{x})\|_{2}^{2}]\\&=E_{\mathbf{x}\sim p_{\mathrm{data}}}\left[\frac{1}{2}\|\nabla_{\mathbf{x}}\log p_{\theta}(\mathbf{x})\|_{2}^{2}+\mathrm{tr}(\underbrace{\nabla_{\mathbf{x}}^{2}\log p_{\theta}(\mathbf{x})}_{\text{Hessian of }\log p_{\theta}(\mathbf{x})})\right]+\mathrm{const.}\end{aligned}$$

每次，从训练集里采一组batch，计算Fisher散度的损失，并进行随机梯度下降优化$\theta$。

$$\begin{aligned}&\frac1n\sum_{i=1}^n\left[\frac12\|\nabla_\mathbf{x}\log p_\theta(\mathbf{x}_i)\|_2^2+\mathrm{tr}(\nabla_\mathbf{x}^2\log p_\theta(\mathbf{x}_i))\right]\\&=\frac1n\sum_{i=1}^n\left[\frac12\|\nabla_\mathbf{x}f_\theta(\mathbf{x}_i)\|_2^2+\mathrm{tr}(\nabla_\mathbf{x}^2f_\theta(\mathbf{x}_i))\right]\end{aligned}$$

一维情况下的证明：

$$\begin{aligned} &\int\left|\left|\nabla_{x}\log p_{data}(x)-\nabla_{x}f_{\theta}(x)\right|\right|^{2}p_{data}(x)dx \\ &=\int\left(\nabla_{x}f_{\theta}(x)^{2}-2\nabla_{x}\log p_{data}(x)\nabla_{x}f_{\theta}(x)\right)p_{data}(x)dx+const && \mathrm{p_{data}的一项与\theta优化无关，直接作为常数} \\ &=-2\int\nabla_{x}\log p_{data}(x)\nabla_{x}f_{\theta}(x) p_{data}(x)dx+E_{p}((\nabla_{x}f_{\theta}(x))^{2})+const && \mathrm{提出来第一个平方项化成期望} \\ &=-2\int\nabla_{x}f_{\theta}(x) \nabla_{x}p_{data}(x)dx+E_{p}(\nabla_{x}f_{\theta}(x))^{2}+const && \mathrm{对logP求导，消掉一个P} \\ & =-2 (\nabla_{x}f_{\theta}(x) p_{data}(x)|_{-\infty}^{+\infty}-E_{p_{data}}\nabla_{x}^{2}f_{\theta}(x))+E_{p_{data}}(\nabla_{x}f_{\theta}(x))^{2}+const && \mathrm{分步积分}\\ &=2E_{p_{data}}(\nabla_{x}^{2}f_{\theta}(x)+\frac{1}{2}(\nabla_{x}f_{\theta}(x))^{2})+const && \mathrm{无穷远处概率为0} \end{aligned}$$

分步积分公式：

$$\begin{aligned} \int_{a}^{b}u(x)v^{\prime}(x) dx& =\left[u(x)v(x)\right]_a^b-\int_a^bu'(x)v(x) dx \\ &=u(b)v(b)-u(a)v(a)-\int_a^bu^{\prime}(x)v(x) dx. \end{aligned}$$

Score based model

不再间接的定义能量函数，而是直接使用神经网络建模score function。

由于score function是一个从$R^N$到$R^N$的一个映射过程，在网络结构设计中，U-Net这种具有远程连接的模型往往会有更好的效果。

$$L(\theta)=\operatorname{E}_{p_{D}(x)}\frac{1}{2}|| s_{\theta}(x)-\nabla_{x}\log p_{D}(x) ||_{2}^{2}$$

但是，存在一个问题,在数据分布的概率密度低的地方，往往由于没有采样到过这些点，也没有对这些地方进行参数更新，而导致这些地方的score function的估计并不准确。也是深度生成模型存在的通病（对data分布之外的地方建模和优化过少），但是由于其他模型大多数祖先采样，而能一定程度上避免这个问题带来的影响。

不完美的解决方案：对真实数据加噪，使得样本逐渐到弥散整个空间，而整个空间的score function都能得到比较好的优化。与此同时，虽然加噪能使模型更容易学习，但是由于加噪会导致模型采样出的样本产生模糊。

进一步的解决方案，定义一系列的噪声等级，采样时先依据高噪声的score function进行采样，在此基础上再依据噪声更小的score function进行更新采样。

据此，模型的损失：

$$L(\theta)=\sum_{i=1}^{L}\lambda_{i} \mathrm{E}_{p_{\sigma_{i}}(x)}\frac{1}{2}\parallel s_{\theta}(x,i)-\nabla_{x}\log p_{\sigma_{i}}(x)\parallel_{2}^{2}$$

实际实现时，不会使用多个神经网络参数化多个score function，而是会将输入多加一维，用于代表当前的噪声等级，进行一个参数共享。

Denoise score matching

普通的score matching方法需要计算二阶导，计算负担大。

首先，不知道$p_{data}(x)$，并且类似于上文所述，可以在真实数据上加噪声构造能覆盖更大样本空间的操作一样，这里尝试对带噪的fisher 散度进行建模。

$$\frac12E_{\tilde{\mathbf{x}}\sim p_{\mathrm{data}}}[\|\boldsymbol{s}_\theta(\tilde{\mathbf{x}})-\nabla_{\tilde{\mathbf{x}}}\log q_\sigma(\tilde{\mathbf{x}})\|_2^2]$$

也就是，优化神经网络描述的分数函数场，近似匹配真实数据加噪情况的分数函数场。

而此时加噪真实数据的梯度仍然是未知的，这里通过建模$q_\sigma(\tilde{x}|x)\sim N(\tilde{x};x,\sigma^2I)$，这样一个已经预定义好的分布来代替原有的概率密度（下面这个等式的证明见Post1）

$$\mathbb{E}_{q_\sigma(\tilde{x})}\left[||s_\theta(\tilde{x})-\frac{\partial log(q_\sigma(\tilde{x}))}{\partial\tilde{x}}||^2\right]=\mathbb{E}_{q_\sigma(\tilde{x}|x)p_{data}(x)}\left[||s_\theta(\tilde{x})-\frac{\partial log(q_\sigma(\tilde{x}|x))}{\partial\tilde{x}}||^2\right]$$

这样，优化目标即可变为后者。后者期望中第二项对$\tilde x$求偏导的过程可以得到，计算变得非常简单：

$$\frac{\partial log(p(\tilde x|x))}{\partial\tilde x}=-\bigl(\frac{\tilde x-x}{\sigma^2}\bigr)=-\frac{\epsilon}{\sigma}$$

其中，

$$\tilde x = x+\epsilon, \epsilon \sim N(0,1)$$

这里，需要注意，由于$s_{\theta}$近似的是加噪后的数据分布，因此需要保证$\sigma$不能太大，否则拟合出的数据分布会严重偏离真实数据分布（实际上通过多级加噪策略，应该还好？）

Post1：

这里的ESM和DSM分别代表显示分数匹配和去噪分数匹配。

$$\begin{aligned} &\mathrm{ESM}_{q_{\sigma}(\theta)}=\mathbb{E}_{q_{\sigma(\tilde{x})}}\left[||s_{\theta}(\tilde{x})-\frac{\partial log(q_{\sigma}(\tilde{x}))}{\partial\tilde{x}}||^{2}\right] \\ &=\mathbb{E}_{q_{\sigma(\tilde{x})}}\left[\left|\left|s_{\theta}(\tilde{x})\right|\right|^{2}\right]-2\mathbb{E}_{q_{\sigma(\tilde{x})}}\left[\left\langle s_{\theta}(\tilde{x}),\frac{\partial log(q_{\sigma}(\tilde{x}))}{\partial\tilde{x}}\right\rangle\right]+C_{1}, \\ &C_1=\mathbb{E}_{q_{\sigma(\tilde{x})}}\left[||\frac{\partial log(q_{\sigma}(\tilde{x}))}{\partial\tilde{x}}||^2\right] \end{aligned}$$

$C_1$与$\theta$无关，因此可以作为优化过程中的一个常数。

$$\begin{aligned} &\mathrm{DSM}_{q_{\sigma}(\theta)}=\mathbb{E}_{q_{\sigma}(\tilde{x},x)}\left[||s_{\theta}(\tilde{x})-\frac{\partial log(q_{\sigma}(\tilde{x}|x))}{\partial\tilde{x}}||^{2}\right],\quad q_{\sigma}(\tilde{x},x)=q_{\sigma}(\tilde{x}|x)p(x) \\ &=\mathbb{E}_{q_{\sigma(\tilde{x},x)}}\left[\left|\left|s_\theta(\tilde{x})\right|\right|^2\right]-2\mathbb{E}_{q_{\sigma(\tilde{x},x)}}\left[\left\langle s_\theta(\tilde{x}),\frac{\partial log(q_\sigma(\tilde{x}|x))}{\partial\tilde{x}}\right\rangle\right]+C_2, \\ &C_2=\mathbb{E}_{q_{\sigma(\tilde{x},x)}}\left[||\frac{\partial log(q_{\sigma}(\tilde{x}|x))}{\partial\tilde{x}}||^2\right] \end{aligned}$$

这里的$C_2$也是优化的常数。

$$\begin{aligned} &\mathbb{E}_{q_{\sigma(\tilde{x})}}\left\lfloor\left|\left|s_{\theta}(\tilde{x})\right|\right|^{2}\right\rfloor=\int_{\tilde{x}}q_{\sigma}(\tilde{x})\left|\left|s_{\theta}(\tilde{x})\right|\right|^{2}d\tilde{x} \\ &\xrightarrow{q_\sigma(\tilde{x})=\int_xq_\sigma(\tilde{x}|x)p(x)dx}\int_{\tilde{x}}\int_xq_\sigma(\tilde{x}|x)p(x)dx||s_\theta(\tilde{x})||^2d\tilde{x} \\ &=\int_{\tilde{x}}\int_xq_\sigma(\tilde{x}|x)p(x)||s_\theta(\tilde{x})||^2dxd\tilde{x}&& \mathrm{||s_{\theta(x_{\theta}})||^2}项与x无关，直接提进去 \\ &=\int_{\tilde{x}}\int_xq_\sigma(\tilde{x},x)||s_\theta(\tilde{x})||^2dxd\tilde{x} \\ &=\mathbb{E}_{q_{\sigma}(\tilde{x},x)}\left[\left|\left|s_{\theta}(\tilde{x})\right|\right|^{2}\right] \end{aligned}$$

这样，即简单的证明了ESM和DSM中的第一项是完全等价的。

下面证明了ESM和DSM的第二项也是完全等价的。下面的尖括号代表内积。

$$\begin{aligned} &\mathbb{E}_{q_{\sigma(\tilde{x})}}\left[\left\langle s_{\theta}(\tilde{x}),\frac{\partial log(q_{\sigma}(\tilde{x}))}{\partial\tilde{x}}\right\rangle\right]=\int_{\tilde{x}}q_{\sigma}(\tilde{x})\left\langle s_{\theta}(\tilde{x}),\frac{\partial log(q_{\sigma}(\tilde{x}))}{\partial\tilde{x}}\right\rangle d\tilde{x} \\ &=\int_{\tilde{x}}q_{\sigma}(\tilde{x})\left\langle s_{\theta}(\tilde{x}),\frac{1}{q_{\sigma}(\tilde{x})}\frac{\partial q_{\sigma}(\tilde{x})}{\partial\tilde{x}}\right\rangle d\tilde{x} && \mathrm{log函数复合求导}\\ &=\int_{\tilde{x}}\left\langle s_{\theta}(\tilde{x}),\frac{\partial q_{\sigma}(\tilde{x})}{\partial\tilde{x}}\right\rangle d\tilde{x} && \mathrm{内外两个q抵消}\\ &\xrightarrow{q_\sigma(\tilde{x})=\int_xq_\sigma(\tilde{x}|x)p(x)dx}\int_{\tilde{x}}\left\langle s_\theta(\tilde{x}),\frac{\partial}{\partial\tilde{x}}\int_xq_\sigma(\tilde{x}|x)p(x)dx\right\rangle d\tilde{x} \\ &=\int_{\tilde{x}}\left\langle s_\theta(\tilde{x}),\int_x\frac{\partial q_\sigma(\tilde{x}|x)}{\partial\tilde{x}}p(x)dx\right\rangle d\tilde{x} && \mathrm{对\tilde x求导，直接带进去（p(x)与\tilde x无关）}\\ &=\int_{\tilde{x}}\left\langle s_{\theta}(\tilde{x}),\int_{x}\frac{\partial log(q_{\sigma}(\tilde{x}|x))}{\partial\tilde{x}}q_{\sigma}(\tilde{x}|x)p(x)dx\right\rangle d\tilde{x} && \mathrm{对log求导，会多产生一个q的倒数抵消}\\ &=\int_{\tilde{x}}\int_{x}p(x)q_{\sigma}(\tilde{x}|x)\left\langle s_{\theta}(\tilde{x}),\frac{\partial log(q_{\sigma}(\tilde{x}|x))}{\partial\tilde{x}}\right\rangle dxd\tilde{x} && \mathrm{s_{\theta}与x积分无关}\\ &=\int_{\tilde{x}}\int_{x}q_{\sigma}(\tilde{x},x)\left\langle s_{\theta}(\tilde{x}),\frac{\partial log(q_{\sigma}(\tilde{x}|x))}{\partial\tilde{x}}\right\rangle dxd\tilde{x} && \mathrm{条件乘边缘得联合}\\ &=\mathbb{E}_{q_{\sigma}(\tilde{x},x)}\left[\left\langle s_{\theta}(\tilde{x}),\frac{\partial log(q_{\sigma}(\tilde{x}|x))}{\partial\tilde{x}}\right\rangle\right] \end{aligned}$$

这样，就能得到：

$$\begin{aligned}&\mathrm{ESM}_{q_{\sigma}(\theta)}=\mathrm{DSM}_{q_{\sigma}(\theta)}-C_{2}+C_{1}\\&\Rightarrow\mathrm{ESM}_{q_{\sigma}(\theta)}\equiv\mathrm{DSM}_{q_{\sigma}(\theta)}\end{aligned}$$

也就是说ESM和DSM从 优化$\theta$的角度 来说是等价的。