zehua

Zehua

论文阅读 - DiffuseVAE 详解

2024-11-27

背景问题

  • 扩散模型:生成质量高,但存在两个缺点:① 缺乏低维、可解释的潜在空间;② 生成速度较慢。

  • VAE:具有低维潜在空间,可以解释和操控生成,但生成图像质量较差。

因此提出了 DiffuseVAE,将VAE嵌入到扩散模型框架中,并设计了新的条件参数化方法,使扩散模型能够利用VAE推导的低维潜在编码。在标准基准(如 CIFAR-10 和 CelebA-64)上,生成质量优于大多数基于 VAE 的方法,并与最先进模型接近,并独特地保留了对低维潜在空间的访问能力,这是其他模型没有的。

 

VAE和扩散模型的优缺点

  • VAE 的特点

    • 显式似然生成模型:以概率分布的方式描述数据生成过程。

    • 低维潜在表示:通过学习隐空间表示,可以捕捉数据的结构和特征。

    • 灵活性:适用于多种任务,包括:

      • 解耦表示学习(分离数据中的不同特征)。

      • 半监督学习(结合少量有标注数据进行学习)。

      • 异常检测(识别偏离正常分布的样本)。

  • VAE 的局限性

    • 模糊的生成样本:生成图像往往缺乏高频细节和清晰度。

    • 结构复杂性:最近的一些改进(如层次化潜在结构)尽管提高了生成质量,但增加了模型复杂度(需要大规模的潜在代码层次化结构)。

    • 与GAN的差距:相比隐式生成模型(如GAN),VAE在样本质量上仍显逊色,特别是在视觉感知方面。

 

  • 扩散模型(DDPM)的优势和局限性

    • 生成质量高:在多个图像合成任务中表现优异,甚至在某些基准上超越了GANs。

    • 高昂的计算成本:生成样本时需要多次迭代的逐步采样过程,计算开销大,生成速度慢。

    • 缺乏低维潜在表示:无法像VAE那样学习和利用低维潜在空间,这限制了其在某些下游任务(如表示解耦、异常检测、可控生成)中的应用。

 

VAE和DDPM复习回顾

变分自编码器 (VAE )

VAE 是一种利用深度神经网络进行变分推断的生成模型,使用编码器将数据x 映射到潜在空间表示z,再通过解码器从z 重建x。它通过最大化数据对数似然\log p(x) 来学习模型,但由于直接计算对数似然通常不可行,VAE 使用 证据下界(ELBO) 来近似优化。

VAE 的损失函数(ELBO)定义如下:

\mathcal{L}(\theta, \phi) = \underbrace{\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]}_{\text{重建误差}} - \underbrace{D_{KL}[q_\phi(z|x) \| p(z)]}_{\text{KL 散度正则化}}

为了使模型的优化过程可微分,VAE 引入了重参数化技巧,将采样过程从随机变量z 中分离。例如,将随机变量z 表示为z = \mu + \sigma \cdot \epsilon,其中\epsilon 是从标准正态分布采样的噪声。默认先验分布p(z) 通常设为标准高斯分布,但可以扩展到更复杂的分布(如混合分布或归一化流)以提升模型的表达能力。

更详细内容请看 VAE

DDPM

DDPM 是一种基于潜变量的生成模型, 利用正向过程将数据逐步加噪到高斯分布,再通过逆向过程从噪声中逐步还原数据。包含两个过程:

正向加噪过程

通过一个固定的高斯马尔可夫链,逐步对数据x_0 加噪,使其最终接近各向同性高斯分布。正向过程的核心是逐步破坏数据结构,每一步的加噪过程使用高斯分布:

q(x_{1:T}|x_0) = \prod_{t=1}^{T} q(x_t|x_{t-1})

这意味着当前状态x_t 仅依赖于前一个状态x_{t-1},整个加噪过程的联合分布可以通过各步的条件分布相乘得到,每一步的加噪过程由高斯分布建模:

q(x_t|x_{t-1}) = \mathcal{N}(\sqrt{1-\beta_t}x_{t-1}, \beta_t\mathbf{I})
  • 其中\beta_t 是噪声调度参数,决定每一步加噪的程度。

通过数学推导,可以从原始数据x_0 直接生成第t 步数据:

q(x_t|x_0) = \mathcal{N}(\sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t)\mathbf{I})
  • 其中\bar{\alpha}_t = \prod_{i=1}^t (1-\beta_i) 是累积衰减系数。

正向过程是固定的,不需要学习,最终会将数据x_0 转化为接近高斯分布的x_T

逆向去噪过程

逆向过程试图逐步还原数据,采用可学习的高斯马尔可夫链:

p_\theta(x_{t-1}|x_t) = \mathcal{N}(\mu_\theta(x_t, t), \Sigma_\theta(x_t, t))
  • 其中均值\mu_\theta 和协方差\Sigma_\theta 是通过神经网络学习的。

采样时,先从p(x_T)(通常选为各向同性高斯分布)中采样一个噪声向量,再运行逆向过程逐步去噪,生成最终样本x_0

噪声调度参数\beta_t 是正向加噪过程的核心,影响生成样本的质量和采样速度。可以是固定值,也可以通过优化学习得到。整个模型通过变分推断优化,目标是最小化正向过程和逆向过程的分布差异。

DDPM 主要优势是生成样本质量高,但缺点是采样速度较慢(需要多次迭代)。这是许多扩散模型后续改进的基础,例如提高采样效率的 DDIM。

 

DiffuseVAE: VAE与扩散模型的结合

提出的 DiffuseVAE 将 VAE 和 DDPM 结合,使生成过程既拥有 VAE 提供的低维潜在空间,又有 DDPM 生成高质量细节的能力。框架中通过两阶段建模,先利用 VAE 的低维潜在空间对样本进行全局特征建模生成粗略样本,再用 DDPM 细化提高质量。通过两阶段建模解决了 VAE 样本模糊和 DDPM 缺乏低维表示的问题。

算法框架

  • 该框架可归纳为生成器-细化器框架

    • 生成器:第一阶段通过VAE 对训练数据x 拟合,生成重构样本\hat{x}

    • 细化器:第二阶段通过条件 DDPM 对 VAE 的重构样本\hat{x} 进行细化建模,从而生成更高质量的样本。

  • 第一阶段 :VAE 建模

    给定训练数据(x_0, y),其中x_0 是高分辨率数据(例如图像),$y$ 是辅助条件信号(例如一些特征描述信息)。

    • 编码器:从(x_0, y) 中提取潜在表示z,即q_\psi(z|y, x_0)。这是一种条件后验分布,因为在给定(y,x_0) 的条件下,我们对z 的分布进行估计。

    • 解码器:在给定z 的前提下生成y,即p_\theta(y|z)。这里y 是由潜在表示z 解码而来的。

    注意,这里 VAE 最后的输出不是直接生成x_0 ,而是生成 条件辅助信号y ,随后我们会在第二阶段使用扩散模型来从yz 再还原出x_0 ,即生成x_0 是扩散模型DDPM算法的任务。形象一点来讲,VAE 的作用是给DDPM提供一个充满了信息(潜在特征)的藏宝图(低维潜在空间),真正一步步(逐渐加噪去噪)找到宝藏(恢复图片)还是要靠DDPM算法来实现

  • 第二阶段:DDPM 建模

    在有了 VAE 产生的条件信号y 以及对应的潜在表示z 后,我们利用条件扩散模型(DDPM)从y,z 出发,对x_0 进行细化生成。

    这里和单纯扩散模型DDPM算法过程类似,依旧是加噪和去噪过程

    • 前向加噪过程

      从真实数据x_0 开始,不断向其添加噪声,经过T 步得到一系列加噪样本x_{1:T}。这一过程用q(x_{1:T}|y,z,x_0) 表示,该分布是已知的(通常是固定高斯加噪过程)。

    • 逆向去噪过程

      扩散模型要学习从纯噪声逐步去噪回到x_0 的逆向分布p_\phi(x_{0:T}|y,z)。这里p_\phi 由参数\phi 决定,需要通过训练来拟合。在条件y 和潜在表示z 下逐步还原出x_0

这里需要从真实数据x_0 开始加噪,然后得到x_{1:T} ,这里从真实数据开始而不是直接从y,z 开始的原因是;真实数据提供了逆向去噪的正确答案,即当我们用真实数据加噪后,逆向去噪应该能从x_T 最终得到原始的x_0,那么就说明去噪过程是正确的。

 

DiffuseVAE 的低维潜在空间能够捕捉生成样本的全局特征(内容、形状或结构),这就意味着用户可以通过修改这个潜在空间中的变量来直接影响生成样本的主要特性。并且 DiffuseVAE 在速度和质量之间取得了更好的平衡。例如,仅需 10 步采样就能生成质量合理的样本,比传统 DDPM 的效率更高,且生成效果优于相同采样次数的其他方法(如 DDIM)。DiffuseVAE 经过预训练后,对条件信号中的不同噪声类型表现出较好的泛化能力,证明其框架不仅适合特定场景,还能应对多种不确定性条件。

 

数学理论

元素表达

  • x_0:我们想生成的高分辨率图像;

  • y:通过 VAE 建模的辅助条件信号;

  • z:与y 相关联的潜在表示;

  • x_{1:T}:扩散模型学习到的T 个加噪表示。

  • \hat{x}_0 : VAE 的重构图像

联合分布的分解

p(x_{0:T}, y, z) = p(z)p_\theta(y|z)p_\phi(x_{0:T}|y, z)
  • p(z) 是潜在变量的先验分布,通常设为标准高斯分布;

  • p_\theta(y|z) 表示在潜在变量z 下生成条件信号y 的概率,这是由 VAE 的解码器学习到的;

  • p_\phi(x_{0:T}|y, z) 是扩散模型的逆向去噪过程,用来从条件信号y 和潜在表示z 逐步恢复数据。

后验分布的近似

在给定数据(x_0, y) 的情况下,潜在表示z 和中间加噪样本x_{1:T} 的后验分布:

p(x_{1:T}, z|y, x_0)

该后验分布没有显式解析解,于是我们用一个近似分布q(x_{1:T}, z|y, x_0) 来替代:

q(x_{1:T}, z|y, x_0) = q_\psi(z|y, x_0)q(x_{1:T}|y, z, x_0)

这里的近似分布由两部分构成:

  • q_\psi(z|y, x_0) 是 VAE 的识别网络,它负责从yx_0 中推断潜在变量z

  • q(x_{1:T}|y, z, x_0) 是扩散模型的前向加噪过程,从x_0 开始加噪到x_{1:T}。这个分布是已知且不需要参数化,因为前向过程是人为设定的加噪机制。

因此我们将复杂的后验分布分解成了两部分:一部分是 VAE 的后验(学习z),另一部分是一个确定的加噪过程(生成x_{1:T}

数据对数似然

我们的最终目标是最大化数据(x_0, y) 的对数似然\log p(x_0, y)

\log p(x_0, y) = \log \int p(x_{0:T}, y, z) dx_{1:T} dz

这表示,我们需要对联合分布p(x_{0:T}, y, z) 进行积分,边缘化掉x_{1:T}z。但是,和 VAE 情况类似,在实际中非常困难计算,为此,我们引入了 证据下界(ELBO) 来对\log p(x_0, y) 进行下界估计,以便通过优化 ELBO 来最大化\log p(x_0, y)

证据下界(ELBO)

ELBO 的定义是:

\begin{aligned} \log p(x_0, y) \geq & \underbrace{\mathbb{E}_{q_\psi(z|y,x_0)} \big[ p_\theta(y|z) \big] - \mathcal{D}_{KL} \big(q_\psi(z|y,x_0) \| p(z)\big)}_{\mathcal{L}_{\text{VAE}}} + \mathbb{E}_{z \sim q(z|y,x_0)} \left[\underbrace{ \mathbb{E}_{q(x_{1:T}|y,z,x_0)} \left[ \frac{p_\phi(x_{0:T}|y,z)}{q(x_{1:T}|y,z,x_0)} \right] }_{\mathcal{L}_{\text{DDPM}}}\right] \end{aligned}

证明

我们希望最大化数据(x_0,y) 的对数似然:

\log p(x_0,y) = \log \int p(x_{0:T}, y, z) \,dx_{1:T}\,dz

这里x_{1:T}z 是潜在变量(未观测量),直接求解该对数似然通常非常困难。

步骤1:引入近似分布

为了对上述不可直接计算的积分进行处理,我们引入两个近似后验分布:

  • q_\psi(z|y,x_0):用于逼近真实后验p(z|x_0,y)

  • q(x_{1:T}|y,z,x_0):用于逼近真实加噪过程的后验(在此情况中,它往往是定义好的正向加噪过程,而非需近似的;这里更多是从数学形式上刻画)。

有了这两个分布,我们可在对数似然中乘以1的巧妙形式:

1 = \frac{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)}{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)}

将此代入后得到:

\log p(x_0,y) = \log \int p(x_{0:T}, y, z) \frac{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)}{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)} \,dx_{1:T}\,dz

步骤2:期望形式表示

将积分形式转换为期望的概率表示:

\log p(x_0,y) = \log \mathbb{E}_{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)}\left[ \frac{p(x_{0:T}, y, z)}{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)} \right]

现在我们有期望的结构,可以使用 Jensen’s inequality(詹森不等式)对对数的期望下界化。

步骤3:应用 Jensen’s inequality 获得ELBO下界

对于任意随机变量X 和凹函数\log(\cdot),詹森不等式有:

\log \mathbb{E}[X] \geq \mathbb{E}[\log X]

将其应用于上式:

\log p(x_0,y) \geq \mathbb{E}_{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)}\left[\log \frac{p(x_{0:T}, y, z)}{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)}\right]

这就得到了所谓的ELBO(下界)。

步骤4:分解联合分布p(x_{0:T}, y, z)

对联合分布使用给定模型的分解方式:

p(x_{0:T}, y, z) = p(z)p_\theta(y|z)p_\phi(x_{0:T}|y,z)

代入后:

\log p(x_0,y) \geq \mathbb{E}_{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)}\left[\log \frac{p(z)p_\theta(y|z)p_\phi(x_{0:T}|y,z)}{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)}\right]

将对数展开为加法:

= \mathbb{E}_{q_\psi(z|y,x_0)q(x_{1:T}|y,z,x_0)}\left[\log p_\theta(y|z) + \log p_\phi(x_{0:T}|y,z) + \log p(z) - \log q_\psi(z|y,x_0) - \log q(x_{1:T}|y,z,x_0)\right]

把联合期望写成双重期望的形式,我们得到:

\log p(x_0,y) \geq \mathbb{E}_{q_\psi(z|y,x_0)}\mathbb{E}_{q(x_{1:T}|y,z,x_0)}[\log p_\theta(y|z) + \log p_\phi(x_{0:T}|y,z) + \log p(z) - \log q_\psi(z|y,x_0) - \log q(x_{1:T}|y,z,x_0)]

 

步骤5:分离VAE部分与DDPM部分

观察上述期望,我们可将其分为两部分:

  1. 仅与z 有关的项(VAE相关项):

    \mathbb{E}_{q_\psi(z|y,x_0)}[\log p_\theta(y|z) + \log p(z) - \log q_\psi(z|y,x_0)]

    此处用标准的VAE推导可知,这部分是VAE的ELBO形式:

    • \log p_\theta(y|z) 是重建项

    • -\log q_\psi(z|y,x_0)\log p(z) 组合后给出KL散度约束项( -D_{KL}[q_\psi(z|y,x_0)||p(z)]

    整合后,这部分为:

    \mathcal{L}_{\text{VAE}} = \mathbb{E}_{q_\psi(z|y,x_0)}[\log p_\theta(y|z)] - D_{KL}[q_\psi(z|y,x_0) \| p(z)]
  2. 与扩散过程相关的项(DDPM相关项):

    \mathbb{E}_{z \sim q_\psi(z|y,x_0)}\left[\mathbb{E}_{q(x_{1:T}|y,z,x_0)}[\log p_\phi(x_{0:T}|y,z) - \log q(x_{1:T}|y,z,x_0)]\right]

    注意到\log p_\phi(x_{0:T}|y,z) - \log q(x_{1:T}|y,z,x_0) 可以进一步写成对数比:

    \mathbb{E}_{q(x_{1:T}|y,z,x_0)}\left[\log \frac{p_\phi(x_{0:T}|y,z)}{q(x_{1:T}|y,z,x_0)}\right]

    将此部分定义为扩散模型相关的ELBO项:

    \mathcal{L}_{\text{DDPM}} = \mathbb{E}_{q(x_{1:T}|y,z,x_0)}\left[\log \frac{p_\phi(x_{0:T}|y,z)}{q(x_{1:T}|y,z,x_0)}\right]

因此,最终得到的ELBO表达式为:

\log p(x_0,y) \geq \underbrace{\mathbb{E}_{q_\psi(z|y,x_0)}[\log p_\theta(y|z)] - D_{KL}[q_\psi(z|y,x_0) \| p(z)]}_{\mathcal{L}_{\text{VAE}}} + \mathbb{E}_{z \sim q_\psi(z|y,x_0)}\left[ \underbrace{\mathbb{E}_{q(x_{1:T}|y,z,x_0)}\left[\log \frac{p_\phi(x_{0:T}|y,z)}{q(x_{1:T}|y,z,x_0)}\right]}_{\mathcal{L}_{\text{DDPM}}}\right]

证毕

上式可以分为两大部分:

  1. VAE 部分 ( \mathcal{L}_{\text{VAE}})

\mathcal{L}_{\text{VAE}} = \mathbb{E}_{q_\psi(z|y,x_0)} \big[ p_\theta(y|z) \big] - \mathcal{D}_{KL} \big(q_\psi(z|y,x_0) \| p(z)\big)
  • 这是标准的 VAE 损失项,包括一个重构项(期望条件下zy 的重构对数似然)以及 KL 散度项(约束q_\psi(z|y,x_0) 接近先验p(z))。

  • VAE 部分的优化确保了潜在空间z 有意义,同时能让解码器在给定z 的情况下生成合理的y

  1. DDPM 部分 ( \mathcal{L}_{\text{DDPM}})

\mathbb{E}_{z \sim q(z|y,x_0)} \left[\underbrace{ \mathbb{E}_{q(x_{1:T}|y,z,x_0)} \left[ \frac{p_\phi(x_{0:T}|y,z)}{q(x_{1:T}|y,z,x_0)} \right] }_{\mathcal{L}_{\text{DDPM}}}\right]

要理解这个公式,首先我们需要回顾 DDPM 的背景:

在 DiffuseVAE 的框架中, x_{1:T} 是扩散模型引入的中间加噪步骤,也就是从x_0 不断加噪直到几乎变成高斯噪声的过程。这里有两个分布很重要:

(a) 正向加噪过程q(x_{1:T}|y,z,x_0)

这个过程是已知且固定的(通常我们设计一个固定的加噪策略)。给定真实数据x_0、潜在变量z 和条件信息y,我们能确定如何一步步往数据里加噪得到x_{1:T}。该分布不需要学习,因为加噪规则是人定义的。

(b) 逆向去噪过程p_\phi(x_{0:T}|y,z)

这个过程是模型需要学习的,即从纯噪声(或高度加噪的x_T)一步步“去噪”回x_0 的概率分布。该分布由参数\phi 控制,通过神经网络拟合。如果参数\phi 拟合的好,说明p_\phi 完美匹配了真正的逆向分布,那么说明可以从噪声逆推回真实数据x_0

再重新看这个公式,它分为内层期望和外层期望两部分

  • 内层期望

    \mathbb{E}_{q(x_{1:T}|y,z,x_0)} \left[ \frac{p_\phi(x_{0:T}|y,z)}{q(x_{1:T}|y,z,x_0)} \right]
    • 分子p_\phi(x_{0:T}|y,z):希望从噪声出发还原数据的模型分布。

    • 分母q(x_{1:T}|y,z,x_0):这是数据真实加噪时的分布。

    如果他们两个的比值期望大,那么说明逆向过程越好(因为真实加噪过程和真实逆向过程是正好相反的,如果模拟逆向过程等于真实逆向过程,那么这个差值最大期望最大),相反,如果p_\phiq 不匹配,这个期望值就会偏小。换句话说,这个期望项在测量p_\phiq 的吻合程度

  • 外层期望

    在上面讨论的内层期望中,我们是固定z 来看问题。但是z 本身是由q_\psi(z|y,x_0) 给定的(VAE的后验分布)。我们希望无论z 如何变化,模型都能保证逆向过程p_\phi 与正向过程q 匹配。

    因此我们对z 的分布再求一个期望,这样就可以确保扩散模型在整个潜在空间分布下(即对各种可能的全局特征)都能完成从噪声到x_0 的还原。这保证了模型的普适性,即对各类潜在表征都能有效地还原数据。

 

综上所述,内层期望是在问:如果按照正向过程q 的方式对x_0 加了噪声,那么逆向过程p_\phi 能否以高概率还原这个加噪后的样本到原始数据x_0?如果能,这个比值就高,期望就大。外层期望再考虑对所有可能的z 做平均,以确保无论潜在表征是什么,模型都能完成这个逆向还原。因此,通过优化这个期望,我们就能让p_\phi 更好地匹配q 的逆过程,从而在训练中“教会”扩散模型如何去噪回到x_0

最后,ELBO 即为两部分之和:

\log p(x_0, y) \geq \mathcal{L}_{\text{VAE}} + \mathbb{E}_{z \sim q(z|y,x_0)}[\mathcal{L}_{\text{DDPM}}]

这为联合训练提供了一个可优化的目标函数。

总的来说,DiffuseVAE 的训练目标通过联合分布、后验分布近似以及 ELBO 分解,将 VAE 和扩散模型的损失函数结合起来。这个设计不仅让模型能够生成高质量的样本,还利用了低维潜在空间,提升了生成的控制性和效率。训练过程中,VAE 负责潜在表示的学习,扩散模型负责生成过程的细节还原,两者相辅相成。

 

简化设计

条件信号y 的简化

在原始 DiffuseVAE 的设计中, y 是一个与x_0​ 有关但不完全等同的条件信号。这样做的目的是为扩散模型提供附加条件,从而在逆向去噪过程中能够利用这些信息对生成样本进行更精细的控制。然而,这种设计在实际实现中可能会变得复杂:

  • 模型需要在去噪的每个步骤处理(y,z) 的条件信息,这增加了运算和实现上的难度。

  • y 的选取和设计也会影响最终模型的性能与适用性,开发者需要思考如何选择合适的y

因此做了一个重要简化,假设y = x_0,也就是说,条件信号直接等于目标图像本身。我们简单来解释他,我们最开始的假想是在去噪的过程中额外输入一个与x_0 有关但不同的y 作为去噪的补充信息,比如说,VAE 生成了一个图像,模模糊糊的,我们现在告诉他,你要生成的应该是一只猫,现在我们做了这个简化,相当于去掉了补充信息,而直接将原图像作为补充信息,换句话说,直接给去噪过程标准答案,告诉他,你应该生成的就是这个猫图像类似的新图像。相当于压根就没有什么y ,研究到最后发现y 这个创新点没用,但是提都提了,就这么草率的把输入图像作为补充信息了。其实本质上就是用了两次输入图像,一份用来给 VAE 提取特征,一份给DDPM 作为参考答案

 

潜在表示z 的简化

潜在表示$z$ 是由 VAE 编码器从(x_0) 中获得的潜在表示。在原本的框架中,扩散模型可能需要同时考虑(y,z) 的条件输入,这意味着扩散模型输入要依赖潜在空间的表达,但是把一个潜在空间作为输入怎么做?没法做! 这就是为什么对潜在空间z 进行了简化。

简化思路是:既然z 是从x_0 映射得到的,那么 VAE 的解码器在给定z 时生成的\hat{x}_0 已经是z 的一个确定性函数。也就是说, \hat{x}_0 完整地代表了z 中包含的信息。这就使得扩散模型可以直接使用\hat{x}_0 作为条件参考,而无需显式地输入z

我们还是简单的来解释他,其实思路是非常漂亮的,通过 VAE 得到输入,再传给 Diffusion 相当于引入了潜在空间,让 diffusion 充分利用潜在空间,完美的符合了DiffuseVAe 的思想,但是实际实现不出来,因为潜在空间不不是一个显式的数据集,是一个难啃的生数据,我们在 Diffusion 中用不了它,要想用,还是得在 VAE 中把z 转换成\hat{x}_0 ,这样扩散模型接收的依旧是“图像形式”的条件,而不是抽象的潜在矢量。本质上是什么,本质上就是把 VAE 得到的结果输入给 DDPM 就是这么简单,DDPM中根本就没用潜在空间。

两阶段独立训练的简化

原本的想法是联合训练 VAE 部分(负责z 的学习和y 的重构)以及扩散部分(负责对重构样本做细化),这样可以在理论上实现一个端到端训练:模型同时学会提取潜在表示和利用扩散模型从潜在表示中生成高质量图像。但是! 联合训练在实际实现中非常复杂,参数多,计算量大,互为条件环路,即VAE 的输出是扩散模型的输入,扩散模型的性能又影响 VAE 的参数搜索空间,非常容易导致训练不稳定

为了简化训练过程,我们采用了顺序的两阶段训练方法:

  1. 第一阶段只训练 VAE,学会从x_0 中提取潜在表示z 并重构出\hat{x}_0。这一步本质是一个标准的 VAE 训练任务,优化相对简单。

  2. 第二阶段固定 VAE 的参数(不再训练),只使用 VAE 训练好的模型表示。然后再单独运行扩散模型,这样就不会被 VAE 的动态变化扰动。

说明了什么,说明了之前说的端对端联合优化没能实现,还是只能使用传统思想,固定一个,训练另一个,VAE 提供一个初步结果,而 DDPM 专注于细化这些结果,从而结合两者的优势。

总结一下这三点简化

  • 条件信号y 被简化为目标数据x_0

  • 第二阶段扩散模型的输入是 VAE 重构结果\hat{x}_0

  • 使用顺序的两阶段训练方式,固定一个,训练另一个

从这里我们可以看出,理论很美好,实践很困难啊!

 

VAE 参数化的选择

在 VAE 的选择上,论文决定只使用最基本、最传统的 VAE 构型(单随机层的简单 VAE),而不是使用更复杂的多层变种:

  • 使用复杂的 VAE(比如多层潜在变量、多级别特征抽取)可以潜在地获得更强大的表示能力,但也会使潜在空间更难解释、更不稳定,从而不利于对潜在表示z 的直观控制和理解。

  • 简单的 VAE 虽然表达能力或许弱一些,但潜在空间更易于理解和利用,符合 DiffuseVAE 利用潜在空间来控制全局特性的初衷。

DDPM扩散模型两种设计假设

在将 VAE 重构的\hat{x}_0 与扩散模型相结合时,有两种主要的设计思路(两种公式假设),分别在正向与逆向过程的依赖关系上做出了不同的简化。

第一种简化假设:

  1. 正向过程的条件独立性

    假设正向加噪过程q(x_{1:T}|x_0,z)z\hat{x}_0 条件独立,也就是正向过程不需要使用z\hat{x}_0的信息,只从原始数据x_0 通过标准扩散加噪就行:

    q(x_{1:T}|x_0, z) \approx q(x_{1:T}|x_0)

    也就是说,加噪过程与 VAE 相关信息无关,依旧是传统的 DDPM 算法的扩散加噪,从原图像x_0 开始,每一步只是简单地注入噪声,潜在空间一点没用上。

  2. 逆向过程的依赖性

    逆向去噪过程中,扩散模型使用VAE 的重构结果\hat{x}_0 作为参考信息,而不是直接依赖潜在表示z,即:

    p(x_{0:T}|z) \approx p(x_{0:T}|\hat{x}_0)

    在每个去噪步骤中,把 \hat{x}_0 与当前带噪声的图像 x_t 一并拼接后[\,x_t,\, \hat{x}_0\,] 输入到扩散网络中,让网络预测噪声或直接预测x_{t-1} \hat{x}_0 提供“目标图像应该长什么样”的参考,而 x_t 是当前带噪声的图像,网络学着如何将 x_t 去噪并向 \hat{x}_0 靠拢。

     

第二种简化假设:

  1. 正向过程的依赖性

    这里正向过程被设计为显式依赖于 VAE 重构结果\hat{x}_0

    q(x_{1:T}|x_0, z) \approx q(x_{1:T}|\hat{x}_0, x_0)

    这意味着在每一步加噪过程中,模型会直接结合x_0\hat{x}_0 的信息,即把 \hat{x}_0 也作为输入条件,使加噪分布围绕着 \hat{x}_0 进行。

  2. 逆向过程的依赖性

    和第一种简化假设相同:

    p(x_{0:T}|z) \approx p(x_{0:T}|\hat{x}_0)

正向过程的具体设计

在第二种假设中,主要针对了正向过程中的输入条件,把 VAE 重构结果\hat{x}_0 也包含在加噪过程中,我们先和标准的DDPM比较一下:

  • 在标准 DDPM 中:如果不断加大噪声,到达x_T 时,理想情况下x_T 就分布于一个各向同性的高斯噪声\mathcal{N}(0, \mathbf{I})

  • 在第二种假设中:t \to T 且噪声足够大时, q(x_T \mid x_0, \hat{x}_0) \approx \mathcal{N}\bigl(\hat{x}_0, \mathbf{I}\bigr) ,我们先给出结论,后续给出推导过程

也就是说,最终的 带噪样本分布是围绕 \hat{x}_0 的一个近似单位方差的高斯,而不是围绕 0,成为了加噪和去噪的中心参考点。这样有一个好处,样本都时刻记得有一个 \hat{x}_0 作为 目标中心,因而在逆向还原时更容易保持与 \hat{x}_0 的结构或风格一致。

但是实际上没有这么高深,其想法很简单,DDPM 算法中输入 A 输出也是 A,在 DiffuseVAE 中输入 VAE 重构结果\hat{x}_0 输出肯定也得是\hat{x}_0 ,然后给出标准答案(原图) x_0 作为参考。

 

正向过程详细数学推导

假设对正向过程每一步都采用变分推断中常用的高斯分布假设:

q(x_t|x_{t-1}, \hat{x}_0) = \mathcal{N}(\sqrt{1-\beta_t}x_{t-1} + (1-\sqrt{1-\beta_t})\hat{x}_0, \beta_t \mathbf{I})
  • x_{t-1}x_t 会注入高斯噪声,方差为\beta_t \mathbf{I}

  • 均值项中不仅有\sqrt{1-\beta_t}\,x_{t-1},还额外加上(1 - \sqrt{1-\beta_t})\,\hat{x}_0

    • \beta_t 很小,则\sqrt{1-\beta_t} \approx 1,此时均值主要受x_{t-1} 控制, \hat{x}_0 占的比重较小

    • 随着\beta_t 增大,每一步加噪时, \hat{x}_0 占的权重会逐渐明显,让 x_t 更靠近 \hat{x}_0

在每个时间步中, x_t 的生成不仅依赖于前一时间步x_{t-1},还结合了\hat{x}_0 的全局信息,即x_t 由原始数据x_0 和 VAE 重构\hat{x}_0 的加权和决定的。

第一步t = 1 的加噪过程可以写为:

q(x_1|x_0, \hat{x}_0) = \mathcal{N}(\sqrt{1-\beta_1}x_0 + \hat{x}_0, \beta_1 \mathbf{I})

这样把原图像x_0 的主要信息留下,同时把 \hat{x}_0 一次性地注入到分布里。

 

正向条件边缘分布

通过连续乘法展开,可得到该正向过程在第t 步时的边缘分布

q(x_t \mid x_0, \hat{x}_0) = \mathcal{N}\Bigl(\sqrt{\bar{\alpha}_t}\, x_0 + \hat{x}_0,\;\bigl(1-\bar{\alpha}_t\bigr)\mathbf{I}\Bigr)

其中

\bar{\alpha}_t = \prod_{i=1}^t (1 - \beta_i)
  • t 逐渐增大, \bar{\alpha}_t 会变得很小(因为(1 - \beta_i) 都小于 1,相乘会衰减),于是\sqrt{\bar{\alpha}_t} \approx 0

  • 故当t 到达T 时, \bar{\alpha}_T \approx 0,即可得到

    q(x_T \mid x_0, \hat{x}_0) \approx \mathcal{N}\bigl(\hat{x}_0, \mathbf{I}\bigr)

 

 

实现

在不同潜在空间进行插值的差异

如前所述,DiffuseVAE 中包含两种潜在表示:

  • 低维潜在向量z_\text{vae}(来自第1阶段 VAE)

  • 高维潜在表示x_{1:T}(来自第2阶段 DDPM 逆向过程)

在实验中,我们分别在z_\text{vae} 空间和x_T 空间(即 DDPM 逆向过程的初始分布采样空间)中进行插值,直观地对比了两种插值方式对合成结果所产生的影响。

z_\text{vae} 空间进行插值

  • 插值公式

    \tilde{z}_\text{vae} \;=\; \lambda z_\text{vae}^{(1)} \;+\; (1 - \lambda)\,z_\text{vae}^{(2)} \quad 0 < \lambda < 1

    将插值得到的\tilde{z}_\text{vae} 输入到 DiffuseVAE 中,先通过 VAE 生成模糊初始图,再进入扩散DDPM细化过程,得到最终清晰结果。

x_T 空间进行插值

  • 插值公式

    \tilde{x}_T \;=\; \lambda\,x_T^{(1)} \;+\; (1 - \lambda)\,x_T^{(2)}\quad 0 < \lambda < 1

    其中x_T^{(1)}x_T^{(2)} 是从 DDPM 的初始分布p(x_T) 中分别采样得到的高维表示(形状与图像相同)。随后, \tilde{x}_T 通过第2阶段逆向过程逐步去噪,生成最终图像。

z_\text{vae} 空间插值用于全局结构的控制, x_T 空间插值则更多地用于细节层面的操作。

 

 

DiffuseVAE 在采样速度与质量权衡中的表现

扩散模型(DDPM) 通常面临“采样步数”与“样本质量”之间的两难:

  • 采样步数越多(如 1000 步),生成样本质量往往越好,但速度更慢;

  • 采样步数越少(如 10、25、50 步),采样速度更快,但生成质量低。

DiffuseVAE 在低采样步数时更优

  • T 从 10 到 100 时,DiffuseVAE 的 FID 分数始终优于无条件 DDPM,特别是在T=25T=50 时优势明显。

无条件 DDPM 在极高步数($T=1000$)时更优

  • 当将采样步数提升到训练时常用的 1000 步,纯 DDPM 的最终 FID 略好于 DiffuseVAE。

  • 论文中推测原因在于 VAE 的先验空洞问题(prior-hole problem):VAE 的先验分布p(z) 与后验q(z) 不匹配,部分潜在区域对应的样本质量太差,DDPM 对这些样本的细化也无能为力,从而整体 FID 拉低。

通过后拟合(Ex-PDE 等)可显著改善

  • 若在训练后,对 VAE 潜在编码分布用更灵活的模型(如高斯混合模型 GMM)进行拟合并从中采样,可以缓解先验空洞问题。结果表示在T=1000 步时其结果大幅改善

与最先进模型的比较

 

在 CIFAR-10 上

接下来我们直接以成绩说话。

VAE Baseline :FID 139.50,IS 3.23。单纯的VAE是当之无愧的垫底。

DiffuseVAE : FID 在 2.62 - 2.95 之间,IS 大约在 9.5 - 9.7 之间,比单纯的 DDPM 效果好一些。在众多算法中排前列,但是我认为这是DDPM的功劳。

 

 

在 CelebA-64x基准上

我们可见 DIffuseVAE 并没有优于 DDPM,这是不应该的。

在 CelebA-256x基准上

还行,主要是靠的 DDPM。