VAE+SDE 双视角解读扩散模型原理

时间:2022-12-10 09:55:02

本文共计5134字,阅读预计需要5分钟,但是5分钟你肯定看不懂,字你都认识,组合起来你不知道我在说啥,老老实实研究吧。

扩散模型对于其他生成模型来讲,从数学角度来看更为优美。扩散模型也因为其背后的数学原理让很多人望而却步。今天就来梳理一下常见的两个理解扩散模型数学原理的角度。


从VEA角度看

之前我们先是讲的VAE模型,详细可以看进来学VAE,VAE都不懂还想懂扩散模型? | Variational Auto-Encoder,扩散模型可以视为多层VAE模型。所以首先我们从VAE角度来看。

生成建模的一个核心问题是模型概率分布的灵活性和易处理性之间的权衡。

扩散模型的基本思想是通过正向扩散过程系统地扰动数据分布中的结构,然后通过学习反向扩散过程来恢复结构,从而产生高度灵活和易处理的生成模型。

去噪扩散概率模型(DDPM)由两个参数化的马尔可夫链组成,并使用变分推理在有限时间步后产生与原始数据匹配的样本。

  • 前向链通过使用预先设计的时间表逐渐添加高斯噪声来扰乱数据分布,直到数据分布收敛到给定的先验,即标准高斯分布。

  • 反向链从给定的先验开始,并使用参数化的高斯转换核,学习逐步恢复未受干扰的数据结构。

形式上:给定一个数据分布 $x_0∼q(x_0)$,我们定义一个前向噪声过程 $q$,通过在时间 $t$添加方差为 $β_t ∈ (0, 1)$的高斯噪声,并产生潜在的$x_1, x_2. . . x_t$到 $x_T$,如下所示: $$ \begin{aligned} q\left(x_1, \ldots, x_T \mid x_0\right) =\prod_{t=1}^T q\left(x_t \mid x_{t-1}\right) &\quad (1) \ \ q\left(x_t \mid x_{t-1}\right) =\mathcal{N}\left(x_t ; \sqrt{1-\beta_t} x_{t-1}, \beta_t \mathbf{I}\right) &\quad (2) \end{aligned} $$

这里我们使用 $q(x_t)$来表示潜变量 $x_t$在前向过程中的分布。等式 (2) 中定义的噪声过程允许我们对直接以输入 $x_0$为条件的噪声潜值的任意步骤进行采样。

当 $α_t = 1 − β_t$且 $\bar{α_t} = \prod^t_{s=0} α_s$ 时,我们可以写为: $$ \begin{aligned} q\left(x_t \mid x_0\right) =\mathcal{N}\left(x_t ; \sqrt{\bar{\alpha}_t} x_0,\left(1-\bar{\alpha}_t\right) \mathbf{I}\right) &\quad (3) \ \ x_t =\sqrt{\bar{\alpha}_t} x_0+\sqrt{1-\bar{\alpha}_t} \epsilon &\quad (4) \end{aligned} $$

当 $\bar{α_T}$接近 0 时,$x_T$实际上与纯高斯噪声无法区分:$p(x_T ) ≈ \mathcal N (x_T ; 0, \mathbf I)$。联合分布 $p_θ(x_{0:T} )$称为逆过程,当 $β_t$足够小时,$q(x_{t−1}|x_t)$近似为高斯分布。因此我们可以将这个反向过程参数化如下,从 $p(x_T) = \mathcal N (x_T ; 0, \mathbf I)$开始:

$$ \begin{aligned} p_\theta\left(\mathbf{x}{0: T}\right) =p\left(\mathbf{x}T\right) \prod{t=1}^T p\theta\left(\mathbf{x}{t-1} \mid \mathbf{x}t\right) &\quad (5) \\ p\theta\left(\mathbf{x}{t-1} \mid \mathbf{x}t\right) =\mathcal{N}\left(\mathbf{x}{t-1} ; \mu_\theta\left(\mathbf{x}t, t\right), \Sigma\theta\left(\mathbf{x}_t, t\right)\right)&\quad (6) \end{aligned} $$

这里我们使用 $p_θ(x_t)$来表示潜在变量 $x_t$ 在逆过程中的分布,或者为简单起见使用 $p(x_t)$。通过优化负对数似然的变分上限来执行训练: $$ \mathbb{E}\left[-\log p_\theta\left(\mathbf{x}0\right)\right] \leq \mathbb{E}q\left[-\log \frac{p\theta\left(\mathbf{x}{0: T}\right)}{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}0\right)}\right]=\mathbb{E}q\left[-\log p\left(\mathbf{x}T\right)-\sum{t \geq 1} \log \frac{p\theta\left(\mathbf{x}{t-1} \mid \mathbf{x}t\right)}{q\left(\mathbf{x}t \mid \mathbf{x}{t-1}\right)}\right]=-L{\mathrm{VLB}}

\quad(7) $$

对(7)右侧进行参数重整化: $\mathbb{E}_{t \sim \mathcal{U}(0, T), \mathbf{x}_0 \sim q\left(\mathbf{x}0\right), \boldsymbol{\epsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\left[\lambda(t)\left|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}\theta\left(\sqrt{\bar{\alpha}_t} \mathbf{x}_0+\sqrt{1-\bar{\alpha}_t} \boldsymbol{\epsilon}, t\right)\right|^2\right]\quad(8)$

这里 $\epsilon_θ$ 是一个神经网络,它使用输入 $x_t$ 和 $t$来预测 $\epsilon$可以生成 $x_t$,类似于在由$t$索引的多个噪声尺度上的去噪分数匹配。$λ(t)$是一个正加权函数。目标函数可以通过 Monte-Carol 算法进行估计,并使用随机优化器进行训练。在生成过程中,我们首先从标准高斯先验中选择一个样本 $x_T$,然后使用反向链顺序绘制样本 $x_t$,直到我们得到一个新的数据 $x_0$。

从SDE角度看

另一个比较常见的角度是SDE。

扩散模型的扩散过程可以看作是基于分数的生成模型。在连续时间设置中,基于分数的生成模型构造一个随机微分方程 (SDE) 以将数据分布平滑地扰乱为已知的先验分布,并构造一个相应的逆时 SDE 以将先验分布转换回数据分布。形式上,前向扩散过程是以下 SDE 的解:

$$ d \mathbf{x}=\mathbf{f}(\mathbf{x}, t) d t+g(t) d \mathbf{w} \quad (9) $$ 这里 $\mathbf w$是标准布朗运动。 $x_0$是原始数据样本,$x_T$是近似标准高斯分布的扰动数据。反向 SDE 的解具有与正向 SDE 相同的边际密度,但它在时间是倒退的 : $$ d \mathbf{x}=\mathbf{f}(\mathbf{x}, t)-g(t)^2 \nabla_{\mathbf{x}} \log q_t(\mathbf{x}) d t+g(t) d \mathbf{w}\quad(10) $$ 逆过程中时间从 $T$倒流到 0,$dt$是一个无穷小的负时间步长。因此,为了逆转扩散过程并生成数据,我们需要的唯一信息是每个时间步 $t$的得分函数:$∇_x \log q_t(x)$。这里我们使用 $q_t(x)$来表示潜变量 $x_t$在前向过程中的分布。 借用分数匹配技术,我们可以训练一个基于时间的分数模型 $s_θ(xt, t)$,通过优化去噪分数匹配目标来估计分数函数: $$ \mathbb{E}_{t, \mathbf{x}_0, \mathbf{x}t}\left[\lambda(t)\left|\mathbf{s}\theta\left(\mathbf{x}t, t\right)-\nabla{\mathbf{x}_t} \log q_t\left(\mathbf{x}_t \mid \mathbf{x}_0\right)\right|^2\right] \quad(11) $$ 证明 DDPM 中的生成过程是反向 SDE 的特殊离散化: $$ d \mathbf{x}=-\frac{1}{2} \beta(t) \mathbf{x} d t+\sqrt{\beta(t)} d \mathbf{w} \quad (12) $$ 其中 $β( \frac i N ) = N β_i$ ,和 (2)中的一样表示方法一样。


这里只是根据北大的综述理论部分进行介绍,相关内容的详细推导可以去看一下苏神的个人网站: