[论文解读] Wuerstchen: An Efficient Architecture for Large-Scale Text-to-Image Diffusion Models

Pablo Pernias, Dominic Rampas+3 more|arXiv (Cornell University)|Jun 1, 2023

Generative Adversarial Networks and Image Synthesis被引用 4

一句话总结

Würstchen 提出了一种三阶段文本到图像扩散架构，通过使用高度压缩的 42:1 潜在空间进行条件控制，实现了 SOTA 图像生成质量，同时将训练成本降低 8 倍（24,602 个 A100-GPU 小时 vs. Stable Diffusion 2.1 的 200,000 个），实现了更快的推理速度和更低的碳足迹，且未牺牲保真度或对齐性。

ABSTRACT

We introduce Würstchen, a novel architecture for text-to-image synthesis that combines competitive performance with unprecedented cost-effectiveness for large-scale text-to-image diffusion models. A key contribution of our work is to develop a latent diffusion technique in which we learn a detailed but extremely compact semantic image representation used to guide the diffusion process. This highly compressed representation of an image provides much more detailed guidance compared to latent representations of language and this significantly reduces the computational requirements to achieve state-of-the-art results. Our approach also improves the quality of text-conditioned image generation based on our user preference study. The training requirements of our approach consists of 24,602 A100-GPU hours - compared to Stable Diffusion 2.1's 200,000 GPU hours. Our approach also requires less training data to achieve these results. Furthermore, our compact latent representations allows us to perform inference over twice as fast, slashing the usual costs and carbon footprint of a state-of-the-art (SOTA) diffusion model significantly, without compromising the end performance. In a broader comparison against SOTA models our approach is substantially more efficient and compares favorably in terms of image quality. We believe that this work motivates more emphasis on the prioritization of both performance and computational accessibility.

研究动机与目标

• 解决大规模文本到图像扩散模型训练与推理的高计算成本问题。
• 在保持或提升图像质量与对齐性的同时，降低训练与推理成本。
• 实现在消费级硬件上更高效、可扩展且可访问的高保真生成模型部署。
• 证明极端压缩的潜在表示可保留语义细节，并指导高保真图像生成。
• 优先考虑计算效率而不牺牲性能，推动可持续生成式 AI 的发展。

提出的方法

• 提出三阶段架构：阶段 C 使用文本条件潜在扩散模型（LDM）生成高度压缩的 42:1 潜在表示。
• 阶段 B 利用该压缩潜在表示，通过基于 VQGAN 的 VQ-VAE 条件化第二个 LDM，重建更高维的 4:1 压缩潜在空间。
• 阶段 A 使用 VQGAN 解码器将最终的潜在表示解码为全分辨率图像。
• 按逆序训练：首先训练 VQGAN（阶段 A），然后训练阶段 B 的 LDM，最后训练阶段 C 的 LDM。
• 使用语义压缩器（高比率编码器）将图像投影到紧凑潜在空间，保留详细的语义信息。
• 利用无分类器指导（CFG）和基于 CLIP 的文本编码，实现文本提示对图像生成的条件控制。

实验结果

研究问题

• RQ1 高度压缩的潜在空间（42:1）是否能保留足够的语义细节，以在扩散模型中指导高保真图像生成？
• RQ2 在如此紧凑的潜在空间上训练文本条件扩散模型，是否能降低训练成本，同时保持或提升图像质量，相比标准 LDM？
• RQ3 将高分辨率生成与条件建模解耦为三阶段，是否能显著降低推理时间与计算开销？
• RQ4 在图像保真度、对齐性与用户偏好方面，该模型与 SOTA 模型（如 Stable Diffusion 2.1）相比表现如何？
• RQ5 该方法在多大程度上可通过降低硬件与能效需求，使高质量文本到图像生成更具可及性？

主要发现

• Würstchen 仅使用 24,602 个 A100-GPU 小时的训练时间，便实现了 SOTA 性能，相比 Stable Diffusion 2.1 的 200,000 个 GPU 小时，降低了 8 倍。
• 在 MS-COCO 上 FID 得分为 2.53，在 Parti-prompts 上为 2.61，PickScore 在 MS-COCO 上为 1.48，在 Parti-prompts 上为 1.52，与 SD 2.1 相比达到相当或更优水平。
• 人工偏好研究表明，用户对 Würstchen 的偏好显著更高：在 Parti-prompts 上 62% 的用户偏好 Würstchen，在 MS-COCO 上为 54%（仅考虑高参与度用户）。
• 由于潜在空间紧凑，Würstchen 的推理速度比 SD 2.1 快两倍以上，显著降低了延迟与碳足迹。
• 该模型所需训练数据更少，且在对齐性与视觉质量方面表现更优，经自动化指标与人工评估均得到验证。
• 整个流程（包括训练与推理代码、模型权重）已公开发布于 GitHub，以支持可复现性与社区使用。