( stable diffusion的开篇之作)

Abstract

    传统扩散模型计算成本高昂,直接在像素空间操作,训练与推理过程都缓慢昂贵。为解决这个问题,本文将扩散过程从像素空间转移到潜在空间,通过引入一个预训练的自编码器的办法。与之前在潜在空间的工作相比,这个成功降低复杂度并保留了细节,并引入交叉注意力机制,允许将任意类型的条件输入编码成一个中间表示从而精确地指导生成过程。

    LDMs 在多种图像生成任务上都取得了顶尖或极具竞争力的表现,包括图像修复(inpainting)、类别条件生成、文生图、无条件生成以及超分辨率等,并显著降低了计算需求。

Introduction 

    图像合成技术发展迅速,但其计算成本,特别是在高分辨率、复杂场景的合成上,极为高昂。

    如今主流的生成模型有:

1 AR:参数大,推理速度慢

2 GAN :训练不稳定(容易出现梯度消失或爆炸),mode-collapse(难以学习  数据分布的全部模式)

3 DM :性能最优,潜力最大,但计算成本高昂。它投入大量计算能力去建模数据中对人类感知而言几乎不可察觉的高频细节 ,并直接在像素空间中进行训练和推理,导致了巨大的计算开销。因此,迫切需要一种方法来降低训练和采样的计算复杂度,实现democratizing。

    作者通过分析像素空间中DMs的率失真权衡 (rate-distortion trade-off),发现其学习过程可以分为两个阶段:

1 感知压缩阶段 (Perceptual Compression):早期阶段,模型主要学习去除高频噪声和细节,类似于一种有损压缩。此阶段对图像的语义内容学习较少。

2 语义压缩阶段 (Semantic Compression):后期阶段,模型开始学习数据中的语义和概念,即图像的真正内容。

    既然DMs本身的学习过程就隐含了一个压缩阶段,那么可以主动地将这个过程分离出来。本文提出,首先使用一个独立的模型将图像压缩到一个在感知上等价 (perceptually equivalent) 但计算上更高效 (computationally more suitable) 的空间,即潜在空间 (Latent Space)。然后,让扩散模型专注于在这个低维的潜在空间中学习语义。

    LDM分为两个训练阶段:训练自编码器(通用)和潜在扩散模型。前者能将高维图像编码为低维潜在表示,并能从该表示中高质量地解码回原始图像;后者则在这个低维潜在空间(保留了丰富的空间信息)中训练扩散模型。并设计了将Transformer与U-Net骨干网络连接的架构。通过交叉注意力机制,可以将任何基于token的条件信息集成到DM的去噪过程。

    贡献总结:LDM可高效处理高分辨率图像;性能强训练开销少;采用解耦的两阶段训练;U-Net可作为一个整体以卷积方式 用于任意大小的输入;交叉注意力机制实现了文生图等多种跨模态生成任务;开源。

Related work

Generative Models for Image Synthesis:

    GAN:高效采样,生成具有良好(perceptual quality)的高分辨率图像。但优化困难,模式覆盖不足难以捕捉到完整的数据分布。

    Likelihood-based Methods:

    VAEs和 Flow-based Models:同样能高效合成高分辨率图像,但生成的样本质量通常不如GANs。

    ARMs:能很好地拟合数据分布,但架构(通常是Transformer)计算成本高,逐像素进行速度极慢,只适用于低分辨率图像。

    似然模型核心问题是,直接在像素空间进行最大似然训练,会使得模型将不成比例的计算能力浪费在建模那些人类几乎无法察觉的高频细节上。

    而扩散模型取得了很高的成果,契合图像数据的inductive biases,它使用加权重的目标函数,允许在图像质量和压缩能力之间进行权衡。但速度慢,训练成本高。

Two-Stage Image Synthesis:

    此前,有基于自回归模型的两阶段方法。VQ-VAE使用自回归模型在一个离散的 潜在空间上学习一个表达能力强的先验分布;VQGAN对VQ-VAE进行改进。它的第一阶段(自编码器)引入了对抗性损失和感知损失,从而允许自回归Transformer扩展到更大的图像。

    但这些方法为实现自回归需要采用高压缩率,但为了保真度还要降低压缩率(计算成本升高)。

    LDM则避免了这种权衡。LDM的骨干网络是卷积的(即U-Net),它对更高维的潜在空间具有更好的尺度适应性。这意味着我们不必为了计算效率而过度压缩图像。因此我们可以自由地选择一个最佳的压缩水平。

Method

    现有的扩散模型虽然可以通过损失加权来减少对图像微小细节的关注,但因为它们依然在像素空间进行运算,所以计算成本依旧高昂。

    我们的解决方案是,先用一个自编码器把图像压缩到一个更小但信息完整的潜在空间,然后再让扩散模型在这个低维空间里学习生成。这样就提高了计算效率。

    我们利用了扩散模型(特别是其 U-Net 架构)天生善于处理具有空间结构的数据这一归纳偏置,获得一个通用的训练好的自编码器。例如,可以在同一个潜在空间上训练一个无条件的DM,一个文生图的DM,一个类别条件的DM等。

3.1 Perceptual Image Compression:

    我们使用一个结合了“感知损失”和“对抗性损失”的自编码器,这种训练方式能生成细节丰富、不模糊的重建图像,效果远超传统的L1/L2损失。该自编码器将 HxW 的图像压缩成一个 hxw 的潜在张量,压缩因子为 f,这个潜在张量保留了空间结构。

    为了让潜在空间更容易学习(不加约束会high-variance),我们用了两种方法来约束它:一种是类似VAE的KL散度惩罚,另一种是类似VQGAN的向量量化(Vector Quantization)。

    因为后续的扩散模型 (DM) 是基于 U-Net 的,它天生就擅长处理具有二维空间结构的数据。潜在空间 z 保留了其二维结构,就不需要太大的f压缩值,就保留更多的图像信息,从而实现非常好的重建质量。而之前的工作第二阶段使用的是自回归模型,如 Transformer(将二维的潜在表示 z 展平成一个一维序列),忽略了 z 固有的二维空间结构。

3.2 Latent Diffusion Model

    标准扩散模型通过训练一个神经网络来预测施加在图像上的噪声,从而学会从纯噪声中逐步生成清晰图像。

    而我们不直接在像素上做扩散,而是在编码器压缩后的那个小而精的潜在空间里做,这样模型可以更专注于学习内容,并且计算成本大大降低。  

   

     核心区别在于:标准 DM 的输入是带噪的图像 xt,而我们的 LDM 的输入是带噪的潜在表示 zt。我们的模型 θ 是一个时间条件的 U-Net。

    z = E(x),对 z 添加噪声 ε,得到带噪的潜在表示 zt。使用训练好的 U-Net 模型 θ,从 t=T 到 t=1 逐步去噪,得到一个干净的潜在表示 z0。这个过程就是 p(z) 的采样。然后用解码器一次性将其“放大”回高质量图像。

3.3 Conditioning Mechanisms

    已有的条件扩散模型大多局限于简单的条件。而我们希望让扩散模型能按指令生成图片,比如根据文字描述。

    于是我们使用了交叉注意力机制,Query (Q):来自 U-Net 内部的图像特征。具体来说,是 U-Net 的第 i 个中间层特征 φi(zt) 经过一个线性变换 W(i)_Q  得到的。

 

      Key (K) 和 Value (V)都是来自条件 y 的编码,这个query发问后,与条件 y 的所有部分(Keys)进行匹配(计算相似度),匹配度高的部分,其对应的信息(Values)就会被更多地“注入”到图像特征中。 

    最终的目标函数。在训练过程中,U-Net θ 和条件编码器 τθ 的参数是一起优化的 (jointly optimized)。这个框架非常灵活,因为条件编码器 τθ 是可插拔的,可以根据y的类型选择。

Experiments

    实验结构可以分为以下几个主要部分:

  1. 感知压缩的权衡分析 (4.1):找到最佳的压缩率。
  2. 无条件图像生成 (4.2):在标准生成任务上展示模型能力。
  3. 条件图像生成 (4.3):展示模型的可控性和灵活性,这是 LDM 的一大亮点。
  4. 具体应用:超分辨率 (4.4):展示 LDM 在图像到图像转换任务上的应用。
  5. 具体应用:图像修复 (4.5):进一步展示其在复杂条件生成任务上的 SOTA 性能。

4.1 On Perceptual Compression Tradeoffs

目标: 探索第一阶段自编码器的最佳压缩因子 f。f 太小,计算优势不明显;f 太大,会损失过多信息,影响最终图像质量。目标是找到计算效率和生成质量的最佳平衡点。

    定义 LDM-f,其中 f 是空间下采样因子(如 f=4 表示将 256x256 图像压缩到 64x64 的潜空间)。特别地,LDM-1 就是在像素空间操作的标准扩散模型,作为基线。在固定的计算预算(一台 NVIDIA A100 GPU)和相同的训练步数下,训练不同 f 值的 LDM 模型。在 ImageNet 数据集上进行类别条件生成任务,比较它们的性能。 

训练效率(上图):低压缩率(LDM-1,2)时训练的特别慢,而高压缩率(LDM-32)丢失了很多重要信息,扩散模型效果一般。 

采样效率与质量(上图):LDM-4 和 LDM-8 不仅在质量上远超像素模型 LDM-1,而且采样速度也快得多。

综上:f=4 和 f=8 是最佳的压缩设置

4.2 Image Generation with Latent Diffusion  

目标: 在标准的无条件图像生成任务上,将 LDM 与当时的其他 SOTA 模型(如 GANs, VAEs, LSGM)进行比较。

在 CelebA-HQ (人脸), FFHQ (人脸), LSUN (场景) 等多个标准数据集上训练无条件的 LDM。衡量 FID 和Precision-and-Recall

得出结论:

在 CelebA-HQ 数据集上,LDM 取得了 5.11 的 FID,创造了新的 SOTA 记录;

 LDM 的性能优于 LSGM(二阶段联合训练);在 Precision 和 Recall 指标上,LDM 持续优于 GANs。

4.3 Conditional Latent Diffusion

4.3.1 Transformer Encoders for LDMs

目标: 证明通过本文提出的交叉注意力机制,LDM可以被塑造成一个通用的可控生成器,能处理文本、语义布局等多种复杂的条件输入。

将条件信息转化为中间表示,示通过交叉注意力机制注入到 U-Net 的多个层中。 

其性能与当时 SOTA 的大型自回归模型(Make-A-Scene)和扩散模型(GLIDE)不相上下,但参数量显著更少。 

LDM 计算需求和参数量都远低于 ADM

4.3.2 Convolutional Sampling Beyond 256^2

     目标: 展示 LDM 的一个独特优势:对于空间对齐的条件任务,可以生成远大于训练分辨率的图像。 

    一个在 256x256 图像上训练的模型,可以成功生成 512x1024 甚至更大的、内容一致的、高质量的连贯图像。(因为 U-Net 主要是卷积结构,可以以“全卷积”的方式处理更大的潜在空间输入)

    这是一个巨大的优势,意味着不需要专门为高分辨率训练模型。

4.4. Super-Resolution with Latent Diffusion

目标: 将 LDM 应用于实际的图像到图像转换任务——超分辨率,并评估其性能。

它将低分辨率图像作为条件,通过拼接方式输入到 U-Net。训练一个 4x  超分模型 (LDM-SR),并与 SOTA 模型 SR3 进行比较。直接进行用户研究 ,因为传统的 PSNR/SSIM 指标不能很好地反映人的感知质量。

还训练了一个更通用的模型 LDM-BSR

LDM-SR 在 FID 指标和主观感知质量上均优于 SR3。

4.5. Inpainting with Latent Diffusion

目标: 在图像修复(特别是大面积缺失)这一更具挑战性的任务上,评估 LDM 的性能。 

将带掩码的图像和掩码本身作为条件输入

LDM 在 FID 指标上优于 LaMa。作者提到,LDM 的 LPIPS 分数稍高可能是因为它能生成更多样化的修复结果,而不是像 LaMa 那样倾向于生成单一的“平均”结果。

Limitations & Societal Impact

局限性:LDM (序列化采样)的生成速度虽然有改进,但仍比 GAN (单步前向传播)慢,这在需要快速生成大量样本的场景下是个劣势。且由于初始的图像压缩是有损的,LDM 可能不适用于那些对像素级精度有极端要求的任务。

正面影响:LDM 通过降低训练和推理成本,有潜力促进技术的普及。

负面影响:虚假信息伪造数据;数据泄露;放大偏见。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐