生成模型SDE

平丘月初

已于 2024-08-08 17:11:49 修改

阅读量823

点赞数 23

分类专栏：算法工作笔记文章标签：机器学习深度学习人工智能

于 2024-08-08 17:11:07 首次发布

本文链接：https://blog.csdn.net/u011994454/article/details/141031749

版权

算法工作笔记专栏收录该内容

3 篇文章 0 订阅

订阅专栏

SDE

Stochastic Differential Equation，在多个噪声尺度上扰动数据是SMLD和DDPM成功的关键，SDE将其泛化到连续时间域，也即无限噪声尺度。

前向过程

用于数据扰动，训练得分模型。
$d x = f (x, t) d t + g (t) d t$

f是x(t)的drift系数，g是x(t)的diffusion系数。

逆向过程

初始化噪声，采样生成样本。
$g(t)^2\Delta_xlogp_t(x)]dt + g(t)d\bar{w}$

目标函数

$\theta^{*} = \mathop{argmin}\limits_{\theta}E_t\{\lambda(t)E_{x(0)}E_{x(t)|x(0)}[||s_\theta(x(t), t)-\Delta_{x(t)}log_{p_\sigma(t)}(x(t)|x(0))||_2^2]\}$

$\begin{aligned} \Delta_{\widetilde{x}}log q(\widetilde{x}|x) &= \Delta_{\widetilde{x}} log\frac{1}{\sqrt{2\pi \sigma^2}}exp\{-\frac{||\widetilde{x}-x||^2}{2\sigma^2}\} \\ &= \Delta_{\widetilde{x}} (log\frac{1}{\sqrt{2\pi \sigma^2}} - \frac{||\widetilde{x}-x||^2} {2\sigma^2})\\ &=-\frac{\widetilde{x} - x}{\sigma^2} \\ &=-\frac{z}{\sigma} \end{aligned}$
所以SDE的loss代码实现为：

if not likelihood_weighting:
	losses = torch.square(score * std[:, None, None, None, None] + z) # 用\sigma^2加权loss
	losses = torch.mean(losses.reshape(losses.shape[0], -1), dim=-1)# [N,]
else:
	g2 = sde.sde(torch.zeros_like(batch, t)[1] ** 2
	losses = torch.square(score + z / std[:, None, None, None])
	losses = torch.mean(losses.reshape(losses.shape[0], -1), dim=-1) * g2
losses = torch.mean(losses)

if中用的SDE推荐的loss加权方式，else中参考了Maximum Likelihood Training of Score-Based Diffusion Models这篇工作中的loss加权方式，默认使用的是else中的加权方式。

逆向过程求解

$s_{\theta}$ 训好后，我们用它构建逆向SDE，通过数值方法生成符合 $p_0$ 数据分布的样本。

通用数值ODE求解器

数值求解器提供了SDE的近似轨迹。存在许多通用的SDE数值求解方法，比如Euler-Maruyama和随机Runge-Kutta方法。
在CIFAR-10数据集上，reverse diffusion samplers比SMLD和DDPM中用的ancestral sampling要略好一些。

Predictor-Corrector采样器

有点复杂，而且巨慢，有需要再看。

Probability Flow ODE

对于所有的扩散过程，存在一个对应的确定性过程，该过程的轨迹和SDE共享边际概率密度 ${p_t(x)\}_{t=0}^T$ 。该确定性过程满足一个ODE：
$\frac{1}{2}g(t)^2\Delta_xlogp_t(x)]dt$
我们称该ODE为probability flow ODE。

似然计算

操纵隐式表征

高效采样

VE, VP and SUB-VP SDEs推导过程

SMLD和DDPM的噪声扰动分别是Variance Exploding(VE)和Variance Preserving(VP) SDEs的离散化形式。sub-VP SDEs做为VP SDEs的修正，常常可以使生成的样本质量和似然表现更好。

VE SDEs

假设总计有N个噪声尺度，SMLD的每个扰动核 $p_{\sigma_i}(x_i|x_0)$ 可从下述马尔可夫链推导出来：

$x_i = x_{i-1} + \sqrt{\sigma^2_i - \sigma^2_{i-1}}z_{i-1}, i = 1, ..., N \tag{20}$
where, $z_{i-1}\sim N(\bold{0}, \bold{I)}, x_0 \sim p_{data}, \sigma_0=0$ , 当 $N$ -> $\infty$ 时，马尔可夫链 ${x_i\}_{i=1}^N$ 变为连续随机过程 $\{x(t)\}^1_{t=0}$ , $\{\sigma_i\}_{i=1}^N$ 变为函数 $\sigma(t)$ , $z_i$ 变为 $z (t)$ , 这里我们使用连续时间变量 $t\in[0, 1]$ 做索引，而不是整数 $i\in\{1, 2, ..., N\}$ 。 $x(\frac{i}{N})=x_i$ , $\sigma(\frac{i}{N})=\sigma_i$ , $z(\frac{i}{N})=z_i$ for i = 1, 2, …, N。可用 $\Delta t=\frac{1}{N}$ 和 $\in\{0, \frac{1}{N}, ..., \frac{N-1}{N}\}$ 重写上式(20)：
$x(t+\Delta t) = x(t) + \sqrt{\sigma^2(t+\Delta t) - \sigma^2(t)} z(t) \approx x(t) + \sqrt{\frac{d[\sigma^2(t)]}{dt}\Delta t} z(t)$
这里的增量 $\sqrt{\Delta t}z(t)\sim N(0, \Delta t)$ 所构成的随机过程，天然满足维纳过程，进而可得：
$\sqrt{\frac{d[\sigma^2(t)]}{dt}} dw \tag{21}$
即不存在drift量。

VP SDEs

DDPM中用到的扰动核 $\{p_{\sigma_i}(x_i|x_0)\}_{i=1}^N$ ，离散马尔可夫链为：
$x_i = \sqrt{1-\beta_i}x_{i-1} + \sqrt{\beta_i}z_{i-1}, i = 1,..., N \tag{22}$
where, $z_{i-1}\sim N(\bold{0}, \bold{I)}$ . 为了获得N-> $\infty$ 时, 该马尔可夫链的极限，我们定义噪声尺度的辅助集合 $\{\bar{\beta}_i = N\beta_i\}_{i=1}^N$ ，并重写式(22)如下：
$x_i = \sqrt{1-\frac{\bar{\beta}_i}{N}}x_{i-1} + \sqrt{\frac{\bar{\beta}_i}{N}}z_{i-1}, i = 1,..., N \tag{23}$
当N趋于无穷大时， $\{\bar{\beta}_i = N\beta_i\}_{i=1}^N$ 成为以 $t\in[0, 1]$ 为索引的函数 $\beta(t)$ ， $\beta(\frac{i}{N})=\bar{\beta}_i$ ， $x(\frac{i}{N})=x_i$ , , $z(\frac{i}{N})=z_i$ ，可用 $\Delta t=\frac{1}{N}$ 和 $\in\{0, \frac{1}{N}, ..., \frac{N-1}{N}\}$ 重写上式(23):
$\begin{aligned} x(t+\Delta t) &= \sqrt{1-\beta(t+\Delta t)\Delta t} x(t) + \sqrt{\beta(t+\Delta t)\Delta t}z(t)\\ &\approx x(t) - \frac{1}{2}\beta(t+\Delta t)\Delta t x(t) + \sqrt{\beta(t+\Delta t)\Delta t}z(t)\\ &\approx x(t) - \frac{1}{2}\beta(t)\Delta tx(t) + \sqrt{\beta(t)\Delta t}z(t) \tag{24} \end{aligned}$

上式用到了泰勒近似，当 $\Delta t$ << 1时，上式中的近似相等成立。因此当 $\Delta$ -> 0时，式(24)收敛到下述VP SDE：
$\frac{1}{2}\beta(t)xdt + \sqrt{\beta(t)}dw \tag{25}$

当 $t$ -> $\infty$ 时，VE SDE是个方差爆炸的过程。相比之下，VP SDE过程的方差是有界的，此外，当 $p (x (0))$ 是单位方差时，对于所有 $t\in[0, \infty)$ ，该过程是固定的单位方差。

根据数理基础可得:
$\frac{d\sum\nolimits_{VP}(t)}{dt} = \beta(t)(I - \sum\nolimits_{VP}(t))$
$\sum\nolimits_{VP}(t)$ 是VP SDE $x (t)$ 的协方差，解这个ODE可得：
$\sum\nolimits_{VP}(t) = I + e^{\int_0^t-\beta(s)ds}(\sum\nolimits_{VP}(0) - I)$
据此可见，给定 $\sum\nolimits_{VP}(0)$ ， $\sum\nolimits_{VP}(t)$ 总是有界的，此外，当 $\sum\nolimits_{VP}(0)=I$ 时， $\sum\nolimits_{VP}(t)$ 恒等于 $I$ 。

受VP-SDE启发，提出了新的SDE叫sub-VP SDE：
$\frac{1}{2}\beta(t)xdt + \sqrt{\beta(t)(1-e^{-2\int_0^t\beta(s)ds})}dw \tag{25}$
VP和sub-VP SDE的期望 $E [x (t)]$ 是相同的。方差函数不同：
$\sum\nolimits_{sub-VP}(t) = I + e^{-2\int_0^t\beta(s)ds}I+ e^{-\int_0^t\beta(s)ds}(\sum\nolimits_{sub-VP}(0) - 2I)$
可以有如下发现：

当 $\sum\nolimits_{sub-VP}(0) = \sum\nolimits_{VP}(0)$ 并共享 $\beta(s)$ 时，对于所有的 $t\geq0$ ， $\sum\nolimits_{sub-VP}(t) \leq \sum\nolimits_{VP}(t)$ 。
如果 $lim_{t->\infty}\int_0^t\beta(s)ds=\infty$ ， $lim_{t->\infty}\sum\nolimits_{sub-VP}(t) = lim_{t->\infty}\sum\nolimits_{VP}(t) = I$ 。
从1可知，我们为什么称之为sub-VP SDE，因为其方差总是被VP SDE限定。

这三个SDE的扰动核 $p_{\sigma_t}(x(t)|x(0))$ 如下：

在这里插入图片描述
其实这三个就是代码中的marginal_prob $p_t(x)$ ，VE SDE计算mean和std非常简单，不做进一步解释。这里讲下(sub-)VP SDE中的积分项，关于 $\beta(t)$ ，我们定义了其两端值(0, 1e-4), （1, 2e-2)，并且需要满足递增性，这样的函数形式有无数种，代码中选择了最简单的线性函数形式， $\beta(t)=(\beta_1-\beta_0)*t + \beta_0$ ，则其在t时刻的积分值为 $\frac{1}{2}t^2(\beta_1-\beta_0) + \beta_0 t$ ，据此就能计算出(sub-)VP SDE的mean和std，用于训练时的数据扰动。