ICLR 2024 spotlight WURSTCHEN - 一个高效的大规模文本到图像扩散模型的架构

ICLR 2024 spotlight WURSTCHEN - 一个高效的大规模文本到图像扩散模型的架构¶

1. 论文信息¶

2. 引言¶

这篇论文探讨了一个当前在图像合成领域中的核心问题：如何在保持生成图像质量的同时，减少计算资源的消耗。目前的先进扩散模型，如Stable Diffusion和DALL·E 2，虽然能够生成接近真实的高质量图像，但它们对计算资源的需求非常高，这在一定程度上限制了它们的应用范围和可达性。例如，Stable Diffusion 1.4版本的训练就耗费了150,000 GPU小时。

虽然存在一些更经济的文本到图像的模型，它们在图像分辨率和美学特征方面的表现通常不如高资源消耗的模型。这是因为，提高图像的分辨率会增加视觉复杂性和计算成本，从而使图像合成变得更加昂贵和数据密集。

为了解决这个问题，论文提出了一种名为“Würstchen”的新颖三阶段架构。这种架构通过在非常低维的潜在空间上训练扩散模型，实现了高压缩比（42:1），从而大幅降低了计算需求，同时还保持了与竞争模型相媲美的性能。这个非常低维的潜在空间用于条件化第二个生成潜在模型，有效地帮助它导航一个更高维的潜在空间，该空间由一个具有4:1压缩比的VQGAN操作。简而言之，这种方法通过三个不同阶段的图像合成（文本条件LDM生成的低维潜在图像表示，然后是另一个LDM生成的更高维度的潜在图像，最后由VQGAN解码器解码以产生全分辨率输出图像）来实现。

训练过程与推理过程相反，先从Stage A的VQGAN开始，以创建一个紧凑的潜在空间，然后是Stage B的潜在扩散过程，最后是Stage C的强压缩潜在空间的文本条件LDM训练。这一逆向训练策略使得在大幅降低计算资源需求的同时，还能保持高质量图像的生成。

这篇论文的贡献在于提出了一个能够显著降低计算需求而不牺牲模型性能的新架构，这对于文本到图像合成领域来说是一个重大进步。通过使用文本条件扩散模型在一个强压缩的潜在空间中，作者证明了他们的模型在大幅降低训练成本和提高推理速度的同时，仍能达到最先进的模型性能。此外，论文还通过自动化指标和人类反馈提供了模型效能的全面实验验证，并公开了源代码和所有模型权重，这为未来在高效、高质量生成模型领域的研究开辟了新的道路。

3. 方法¶

在这篇文章中，Stage A和Stage B是减少计算负担的重要环节，通过将数据压缩成较小的表示来实现。让我们具体来看每个阶段：

Stage A¶

在Stage A中，我们采用了VQGAN模型来将原始图像（尺寸为 $ 1024 \times 1024 \times 3 $）映射到一个离散的潜在表示，这个表示的分辨率为 $ 256 \times 256 $，从而实现 $ 4:1 $ 的空间压缩。这是通过学习的编码器函数 $ f_{\Theta} $ 完成的，输出的离散编码为 $ X_q $。编码器和解码器（$ f_{\Theta}^{-1} $）被训练为可以从这些压缩的潜在表示中重建原始图像。

Stage B¶

在Stage B，我们使用一个名为“语义压缩器”的编码器网络来创建更强压缩率的潜在表示（尺寸为 $ 16 \times 24 \times 24 $），这些表示被用作扩散过程的引导。为此，我们将图像从 $ 1024 \times 1024 $ 调整大小为 $ 786 \times 786 $，以充分利用“语义压缩器”的参数，然后进一步压缩。Stage B的模型被训练为在Stage A的未量化潜在空间中重建图像，同时考虑到由语义压缩器提供的高度压缩的视觉嵌入和文本条件。Stage B的解码器接受被噪声化的潜在表示、语义压缩器的视觉嵌入、文本条件和时间步长作为输入，并且要预测原始的离散编码。这样，我们就能够实现总体 $ 42:1 $ 的空间压缩。

Stage C¶

最后的Stage C是文本条件的阶段。在这个阶段，我们使用16个没有下采样的ConvNeXt块，并在每个块之后通过交叉注意力应用文本和时间步条件。这里遵循的是标准的扩散过程，应用于经过微调的“语义压缩器”的潜在空间。噪声化过程遵循一个特定的公式，并使用余弦时间表来生成 $ \bar\alpha_t $ 值，并使用连续的时间步。扩散模型接收噪声化的嵌入、文本条件和时间步长作为输入，并返回对噪声的预测。

在这个阶段，目标函数被特别设计为稳定训练过程。通过这种方式，模型在初始时候趋向于预测接近于零的输出，这对于含有大量噪声的时间步来说，损失会比较小。损失函数是预测的噪声和实际噪声之间的均方误差，且加上了 $ p_2(t) $ 的权重，该权重随噪声水平的提高而增加，使得模型更注重高噪声水平下的性能。此外，为了实现无分类器引导（classifier-free guidance），文本条件有5%的时间会被随机丢弃并被替换为一个空标签。

通过上述三个阶段，文章的模型成功地在保持图像质量的同时，显著降低了图像合成过程的计算成本。

在论文的上下文中，“图像生成（采样）”和“模型决策”部分详细介绍了如何从随机噪声开始，经过三个阶段的逆向扩散过程，最终生成高分辨率的图像。

图像生成（采样）¶

图像生成（采样）的过程从Stage C开始。Stage C的任务是利用文本嵌入条件生成语义压缩器的潜在表示。这一过程以初始随机噪声 $X_{ {sc}, \tau_{C}} = {N}(0, {I})$ 开始，并通过以下操作进行 $\tau_{C}$ 步迭代：

该操作的结果 $\bar{X}_{sc}$ 具有 $ 16\times24\times24 $ 的形状，并被展平成 $576\times16$ 形状，然后与用来采样 $\bar{X}_{sc}$ 的相同文本嵌入一起，作为条件输入到Stage B。Stage B在 $4\times256\times256$ 的未量化VQGAN潜在空间操作，并且使用标准的潜在扩散模型（LDM）方案进行 $\tau_{B}$ 步迭代采样。最终，使用VQGAN的解码器 $f_\Theta^{-1} $ 将 $\tilde{X}$ 投影回像素空间，生成最终的图像 $\bar{X}$。

模型决策¶

理论上，任何特征提取器都可以用作语义压缩器的主干网络。然而，选择一个已经能够很好地表示各种图像特征的网络会更有利于模型的性能。此外，一个小型的语义压缩器能使得Stage B和C的训练更快。特征维数也至关重要：如果过小，可能无法捕获足够的图像细节，或者没有充分利用参数；如果过大，可能不必要地增加计算要求并延长训练时间。

基于以上原因，作者选择了预训练在ImageNet1k上的EfficientV2（S）作为语义压缩器的主干网络，因为它在高压缩、良好的特征表示和计算效率方面达到了良好的平衡。

在Stage C，作者偏离了标准的U-Net架构。由于图像已经通过42的因子压缩，进一步压缩将对模型质量有害。因此，作者采用了一系列不进行下采样的16个ConvNeXt块。每个块之后应用时间和文本条件。通过这种设计，Stage C能够更有效地进行图像的潜在空间扩散过程，并最终生成具有高分辨率和高质量的图像。

4. 实验¶

根据提供的实验结果，我们可以对“Würstchen”模型的图像质量进行评估。这些实验使用了PickScore来比较不同模型基于相同文本描述生成的图像。在PickScore中，较高的百分比表示相对于另一个模型，更多的图片被选择为质量更好的图像。

从表格中可以看出：

• Würstchen 在三种不同数据集上（MS-COCO-30k, Localized Narratives-COCO-5K, Parti-prompts）都得到了极高的PickScore百分比（96.5%, 96.6%, 和 98.6%）。这表明在文本到图像的生成任务中，相比于其他模型，Würstchen生成的图像更受偏好。

• 在与Baseline LDM（其训练成本约为25,000 GPU小时）的比较中，Würstchen模型表现出显著的优势，PickScore超过了96%。这意味着在几乎所有情况下，Würstchen生成的图像都被评价为更好。

• 对比SD（Stable Diffusion）1.4和SD 2.1，Würstchen的表现同样出色。尽管SD 1.4和SD 2.1的训练成本分别为150,000和200,000 GPU小时，远高于Würstchen的24,602 GPU小时，但Würstchen的图像选择率却更高。

• 相比于SD XL，Würstchen的PickScore是其近三倍，这表明Würstchen在图像质量上具有显著的优势。

这些结果表明Würstchen模型不仅在训练成本上更为经济，在图像质量上也优于或至少可与使用了更多计算资源的现有模型相媲美。这些结果证明了Würstchen模型在效率和质量上的双重成功，这对于图像合成领域来说是一个重要的进展。

在效率方面，论文中提出的Würstchen模型表现卓越。它在保持与先进模型相媲美的图像生成质量的同时，显著减少了所需的训练成本和步骤。具体而言，它以远低于Stable Diffusion模型的GPU小时数达成了相似的FID得分，并且在Inception Score上超越了所有比较模型，表明其生成图像的多样性和吸引力。此外，Würstchen使用的训练样本量较少，其碳足迹估计也低于其他高性能模型，这些都是其效率高的明证。因此，Würstchen模型在资源效率、步骤效率和环境效率方面都表现出色，是一个在计算资源有限的情况下进行高质量图像合成的理想选择。

5.讨论¶

这篇论文中提出的Würstchen模型具备明显的优势，特别是在效率方面。它利用了一个层次化的生成框架，通过压缩的潜在空间来降低了训练和生成图像所需的计算量。尽管参数数量与高性能的Baseline LDM相似，但在保持了相似的图像质量的同时，大幅减少了训练过程的GPU小时数和所需的采样步骤，这使得模型更为快速且节能。特别值得注意的是，Würstchen在图像质量评价指标如FID和IS上的表现与更耗资源的模型如Stable Diffusion相媲美，或甚至更好。这表明了它在生成多样化和高质量图像方面的能力，同时保持了资源的有效利用。

此外，模型减少了训练过程中的碳排放量，这是目前深度学习社区日益关注的一个重要议题。通过使用更少的训练样本并减少所需的能源消耗，Würstchen在环境可持续性方面也展现出了优势。

然而，这项工作可能也有其局限性。虽然减少了资源的使用，但Würstchen可能需要特定的结构优化和调参来达到如此高效的表现。这可能意味着模型在应对未见过的数据或更复杂任务时，可能需要进一步的调整或优化。此外，尽管它减少了训练资源的需求，但在部署和应用阶段仍然需要相对较新的硬件来执行这些复杂的模型。这些潜在的缺点，并不削弱Würstchen作为一个有前途的、高效的图像生成模型的地位，但它们提醒我们在实践中可能需要平衡不同的技术考量。

6. 结论¶

本文提出的Würstchen模型在图像合成领域取得了重要进展，实现了计算效率和图像质量之间的优秀平衡。该模型采用分层扩散过程，在压缩的潜在空间中运作，显著减少了训练时间和碳排放，同时在生成的图像质量上仍然与顶尖模型相媲美。通过定量指标和定性评估所反映出的令人鼓舞的结果表明，Würstchen模型为高效、高质量的图像生成树立了新的标杆。它不仅为未来研究提供了一个有力的示范，即创造更可持续、更普遍可及的人工智能驱动的图像合成工具，同时通过开源模型权重和代码库，进一步推动了高级生成模型的民主化，允许更广泛的实验和研究社区的进一步改进。尽管模型表现出色，但对于其在更广泛和更复杂任务中的应用范围，可能还需进一步探索和调整。