【论文阅读笔记】Würstchen: AN EFFICIENT ARCHITECTURE FOR LARGE-SCALETEXT-TO-IMAGE DIFFUSION MODELS

这篇文章提出了一个高效的用于文本到图像生成模型架构，整体思路比较直白，在不损失图像生成质量的情况下，相比于现有T2I模型（SD1.4，SD2.1等）大大节约了成本。附录部分给了一些有趣的东西，比如FID的鲁棒性

整篇文章还有点疑惑，比如阶段B的训练，使用的模型；节省成本主要是在说C阶段？那A和B呢；256的潜在空间训练，不应该比SD在64的潜在空间训练更加成本高昂吗？看hf的权重大小，STAGE-A 296M，STAGE-B 3.4G， STAGE-C 3.97G好像和SD2.1的也差不多

更新1：附录D提到**“在我们的工作中，我们将阶段 C 视为模型的主要工作部分，当涉及到从文本中生成图像时”**。
更新2：一定要看附录D！

1. Würstchen是一种通过将计算成本高昂的文本条件阶段移动到高度压缩的潜在空间来训练文本条件模型的新框架。常见的方法使用单级压缩，而Würstchen引入了另一个引入更多压缩的阶段。

2. 框架包含负责压缩图像的阶段A和B，以及在低维潜在空间中学习文本条件部分的阶段C。

3. Würstchen实现了42倍的压缩因子，同时仍然忠实地重建图像。这使得阶段C的训练能够快速且计算成本低廉。

paper：https://arxiv.org/abs/2306.00637
code：https://github.com/dome272/wuerstchen
publication：ICLR 2024 Oral

摘要

开发了一种潜在的扩散技术，其中我们学习了一个用于指导扩散过程的详细但极其紧凑的语义图像表示。
与语言的潜在表示相比，图像的这种高度压缩的表示提供了更详细的指导，这显着减少了实现最先进结果的计算要求。(直观上理解就是文本的表示和图像的表示tokens一起作为条件引导图像生成)
基于用户偏好研究提高了文本条件图像生成的质量。
与稳定扩散 2.1 的 200,000 GPU 小时相比，我们方法的训练需求由 24,602 A100-GPU 小时组成。我们的方法还需要较少的训练数据来实现这些结果。（C阶段的时间，文本条件和图像条件的转换？）
此外，我们紧凑的潜在表示允许我们对最先进的 (SOTA) 扩散模型的通常成本和碳足迹进行两次推断，而不会损害最终性能。
与 SOTA 模型更广泛的比较中，我们的方法更有效，并且在图像质量方面表现出色。我们相信这项工作更强调性能和计算可访问性的优先级。

贡献

1. 我们提出了一种新的三阶段结构，用于强压缩比的文本图像合成，由两个条件潜在扩散阶段和一个潜在图像解码器组成。
2. 我们表明，通过在强压缩的潜在空间中使用文本条件扩散模型，我们可以在显着降低的训练成本和推理速度的情况下实现最先进的性能。
3. 我们提供了基于自动化指标和人工反馈的模型功效的综合实验验证。
4. 我们公开发布源代码和整套模型权重。

方法

训练

**Stage A：**最初，训练一个 VQGAN。
**StageB：**阶段 B 被训练为阶段 A 的潜在空间内的扩散模型。阶段 B 以文本嵌入和语义压缩器的输出为条件，它产生同一图像的强烈下采样的潜在表示。
**StageC：**阶段 C 在语义压缩器作为文本条件 LDM 的潜在表示上进行训练，有效地以 42 : 1 的压缩比运行。（只有这个阶段需要重头训练）

推理

1. 使用文本条件 LDM（阶段 C）以强压缩比生成潜在图像。(Text conditioning is applied on Stage C using CLIP-H )
2. 这种表示通过负责这种重建的次要模型转换为压缩较少的潜在空间（阶段 B）。
3. 对该中间分辨率中包含潜在图像的标记被解码以产生输出图像（阶段 A）。
   

实验

表 2 显示了与原始 SD 1.4 和 2.1 相比训练 W ̈urstchen 的计算成本。基于 4.1 节中的评估，可以看出所提出的将高分辨率图像投影与实际文本条件生成的解耦设置可以更多地利用，同时仍然保持相当或优于质量、保真度和对齐。与 SD 2.1 的 200,000 个 GPU 小时相比，阶段 C 是从头开始训练的最昂贵的阶段，只需要 24,602 个 GPU 小时，使其成为 8 倍的改进。此外，SD 1.4 和 2.1 处理的图像样本明显更多。后一个指标基于所有训练和微调的步骤总数，并与各自的批量大小相乘。即使考虑到用于训练阶段 B 的 11,000 个 GPU 小时和 318M 训练样本，W ̈urstchen 也比 SD 模型更有效地训练。此外，虽然需要同时使用阶段 A 和 B 进行采样来生成 VQGAN 潜在 ̄xq ，但总推理仍然比 SD 2.1 和 XL 快得多（见图 4）。

使用 PickScore (Kirstain et al., 2023) 评估 MS-COCO 和本地化叙事 (Pont-Tuset et al., 2020) 上的图像质量，通过两个不同的模型从同一字幕生成的图像。Ẅurstchen 优于所有大小相等的模型，尽管使用明显更高计算预算的稳定扩散模型。

结论

在这项工作中，提出了我们的文本条件图像生成模型 Ẅurstchen，该模型采用了三个阶段的过程，将文本条件图像生成与高分辨率空间解耦。该过程能够有效地训练大规模模型，大大减少计算需求，同时提供高保真图像。我们训练的模型实现了与使用更多计算资源训练的模型相当的性能，说明了这种方法的可行性，并提出了对更大模型参数的潜在有效可扩展性。我们希望我们的工作可以作为进一步研究生成 AI 的更可持续和计算效率更高的领域的起点，并为消费者硬件上的训练、微调和部署大规模模型开辟了更多可能性。

附录

附录A

一些额外可视化

附录B

FID评价指标：尽管在imagenet数据集上训练，但能够泛化到其他数据集；并给出了一些常见操作对FID的影响

附录C

关于用户研究指标的设置描述

附录D

解决了疑问：

1. 在我们的工作中，我们将阶段 C 视为模型的主要工作部分，当涉及到从文本中生成图像时。
2. 简要讨论阶段 B 和阶段 C 如何共享图像生成的工作量。通过这样做，我们证明了阶段 C 负责图像的内容，而阶段 B 充当细化模型，添加细节并提高分辨率，但最终不会以语义上有意义的方式更改图像。
   1. 为了研究，我们训练了一个小（3.9M 参数）解码器来从阶段 C 产生的潜伏期重建图像，并将重建与以阶段 C 为条件的阶段 B 的重建进行了比较。图 17, 18, 19 和 20 中的结果表明，阶段 C 生成的图像与阶段 B 和 C 组合生成的图像非常相似。
   2. 由此我们得出结论，阶段 C 是将文本转换为图像时的主要因素。这进一步得到了以下事实的支持：在替代训练机制上进行的简短实验表明，阶段 B 的文本调节并不能提高图像的质量，并且可以在未来几代中丢弃我们的模型。

附录E

关于网络架构的详细描述