VDM Variational Diffusion Models

VDM: Variational Diffusion Models

导语

Diffusion 模型很强，但它能否被用来计算似然呢？本文给出了肯定的回答，并提出了一系列基于 Diffusion 的生成模型，达到了 SOTA 的似然计算性能。不同于其他的扩散模型，本文所提出的模型可以将扩散过程的 noise schedule 和网络的参数一起进行优化。本文还展示了变分下界（VLB）可以从信噪比（SNR）的角度，化简为一个相当简洁的形式，并借此将之前文献中的多种扩散模型统一起来。

本文表明，连续时间情形下，VLB 除了端点处的信噪比外，关于 noise schedule 是不变的。这使得我们能够调整中间的 noise schedule，最小化生成的 VLB 估计器的方差，从而实现更快的优化。再加上网络架构的改进，本文在图像密度估计基准上获得了 SOTA 的似然计算性能，优于多年来一直主导这些基准的自回归模型，并且优化速度通常要快得多。

本文通过直接对边缘分布 $q(z_t|z_0)$ 的均值和方差进行参数化，将前向过程 noise schedule 也和模型参数一同进行优化，这使得整个模型又成了 VAE（前向编码器和反向解码器都需要学习）。而之前的方法都是对前向过程的每一个单独的时间步 $q(z_{t+ϵ}|z_t)$ 进行参数化，前向过程都是使用固定的 noise schedule。

方法

前向过程

最基础的生成任务，假设我们有一个数据集 $\mathbf{x}$，我们的任务是估计边缘分布 $p(\mathbf{x})$。对于真实数据 $\mathbf{x}$，我们添加一系列逐步增强的噪声得到一组隐变量 ${\mathbf{z}}_t$，$t$ 的取值范围从 0 （ ${\mathbf{z}}_0$ 噪声最小）到 1（${\mathbf{z}}_1$ 噪声程度最大），隐变量 ${\mathbf{z}}_t$ 在给定 $\mathbf{x}$ 下的条件分布为：
$$q({\mathbf{z}}_t|\mathbf{x})=\mathcal{N}({\alpha }_t\mathbf{x},{\sigma }_t^2\mathbf{I})$$
其中 ${\alpha }_t,{\sigma }_t^2$ 是关于 $t$ 的标量函数。我们进一步定义信噪比（Signal-to-Noise Ratio，SNR）为：
$$\text{SNR}(t)={\alpha }_t^2/{\sigma }_t^2$$
在扩散模型中，随着时间 $t$ 的增大，${\mathbf{z}}_t$ 中噪声的占比不断增大，因此信噪比 $\text{SNR}(t)$ 关于 $t$ 一般是严格单调递减的。另外我们假设 ${\alpha }_t$ 和 ${\sigma }_t^2$ 都是光滑的，即他们关于时间 $t$ 都是可导的。

常见扩散模型的具体规格（规格，即 specification，这里是指 formalize 一个具体扩散模型的参数化，包括 ${\alpha }_t$ 和 ${\sigma }_t^2$）包括 variance preserving（其 ${\alpha }_t=\sqrt{1-{\sigma }_t^2}$ ）和 variance exploding （其${\alpha }_t^2=1$ ）两种，本文主要研究 variance preserving。

这里还有几个推论：

1. 对任意的 $t>s$，分布 $q({\mathbf{z}}_t|{\mathbf{z}}_s)$ 也是高斯分布
2. 对于时间步的任意子序列 $0\le s<t<u\le 1$ ，隐变量 $({\mathbf{z}}_s,{\mathbf{z}}_t,{\mathbf{z}}_u)$ 的联合概率分布是马尔可夫的：$q({\mathbf{z}}_u|{\mathbf{z}}_t,{\mathbf{z}}_s)=q({\mathbf{z}}_u|{\mathbf{z}}_t)$
3. 根据上述分布及贝叶斯公式，可知对于任意 $0\le s<t\le 1$ ，$q({\mathbf{z}}_s|{\mathbf{z}}_t,\mathbf{x})$ 同样是高斯分布

Noise Schedule

在之前的扩散模型中，noise schedule 都是人工设定的固定的一组超参数。而本文中，noise schedule 是可学习的，其参数化为：
$${\sigma }_t^2=\text{sigmoid}({\gamma }_{\eta }(t))$$
其中 ${\gamma }_{\eta }(t)$ 是一个可学习的神经网络，参数为 $\eta$。

考虑到不同扩散模型的等价性（后面会详细介绍），本文的离散时间模型和连续时间模型均采用 variance preserving（也就是 DDPM）的 ${\alpha }_t=\sqrt{1-{\sigma }_t^2}$ 。可以推出 ${\alpha }_t^2$ 和 $\text{SNR}(t)$ 都是关于 ${\gamma }_{\eta }(t)$ 的函数，化简有：
$$\begin{gathered} {\alpha }_t^2=\text{sigmoid}(-{\gamma }_{\eta }(t)) \\ \text{SNR}(t)=\text{exp}(-{\gamma }_{\eta }(t)) \end{gathered}$$

反向过程

生成过程即为前向扩散过程的反过程，时间 $t$ 由 1 到 0，由多级生成模型，采样并生成出一个隐变量序列 ${\mathbf{z}}_t$ 。我们同时考虑有限总步数 $T$ 的离散时间情形和 $T\to \infty$ 的连续时间情形。我们将首先讨论反向过程，然后分别讨论离散时间情形和连续时间情形。

给定有限的时间步 $T$，我们将时间均匀地分为 $T$ 个时间步，每步的宽度为 $\tau =1/T$ 。定义相邻步 $s(i)=(i-1)/T$，$t(i)=i/T$ ，则我们的多级生成模型为：
$$p(\mathbf{x})={\int }_{\mathbf{z}}p({\mathbf{z}}_1)p(\mathbf{x}|{\mathbf{z}}_0)\prod _{i=1}^{T}p({\mathbf{z}}_{s(i)}|{\mathbf{z}}_{t(i)})$$
（这里相当于把 z 从联合概率分布里积分掉，得到 p(x)）

在 variance preserving 扩散模型中，$\text{SNR}(1)$ 足够小，则有 $q({\mathbf{z}}_1|\mathbf{x})\approx \mathcal{N}({\mathbf{z}}_1;0,\mathbf{I})$。从而我们将 ${\mathbf{z}}_1$ 的边缘分布建模为一个球形高斯：
$$p({\mathbf{z}}_1)=\mathcal{N}({\mathbf{z}}_1;0,\mathbf{I})$$
我们想要训练一个模型 $p(\mathbf{x}|{\mathbf{z}}_0)$ 来拟合未知的 $q(\mathbf{x}|{\mathbf{z}}_0)$。记 $x_i,z_{0,i}$ 分别为 $\mathbf{x},{\mathbf{z}}_0$ 中的第 $i$ 个元素，则有分解分布的形式：
$$p(\mathbf{x}|{\mathbf{z}}_0)=\prod _{i}p(x_i|z_{0,i})$$
我们选择 $p(x_i|z_{0,i})\propto q(z_{0,i}|x_i)$ ，通过对 $x_i $ 所有可能的离散取值（比如对 8bit 图像来说，共 256 个取值）求和来进行归一化。又由于 $\text{SNR}(0)$ 足够小，未知的数据分布 $q(\mathbf{x})$ 的影响被似然 $q({\mathbf{z}}_0|\mathbf{x})$ 吞没， 所以这会是对真实的 $q(\mathbf{x}|{\mathbf{z}}_0)$ 的非常接近的近似。

最后，我们选用条件分布：
$$p({\mathbf{z}}_s|{\mathbf{z}}_t)=q({\mathbf{z}}_s|{\mathbf{z}}_t,\mathbf{x}={\hat{\mathbf{x}}}_{\theta }({\mathbf{z}}_t;t))$$
其与 $q({\mathbf{z}}_s|{\mathbf{z}}_t,\mathbf{x})$ 相同，但将原始数据 $\mathbf{x}$ 替换成去噪模型根据 ${\mathbf{z}}_t$ 预测出的 ${\hat{\mathbf{x}}}_{\theta }({\mathbf{z}}_t;t)$ 。注意，我们实际上是将去噪模型参数化为一个噪声预测模型（而不是直接预测去噪后的隐变量），这就跟 DDPM 一致了。$p({\mathbf{z}}_s|{\mathbf{z}}_t)$ 的均值和方差还可以被大幅简化，略。

噪声预测模型和傅里叶特征

本文将去噪模型参数化为一个噪声预测模型 ${\widehat{ϵ}}_{\theta }({\mathbf{z}}_t;t)$：
$${\hat{\mathbf{x}}}_{\theta }({\mathbf{z}}_t;t)=({\mathbf{z}}_t-{\sigma }_t{\widehat{ϵ}}_{\theta }({\mathbf{z}}_t;t))/{\alpha }_t$$
其中 ${\widehat{ϵ}}_{\theta }({\mathbf{z}}_t;t)$ 是参数化的神经网络。相当于预测出噪声，再将噪声出当前噪声图中减去，得到前一步的噪声图，而不是直接预测前一步的噪声图。本文使用的噪声预测模型与 DDPM 基本一样，区别在于 DDPM 仅使用了原始的像素空间输入，而本文还加入了傅里叶特征。

之前的扩散模型主要聚焦在生成图片的感知质量，比较强调粗尺度模式和生成图像的全局一致性。而本文方法是直接优化似然，因此对细粒度的细节和每个像素的精确值比较敏感。为了捕捉数据细粒度的细节信息，本文在噪声预测模型的输入中添加了一组傅里叶特征。真实图片 $\mathbf{x}$ 和隐变量 $\mathbf{z}$ 的数值范围均为 $[0,1]$，我们添加两组新的通道，分别为 $\sin (2^n\pi \mathbf{z})$ 和 $\cos (2^n\pi \mathbf{z})$ ，其中 $n$ 为一组整数值 { n m i n , … , n m a x } \{n_{min},\dots,n_{max}\} {nmin​,…,nmax​} 。这些特征是高频特征，可以放大 ${\mathbf{z}}_t$ 中的微小变动。引入这种傅里叶特征对噪声预测模型的性能有很大的提升，特别是配合上本文提出的可学习的 SNR 函数，提升更为显著。在 AR 模型中引入傅里叶特征并没有如此显著的提升。

变分下界

我们的优化目标是边缘似然的变分下界（VLB）：
$$-\log p(\mathbf{x})\le \text{VLB}(\mathbf{x})=\underset{\text{Prior Loss}}{\underbrace{D_{KL}(q({\mathbf{z}}_1|\mathbf{x})||p({\mathbf{z}}_1))}}+\underset{\text{Reconstruction Loss}}{\underbrace{{\mathbb{E}}_{q({\mathbf{z}}_0|\mathbf{x})}[-\log p(\mathbf{x}|{\mathbf{z}}_0)]}}+\underset{\text{Diffusion Loss}}{\underbrace{{\mathcal{L}}_T(x)}}$$
其中先验损失 Prior Loss 和重构损失 Reconstruction Loss 比较常规，计算方式与 VAE 中类似。扩散损失 Diffusion Loss 比较复杂，他还与超参数 $T$ 的值，即扩散模型的深度，有关。

离散时间模型

当 $T$ 为有限值时，为离散时间模型，使用 $s(i)=(i-1)/T$，$t(i)=i/T$ ，则扩散损失为：
$${\mathcal{L}}_T(\mathbf{x})=\sum _{i=1}^T{\mathbb{E}}_{q({\mathbf{z}}_{t(i)}|\mathbf{x})}{D}_{KL}[q({\mathbf{z}}_{s(i)}|{\mathbf{z}}_{t(i)},\mathbf{x})||p({\mathbf{z}}_{s(i)}|{\mathbf{z}}_{t(i)})]$$
结合我们定义的信噪比函数 $\text{SNR}(t)$，上式可以被大幅化简：
$${\mathcal{L}}_T(\mathbf{x})=\frac{T}{2}{\mathbb{E}}_{ϵ\sim \mathcal{N}(0,\mathbf{I}),i\sim U(1,T)}[\text{SNR}(s)-\text{SNR}(t)||\mathbf{x}-{\hat{\mathbf{x}}}_{\theta }({\mathbf{z}}_t;t)|{|}_2^2]$$
其中 $U(1,T)$ 是指对 1 到 T 每个时间步均匀采样，${\mathbf{z}}_t={\alpha }_t\mathbf{x}+{\sigma }_tϵ$ 。上式是对于任意的 $({\sigma }_t,{\alpha }_t)$ 规格的扩散模型的扩散损失的一般形式。对于我们本文讨论的 variance preserving 的 $({\sigma }_t,{\alpha }_t)$ 以及参数化去噪网络 ${\hat{\mathbf{x}}}_{\theta }({\mathbf{z}}_t;t)$ 来说，有：
$${\mathcal{L}}_T(\mathbf{x})=\frac{T}{2}{\mathbb{E}}_{ϵ\sim \mathcal{N}(0,\mathbf{I}),i\sim U(1,T)}[(\exp ({\gamma }_{\eta }(t)-{\gamma }_{\eta }(s))-1)||ϵ-{\widehat{ϵ}}_{\theta }({\mathbf{z}}_t;t)|{|}_2^2]$$
其中 ${\mathbf{z}}_t=\text{sigmoid}(-{\gamma }_{\eta }(t))\mathbf{x}+\text{sigmoid}({\gamma }_{\eta }(t))ϵ$ 。在离散时间情形下，我们对上式进行蒙特卡洛估计，最大化 VLB，来优化参数 $\theta$ 和 $\eta$ 。直观地讲，离散时间扩散损失是严格递减函数积分的上黎曼和近似，这意味着更精细的近似，即更大的总步数 $T$ ，会得到更低的扩散损失。

连续时间模型

更多的步数会得到更好的 VLB 优化结果，我们现在来看一下当 $T\to \infty$ 时，即连续时间情形下的模型。在宋飏等人的研究中，使用 SDE 来描述连续时间情形下的扩散模型 $p({\mathbf{z}}_t)$ 。本文推导出，在连续时间模型下，扩散损失 ${\mathcal{L}}_T(\mathbf{x})$ 可以进一步近似：
$$\begin{array}{rl}{\mathcal{L}}_{\infty }(\mathbf{x})& =-\frac{1}{2}{\mathbb{E}}_{ϵ\sim \mathcal{N}(0,\mathbf{I})}{\int }_0^1{\text{SNR}}^{\prime }(t)||\mathbf{x}-{\hat{\mathbf{x}}}_{\theta }({\mathbf{z}}_t;t)|{|}_2^{2}dt\\ & =-\frac{1}{2}{\mathbb{E}}_{ϵ\sim \mathcal{N}(0,\mathbf{I}),t\sim \mathcal{U}(0,1)}[{\text{SNR}}^{\prime }(t)||\mathbf{x}-{\hat{\mathbf{x}}}_{\theta }({\mathbf{z}}_t;t)|{|}_2^2]\end{array}$$
其中 ${\text{SNR}}^{\prime }(t)=d\text{SNR}(t)/dt$ ，${\mathbf{z}}_t={\alpha }_t\mathbf{x}+{\sigma }_tϵ$ 。上式同样是一般的通用形式，在本文的规格下，进一步有：
$${\mathcal{L}}_{\infty }(\mathbf{x})=\frac{1}{2}{\mathbb{E}}_{ϵ\sim \mathcal{N}(0,\mathbf{I}),t\sim \mathcal{U}(0,1)}[{\gamma }_{\eta }^{\prime }(t)||ϵ-{\widehat{ϵ}}_{\theta }({\mathbf{z}}_t;t)|{|}_2^2]$$
其中 ${\gamma }_{\eta }^{\prime }(t)=d{\gamma }_{\eta }(t)/dt$ 。同样使用蒙特卡洛估计来进行计算和优化。

连续时间情形下不同扩散模型的等价性

由于在定义时我们假设 $\text{SNR}(t)$ 是单调的，因此它是可逆的，从而我们可以进行变量变换，将所有原来关于 $t$ 的函数都表示为关于 $v\equiv \text{SNR}(t)$ 的函数。从而可在 $t={\text{SNR}}^{-1}(v)$ 下分别求得有 ${\alpha }_v,{\sigma }_v$ ，进而 ${\mathbf{z}}_v={\alpha }_v\mathbf{x}+{\sigma }_vϵ$ ，以及 ${\tilde{\mathbf{x}}}_{\theta }(\mathbf{z};v)={\hat{\mathbf{x}}}_{\theta }(\mathbf{z};{\text{SNR}}^{-1}(v))$ 。经过变量变换之后，连续时间情形下的扩散损失可写为：
$${\mathcal{L}}_{\infty }(\mathbf{z})=\frac{1}{2}{\mathbb{E}}_{ϵ\sim \mathcal{N}(0,\mathbf{I})}{\int }_{{\text{SNR}}_{\text{min}}}^{{\text{SNR}}_{\text{max}}}||\mathbf{x}-{\tilde{\mathbf{x}}}_{\theta }({\mathbf{z}}_v;v)||dv$$

从之前关于 $t$ 的积分变成了关于 $v$ 的积分。

仔细观察上式，我们发现，扩散模型的规格 $\alpha (t),\sigma (t)$ 以及 $\text{SNR}(t)={\alpha }_t^2/{\sigma }_t^2$ 对该扩散损失的影响仅体现在积分的区间端点上，即 ${\text{SNR}}_{\text{min}}$ 和 ${\text{SNR}}_{\text{max}}$ ，而这两个值 ${\text{SNR}}_{\text{min}}=\text{SNR}(1)$ 、${\text{SNR}}_{\text{max}}=\text{SNR}(0)$ 又仅与扩散模型时间步的 $t=0$ 和 $t=1$ 两端点有关。也就是说 ${\text{SNR}}_{\text{min}}$ 和 ${\text{SNR}}_{\text{max}}$ 一旦给定，扩散损失就已经确定下来了，而与 $t=0$ 和 $t=1$ 之间的信噪比函数 $\text{SNR}(t)$ 的形状无关。VLB 受到信噪比函数 $\text{SNR}(t)$ 的影响仅体现在两端点处。

进一步我们发现，生成模型定义的数据分布 $p(\mathbf{x})$ 也与扩散模型的规格无关。具体来说，如果我们在扩散模型规格 { α v A , σ v A , x ~ θ A ( z v ; v ) } \{\alpha^A_v,\sigma^A_v,\tilde{\mathbf{x}}^A_\theta(\mathbf{z}_v;v)\} {αvA​,σvA​,x~θA​(zv​;v)} 下定义一个分布 $p^A(\mathbf{x})$ ，然后在不同的扩散模型规格 { α v B , σ v B , x ~ θ B ( z v ; v ) } \{\alpha^B_v,\sigma^B_v,\tilde{\mathbf{x}}^B_\theta(\mathbf{z}_v;v)\} {αvB​,σvB​,x~θB​(zv​;v)} 下定义一个分布 $p^B(\mathbf{x})$ ，作者证明了，如果两种规格下的 ${\text{SNR}}_{\text{min}},{\text{SNR}}_{\text{max}}$ 相同，且对模型乘上一个尺度定值 ${\tilde{\mathbf{x}}}_{\theta }^B(\mathbf{z};v)\equiv {\tilde{\mathbf{x}}}_{\theta }^B(({\alpha }_v^A/{\alpha }_v^B)\mathbf{z};v)$，那么就有 $p^A(\mathbf{x})=p^B(\mathbf{x})$。并且所有隐变量 ${\mathbf{z}}_v$ 的特征的分布也是相同的，顶多需要乘一个常系数。也就是说，任何两个扩散模型，只要满足我们在 ”前向过程“ 中给出的一点轻微约束，那么他们在连续时间情形下就是等价的。

加权扩散损失

即使优化目标不是 VLB，而是一个加权的扩散损失，不同扩散模型规格下的等价性也是成立的。对于加权扩散损失：
$${\mathcal{L}}_{\infty }(\mathbf{z})=\frac{1}{2}{\mathbb{E}}_{ϵ\sim \mathcal{N}(0,\mathbf{I})}{\int }_{{\text{SNR}}_{\text{min}}}^{{\text{SNR}}_{\text{max}}}w(v)||\mathbf{x}-{\tilde{\mathbf{x}}}_{\theta }({\mathbf{z}}_v;v)||dv$$

其中 $w(v)$ 是一个关于 $v$ 的（自然也就是关于 $t$）的加权函数。一般来说，噪声越大，加的权重也会越大，从而可以提高 FID/IS 等指标度量出的感知质量。加入损失加权之后，就能囊括宋飏 SDE 论文中讨论的所有不同的目标函数了。

最小化方差

降低连续时间损失的蒙特卡洛估计量的方差通常会提高优化效率。作者发现，使用低差异采样器（可显著降低方差。此外，由于连续时间情形下的等价性，端点之间的 noise schedule 对扩散损失没有影响，我们可以对其进行调整优化，以最小化我们的损失估计量的方差。

总结

本文定义了信噪比函数 $\text{SNR}(t)={\alpha }_t^2/{\sigma }_t^2$，从变分下界的角度推导了扩散模型的损失函数。并对其中的扩散损失项在离散时间情形和连续时间情形下进行了推导，结合信噪比函数得出了一个极简洁的形式。进一步地，作者推导出在连续时间情形下，扩散模型规格（以 SNR 的形式）对扩散损失的影响仅体现在 $t=0$ 和 $t=1$ 两处端点上，而与中间的 ${\alpha }_t,{\sigma }_t$ 的形状无关，从而推导出不同的扩散模型规格在连续时间情形下的等价性（宋飏连续时间 SDE 的推导也有此结论）。