(ICLR-2024)SDXL：改进高分辨率图像合成的潜在扩散模型

SDXL：改进高分辨率图像合成的潜在扩散模型

Paper Title：SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis

Paper是Stability AI发表在ICLR 2024的工作

Paper地址

Code地址

Abstract

我们提出了 SDXL，一种用于文本到图像合成的潜在扩散模型。与之前版本的 Stable Diffusion 相比，SDXL 利用了三倍大的 UNet 主干：模型参数的增加主要是由于更多的注意力块和更大的交叉注意力上下文，因为 SDXL 使用第二的文本编码器。我们设计了多种新颖的条件方案并在多种纵横比上训练 SDXL。我们还引入了一个细化模型，该模型用于提高使用事后图像到图像技术由 SDXL 生成的样本的视觉保真度。
我们证明，与之前版本的 Stable Diffusion 相比，SDXL 的性能有了显著提高，并且取得了与黑盒最先进的图像生成器相媲美的结果。本着促进开放研究和促进大型模型训练和评估透明度的精神，我们提供代码和模型权重的访问权限。

1 Introduction

去年，深度生成模型在各个数据领域取得了巨大飞跃，例如自然语言 [50]、音频 [17] 和视觉媒体 [38、37、40、44、15、3、7]。在本报告中，我们重点关注后者，并推出了 SDXL，这是 Stable Diffusion 的大幅改进版本。Stable Diffusion 是一种潜在的文本到图像扩散模型 (DM)，它是一系列最新进展的基础，例如 3D 分类 [43]、可控图像编辑 [54]、图像个性化 [10]、合成数据增强 [48]、图形用户界面原型设计 [51] 等。值得注意的是，应用范围非常广泛，涵盖了从音乐生成 [9] 到从 fMRI 脑部扫描重建图像 [49] 等各种领域。

用户研究表明，SDXL 始终以显著优势超越所有先前版本的稳定扩散 (见图 1)。在本报告中，我们介绍了导致性能提升的设计选择，包括 i) 与先前的稳定扩散模型相比，UNet 主干大 3 倍 (第 2.1 节)，ii) 两种简单但有效的附加调节技术 (第 2.2 节)，不需要任何形式的额外监督，以及 iii) 一个单独的基于扩散的细化模型，该模型将噪声去噪过程 [28] 应用于 SDXL 生成的潜在数据，以提高其样本的视觉质量 (第 2.5 节)。

图 1：左图：比较 SDXL 与 Stable Diffusion 1.5 和 2.1 之间的用户偏好。虽然 SDXL 的表现已经明显优于 Stable Diffusion 1.5 和 2.1，但添加额外的细化阶段可以提高性能。右图：两阶段管道的可视化：我们使用 SDXL 生成大小为 128 × 128 的初始潜在数据。之后，我们使用专门的高分辨率细化模型，并使用相同的提示对第一步生成的潜在数据应用 SDEdit [28]。SDXL 和细化模型使用相同的自动编码器。

视觉媒体创作领域的一个主要问题是，虽然黑盒模型通常被认为是最先进的，但其架构的不透明性阻碍了对其性能进行如实地评估和验证。这种缺乏透明度的情况会阻碍可重复性，扼杀创新，并阻止社区在这些模型的基础上进一步推动科学和艺术的进步。此外，这些闭源策略使得以公正客观的方式评估这些模型的偏差和局限性变得具有挑战性，而这对于负责任和合乎道德的部署至关重要。借助 SDXL，我们发布了一个开放模型，该模型的性能可与黑盒图像生成模型相媲美（见图 10 和图 11）。

2 Improving Stable Diffusion

在本节中，我们将介绍对稳定扩散架构的改进。这些改进是模块化的，可以单独使用或一起使用来扩展任何模型。虽然以下策略是作为潜在扩散模型 (LDM) [38] 的扩展实现的，但其中大多数也适用于像素空间对应模型。

2.1 Architecture & Scale

---

从 Ho 等人 [14] 和 Song 等人 [47] 的开创性研究开始，这些研究证明了 DM 是用于图像合成的强大生成模型，卷积 UNet [39] 架构一直是基于扩散的图像合成的主导架构。然而，随着基础 DM [40, 37, 38] 的发展，底层架构也在不断发展：从添加自注意力和改进的上采样层 [5]，到用于文本到图像合成的交叉注意力 [38]，再到纯粹基于 Transformer 的架构 [33]。

我们顺应这一趋势，并按照 Hoogeboom 等人 [16] 的做法，将大部分 Transformer 计算转移到 UNet 中的较低级别特征上。具体来说，与原始的 Stable Diffusion 架构相比，我们在 UNet 中使用了异构分布的 Transformer 块：出于效率原因，我们省略了最高特征级别的 Transformer 块，在较低级别使用 2 和 10 个块，并完全删除 UNet 中的最低级别（8 倍下采样）——请参阅表 1，了解 Stable Diffusion 1.x & 2.x 与 SDXL 架构之间的比较。
我们选择了一种更强大的预训练文本编码器来进行文本调节。具体来说，我们将 OpenCLIP ViT-bigG [19] 与 CLIP ViT-L [34] 结合使用，其中我们沿通道轴连接倒数第二个文本编码器输出 [1]。除了使用交叉注意层来根据文本输入调节模型外，我们还遵循 [30]，并根据 OpenCLIP 模型中的池化文本嵌入来调节模型。这些变化导致 UNet 中的模型大小为 2.6B 参数，见表 1。文本编码器的总参数大小为 817M。

2.2 Micro-Conditioning

---

根据图像大小调节模型 LDM 范式 [38] 的一个臭名昭著的缺点是，由于其两阶段架构，训练模型需要最小图像大小。解决这个问题的两种主要方法是丢弃所有低于某个最小分辨率的训练图像（例如，Stable Diffusion 1.4/1.5 丢弃所有尺寸低于 512 像素的图像），或者将太小的图像放大。但是，根据所需的图像分辨率，前一种方法可能会导致大量训练数据被丢弃，这可能会导致性能下降并损害泛化。我们在图 2 中可视化了 SDXL 预训练数据集的这种影响。对于这个特定的数据选择，丢弃所有低于我们预训练分辨率 2562 像素的样本将导致丢弃 39% 的数据。另一方面，第二种方法通常会引入放大伪影，这些伪影可能会泄漏到最终的模型输出中，例如导致样本模糊。

图 2：我们预训练数据集的高度与宽度分布。如果没有建议的尺寸条件，39% 的数据将由于边缘长度小于 256 像素而被丢弃，如虚线黑线所示。每个可视化单元中的颜色强度与样本数量成正比。

相反，我们建议在原始图像分辨率的条件下训练 UNet 模型，这在训练期间是容易获得的。具体来说，我们提供图像的原始高度和宽度（即，在任何缩放之前）作为对模型的额外条件 ${\mathbf{c}}_{\text{size }}=\left(h_{\text{original }},w_{\text{original }}\right)$。每个组件分别使用傅里叶特征编码进行嵌入，这些嵌入被连接成一个向量，通过将其添加到时间步嵌入来输入模型 [5]。

在推理时，用户可以通过此尺寸调节来设置图像所需的表观分辨率。显然（见图 3），该模型已经学会将调节$c_{size}$与依赖于分辨率的图像特征相关联，可以利用这些特征来修改与给定提示相对应的输出的外观。请注意，对于图 3 中所示的可视化，我们可视化了由 512 × 512 模型生成的样本（详情见第 2.5 节），因为在我们用于最终 SDXL 模型的后续多纵横比（比例）微调之后，尺寸调节的效果不太明显。

图 3：不同尺寸条件的效果：我们从 SDXL 中用相同的随机种子绘制 4 个样本，并根据每列上方显示的内容变化尺寸条件。图像质量在条件为较大图像尺寸时明显提高。样本来自 $512^2$ 模型，参见第 2.5 节。注意：对于此可视化，我们使用 $512\times 512$ 像素的基础模型（参见第 2.5 节），因为在 $1024\times 1024$ 微调之前，尺寸条件的效果更清晰。建议放大观看。

我们通过在空间大小为$512^2$的类条件 ImageNet [4] 上训练和评估三个 LDM，定量评估了这种简单但有效的条件技术的效果：对于第一个模型（仅限 CIN-512），我们丢弃所有至少有一条边小于 512 像素的训练示例，从而导致训练数据集只有 70k 张图像。对于 CIN-nocond，我们使用所有训练示例，但不进行大小条件处理。这种额外的条件处理仅用于 CIN-size-cond。训练后，我们为每个模型生成具有 50 个 DDIM 步骤 [46] 和（无分类器）指导尺度为 5 [13] 的 5k 个样本，并计算 IS [42] 和 FID [12]（针对完整验证集）。对于 CIN-size-cond，我们生成始终以 csize =（512, 512）为条件的样本。
表 2 总结了结果并验证了 CIN-size-cond 在两个指标上都比基线模型有所改进。我们认为 CIN-512-only 的性能下降是由于在小型训练数据集上过度拟合导致的泛化效果不佳，而 CIN-nocond 样本分布中的模糊样本模式的影响导致 FID 分数降低。请注意，虽然我们发现这些经典的定量分数不适合评估基础（文本到图像）DM 的性能 [40, 37, 38]（参见附录 F），但它们在 ImageNet 上仍然是合理的指标，因为 FID 和 IS 的神经主干已经在 ImageNet 本身上进行了训练。

根据裁剪参数对模型进行条件调节 图 4 的前两行说明了以前的 SD 模型的典型失败模式：合成的对象可以被裁剪，例如 SD 1-5 和 SD 2-1 的左侧示例中被切掉的猫头。这种行为的一个直观解释是在模型训练期间使用随机裁剪：由于在 PyTorch [32] 等深度学习框架中整理一批数据需要相同大小的张量，因此典型的处理流程是 (i) 调整图像大小，使最短尺寸与所需的目标尺寸匹配，然后 (ii) 沿长轴随机裁剪图像。虽然随机裁剪是一种自然的数据增强形式，但它可能会泄漏到生成的样本中，导致上面显示的恶意影响。

图 4：SDXL 的输出与 Stable Diffusion 先前版本的比较。对于每个提示，我们展示了 DDIM 采样器 [46] 和 cfg-scale 8.0 [13] 的 50 步的相应模型的 3 个随机样本。

为了解决这个问题，我们提出了另一种简单但有效的调节方法：在数据加载期间，我们均匀地采样裁剪坐标$c_{\text{top }}$和$c_{\text{left }}$（分别指定沿高度和宽度轴从左上角裁剪的像素数量的整数），并通过傅里叶特征嵌入将它们作为调节参数输入模型，类似于上面描述的尺寸调节。然后使用连接的嵌入 $c_{crop}$ 作为附加调节参数。我们强调，这种技术不仅限于 LDM，可以用于任何 DM。请注意，裁剪和尺寸调节可以很容易地组合在一起。在这种情况下，我们沿通道维度连接特征嵌入，然后将其添加到 UNet 中的时间步嵌入中。算法 1 说明了如果应用这种组合，我们如何在训练期间对$c_{crop}$和$c_{size}$进行采样。

鉴于我们经验中大型数据集平均而言是以对象为中心的，我们在推理过程中设置 $\left(c_{\text{top }},c_{\text{left }}\right)=(0,0)$，从而从训练模型中获得以对象为中心的样本。参见图 5 进行说明：通过调整 $\left(c_{\text{top }},c_{\text{left }}\right)$，我们可以成功模拟推理过程中裁剪的程度。这是一种条件增强形式，已经在各种自回归 [20] 模型中使用，最近也在扩散模型中使用 [21]。

虽然数据分桶 [31] 等其他方法成功解决了相同的任务，但我们仍然从裁剪引起的数据增强中受益，同时确保其不会泄露到生成过程中——我们实际上利用它来获得对图像合成过程的更多控制。此外，它易于实现，并且可以在训练过程中以在线方式应用，而无需额外的数据预处理。

图 5：如第 2.2 节所述，改变裁剪条件。参见图 4 和图 14，其中 SD 1.5 和 SD 2.1 中的样本没有提供对此参数的明确控制，因此引入了裁剪伪影。来自$512^2$模型的样本，参见第 2.5 节。

2.3 Multi-Aspect Training

---

现实世界的数据集包含各种尺寸和纵横比的图像（参见图 2）。虽然文本到图像模型的常见输出分辨率是 $512\times 512$ 或 $1024\times 1024$ 像素的方形图像，但考虑到广泛分布和使用的横向（例如，16:9）或纵向格式的屏幕，我们认为这是一个相当不自然的选择。

受到此启发，我们对模型进行微调，使其能够同时处理多个纵横比：我们遵循常见做法 [31]，将数据划分为不同纵横比的桶，其中我们保持像素数尽可能接近 $1024^2$ 像素，并相应地按 64 的倍数调整高度和宽度。训练中使用的所有纵横比的完整列表见附录 I。优化过程中，一个训练批次由同一个桶中的图像组成，并在每个训练步骤交替不同的桶大小。此外，模型接收桶大小（或目标大小）作为条件，表示为整数元组 ${\mathbf{c}}_{\mathrm{ar}}=\left(h_{\mathrm{tg}},w_{\mathrm{tgt}}\right)$，这些元组嵌入到傅里叶空间中，类似于上述尺寸和裁剪条件。

实际上，我们在固定纵横比和分辨率下预训练模型后，将多纵横比训练应用于微调阶段，并结合第 2.2 节中介绍的条件技术，通过在通道轴上连接的方式进行。附录 J 的图 16 提供了这一操作的 Python 代码。请注意，裁剪条件和多纵横比训练是互补的操作，裁剪条件仅在桶边界内工作（通常为 64 像素）。然而，为了实现方便，我们选择保留这一控制参数用于多纵横比模型。

2.4 Improved Autoencoder

---

稳定扩散是一种 LDM，在自动编码器的预训练、学习（和固定）潜在空间中运行。虽然大部分语义组合是由 LDM [38] 完成的，但我们可以通过改进自动编码器来改善生成图像中的局部高频细节。为此，我们以更大的批处理大小（256 vs 9）训练用于原始稳定扩散的相同自动编码器架构，并使用指数移动平均线跟踪权重。生成的自动编码器在所有评估的重建指标中均优于原始模型，见表 3。我们在所有实验中都使用此自动编码器。

2.5 Putting Everything Together

---

我们在多阶段过程中训练最终模型 SDXL。SDXL 使用第 2.4 节中的自动编码器和具有 1000 步的离散时间扩散调度 [14,45]。首先，我们在内部数据集上预训练一个基础模型（见表 1），其高度和宽度分布在图 2 中可视化，进行 600000 次优化步骤，分辨率为 $256\times 256$ 像素，批量大小为 2048，使用第 2.2 节描述的尺寸和裁剪条件。我们继续在 $512\times 512$ 像素图像上进行另外 200000 次优化步骤，最后利用多纵横比训练（第 2.3 节）结合 0.05 的偏移噪声 [11,25] 水平，在 $\sim 1024\times 1024$ 像素面积的不同纵横比上训练模型（第 2.3 节，附录 I）。

细化阶段 实验证明，生成的模型有时会产生局部质量较低的样本，见图 6。为了提高样本质量，我们在相同潜在空间中训练一个专门处理高质量、高分辨率数据的独立 LDM，并在基础模型的样本上使用 SDEdit [28] 引入的加噪-去噪过程。我们遵循 [1] 并将这个细化模型专门用于前 200 个（离散）噪声尺度。在推理过程中，我们从基础 SDXL 中渲染潜在变量，并使用细化模型在潜在空间中直接扩散和去噪，使用相同的文本输入。我们注意到这一步是可选的，但它可以提高详细背景和人脸的样本质量，如图 6 和图 13 所示。

为了评估我们模型的性能（有无细化阶段），我们进行了用户研究，让用户从以下四个模型中选择他们最喜欢的生成：SDXL、SDXL（带细化器）、Stable Diffusion 1.5 和 Stable Diffusion 2.1。结果表明，带有细化阶段的 SDXL 是评分最高的选择，并显著优于 Stable Diffusion 1.5 和 2.1（获胜率：SDXL 带细化：48.44%，SDXL 基础：36.93%，Stable Diffusion 1.5：7.91%，Stable Diffusion 2.1：6.71%）。参见图 1，其中还提供了整个流程的概览。然而，当使用经典的性能指标如 FID 和 CLIP 分数时，SDXL 相对于以前方法的改进并未反映出来，如图 12 所示并在附录 F 中讨论的那样。这与 Kirstain 等人的发现 [23] 一致并进一步支持了这些发现。

图 6：SDXL 中的$1024^2$个样本（带放大图），未使用（左）和使用（右）所讨论的细化模型。提示：“纽约市被海洋淹没的史诗级远距离城市景观照片，以及雨林中杂草丛生的建筑物和丛林废墟，日落时分，电影镜头，细节丰富，8k，金色光线”。有关其他样本，请参见图 13。