526互联

Latent Diffusion Model快读

发布时间 2023-08-06 18:19:51作者: 戈壁与草原

总览

  原始的扩散模型训练及推理是在像素空间进行的，这将耗费大量的GPU和能量资源。为了缓解这一状况，作者提出了一种方案，将高分辨率图像的合成过程分解为两个过程：图像编解码和隐特征空间扩散。

  • 图像编码器将像素空间中图像编码为隐特征latent code（生成过程不需要此步骤）
  • 扩散模型DM在隐特征空间latent space上取样sample进行去噪得到denoised_sample
  • 去噪样本经过图像解码器得到图像

这种方案有这样几个好处：

  • 扩散过程的训练推理，由高维的像素空间降低到低维的隐空间进行。比如原始像素空间512x512x3 --> 隐空间 1024. 维度降低了768倍
  • 利用了从UNet结构继承DM归纳偏置，这将对空间结构的数据十分有效，因此减轻了对先前方法所需的激进的质量降低的压缩级别的需求。
  • 获得了通用压缩模型，其潜在空间可用于训练多个生成模型，也可用于其他下游应用，例如单图像CLIP引导合成。

主体

感知图像压缩 Perceptual Image Compression

  感知压缩模型基于之前工作[1], 包含一个自编码器，其训练采用混合损失函数形式： 感知损失函数[2]和基于图像块的对抗损失函数[1]. 这种方式确保了局部真实，而且避免了仅仅依赖简单像素空间中L1或者L2损失函数带来的模糊问题。
具体是，

  • 给定RGB图像空间中的图像\(x \in R^{H \times W \times 3}\)
  • 编码器\(\mathcal{E}\)将图像\(x\)编码为隐空间表示\(z=\mathcal{E}(x) \in R^{h \times w \times c}\)
  • 解码器\(\mathcal{D}\)从该隐变量值解码/重建出图像 $\bar x = \mathcal{D}(z) $
  • 要求图像的下采样因子必须是2的幂次方，即\(f=\frac{H}{h}=\frac{W}{w}=2^m, m \in R\)

为了避免隐空间较大的方差，引入了KL正则化或者VQ正则化项。

隐扩散模型 Latent Diffusion Models

  相比原始扩散模型在像素空间进行训练推理，LDM是在维度更低的隐空间进行训练推理，其计算量则小了很多。其次利用了基于2D卷积UNet，有利于图像的归纳偏置。这里的UNet条件依赖于时间t。

条件建模机制 Conditioning Mechanisms

  与其它类型的生成模型一样，扩散模型DM也可以对条件分布\(p(z|y)\)进行建模拟合，通过条件去噪自编码器实现\(\epsilon_{\theta}(z_t, t, y)\)，通过条件输入如文本、语义图、图像-图像翻译任务等控制图像的生成。通过使用交叉注意力机制增强底层UNet结构，将DM变成更加灵活的条件图像生成器。这对于各种输入模式的学习注意力模型是有效的。为了处理不同模态的输入如文本输入，引入了一个模态相关的编码器 \(\tau_{\theta}\) ，将原始输入\(y\)，编码为\(\tau_{\theta}(y) \in R^{M \times d}\)，并将该特征输入到UNet的交叉注意力层。具体是

  • 输入1： 条件项特征$\tau_{\theta}(y) \in R^{M \times d} $
  • 输入2： UNet网络结构第i层t时间步特征 \(\psi_i(z_t) \in R^{N \times d}\)
  • 注意力公式 \(Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d}})\cdot V\)
  • \(Q=W_Q^{i}\cdot\psi_i(z_t), K=W_K^{i}\cdot\tau_{\theta}(y), V=W_V^{i}\cdot\tau_{\theta}(y)\)

实验

感知压缩平衡

  本节分析LDM在不同下采样因子情况下的表现，\(f \in \{1，2，4，8，16，32\}\)，简写为\(LDM-f\), \(LDM-1\)则对于像素空间的扩散模型，没有任何下采样。为了可比较，固定一致的计算资源-单张A100， 以及相同的参数量训练相同的迭代步数。

  • 论文中Tab. 8展示了不同采样因子与KL-reg/VQ-reg正则化组合实验与经典VQGAN，DALL-E方法的实验对比结果。整体上不管是KL-reg还是VQ-reg方法，\(f=4,8\)时的\(R-FID\downarrow\)，\(PSIM\downarrow\)指标都优于之前方法；\(f=4,8\)时\(PSNR\uparrow\)，\(SSIM\uparrow\)指标大多数都优于之前的方法。

  • 论文中Fig. 6展示在ImageNet数据集上训练了2M步类别条件模型的取样质量。更小的下采样因子\(f=1,2\)导致较慢的训练过程；而较大的下采样因子\(f=32\)则在训练了有限迭代步数后，生成图像的保真度基本停止不动了。而\(LDM-\{4，16\}\)则在效率和感知真实度上达到了很好的平衡。比如在2M迭代后，\(LDM-1\)的\(FID\downarrow\)指标比\(LDM-8\)的\(FID\downarrow\)指标值差距达到38。表明了\(LDM-8\)生成图像的真实度要高于\(LDM-1\)很多。

  • 论文中Fig. 7展示了分别在CelebA-HQ和ImageNet上训练的LDM模型。基于DDIM取样方法，采用不同采样步数下，采样步数与FID指标之间的关系。\(LDM-\{4,8\}\)则优于其它模型。特别是与像素空间上的\(LDM-1\)模型相比，在达到更低的\(FID\downarrow\)指标时，也能达到更快的样本生成能力。而复杂的数据集如ImageNet，则要求减小压缩率以避免较差的取样质量。从图中可以看到ImageNet在下采样因子\(f=32\)时，\(FID\downarrow\)指标值（棕色线）较高对应生成图像质量差。而CelebA-HQ这个单一的人脸数据集，在下采样因子\(f=32\)时（棕色线）并没有出现与ImageNet类似的情况。

Tab. 8
image

Fig. 6
image

Fig. 7
image

隐空间无条件图像生成

  论文在4个数据集CelebA-HQ, FFHQ, LSUN-Churches, LSUN-Bedrooms上训练无条件生成256分辨率图像，通过生成图片的生成质量，见Fig. 4；数据流形覆盖度用\(FID\downarrow\)和Precision-Recall指标来衡量LDM性能，见Tab. 1。

Fig. 4
image

Tab. 1
image

隐空间有条件图像生成

用于LDM输入的Transformer Encoder

  通过引入交叉注意力用于LDM的条件建模，为各种模态的条件依赖打开了一条道路。对于文生图的图像建模，论文在LAION-400M数据集上，训练了1.45B参数量的KL正则化的LDM模型。采用bert-tokenizer将文本信息token化，用transfomer实现\(\tau_{\theta}\)，将文本信息最终编码输入到UNet网络中。这种领域特定的语言表示与视觉合成产生了一个强大的模型，可以很好地推广到复杂的、用户定义的文本提示。见Fig. 5, Fig. 8

Fig. 5
image

Fig. 8
image

  定量的图像分析，依照之前的工作，在MS-COCO验证集上评估文生图的质量。
Fig.
image

用论文中最好的在ImageNet上训练的类别条件模型\(LDM-\{4,8\}\)与之前工作进行对比

Tab. 3
image

卷积取样256之外

  通过将空间对齐的条件信息聚合到扩散模型输入，LDM模型可以有效适用于一般目的的图像-图像迁移任务。比如语义合成、超分辨率合成、图像修复等。

基于LDM的超分辨合成

  LDM可以通过聚合低分辨率的图像有效训练出超分辨率生成模型，基于之前提到的条件建模机制。在第一个实验中，论文依照SR3论文中数据处理方法，采用双三次插值将训练图像进行4倍的下采样。在OpenImage数据集上训练出模型\(LDM-4\)（VQ-reg正则化），直接将低分辨率图像输入给UNet网络结构即\(\tau\)是恒等变换。定性和定量分析结果参考Tab.5 和 Fig. 10.

Tab.5
image

Fig. 10
image

基于LDM的图像修复

在图像修复问题上，分析了不同下采样因子对训练速度和样本生成质量进行了定量分析。

  • LDM-4比像素空间LDM-1模型速度上有2.7倍的提升，同时\(FID\downarrow\)指标上有1.6倍的提升

Tab. 6
image

Tab. 7
image

引用数据集

数据集 简介
FFHQ 从Flickr上收集的高质量人脸图片，70000张，1024x1024分辨率，PNG格式。https://github.com/NVlabs/ffhq-dataset
CelebA 网上收集的大规模人脸属性数据集，约1万个名人ID，20万张图片，每张图片都标注了5个关键点及40种二值属性。https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
CelebA-HQ 从CelebA数据集中挑选出3万张分辨率达到1024x1024图片，https://huggingface.co/datasets/huggan/CelebA-HQ
LSUN 场景理解数据集，包含10大场景和20种物体类别。https://www.yf.io/p/lsun
MS-COCO 一个目标检测、分割和语义注释的数据集 https://cocodataset.org/#home
LAION 图文对数据集 https://laion.ai/blog/laion-5b/
OpenImages 图像的标注信息：图像级标签、目标框、物体语义分割图、视觉语义关系、点级标签、多模态描述（语音、文本、图片对应）等 https://storage.googleapis.com/openimages/web/factsfigures_v7.html

参考文献

[1] Taming Transformers for High-Resolution Image Synthesis
[2] The Unreasonable Effectiveness of Deep Features as a Perceptual Metric