SwiftBrush: One-Step Text-to-Image Diffusion Model with Variational Score Distillation

📄 arXiv: 2312.05239v7 📥 PDF

作者: Thuan Hoang Nguyen, Anh Tran

分类: cs.CV

发布日期: 2023-12-08 (更新: 2024-11-16)

备注: Accepted to CVPR 2024; Github: https://github.com/VinAIResearch/SwiftBrush

💡 一句话要点

SwiftBrush：一种基于变分分数蒸馏的单步文本到图像扩散模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱三：空间感知与语义 (Perception & Semantics)

关键词: 文本到图像生成 扩散模型 模型蒸馏 单步生成 变分分数蒸馏

📋 核心要点

  1. 文本到图像扩散模型生成高质量图像，但迭代采样速度慢，模型蒸馏是加速方法，但现有蒸馏方法需要大量图像数据且难以保持生成质量。
  2. SwiftBrush借鉴文本到3D合成的思路，利用专门的损失函数，将多步扩散模型蒸馏为单步模型，无需任何训练图像数据。
  3. SwiftBrush在COCO-30K上取得了16.67的FID和0.29的CLIP分数，性能与Stable Diffusion相当，甚至超越了现有蒸馏技术。

📝 摘要（中文）

文本到图像的扩散模型虽然能够从文本提示生成高分辨率和多样化的图像，但通常面临迭代采样过程缓慢的问题。模型蒸馏是加速这些模型最有效的方向之一。然而，以往的蒸馏方法在需要大量真实或教师模型合成的图像进行训练时，无法保持生成质量。为了解决这个限制，我们提出了一种名为$ extbf{SwiftBrush}$的全新无图像蒸馏方案。受到文本到3D合成的启发，该方法通过一个专门的损失函数，从2D文本到图像的扩散先验中获得与输入提示对齐的3D神经辐射场，而无需任何3D数据真值。我们的方法将相同的损失函数重新用于将预训练的多步文本到图像模型蒸馏到学生网络，该网络只需一步推理即可生成高保真图像。尽管其简单性，我们的模型是首批能够生成与Stable Diffusion质量相当的图像，且不依赖任何训练图像数据的单步文本到图像生成器之一。值得注意的是，SwiftBrush在COCO-30K基准测试上实现了$ extbf{16.67}$的FID分数和$ extbf{0.29}$的CLIP分数，取得了具有竞争力的结果，甚至大大超过了现有的最先进的蒸馏技术。

🔬 方法详解

问题定义：论文旨在解决文本到图像扩散模型推理速度慢的问题。现有的模型蒸馏方法通常需要大量的图像数据进行训练，无论是真实数据还是由教师模型生成的合成数据，这增加了训练成本和复杂度，并且难以保证蒸馏后的模型能够保持原有的生成质量。

核心思路：SwiftBrush的核心思路是借鉴文本到3D合成中的技术，利用变分分数蒸馏（Variational Score Distillation, VSD）的思想，将预训练的多步文本到图像扩散模型蒸馏成一个单步模型。关键在于，它不需要任何真实的图像数据进行训练，而是直接利用预训练模型的梯度信息进行学习。

技术框架：SwiftBrush的整体框架包括以下几个主要步骤：1) 使用预训练的多步文本到图像扩散模型作为教师模型。2) 定义一个单步学生模型，通常是一个生成对抗网络（GAN）或类似的结构。3) 使用变分分数蒸馏损失函数，该损失函数鼓励学生模型的输出与教师模型在特定噪声水平下的分数匹配。4) 通过优化学生模型，使其能够一步生成与文本提示对齐的高质量图像。

关键创新：SwiftBrush最重要的技术创新在于其无图像蒸馏方案。与传统的蒸馏方法不同，它不需要任何真实的图像数据或教师模型生成的图像数据进行训练。这大大降低了训练成本和复杂度，并且避免了合成数据可能引入的偏差。此外，它将文本到3D合成中的技术成功地应用于文本到图像的蒸馏任务中。

关键设计：SwiftBrush的关键设计包括：1) 变分分数蒸馏损失函数的设计，该损失函数需要能够有效地衡量学生模型和教师模型之间的差异。2) 学生模型的网络结构设计，需要能够一步生成高质量的图像。3) 噪声水平的选择，不同的噪声水平可能会影响蒸馏的效果。具体来说，论文可能采用了特定的噪声调度策略，并对损失函数中的权重进行了精细调整。

📊 实验亮点

SwiftBrush在COCO-30K数据集上取得了显著的成果，FID分数为16.67，CLIP分数为0.29。该结果与Stable Diffusion等先进的多步扩散模型相比具有竞争力，并且优于现有的单步图像生成模型。值得注意的是，SwiftBrush在没有任何图像数据训练的情况下实现了这一性能。

🎯 应用场景

SwiftBrush具有广泛的应用前景，包括快速图像生成、实时内容创作、个性化图像定制等。该技术可以应用于游戏开发、广告设计、虚拟现实等领域，显著提升图像生成效率，降低计算成本。未来，该技术有望进一步发展，实现更高质量、更可控的图像生成。

📄 摘要（原文）

Despite their ability to generate high-resolution and diverse images from text prompts, text-to-image diffusion models often suffer from slow iterative sampling processes. Model distillation is one of the most effective directions to accelerate these models. However, previous distillation methods fail to retain the generation quality while requiring a significant amount of images for training, either from real data or synthetically generated by the teacher model. In response to this limitation, we present a novel image-free distillation scheme named $\textbf{SwiftBrush}$. Drawing inspiration from text-to-3D synthesis, in which a 3D neural radiance field that aligns with the input prompt can be obtained from a 2D text-to-image diffusion prior via a specialized loss without the use of any 3D data ground-truth, our approach re-purposes that same loss for distilling a pretrained multi-step text-to-image model to a student network that can generate high-fidelity images with just a single inference step. In spite of its simplicity, our model stands as one of the first one-step text-to-image generators that can produce images of comparable quality to Stable Diffusion without reliance on any training image data. Remarkably, SwiftBrush achieves an FID score of $\textbf{16.67}$ and a CLIP score of $\textbf{0.29}$ on the COCO-30K benchmark, achieving competitive results or even substantially surpassing existing state-of-the-art distillation techniques.