CVPR2022 | 扩散自编码器

2023-10-24 10:45:31 浏览数 (1)

简介

基于扩散和基于分数的生成模型最近在合成逼真的高分辨率图像方面取得了成功,相比 GAN 取得了更好的结果。这两种模型紧密相关,优化的是相似的目标。而本文考虑的是 DPMs 是否能够成为一个良好的表示学习器.具体而言,本文试图提取一个包含高层语义且允许近乎精确重构图像的有意义且可解码的表示。本文的探索重点是扩散模型,但这些贡献也适用于基于分数的模型。

寻找一个有意义且可解码的表示需要捕捉高层语义和低层随机变化。本文的关键思想是利用可学习的编码器来学习这两个级别的表示,用于发现高层语义,并利用 DPM 进行解码和建模随机变化。具体而言,本文使用了一种条件变种的去噪扩散隐式模型(DDIM)作为解码器,并将潜在编码分为两个子编码。第一个语义子编码由 CNN 编码器推断得到,较为紧凑;第二个随机子编码通过反转 DDIM 变体的生成过程,以语义子编码为条件进行推断。DDIM 与其他 DPM 不同的地方在于它将前向过程修改为一个非马尔可夫过程,同时保持 DPM 的训练目标。这种修改使得 DDIM 可以确定地将图像编码为相应的初始噪声,这就是本文中的随机子代码。

这个框架的含义有两个方面。首先,通过以目标输出的语义信息作为 DDIM 的条件,去噪变得更容易、更快速。其次,这种设计产生了一种线性的、语义上有意义的、可解码的表示——这是 DPM 潜在变量的一个新属性。这个重要的属性允许利用 DPM 来完成许多任务,包括那些对任何基于 GAN 的方法都具有极高挑战性的任务,例如在真实图像上进行插值和属性操作。与依赖错误率高的反演才能在真实图像上操作的 GAN 不同,本文的方法不需要优化来编码输入,并且会产生保留原始细节的高质量输出。

本文提出的框架还可以用于图像生成任务。尽管它是自编码器,通常不用于无条件生成,但我们的框架可以通过适应另一个 DPM 到语义子编码码分布来生成图像样本。这种组合在无条件生成方面的 FID 得分与纯 DPM 相比具有竞争力。此外,从我们的紧凑、有意义的潜在空间中进行采样的能力还支持少样本条件生成。与其他基于 DPM 的技术在少样本设置下所使用的对比学习不同,我们的方法仅使用少量标记示例即可产生令人信服的结果。

方法

为了得到有意义的潜在编码,本文设计了一个条件 DDIM 解码器

p({rm x}_{t-1}|{rm x}_t, rm z_{sem})

,以额外的隐变量

rm z_{sem}

作为条件,并设计了一个语义编码器

rm z_{sem}=Enc_{phi}(x_0)

学习如何将输入图像

rm x_0

映射到在语义上有意义的

rm z_{sem}

。这个条件 DDIM 解码器的输入表示为隐变量

{rm z}=({rm z_{sem}}, {rm x}_T)

,包含高层次的语义子编码

rm z_{sem}

和通过 DDIM 生成过程的反向过程推断得到的低层次的随机子编码

{rm x}_T

。在这个框架中,DDIM 既是解码器也是随机编码器。整体构成如图 1 所示。

图1 Diff-AE框架示意图

与其他使用空间条件变量(例如 2D 潜在图)的条件 DPM 不同,本文中的

rm z_{sem}

是一个非空间的维度为

d=512

的向量,类似于StyleGAN中的风格向量(style vector),它允许模型编码全局语义而不是特定于任何空间区域。本文的一个目标是学习一个语义丰富的潜在空间,使得可以进行平滑插值,类似于GAN所学习的那种能力,同时保持扩散模型在重构方面的优越性。

基于扩散的解码器

本文中的条件 DDIM 解码器以

{rm z}=({rm z_{sem}}, {rm x}_T)

为输入产生图像。该解码器用于建模

p_{theta}({rm x}_{t-1}|{rm x}_t, rm z_{sem})

,以匹配扩散模型中定义的推理分布

q({rm x}_{t-1}|{rm x}_t, rm x_0)

。其生成过程如公式 1 和 2 所示。

p_{theta}({rm x}_{0:T}|{rm z_{sem}})=p({rm x}_T)prod_{t=1}^Tp_{theta}({rm x}_{t-1}|{rm x}_t, rm z_{sem}) quad (1)
p_{theta}({rm x}_{t-1}|{rm x_t}, rm z_{sem})= begin{cases} cal{N}({rm f}_{theta}({rm x_1}, {rm 1}, {rm z_{sem}}), {rm 0})~~~~~~~~~{rm{if}}~~t=rm 0\ q({rm x}_{t-1}|{rm x}_t, {rm f}_{theta}({rm x}_t, t, rm z_{sem})) ~~otherwise end{cases} quad (2)

然后将公式 2 中的

{rm f}_{theta}

参数化为一个噪声预测网络

epsilon_{theta}({rm x}_t, t, rm z_{sem})

,如公式 3 所示。

{rm f}_{theta}({rm x}_t, t, {rm z_{sem}})=frac{1}{sqrt{alpha_t}}({rm x}_t-sqrt{1-alpha_t}epsilon_{theta}({rm x}_t, t, {rm z_{sem}})) quad (3)

这个网络是一个改进版本的 UNet。训练的目标函数如公式 4 所示。

L_{simple}=sum_{t=1}^T{Bbb E}_{{rm x_0},epsilon_t}[|epsilon_{theta}({rm x}_t, t, {rm z_{sem}})-epsilon_t|_2^2]quad (4)

其中

epsilon_tinBbb R^{3times h times w}simcal N({bf 0,I}),{rm x}_t=sqrt{alpha_t}{rm x}_0 sqrt{1-alpha_t}epsilon_t

T

设定为一个较大的数,如1000。在训练时不需要随机子编码

{rm x}_t

。本文使用 AdaGN 对 UNet 进行调节,它通过在归一化的特征图上应用通道级别的缩放和偏移来扩展组归一化。AdaGN 以

t

{rm z_{sem}}

为条件。

{rm AdaGN}({rm h},t,{rm z_{sem}})={rm z}_s({rm t}_s{rm GroupNorm}({rm h}) {rm t}_b)quad (5)

其中

{rm z}_sin{Bbb R}^c={rm Affline}({rm z_{sem}})

({rm t}_s,{rm t}_b)in{Bbb R}^{2times c}={rm MLP}(psi(t))

是具有正弦编码函数的多层感知器的输出。这些层在整个 UNet 中都被使用。

语义编码器

语义编码器

{rm Enc}({rm x}_0)

的目标是将输入图像总结成一个描述性向量

{rm z_{sem}}={rm Enc}({rm x}_0)

,使其具有必要的信息来帮助解码器

p_{theta}({rm x}_{t-1}|{rm x_t}, rm z_{sem})

去噪和预测输出图像。本文没有对这个编码器做任何特定的架构假设;然而,在实验中,这个编码器与 UNet 解码器的前半部分有相同的架构。使用信息丰富的

{rm z_{sem}}

来调节 DDIM 的一个好处是使去噪过程更加高效。

随机编码器

除了解码,条件 DDIM 还可以通过将其生成过程反向来将输入图像

{rm x}_0

编码为随机子编码

{rm x}_T

,如公式 6 所示。

{rm x}_{t 1}=sqrt{alpha_{t 1}}{rm f}_{theta}({rm x}_t, t, {rm z_{sem}}) sqrt{1-alpha_{t 1}}epsilon_{theta}({rm x}_t, t, {rm z_{sem}}) quad (6)

这个过程可以看作一个随机编码器。它只编码由

{rm z_{sem}}

剩下的信息,因为

{rm z_{sem}}

对压缩随机细节的容量有限。利用语义和随机编码器,该自编码器可以详尽地捕获输入图像中的细节,同时为下游任务提供高级别的表示

{rm z_{sem}}

通过扩散自编码器采样

{rm z_{sem}}

调节解码器时,扩散自动编码器将不再是生成模型。因此,为了实现从自编码器中进行采样,需要一个额外的机制来从潜在分布中对

{rm z_{sem}}in{Bbb R}^d

进行采样。虽然 VAE 是这个任务的一个吸引人的选择,但在保留潜在代码中丰富的信息和维护 VAE 的采样质量之间取得平衡很困难。GAN 是另一个选择,但它使训练稳定性复杂化,而这是 DPM 的一个主要优势。这里本文选择拟合另一个 DDIM,称为潜在 DDIM

p_{omega}({rm z}_{{rm sem},t-1}|{rm z}_{{rm sem},t})

,到

{rm z_{sem}=Enc_{phi}(x_0)}, {rm x_0}sim p({rm x_0})

的潜在分布中。与公式 3 和 4 类似,训练过程主要通过优化与

omega

相关的

L_{latent}

来实现。

L_{latent}=sum_{t=1}^TBbb E_{{rm z}_{{rm sem}},epsilon_t}[|epsilon_{omega}({rm z}_{{rm sem},t},t)-epsilon_t|_1] quad (7)

其中

epsilon_tinBbb R^dsimcal N({bf 0,I}),{rm z}_{{rm sem},t}=sqrt{alpha_t}{rm z_{sem}} sqrt{1-alpha_t}epsilon_t

T

与 DDIM 解码器中的相同。对于

L_{latent}

,本文发现

L_1

损失函数比

L_2

损失函数的表现更好。与1D/2D图像不同,非空间数据没有一个被广泛接受的DPM架构,但我们发现具有跳跃连接的深度MLP(10-20层)表现得相当好。本文首先训练语义编码器(

phi

)和图像解码器(

theta

),直到收敛,然后在固定语义解码器的情况下训练潜在 DDIM。在实践中,由潜在 DDIM 建模的潜在分布首先被归一化为零均值和单位方差。因此,从扩散自编码器中进行无条件采样的过程是:从潜在 DDIM 中采样

{rm z_{sem}}

并进行反归一化,然后采样

{rm x}_Tsim{cal N}({bf 0,I})

,最后使用解码器解码

{rm z}=({rm z_{sem}},{rm x}_T)

。本文选择事后训练潜在 DDIM 有几个实际原因。首先,由于训练潜在 DDIM 只需要完整训练时间的一小部分,事后训练可以在相同的扩散自编码器上进行对不同潜在 DDIM 的快速实验。另一个原因是通过不施加任何约束(如 VAE 中的先验损失),尽可能地保持

{rm z_{sem}}

的表达力,可以保证潜在变量的质量不受影响。

实验

本文的实验部分证明了高层次语义主要在

{rm z_{sem}}

中而基本不被

{rm x}_T

捕捉到,并且展示了 Diff-AE 框架在图像插值、图像属性编辑、图像重建、去噪过程加速、类别条件生成和无条件生成等任务中的表现。

图像插值

图 2 展示了在两张真实图像之间插值的效果。本文将 Diff-AE 与 StyleGAN 和 DDIM 进行了对比,展示了 Diff-AE 在插值的连续性、平滑性以及保持原图像特征方面优于二者。

图2 图像插值效果图

图像重建

表 1 展示了 Diff-AE 与其他一些模型在图像重建任务中的对比。模型使用 FFHQ 数据集进行训练,并在 CelebA-HQ 数据集上测试。结果表明 Diff-AE 与最先进的 NVAE 相比具有竞争力,同时在一个紧凑的 512 维向量

{rm z_{sem}}

中保存了有用的高层次语义信息。

表1 图像重建质量对比

此外,表 2 展示了消融实验结果,即提供随机采样的

{rm x}_T

和不同维度的

{rm z_{sem}}

重建的图像质量的对比。结果表明即使在

{rm x}_T

随机的情况下重建的图像在感知上也与原图相似(参考 LPIPS 指标)。而更高维度的

{rm z_{sem}}

会让重建的图像具有更高的保真度。

表2 消融实验结果

加速去噪过程

图 3 展示了在设定

T=10

时,时间步骤分别为

9,8,7,5,2,0

时 Diff-AE 和 DDIM 模型对输入图像

rm x_0

的预测结果。实验表明通过以

{rm z_{sem}}

作为条件,Diff-AE 可以明显更快地生成图像。

图3 模型预测效果对比

类别条件生成和无条件生成

表 3 展示了类别条件下在 CelebA 64 数据集的 5000 个采样图像和目标子集之间计算的 FID 分数。其中 Binary 代表二元分类器的结果,由 50 个正例和 50 个负例训练。PU 代表正无标签分类器的结果,由 100 个正例和 10,000 个无标签示例训练。Naive 代表所有图像和目标子集之间计算的结果。

表3 类别条件生成结果

表 4 展示了无条件生成的结果对比,采用的数据集有 FFHQ,LSUN Horse&Bedroom 和 CelebA 64。

表4 无条件生成结果

总结

本文提出了扩散自编码器 Diff-AE,可以从输入图像中分别推断出语义和随机信息。与 DPM 和像 NVAE 这样的高保真自编码器相比,Diff-AE 得到的潜在表示允许几乎精确的解码,同时包含对下游任务非常有用的紧凑语义。这些特性使得能够简单地解决各种真实图像编辑任务,而无需使用容易出错的 GAN 和其反演过程。该框架还提高了去噪效率,并保持了 DPM 的有竞争力的无条件采样能力。Diff-AE 仍然存在一些局限性。当对训练分布之外的图像进行编码时,由于来自DDIM的高维随机子代码,Diff-AE 仍然可以很好地重构图像。然而,推断得到的语义和随机子代码可能会落在学习到的分布之外,导致表示质量较差,无法再被解释或插值。尽管使用了非空间潜在代码来学习全局语义,但某些图像和空间推理任务可能需要更精确的局部潜在变量。对于这些任务,引入 2D 潜在映射可能是有益的。在图像生成方面,StyleGAN 具有的一个独特特性是能够控制特定尺度的生成,而 Diff-AE 没有这一能力。且就生成速度而言,Diff-AE 将从 DDIM 中获得高质量样本所需的时间步骤显著减少,但仍落后于只需要单个生成器遍历以生成图像的GANs。

0 人点赞