Diffenc: Variational Diffusion With A Learned Encoder

Abstract

扩散模型可以看作是层次化变分自编码器（VAE），并具有两项改进：生成过程中的条件分布参数共享，以及将损失计算为层次结构中独立项的高效计算。我们对扩散模型进行了两项改进，这些改进在保留上述优势的同时增加了模型的灵活性。

首先，我们在扩散过程中引入了一个依赖于数据和深度的均值函数，从而导致修改后的扩散损失。我们提出的框架 DiffEnc 在 CIFAR-10 数据集上显著提高了似然度。

其次，我们将逆编码器过程和生成过程中的噪声方差比率设置为一个自由权重参数，而不是固定为1。这带来了理论上的见解：对于有限深度的层次结构，证据下界（ELBO）可以用作加权扩散损失方法的目标，并用于专门针对推断优化噪声计划。而对于无限深度的层次结构来说，为了使 ELBO 定义良好，该权重参数必须为1。

1. Introduction

扩散模型可以被视为按时间索引的层次结构，其中潜在变量按顺序生成，并仅依赖于前一步的潜在向量。因此，扩散模型可以理解为分层变分自编码器（VAEs）（Kingma & Welling, 2013; Rezende et al., 2014; Sønderby et al., 2016），但有三个限制：(1) 前向扩散过程（即变分推理中的推断模型）是固定且非常简单的；(2) 生成模型是马尔可夫性的——每一层（按时间索引的）潜在变量仅依赖于前一层生成；(3) 参数共享——生成模型的所有步骤共享相同的参数。

前向过程的简单性（1）和生成模型的马尔可夫性质（2）允许证据下界（ELBO）表示为对随机变量层的期望，即从随机过程角度来看是对时间的期望。由于生成模型中的大量参数共享（3），这种期望可以通过单个蒙特卡罗样本有效估计。这些特性使得扩散模型在高效可扩展性的同时，仍然保留了前述限制的灵活性。

在这项工作中，我们放宽了假设（1），以提高扩散模型的灵活性，同时保持其可扩展性。具体而言，我们不再假设一个恒定的扩散过程，同时仍保持足够的简单性，以便将ELBO表示为对时间的期望。我们引入了一个时间依赖的编码器，该编码器参数化了扩散过程的均值：学习到的去噪模型不再预测原始图像 $x$，而是预测时间 $t$ 的编码图像 $x_t$。关键是，这个编码器仅在训练阶段使用，而在采样过程中不使用。因此，提出的这类扩散模型 DiffEnc 比标准扩散模型更灵活，而不影响采样时间。

为了推导 DiffEnc 的负对数似然损失（Eq. (18)），我们首先展示如何将时间依赖的编码器引入到扩散过程中，并说明如果我们使用生成模型中通常的均值表达式，会在损失中引入一个额外项（第3节）。然后，我们展示如何使用编码器的特定参数化来抵消这个额外项（第4节）。

我们在 MNIST、CIFAR-10 和 ImageNet32 上进行实验，使用了两种不同的编码器参数化方法。实验结果表明，使用可训练编码器时，DiffEnc 在 CIFAR-10 上提高了总似然性，并在所有数据集上改善了潜在损失，而不会损害扩散损失。我们观察到，对于早期和晚期时间步长， $x_t$ 的变化显著不同，显示了编码器的非平凡、时间依赖的行为（见图2）。

此外，我们研究了扩散模型中的一个常见假设的放松：即生成过程的方差 ${\sigma }_P^2$ 等于前向扩散过程反向公式的方差 ${\sigma }_Q^2$。这在扩散损失中引入了一个额外的项，可以解释为一个带时间依赖权重 $w_t$ 的加权损失。我们随后解析地推导了最优的 ${\sigma }_P^2$。虽然这在离散时间（即，具有有限层数）训练或采样时是相关的，但我们证明了在连续时间极限中，当方差相等时，ELBO 达到最大值（事实上，如果方差不相等，ELBO 会发散）。

---

2. Preliminaries On Variational Diffusion Models

首先介绍扩散模型的变分形式（Kingma等，2021）。我们定义一个具有T + 1层潜变量的层次生成模型：

$$p_{\theta }(x,z)=p(x|z_0)p(z_1)\prod _{i=1}^T{p}_{\theta }(z_{s(i)}|z_{t(i)})\text{\hspace{1em}(1)}$$
其中 $x\in X$ 是数据点，$\theta$ 是模型参数，$s(i)=\frac{i-1}{T}$，$t(i)=\frac{i}{T}$，并且 $p(z_1)=N(0,I)$。接下来，我们将省略索引 $i$，假设 $0\le s<t\le 1$。
定义一个扩散过程 $q$，其边际分布为：
$$q(z_t|x)=N({\alpha }_{t}x,{\sigma }_t^{2}I)\text{\hspace{1em}(2)}$$

其中 $t\in [0,1]$ 是时间索引，${\alpha }_t$，${\sigma }_t$ 是 $t$ 的正标量函数。要求方程（2）对任意的 $s$ 和 $t$ 都成立，条件分布变为：

$$q(z_t|z_s)=N({\alpha }_{t|s}{z}_s,{\sigma }_{t|s}^{2}I)$$

其中
$${\alpha }_{t|s}={\alpha }_t/{\alpha }_s,{\sigma }_{t|s}^2={\sigma }_t^2-{\alpha }_{t|s}^2{\sigma }_s^2.$$

利用贝叶斯法则，我们可以反转扩散过程的方向：

$$q(z_s|z_t,x)=N({\mu }_Q,{\sigma }_Q^{2}I)\text{\hspace{1em}(3)}$$

其中

$${\sigma }_Q^2=({\sigma }_{t|s}^2{\sigma }_s^2)/{\sigma }_t^2,{\mu }_Q=\frac{{\alpha }_{t|s}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t|s}^2}{{\sigma }_t^2}x.\text{\hspace{1em}(4)}$$

我们现在可以以与生成模型相似的方式表达扩散过程：

$$q(z|x)=q(z_1|x)\prod _{i=1}^{T}q(z_{s(i)}|z_{t(i)},x)$$

并且可以以与方程（3）相同的功能形式定义生成过程的一步：

$$p_{\theta }(z_s|z_t)=N({\mu }_P,{\sigma }_P^{2}I)\text{\hspace{1em}(5)}$$

其中

$${\mu }_P=\frac{{\alpha }_{t|s}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t|s}^2}{{\sigma }_t^2}{\hat{x}}_{\theta }(z_t,t).\text{\hspace{1em}(6)}$$

在扩散模型中，去噪方差 ${\sigma }_P^2$ 通常选择与反向扩散过程方差相等：${\sigma }_P^2={\sigma }_Q^2$。虽然我们最初不做这一假设，但我们将在连续时间极限中证明这最优。根据VDM，我们通过信噪比（SNR）参数化噪声调度：

$$\text{SNR}(t)\equiv \frac{{\alpha }_t^2}{{\sigma }_t^2}$$

及其对数：${\lambda }_t\equiv \log \text{SNR}(t)$。在所有实验中，我们将使用方差保持形式：${\alpha }_t^2=1-{\sigma }_t^2=\text{sigmoid}({\lambda }_t)$。

所定义模型的证据下界（ELBO）为：

$$\log p_{\theta }(x)\ge E_{q(z|x)}\left[\frac{p_{\theta }(x|z)p_{\theta }(z)}{q(z|x)}\right]\equiv \text{ELBO}(x)$$

损失 $L\equiv -\text{ELBO}$ 是重构损失（$L_0$）、扩散损失（$L_T$）和潜在损失（$L_1$）的总和：

$$\mathcal{L}=L_0+L_T+L_1$$

其中

$$L_0=-E_{q(z_0|x)}[\log p(x|z_0)],L_1=D_{KL}(q(z_1|x)||p(z_1)).$$

由于生成和反向噪声过程的匹配分解（见方程（1）和（5）），以及 $q(z_t|x)$ 以封闭形式可用，因为 $q$ 是马尔可夫和高斯的，扩散损失 $L_T$ 可以写成随机变量层的和或期望：

$$L_T(x)=\sum _{i=1}^T{E}_{q(z_{t(i)}|x)}\left[D_{KL}(q(z_{s(i)}|z_{t(i)},x)\Vert p_{\theta }(z_{s(i)}|z_{t(i)}))\right]\text{\hspace{1em}(7)}$$

= T E i ∼ U { 1 , T } , q ( z t ( i ) ∣ x ) [ D K L ( q ( z s ( i ) ∣ z t ( i ) , x ) ∥ p θ ( z s ( i ) ∣ z t ( i ) ) ) ] , (8) = T E_{i \sim U\{1, T\}, q(z_{t(i)}|x)} \left[ D_{KL}(q(z_{s(i)}|z_{t(i)}, x) \| p_\theta(z_{s(i)}|z_{t(i)})) \right],\tag{8} =TEi∼U{1,T},q(zt(i)​∣x)​[DKL​(q(zs(i)​∣zt(i)​,x)∥pθ​(zs(i)​∣zt(i)​))],(8)

其中 U { 1 , T } U\{1, T\} U{1,T} 是索引 $1$ 到 $T$ 的均匀分布。由于所有分布都是高斯的，KL散度具有封闭形式表达（见附录E）：

$$D_{KL}(q(z_s|z_t,x)\Vert p_{\theta }(z_s|z_t))=d/2(w_t-1-\log w_t)+\frac{w_t}{2{\sigma }_Q^2}\Vert {\mu }_P-{\mu }_Q{\Vert }^2,\text{\hspace{1em}(9)}$$

其中绿色部分是使用 ${\sigma }_P^2\ne {\sigma }_Q^2$ 而不是 ${\sigma }_P^2={\sigma }_Q^2$ 时的差异，我们定义了加权函数：

$$w_t=\frac{{\sigma }_Q^2}{{\sigma }_P^2}$$

且 ${\sigma }_Q^2$ 和 ${\sigma }_P^2$ 对 $s$ 的依赖性留隐式，因为步长 $t-s=\frac{1}{T}$ 是固定的。最优生成方差可以封闭形式计算（见附录F）：

$${\sigma }_P^2={\sigma }_Q^2+\frac{1}{d}{E}_{q(x,z_t)}\left[\Vert {\mu }_P-{\mu }_Q{\Vert }^2\right]$$

3. DiffEnc

DiffEnc的主要组成部分是时间依赖编码器，我们定义为 $x_t\equiv x_{\varphi ({\lambda }_t)}$，其中 $x_{\varphi ({\lambda }_t)}$ 是一个参数为 $\varphi$ 的函数，依赖于通过 ${\lambda }_t\equiv \log \text{SNR}(t)$ 的 $x$ 和 $t$。方程（2）的广义版本为：

$$q(z_t|x)=N({\alpha }_t{x}_t,{\sigma }_t^{2}I)\text{\hspace{1em}(10)}$$

图1 展示了这一扩散过程的变化，详细图示见附录A。要求过程在边际化时保持一致，即：

$$q(z_t|x)=\int q(z_t|z_s,x)q(z_s|x)d{z}_s$$

则会导出以下条件分布（见附录B）：

$$q(z_t|z_s,x)=N({\alpha }_{t|s}{z}_s+{\alpha }_t(x_t-x_s),{\sigma }_{t|s}^{2}I)\text{\hspace{1em}(11)}$$

这里，由深度依赖编码器引入了额外的均值偏移项。与第二节一样，我们可以推导反向过程（见附录C）：

$$q(z_s|z_t,x)=N({\mu }_Q,{\sigma }_Q^{2}I)\text{\hspace{1em}(12)}$$

其中

$${\mu }_Q=\frac{{\alpha }_{t|s}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t|s}^2}{{\sigma }_t^2}{x}_t+{\alpha }_s(x_s-x_t)\text{\hspace{1em}(13)}$$

而 ${\sigma }_Q^2$ 如方程（4）所给。我们将在第4节中展示我们如何参数化编码器。

Infinite-depth limit.

Kingma等（2021）推导了扩散损失的连续时间极限，即在 $T\to \infty$ 的极限下的损失。我们可以将该结果扩展到我们的情况。使用方程（13）中的 ${\mu }_Q$ 和方程（6）中的 ${\mu }_P$，在未加权的情况下，KL散度 $\frac{1}{2{\sigma }_Q^2}\Vert {\mu }_P-{\mu }_Q{\Vert }^2$ 可以重新写为如下形式（见附录G）：

$$\frac{1}{2{\sigma }_Q^2}\Vert {\mu }_P-{\mu }_Q{\Vert }^2=-\frac{1}{2}\frac{\Delta \text{SNR}}{\Delta x}{\left({\hat{x}}_{\theta }(z_t,t)-x_{\varphi ({\lambda }_t)}-\frac{\text{SNR}(s)}{\Delta \text{SNR}}\Delta x\right)}^2$$

其中 $\Delta x\equiv x_{\varphi ({\lambda }_t)}-x_{\varphi ({\lambda }_s)}$，SNR 也是类似处理。在附录G中，我们还展示了当 $T\to \infty$ 时，最优的 ${\sigma }_P$ 表达式趋向于 ${\sigma }_Q$，并且允许 ${\sigma }_P^2\ne {\sigma }_Q^2$ 所引入的扩散损失的额外项趋向于0。这一结果与之前关于随机过程的变分方法的研究相一致（Archambeau等，2007）。我们已证明在连续极限下，ELBO必须是未加权的损失（即 $w_t=1$）。然而，考虑对有限层数的优化加权损失是有意义的，但我们将其留待未来的研究。

扩散损失的无限深度极限 $L_{\infty }(x)\equiv {\lim }_{T\to \infty }{L}_T(x)$ 变为（见附录G）：

$$L_{\infty }(x)=-\frac{1}{2}{E}_{t\sim U(0,1)}{E}_{q(z_t|x)}\left[{\left(\frac{d\text{SNR}(t)}{dt}{\hat{x}}_{\theta }(z_t,t)-x_{\varphi ({\lambda }_t)}-\frac{d{x}_{\varphi ({\lambda }_t)}}{d{\lambda }_t}\right)}^2\right]\text{\hspace{1em}(14)}$$

因此， $L_{\infty }(x)$ 与VDM中的标准连续时间扩散损失非常相似，但多了一个源于均值偏移项的额外梯度。在第四节中，我们将开发一个修改的生成模型，以应对这一额外项。在附录H中，我们推导了描述DiffEnc在无限深度极限下生成模型的随机微分方程（SDE）。

4. Parameterization Of The Encoder And Generative Model

编码器和生成模型的参数化

我们现在转向编码器 $x_{\varphi ({\lambda }_t)}$ 的参数化。重建损失和潜在损失对编码器在潜在变量层次结构两端的行为施加了约束：我们使用的似然性构造使得重建损失在 $x_{\varphi ({\lambda }_0)}=x$ 时最小化。同样，潜在损失在 $x_{\varphi ({\lambda }_1)}=0$ 时最小化。在 $0<t<1$ 的范围内，一个非平凡的编码器可以改善扩散损失。

我们提出两种相关的编码器参数化方式：一个是可训练的 $x_{\varphi }$，另一个是更简单的不可训练的 $x_{nt}$（其中 nt 代表不可训练）。设 $y_{\varphi }(x,{\lambda }_t)$ 是一个带参数 $\varphi$ 的神经网络，简记为 $y_{\varphi }({\lambda }_t)$。我们将可训练编码器定义为：

$$x_{\varphi }({\lambda }_t)=(1-{\sigma }_t^2)x+{\sigma }_t^2{y}_{\varphi }({\lambda }_t)={\alpha }_t^{2}x+{\sigma }_t^2{y}_{\varphi }({\lambda }_t)\text{\hspace{1em}(15)}$$

而不可训练编码器定义为：

$$x_{nt}({\lambda }_t)={\alpha }_t^{2}x\text{\hspace{1em}(15)}$$

对这些参数化方式的更多动机可以在附录I中找到。可训练编码器 $x_{\varphi }$ 在初始化时设置为 $y_{\varphi }({\lambda }_t)=0$，因此在训练开始时其行为类似于不可训练编码器 $x_{nt}$（但与VDM的恒等编码器不同）。

为了更好地适应方程（14）中的无限深度扩散损失，我们定义生成模型 $p_{\theta }(z_s|z_t)$ 的新均值 ${\mu }_P$，这是对方程（6）的修改。具体而言，我们希望在 ${\mu }_P$ 中引入一个对抗项，该项在取连续极限时大致抵消 $\frac{d{x}_{\varphi }({\lambda }_t)}{d{\lambda }_t}$。这个项应以 ${\hat{x}}_{\theta }({\lambda }_t)$ 而非 $x_{\varphi }$ 表示。对于不可训练编码器，我们有：

$$\frac{d{x}_{nt}({\lambda }_t)}{d{\lambda }_t}={\alpha }_t^2{\sigma }_t^{2}x={\sigma }_t^2{x}_{nt}({\lambda }_t)$$

因此，对于不可训练编码器，我们可以用 ${\sigma }_t^2{\hat{x}}_{\theta }({\lambda }_t)$ 来近似 $\frac{d{x}_{nt}({\lambda }_t)}{d{\lambda }_t}$。可训练编码器则更复杂，因为它还包含了 $y_{\varphi }$ 的导数，我们无法简单地用 ${\hat{x}}_{\theta }$ 来表示。因此，我们选择将 $\frac{d{x}_{\varphi }({\lambda }_t)}{d{\lambda }_t}$ 近似为与 $\frac{d{x}_{nt}({\lambda }_t)}{d{\lambda }_t}$ 相同的形式。我们将探索不同的策略来近似该梯度留待未来工作。

通过选择的对抗项，在连续极限下应大致抵消方程（13）中的均值偏移项的影响，新均值 ${\mu }_P$ 定义为：

$${\mu }_P=\frac{{\alpha }_{t|s}{\sigma }_s^2}{{\sigma }_t^2}{z}_t+\frac{{\alpha }_s{\sigma }_{t|s}^2}{{\sigma }_t^2}{\hat{x}}_{\theta }({\lambda }_t)+{\alpha }_s({\lambda }_s-{\lambda }_t){\sigma }_t^2{\hat{x}}_{\theta }({\lambda }_t)\text{\hspace{1em}(17)}$$

类似于上述内容，我们在编码器参数化为方程（15）时推导无限深度扩散损失 $L_{\infty }(x)$，通过取 $L_T$ 的极限 $T\to \infty$ （见附录J）：

$$L_{\infty }(x)=-\frac{1}{2}{E}_{ϵ,t\sim U[0,1]}\left[{\left(\frac{d{\lambda }_t}{dt}{\hat{x}}_{\theta }(z_t,{\lambda }_t)+{\sigma }_t{\hat{x}}_{\theta }(z_t,{\lambda }_t)-x_{\varphi }({\lambda }_t)-\frac{d{x}_{\varphi }({\lambda }_t)}{d{\lambda }_t}\right)}^2\right]\text{\hspace{1em}(18)}$$

其中 $z_t={\alpha }_t{x}_t+{\sigma }_tϵ$ 且 $ϵ\sim N(0,I)$。

在我们的实验中，我们使用 v-参数化（Salimans & Ho, 2022）作为损失，这意味着对于可训练编码器，我们使用的损失为：

$$L_{\infty }(x)=-\frac{1}{2}{E}_{ϵ,t\sim U[0,1]}\left[{\lambda }_t^{\prime }{\alpha }_t^2{\left(v_t-{\hat{v}}_{\theta }+{\sigma }_t\left({\hat{x}}_{\theta }({\lambda }_t)-x_{\varphi }({\lambda }_t)+y_{\varphi }({\lambda }_t)-\frac{d{y}_{\varphi }({\lambda }_t)}{d{\lambda }_t}\right)\right)}^2\right]\text{\hspace{1em}(19)}$$

而对于不可训练编码器，我们使用的损失为：

$$L_{\infty }(x)=-\frac{1}{2}{E}_{ϵ,t\sim U[0,1]}\left[{\lambda }_t^{\prime }{\alpha }_t^2\Vert v_t-{\hat{v}}_{\theta }+{\sigma }_t({\hat{x}}_{\theta }({\lambda }_t)-x_{\varphi }({\lambda }_t)){\Vert }^2\right]\text{\hspace{1em}(20)}$$

方程（19）和（20）的推导见附录K。我们注意到，当使用 v-参数化时，当 ${\alpha }_t$ 趋近于0时，损失变为与 $ϵ$ 预测参数化相同。而当 $t$ 趋近于1时，损失的行为取决于编码器：对于可训练编码器，我们有 ${\hat{v}}_{\theta }\approx \frac{d{y}_{\varphi }({\lambda }_t)}{d{\lambda }_t}$，这表明编码器在原则上可以指导扩散模型。