[论文笔记&随手] Training with Weighted Sum of Denoising Score Matching Objectives

[note] Training with Weighted Sum of Denoising Score Matching Objectives

利用 去噪分数匹配目标的加权和 进行训练，去噪指的是使用sde的方法就不需要自行补充噪声了。

本文的目的是解释如何对原始数据进行扰动。 from https://yang-song.github.io/blog/2021/score/

一、理论

首先，挑选一个随机过程（SDE）对原始数据分布$p_0$进行扰动得到扰动后数据的概率密度分布$p_t$。

本文选择的随机过程为：
$$d\mathbf{x}={\sigma }^{t}d\mathbf{w},t\in [0,1]$$
在这种情况下，扰动后数据的概率密度分布$p_t$，在原始数据下的条件概率分布为：
$$p_{0t}(\mathbf{x}(t)\mid \mathbf{x}(0))=\mathcal{N}(\mathbf{x}(t);\mathbf{x}(0),\frac{1}{2\log \sigma }({\sigma }^{2t}-1)\mathbf{I})$$
关于这个函数的解释是，使用参数$ \frac{1}{2\log \sigma}(\sigma^{2t} - 1) $作为我们的权重函数，即$\lambda(t) = \frac{1}{2 \log \sigma}(\sigma^{2t} - 1)$.

当参数$\sigma$变得非常大的时候，其中的先验分布$p_{t=1}$，也就是最终扰动后的数据分布就可以变成一个正太分布：
$$\int p_0(\mathbf{y})\mathcal{N}(\mathbf{x};\mathbf{y},\frac{1}{2\log \sigma }({\sigma }^2-1)\mathbf{I})d\mathbf{y}\approx \mathbf{N}(\mathbf{x};0,\frac{1}{2\log \sigma }({\sigma }^2-1)\mathbf{I}),$$
直观地说，这个SDE通过一个变种函数$\frac{1}{2log\sigma }({\sigma }^{2t}-1)$帮助我们捕获了高斯扰动的数据变量集合（连续统continuum），即$x(t)$。这个数据变量集合可以帮助我们逐渐将原始数据分布$p_0$变成了一个简单的高斯分布$p_1$，也就是t=1时候的分布。

二、代码实现

1） 对t进行连续采样

 # 对时间特征t进行均匀采样
 random_t = torch.rand(x.x.shape[0]//30, device=device) * (1. - eps) + eps # 防止采样到0

2）定义权重函数

可以看到，这里定义的权重函数就是作者在上面提到的$\lambda (t)$函数。

def marginal_prob_std(t, sigma):
    # t = torch.tensor(t, device=device)
    return torch.sqrt((sigma ** (2 * t) - 1.) / 2. / np.log(sigma))

3）对数据进行扰动

# 表征时间的特征t, 从0到1上进行均匀采样
random_t = torch.rand(batchsize, device=device) * (1. - eps) + eps # 这里的eps是为了防止采样到t=0

# 构造一个与原始数据结构一样的向量，并在[0,1)上进行均匀采样。
z = torch.randn_like(x.x)

# 利用前面均匀采样的时间特征t，求得权重函数的值，这个权重函数的目的就是为了使得t=1时的扰动数据达到一个正太分布的结果。重复30遍的目的是因为一轮训练中设置的batch_size = 30
std = marginal_prob_std_func(random_t).repeat(1, 30).view(-1, 1)

# 这里将噪声与标准差相乘，
perturbed_x = copy.deepcopy(x)
perturbed_x.x += z * std

4）利用扰动的数据进行训练

需要补充一下，为了训练积分函数模型，目前的目标函数变成了下面这个样子：
$${\mathbb{E}}_{t\in u(0,T)}{\mathbb{E}}_{p_t(x)}[\lambda (t)||{\nabla }_{x}log{p}_t(x)-s_{\theta }(x,t)|{|}_2^2]$$

这里是最基本的目标函数的样子：
$${\mathbb{E}}_{p(x)}[{||{\nabla }_{x}logp(x)-s_{\theta }(x)||}_2^2]=\int p(x)||{\nabla }_{x}logp(x)-s_{\theta }(x)|{|}_2^{2}dx.$$
为了估计这个目标函数，需要如下估计，即使用Score Matching的方法进行估计（Hyvärinen 2005）：

可以看到，去估计如下的目标函数是可以达到的。

$${\mathbb{E}}_{p_{data}(x)}[\frac{1}{2}||s_{\theta }(x)|{|}_2^2+trace({\nabla }_x{s}_{\theta }(x))]$$

具体上，体现在代码上，用的是如下的公式：
$$\begin{gathered} \frac{1}{N}\sum _{i=1}^N[\frac{1}{2}||s_{\theta }(x_i)|{|}_2^2+trace({\nabla }_x{s}_{\theta }(x_i))] \\ \approx \frac{1}{N}\sum _{i=1}^N[\frac{1}{2}||s_{\theta }(x_i)|{|}_2^2+trace({\nabla }_x{s}_{\theta }(x_i)) \end{gathered}$$

# 计算积分函数的值
output = model(perturbed_x, random_t)
# score matching的损失函数，与上式不一致的原因在于，本文的目标函数中还有一个参数\lambda(t)，所以表现为如下的形式。
loss_ = torch.mean(torch.sum(((output * std + z)**2).view(batch_size, -1)), dim=-1)
# 一轮训练之后，将score matching的目标函数的结果返回
return loss_

🙋‍♂️ 我有一个问题，这个目标函数是怎么推理得到的呀? 🤔