对 diffusion 结构的改进

Generation Model

RAPHAEL: Text-to-Image Generation via Large Mixture of Diffusion Paths

香港大学，商汤科技，23.5.29，https://raphael-painter.github.io/

1. Research Problem: The research problem addressed in this paper is the generation of highly artistic images that accurately portray text prompts. Existing models often fail to adequately preserve textual concepts within the generated images due to the reliance on a classic cross-attention mechanism for integrating text descriptions into visual representations.

2. Motivation: The motivation behind this research is to improve the precision of alignment between text and image in the generation process. This is aimed at enhancing the quality and aesthetic appeal of the generated images, and enabling them to accurately represent the associated text prompt.

3. Solution: The authors propose a text-conditional image diffusion model, RAPHAEL, which stacks tens of mixture-of-experts (MoE) layers, including both space-MoE and time-MoE layers. This configuration leads to billions of diffusion paths from the network input to the output. Each path can act as a "painter" responsible for rendering a particular concept to an image region at a specific timestep. The authors also propose an edge-supervised learning module to further enhance the image quality and aesthetic appeal of the generated images.

   

4. What's New: The novelty of this research lies in the introduction of RAPHAEL, which outperforms recent cutting-edge models in terms of both image quality and aesthetic appeal. RAPHAEL exhibits superior performance in switching images across diverse styles, such as Japanese comics, realism, cyberpunk, and ink illustration. It also establishes a new state-of-the-art with a zero-shot FID-30k score of 6.61 on the COCO dataset. Furthermore, RAPHAEL is capable of generating images with resolutions up to 4096 × 6144 with rich image contents and details, when combined with a tailor-made SR-GAN model.

Mixture of experts？

一种集成学习(Ensemble Learning) 技术，一个系统中包含多个分开的网络，每个网络去处理全部训练样本的一个子集。这种方式可以看做是把多层网络进行了模块化的转换。

假设我们已经知道数据集中存在一些天然的子集（比如来自不同的domain，不同的topic），那么用单个模型去学习，就会受到很多干扰（interference），导致学习很慢、泛化困难。这时，我们可以使用多个模型（即专家，expert）去学习，使用一个门网络（gating network）来决定每个数据应该被哪个模型去训练，这样就可以减轻不同类型样本之间的干扰。

MoE 将预测建模任务分解为若干子任务，在每个子任务上训练一个专家模型（Expert Model），开发一个门控模型（Gating Model），该模型根据要预测的输入来学习信任哪个专家，并组合预测结果。

Detail

• LAION-5B and some internal datasets
• Multi-scale Training
• 3B
• GPT3.5生成的100个常见的形容词

SnapFusion: Text-to-Image Diffusion Model on Mobile Devices within Two Seconds

Snap Inc.，东北大学，23.6.1，https://snap-research.github.io/SnapFusion/

1. 研究问题：这篇论文解决的研究问题是如何在移动设备上快速运行文本到图像的扩散模型。现有的模型虽然能够生成出色的图像，但由于其复杂的网络架构和大量的去噪迭代，使得它们在计算上非常昂贵，运行速度慢。
2. 动机：这项研究的动机是改进现有模型的运行速度和计算成本。这不仅可以降低运行模型的成本，还可以避免在用户数据被发送到第三方时涉及到的隐私问题。
3. 解决方法：作者通过引入高效的网络架构和改进步骤蒸馏来解决这个问题。他们提出了一个高效的UNet，通过识别原始模型的冗余，并通过数据蒸馏减少图像解码器的计算。此外，他们通过探索训练策略和引入分类器无关指导的正则化来增强步骤蒸馏。
4. 创新之处：这项研究的创新之处在于，它是第一个能够在移动设备上运行文本到图像扩散模型的方法，运行时间少于2秒。此外，他们的模型在8个去噪步骤中比Stable Diffusion v1.5的50个步骤获得了更好的FID和CLIP分数。这项工作通过将强大的文本到图像扩散模型带到用户手中，使内容创建民主化。

为了优化 UNet 结构，我们提出一套 UNet 结构自动评估、进化流程：先对 UNet 进行鲁棒性训练（Robust Training），在训练中随机 drop 一些模块，以此来测试出每个模块对性能的真实影响，从而构建一个 “对 CLIP score 的影响 vs. latency” 的查找表；然后根据该查找表，优先去除对 CLIP score 影响不大同时又很耗时的模块。这一套流程是在线自动进行，完成之后，我们就得到了一个全新的 UNet 结构，称为 Efficient UNet。相比原版 UNet，实现 7.4x 加速且性能不降。

这篇论文中的蒸馏过程主要包括两个部分：图像解码器的蒸馏和步骤蒸馏。

1. 图像解码器的蒸馏：作者提出了一个蒸馏流程，使用合成数据来学习通过通道减少获得的高效图像解码器。他们使用文本提示从SD-v1.5的UNet中获取潜在表示，然后将其转发到他们的高效图像解码器和SD-v1.5的解码器以生成两个图像。然后，通过最小化两个图像之间的均方误差来优化解码器。
2. 步骤蒸馏：步骤蒸馏是一种减少UNet迭代去噪步骤的方法，以实现更快的速度。他们采用了步骤蒸馏的研究方向，其中通过蒸馏教师模型来减少推理步骤。例如，如果教师模型需要32步，那么学生模型可能只需要16步，这样，学生模型的速度就可以提高2倍。

这篇论文还提出了一种新的蒸馏方法，称为CFG-Aware Step Distillation，它在计算损失之前对教师和学生进行分类器无关的指导。这种方法在提高CLIP分数方面表现出显著的效果。

贡献：

（1）通过对现有 UNet 的逐层分析，定位速度瓶颈，提出一种新的高效 UNet 结构（Efficient UNet），可以等效替换原 Stable Diffusion 中的 UNet，实现 7.4x 加速；

（2）对推理阶段的迭代步数进行优化，提出一种全新的步数蒸馏方案（CFG-aware Step Distillation），减少步数的同时可显著提升 CLIP score，实现 6.25x 加速。

不足：

1. SD 模型在多种图像生成场景中都可以使用，本文囿于时间，目前只关注了 text to image 这个核心任务，后期将跟进其他任务（如 inpainting，ControlNet 等等）。
2. 本文主要关注速度上的提升，并未对模型存储进行优化。我们相信所提出的 Efficient UNet 仍然具备压缩的空间，结合其他的高性能优化方法（如剪枝，量化），有望缩小存储，并将时间降低到 1 秒以内，离端上实时 SD 更进一步。