图像生成（Text-to-Image）发展脉络

SoaringPigeon

已于 2024-08-05 18:00:32 修改

阅读量646

点赞数 20

分类专栏：生成文章标签：人工智能生成对抗网络深度学习

于 2024-07-21 23:29:15 首次发布

本文链接：https://blog.csdn.net/weixin_45429089/article/details/140593090

版权

生成专栏收录该内容

1 篇文章 0 订阅

订阅专栏

这篇博客对 图像生成（image generation） 领域的经典工作发展进行了梳理，包括重要的一些改进，目的是帮助读者对此领域有一个整体的发展方向把握，并非是对每个工作的详细介绍。

脉络发展（时间顺序）

GAN [MILA] (2014.06) → VQ-VAE [Google] (2017.11) → DALLE [OpenAI] (2021.02) → DDPM [Google] (2020.06) → DDIM [Stanford] (2020.10) → improved DDPM [Google] (2021.02) → Guided-diffusion [OpenAI] (2021.05) → GLIDE [OpenAI] (2021.12) → DALLE-2 [OpenAI] (2022.04) → Latent Diffusion Model [Heidelberg University] (2022.04) → Stable Diffusion V1 [Stability AI] (2022.08) → Stable Diffusion V1.5 [Stability AI] (2022.10) → Stable Diffusion V2 [Stability AI] (2022.11) → → DiT [Berkeley] (2022.12) → Stable Diffusion XL [Stability AI] (2023.06) → DALLE-3 [OpenAI] (2023.10) → SDXL-Turbo [Stability AI] (2023.11) → Stable Diffusion V3 [Stability AI] (2024.03)

主要发展路线分为自回归模型和扩散模型两个大类，因此我们加下来会按照两大发展脉络来进行梳理具体的技术更新。

GAN

GAN的核心思想是“左右手互搏”，有两个网络，一个生成器G，一个判别器D。它刚提出时，被认为是深度学习领域当时最具创意的思想。

生成器的输入是随机噪声，输出是图像x’。判别器的输入是x‘和真实图像x，输出是二分类，表示x’是否是真实图片。生成器的目标是以假乱真，糊弄判别器。而判别器的目标是练就一双火眼金睛，识别伪造的图片。训练过程G和D不断提升，最后G生成非常逼真的图片。

GAN的目标函数就是为了以假乱真，所以GAN生成的图片保真度非常高。即便人眼也很难区分真假。使用GAN的DeepFake曾经十分火爆。

经过多年的优化，GAN现在很好用，但是它还有一些缺点。首先，训练不稳定。因为它要训练两个网络，不太好平衡。其次，它生成过程的随机性来自初始的随机噪声。这导致GAN生成的图片缺乏多样性和创造性。再者，GAN不是一个概率模型。它的生成都是隐式的，通过一个网络完成的。我们没法知道它具体做了什么，遵循什么分布。GAN在数学上不如后期的VAE、diffusion模型优美。

扩散模型

DDPM

扩散模型早在2015年就提出来了，但是它真正产生好的效果走入人们的视野是2020年DDPM[4]论文之后。

DDPM由Berkeley的三位大佬提出，它算是扩散模型在图像生成领域的开山之作。它的贡献主要有两个：
首先，之前人们在扩散过程中想直接实现X(t)到X(t-1)即图像到图像的转换。DDPM认为直接预测图像比较困难，它转而预测噪声，类似ResNet的思想。
其次，如果要预测正态分布（噪声），只需要学习它的均值和方差。DDPM发现甚至连方差都不用学习（设置成一个常数），只学习均值就能取得很好的效果，再次降低模型优化的难度。

也可以从VAE的角度理解DDPM，只不过有以下差别：

第一，DDPM的编码器是前向扩散过程，不需要学习。

第二，扩散过程每一步中间结果的维度都和输入一样，而VAE的中间结果维度比输入小。

第三，扩散模型有步数，time embedding的概念。每一步的U-Net模型共享参数。

如果想详细学习DDPM建议参考知乎文章和论文原文

DDIM

基于DDPM，DDIM论文主要提出了两项改进：
第一，对于一个已经训练好的DDPM，只需要对采样公式做简单的修改，模型就能在去噪时「跳步骤」，在一步去噪迭代中直接预测若干次去噪后的结果。比如说，假设模型从T=100时刻开始去噪，新的模型可以在每步去噪迭代中预测10次去噪操作后的结果，也就是逐步预测时刻 t=90,80,…,0 的结果。这样，DDPM的采样速度就被加速了10倍。
第二，DDIM论文推广了DDPM的数学模型，从更高的视角定义了DDPM的前向过程（加噪过程）和反向过程（去噪过程）。在这个新数学模型下，我们可以自定义模型的噪声强度，让同一个训练好的DDPM有不同的采样效果。

improved DDPM

改进主要包括：

第一，它没有采用常数方差，而是和均值一样，通过模型学习方差。

第二，改变增加噪声的schedule，从线性schedule变成余弦schedule，和学习率的schedule类似。

第三，尝试大模型，发现扩散模型scale非常好，增加模型参数量可以提升效果，可以大力出奇迹。

Guided Diffusion-“Diffusion Beats GAN”

出了一个新的思想，classifier guidance，引导模型做采样和生成。文章不仅提高了图片的质量，还大大提高了推理速度，只需25步就能完成反向扩散生成图片。

作者受目前GAN方法里通常会使用的类别信息辅助图像生成的原理启发，开发了一个将类别信息引入扩散模型中的方法Classifier Guidance Diffusion，这个方法通俗的说是会训练一个图片分类器，在扩散模型的生成过程中的中间的latend code会通过分类器计算得到一个梯度，该梯度会指导扩散模型的迭代过程。其实这一操作也比较make sense，有一个分类器的存在能更好的告诉U-Net的模型在反向过程生成新图片的时候，当前图片有多像需要生成的物体。有点类似GAN中存在一个判别器的意思。

在论文中提到使用Classifier Guidance的技术能更好的生成逼真的图像，同时能加速图像生成的速度。论文中也提到，通过使用Classifier Guidance的track会牺牲掉一部分的多样性，换取图片的真实性。

classifier guidance仅在扩散模型生成过程中使用，训练过程中不使用

在《Diffusion Models Beat GANs on Image Synthesis》中，通过在生成过程中的近似噪声中加入分类器梯度信息来进行指导：

生成过程：
在这里插入图片描述

然而，人们很快发现使用分类器来做引导存在几个致命的缺点：首先，我们需要相应的分类器来输出置信度。虽然在ImageNet和CIFAR等数据集上我们已经有了不错的分类模型，但仅可以引导特定类别的约束同时要基于用户提供的复杂描述训练一个适用的图片分类器却非常困难。其次，这些分类器都是在正常的图片上进行训练的，而扩散模型的输出信噪比很高，特别是在去噪的前几个阶段。因此，直接使用这些分类器可能会导致输入领域的偏移（Domain Shift），从而引导方向可能不准确。

GLIDE

于是openAI又提出了classifer-free guidance，诞生了GLIDE模型。
GLIDE在引导时不需要额外训练一个classifier。它用什么作为引导信号呢？
它在训练阶段，得到两个输出。一个是有条件的输出，一个是无条件的输出。举例来说，如果训练是图像-文本对，一个输出是有文本作为条件的输出Xy，一个是没有文本作为条件，即将文本随机置成空序列，得到的输出X，保留X与Xy之间的差距作为引导信号。先得到无条件的输出X，再叠加引导信号，可以得到有条件的输出。

classifier-free guidance训练过程中一定几率丢弃condition训练, 生成过程中使用一次去噪过程生成condition和uncondition结果，此处强烈建议看cfg部分源码理解操作。

if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
    e_t = self.model.apply_model(x, t, c)
else:
    x_in = torch.cat([x] * 2)
    t_in = torch.cat([t] * 2)
    c_in = torch.cat([unconditional_conditioning, c])
    e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
    e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)

训练过程：
在这里插入图片描述
生成过程：

在生成过程中，则是同时使用了 conditional model 和 unconditional model 进行计算，在使用一个 w 权重来调整两者的占比。
w越大的时候，conditional model 作用越大，则生成的图像越真实（符合用户需求），但是缺点是损失生成多样性和细节纹理恢复能力。

Latent Diffusion Model

参考：
文生图大模型三部曲：DDPM、LDM、SD 详细讲解！
Stable Diffusion ———LDM、SD 1.0, 1.5, 2.0、SDXL、SDXL-Turbo等版本之间关系现原理详解

主要创新点

LDM提出了cross-attention的方法来实现多模态训练，使得条件图片生成任务也可以实现。论文中提到的条件图片生成任务包括：类别条件图片生成（class-condition）, 文图生成（text-to-image）, 布局条件图片生成（layout-to-image）。这也为日后Stable Diffusion的开发奠定了基础。本文在扩散过程中引入了条件机制（Conditioning Mechanisms），通过cross-attention的方式来实现多模态训练，使得条件图片生成任务得以实现。具体做法是通过训练一个条件时序去噪自编码器ϵ_θ (z_t,t,y)，来通过 y来控制图片合成的过程。为了能够从多个不同的模态预处理 y ，论文引入了一个领域专用编码器τ_θ（在stable diffusion中替代为了预训练冻结的CLIP text encoder），它用来将 y 映射为一个中间表示τ_θ (y) ，这样我们就可以很方便的引入各种形态的条件（文本、类别等等）。最终模型就可以通过一个cross-attention层映射将控制信息融入到UNet的中间层。
DDPM在像素空间上训练模型，需要反复迭代计算，因此训练和推理代价都很高。DLM提出一种在潜在表示空间上进行扩散过程的方法，能够显著减少计算复杂度，同时也能达到十分不错的图片生成效果。
相比于其它空间压缩方法，论文提出的方法可以生成更细致的图像，并且在高分辨率图片生成任务（如风景图生成，百万像素图像）上表现得也很好。

UNet的具体结构图：condition y是在cross-attention中引入，step t是在ResNet的两层卷积之间add进入。
在这里插入图片描述

Stable Diffusion V1

这里正式开启Stable Diffusion发展路线，VAE压缩+去噪模型+条件控制。

概要

Stable diffusion是一种潜在的文本到图像的扩散模型。基于之前的大量工作（如DDPM、LDM的提出），并且在Stability AI的算力支持和LAION的海量数据支持下，Stable diffusion才得以成功。

Stable diffusion在来自LAION- 5B数据库子集的512x512图像上训练潜在扩散模型。与谷歌的Imagen类似，这个模型使用一个冻结的CLIP vit L/14文本编码器来根据文本提示调整模型。

Stable diffusion拥有860M的UNet和123M的文本编码器，该模型相对轻量级，可以运行在具有至少10GB VRAM的GPU上。

主要改进点

Stable diffusion是在LDM的基础上建立的，同时在LDM的基础上进行了一些改进：

数据集：在更大的数据集LAION- 5B上进行训练

条件机制：使用更强大的CLIP模型，代替原始的交叉注意力调节机制

除此之外，随着各种图形界面的出现、微调方法的发布、控制模型的公开，SD进入全新架构SDXL时代，功能更加强大。

模型训练

SD的训练是采用了32台8卡的A100机器（32 x 8 x A100_40GB GPUs），单卡的训练batch size为2，并采用gradient accumulation，其中gradient accumulation steps=2，那么训练的总batch size就是32x8x2x2=2048。训练优化器采用AdamW，训练采用warmup，在初始10,000步后学习速率升到0.0001，后面保持不变。至于训练时间约150,000小时（A100卡时），如果按照256卡A100来算的话，那么大约需要训练25天左右。

SD提供了不同版本的模型权重可供选择：

SD v1.1：在laion2B-en数据集上以256x256大小训练237,000步，上面我们已经说了，laion2B-en数据集中256以上的样本量共1324M；然后在laion5B的高分辨率数据集以512x512尺寸训练194,000步，这里的高分辨率数据集是图像尺寸在1024x1024以上，共170M样本。

SD v1.2：以SD v1.1为初始权重，在improved_aesthetics_5plus数据集上以512x512尺寸训练515,000步数，这个improved_aesthetics_5plus数据集上laion2B-en数据集中美学评分在5分以上的子集（共约600M样本），注意这里过滤了含有水印的图片（pwatermark>0.5)以及图片尺寸在512x512以下的样本。

SD v1.3：以SD v1.2为初始权重，在improved_aesthetics_5plus数据集上继续以512x512尺寸训练195,000步数，不过这里采用了CFG（以10%的概率随机drop掉text）。

SD v1.4：以SD v1.2为初始权重，在improved_aesthetics_5plus数据集上采用CFG以512x512尺寸训练225,000步数。

SD v1.5：以SD v1.2为初始权重，在improved_aesthetics_5plus数据集上采用CFG以512x512尺寸训练595,000步数。

其实可以看到SD v1.3、SD v1.4和SD v1.5其实是以SD v1.2为起点在improved_aesthetics_5plus数据集上采用CFG训练过程中的不同checkpoints，目前最常用的版本是SD v1.4和SD v1.5。

条件控制

SD采用CLIP text encoder来对输入text提取text embeddings，具体的是采用目前OpenAI所开源的最大CLIP模型：clip-vit-large-patch14，这个CLIP的text encoder是一个transformer模型（只有encoder模块）：层数为12，特征维度为768，模型参数大小是123M。对于输入text，送入CLIP text encoder后得到最后的hidden states（即最后一个transformer block得到的特征），其特征维度大小为77x768（77是token的数量），这个细粒度的text embeddings将以cross attention的方式送入UNet中。
值得注意的是，这里的tokenizer最大长度为77（CLIP训练时所采用的设置），当输入text的tokens数量超过77后，将进行截断，如果不足则进行paddings，这样将保证无论输入任何长度的文本（甚至是空文本）都得到77x768大小的特征。在训练SD的过程中，CLIP text encoder模型是冻结的。在早期的工作中，比如OpenAI的GLIDE和latent diffusion中的LDM均采用一个随机初始化的tranformer模型来提取text的特征，但是最新的工作都是采用预训练好的text model。比如谷歌的Imagen采用纯文本模型T5 encoder来提出文本特征，而SD则采用CLIP text encoder，预训练好的模型往往已经在大规模数据集上进行了训练，它们要比直接采用一个从零训练好的模型要好。
下面是SD中使用的条件控制模型CLIP的结构示意图

与其他模型对比

DALL-E2 ：出自OpenAI，其基本原理和SD一样，都是源于最初的扩散概率模型（DDPM），与之不同发是，SD继承了LDM的思想，在潜在空间中进行扩散学习；而DALL-E2是在像素空间中进行扩散学习，所以其计算复杂度较高。

Imagen：由谷歌发布，采用预训练好的文本编码器T5，通过扩散模型，实现文本到低分辨率图像的生成，最后将低分辨率图像进行两次超分，得到高分辨率图像。

Stable Diffusion V1.5

Stable Diffusion 的 V1.5 版本，由 runway 发布，位于代码库 GitHub - runwayml/stable-diffusion: Latent Text-to-Image Diffusion 中。

该版本发布于 2022 年 10 月，主要包含两个模型：

sd-v1-5.ckpt：

复用 sd-v1-2.ckpt，在 LAION-aesthetics v2 5+ 上以 512x512 分辨率继续训练 595k step，使用了 Classifier Free Guidance 技术，以 10% 概率删除文本条件。

sd-v1-5-inpainting.ckpt：

复用 sd-v1-5.ckpt，在 LAION-aesthetics v2 5+ 上以 512x512 分辨率以 inpainting 训练了 440k step，使用 Classifier Free Guidance 技术，以 10% 概率删除文本条件。在 U-Net 的输入中额外加了 5 个 channel，4 个用于 masked 的图像，1 个用于 mask 本身。

对应的 FID 和 CLIP 分数如下图所示，可以看出，v1.5 相比 v1.4 的提升也不是很明显：
在这里插入图片描述

Stable Diffusion V2

Stable Diffusion 的 V2 版本，由 Stability-AI 发布，位于代码库 GitHub - Stability-AI/stablediffusion: High-Resolution Image Synthesis with Latent Diffusion Models 中。

V2 包含三个子版本，分别为 v2.0，v2.1 和 Stable UnCLIP 2.1：

v2.0：

发布于 2022 年 11 月，U-Net 模型和 V1.5 相同，Text encoder 模型换成了 OpenCLIP-ViT/H 中的 text encoder。

SD 2.0-base：分别率为 512x512

SD 2.0-v：基于 2.0-base 微调，分辨率提升到 768x768，同时利用 [2202.00512] Progressive Distillation for Fast Sampling of Diffusion Models 提出的技术大幅降低 Diffusion 的步数。

发布了一个文本引导的 4 倍超分模型。

基于 2.0-base 微调了一个深度信息引导的生成模型。

基于 2.0-base 微调了一个文本信息引导的修复模型。

v2.1：

发布于 2022 年 12 月，模型结构和参数量都和 v2.0 相同。并在 v2.0 的基础上使用 LAION 5B 数据集（较低的 NSFW 过滤约束）微调。同样包含 512x512 分辨率的 v2.1-base 和 768x768 分辨率的 v2.1-v。

Stable UnCLIP 2.1：

发布于 2023 年 3 月，基于 v2.1-v（768x768 分辨率）微调，参考 OpenAI 的 DALL-E 2（也就是 UnCLIP），可以更好的实现和其他模型的联合，同样提供基于 CLIP ViT-L 的 Stable unCLIP-L 和基于 CLIP ViT-H 的 Stable unCLIP-H。

如下图所示为 v2.0 和 v2.0-v 与 v1.5 的对比，可见其都有明显提升：
在这里插入图片描述

DiT

如果只是对 DiT 的结构感兴趣的话，可以去直接通过读 SD3 的源码来学习。DiT模块有两个作用，一个是对特征进行加工，另一个是融合图像的特征和不同模态的条件特征，读 DiT 论文时只需要着重学习 AdaLayerNormZero 模块。

在这里插入图片描述

下面截取部分源码来简单看看：

上下文条件（In-context conditioning）
基于上下文条件的DiT直接将条件特征附加到输入序列中，这个操作类似于在输入序列中添加了一个[CLS] token。DiT的条件编码是通过LabelEmbedder类实现的，具体实现见下面代码片段：

class LabelEmbedder(nn.Module):
  """
  Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
  """
  def __init__(self, num_classes, hidden_size, dropout_prob):
      super().__init__()
      use_cfg_embedding = dropout_prob > 0
      self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
      self.num_classes = num_classes
      self.dropout_prob = dropout_prob
 
  def token_drop(self, labels, force_drop_ids=None):
      """
      Drops labels to enable classifier-free guidance.
      """
      if force_drop_ids is None:
          drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
      else:
          drop_ids = force_drop_ids == 1
      labels = torch.where(drop_ids, self.num_classes, labels)
      return labels
 
  def forward(self, labels, train, force_drop_ids=None):
      use_dropout = self.dropout_prob > 0
      if (train and use_dropout) or (force_drop_ids is not None):
          labels = self.token_drop(labels, force_drop_ids)
      embeddings = self.embedding_table(labels)
      return embeddings

交叉注意力块（Cross-Attention）
我们将时间片特征t和条件特征c拼成一个长度为2的序列，然后将这个序列输入到一个多头交叉注意力模块中和图像特征进行融合。但是DiT采用的不是这种形式，二是adaLN-Zero
自适应层归一化块（Adaptive Layer Normalization，AdaLN）
DiT在模型中尝试了AdaLN，AdaLN的核心思想是使用模型中的一些信息学习scale和shift两个归一化参数。DiT是使用时间片特征t和条件特征y相加后的结果计算这两个参数。此外，DiT在每个残差连接之后还接了一个回归缩放参数scale，它同样是这样计算得到。接下来我们根据下面的代码片段详细介绍DiT的具体结构。

class DiTBlock(nn.Module):
    """
    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
    """
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        approx_gelu = lambda: nn.GELU(approximate="tanh")
        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )
 
    def forward(self, x, c):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x

首先我们观察forword函数的第1行，它使用adaLN_modulation计算了6个变量shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp，这6个变量分别对应了多头自注意力的LN的归一化参数与缩放参数以及MLP的LN的归一化参数与缩放参数。forward函数的第二行是计算多头自注意力以及它的LN，它首先计算的是modulate函数，即相当于使用学习好的 $\beta$ 和 $\gamma$ 对LN进行归一化。接下来再计算的注意力模块，计算方式和Transformer相同。最后在通过乘以gate_msa对注意力计算的结果进行缩放。modulate函数的实现如下：

def modulate(x, shift, scale):
     return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

当对特征加工完之后，我们需要使用FinalLayer模块来将特征还原为与输入相同的尺寸。它是由一个AdaLN和一个线性层组成，具体实现见下面代码片段：

class FinalLayer(nn.Module):
    """
    The final layer of DiT.
    """
    def __init__(self, hidden_size, patch_size, out_channels):
        super().__init__()
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
        )
 
    def forward(self, x, c):
        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
        x = modulate(self.norm_final(x), shift, scale)
        x = self.linear(x)
        return x

adaZero-Block
之前有研究表明使用0初始化网络中的某些参数可以加速模型的训练。例如我们可以将残差网络的残差部分初始化为0，这样初始化后的残差块相当于一个单位映射，可以直接将上一层的特征透传给下一层。我们也可以将BN的归一化因子初始化为0来加速模型的训练。DiT对模型参数的初始化都是在initialize_weights函数中实现的，它的作用是对DiT中的变量进行初始化，我们具体看一下这个函数。

def initialize_weights(self):
    # Initialize transformer layers:
    def _basic_init(module):
        if isinstance(module, nn.Linear):
            torch.nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.constant_(module.bias, 0)
    self.apply(_basic_init)
 
    # Initialize (and freeze) pos_embed by sin-cos embedding:
    pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5))
    self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
 
    # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
    w = self.x_embedder.proj.weight.data
    nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
    nn.init.constant_(self.x_embedder.proj.bias, 0)
 
    # Initialize label embedding table:
    nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
 
    # Initialize timestep embedding MLP:
    nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
    nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
 
    # Zero-out adaLN modulation layers in DiT blocks:
    for block in self.blocks:
        nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
        nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
 
    # Zero-out output layers:
    nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
    nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
    nn.init.constant_(self.final_layer.linear.weight, 0)
    nn.init.constant_(self.final_layer.linear.bias, 0)

逆Patch操作

def unpatchify(self, x):

    x: (N, T, patch_size**2 * C)
    imgs: (N, H, W, C)
    """
    c = self.out_channels
    p = self.x_embedder.patch_size[0]
    h = w = int(x.shape[1] ** 0.5)
    assert h * w == x.shape[1]

    x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
    x = torch.einsum('nhwpqc->nchpwq', x)
    imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
    return imgs

DIT完整forward

 def forward(self, x, t, y):

        x = self.x_embedder(x) + self.pos_embed  # (N, T, D), where T = H * W / patch_size ** 2
        t = self.t_embedder(t)                   # (N, D)
        # time step embedding
        y = self.y_embedder(y, self.training)    # (N, D)
        c = t + y                                # (N, D)
        # 送入上述的DIT-Block中
        for block in self.blocks:
            x = block(x, c)                      # (N, T, D)
        x = self.final_layer(x, c)                # (N, T, patch_size ** 2 * out_channels)
        x = self.unpatchify(x)                   # (N, out_channels, H, W)
        return x

带CFG的forward：
参考：Open-Sora-Plan-v1.0.0

# Create sampling noise:
n = len(class_labels)
z = torch.randn(n, 4, latent_size, latent_size, device=device)
y = torch.tensor(class_labels, device=device)
 
# Setup classifier-free guidance:
z = torch.cat([z, z], 0)
y_null = torch.tensor([1000] * n, device=device)
y = torch.cat([y, y_null], 0)
model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)

def forward_with_cfg(self, x, t, y, cfg_scale):

    half = x[: len(x) // 2]
    combined = torch.cat([half, half], dim=0)
    model_out = self.forward(combined, t, y)
    #确保精确的可重复性，默认只在三个通道上应用无分类器引导
    eps, rest = model_out[:, :3], model_out[:, 3:]
    cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
    #条件控制权重参数
    half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
    eps = torch.cat([half_eps, half_eps], dim=0)
    return torch.cat([eps, rest], dim=1)

Stable Diffusion XL (SDXL)

Stable Diffusion 的 XL 版本，由 Stability-AI 发布，位于代码库 Generative Models by Stability AI。

该版本发布于 2023 年 06 月，主要包含两个模型：

SDXL-base-0.9：基于多尺度分辨率训练，最大分辨率 1024x1024，包含两个 Text encoder，分别为 OpenCLIP-ViT/G 和 CLIP-ViT/L。

SDXL-refiner-0.9：用来生成更高质量的图像，不应直接使用，此外文本条件只使用 OpenCLIP 中的 Text encoder。

2023 年 07 月发布 1.0 版本，同样对应两个模型：

SDXL-base-1.0：基于 SDXL-base-0.9 改进。

SDXL-refiner-1.0：基于 SDXL-refiner-0.9 改进。

2023 年 11 月发表 SDXL-Trubo 版本，也就是优化加速的版本。

SDXL 模型概览

如下图所示，SDXL 相比 SD 主要的修改包括（模型总共 2.6B 参数量，其中 text encoder 817M 参数量）：

增加一个 Refiner 模型，用于对图像进一步地精细化
使用 CLIP ViT-L 和 OpenCLIP ViT-bigG 两个 text encoder
基于 OpenCLIP 的 text embedding 增加了一个 pooled text embedding

微条件（Micro-Conditioning）

（一）以图像大小作为条件
在 SD 的训练范式中有个明显的缺陷，对图像大小有最小长宽的要求。针对这个问题有两种方案：

丢弃分辨率过小的图像（例如，SD 1.4/1.5 丢弃了小于 512 像素的图像）。但是这可能导致丢弃过多数据，如下图 Figure 2 所示为预训练数据集中图像的长、宽分布，如果丢弃 256x256 分辨率的图像，将导致 39% 的数据被丢弃。

在这里插入图片描述
2. 另一种方式是放大图像，但是可能会导致生成的样本比较模糊。针对这种情况，作者提出将原始图像分辨率作用于 U-Net 模型，并提供图像的原始长和宽（csize = (h, w)）作为附加条件。并使用傅里叶特征编码，然后会拼接为一个向量，把它扩充到时间步长 embedding 中并一起输入模型。如下图所示，在推理时指定不同的长宽即可生成相应的图像，（64,64）的图像最模糊，（512, 512）的图像最清晰：
在这里插入图片描述

（二）以裁剪参数作为条件
此外，以前的 SD 模型存在一个比较典型的问题：生成的物体不完整，像是被裁剪过的，如下图 SD1.5 和 SD 2.1 的结果。作者猜测这可能和训练阶段的随机裁剪有关，考虑到这个因素，作者将裁剪的左上坐标（top, left）作为条件输入模型，和 size 类似。如下图 Figure 4 中 SDXL 的结果，其生成结果都更加完整：
在这里插入图片描述
如下图 Figure 5 所示，在推理阶段也可以通过裁剪坐标来控制位置关系：

（三）多分辨率训练
真实世界的图像会包含不同的大小和长宽比，而文本到模型生成的图像分辨率通常为 512x512 或 1024x1024，作者认为这不是一个自然的选择。受此启发，作者以不同的长宽比来微调模型：首先将数据划分为不同长宽比的桶，其中尽可能保证总像素数接近 1024x1024 个，同时以 64 的整数倍来调整高度和宽度。如下图所示为作者使用的宽度和高度。在训练过程中，每次都从同样的桶中选择一个 batch，并在不同的桶间交替。此外，和之前的 size 类似，作者会将桶的高度和宽度（h, w）作为条件，经傅里叶特征编码后添加到时间步 embedding 中：
在这里插入图片描述
（四）训练
SDXL 模型的训练包含多个步骤：

基于内部数据集，以 256x256 分辨率预训练 6,000,000 step，batch size 为 2048。使用了 size 和 crop 条件。

继续以 512x512 分辨率训练 200,000 step。

最后使用多分辨率（近似 1024x1024）训练。

根据以往的经验，作者发现所得到的的模型有时偶尔会生成局部质量比较差的图像，为了解决这个问题，作者在同一隐空间训练了一个独立的 LDM（Refiner），该 LDM 专门用于高质量、高分辨率的数据。在推理阶段，直接基于 Base SDXL 生成的 Latent code 继续生成，并使用相同的文本条件（当然，此步骤是可选的），实验证明可以提高背景细节以及人脸的生成质量。

（五）实验结果
如下图所示，作者基于用户评估，最终带有 Refiner 的 SDXL 获得了最高分，并且 SDXL 结果明显优于 SD 1.5 和 SD 2.1。
在这里插入图片描述
如下图 Figure 10 所示为 SDXL（没有 Refiner）和 Midjourney 5.1 的对比结果，可见 SDXL 的结果略胜一筹：
如下图 Figure 11 所示为 SDXL（带有 Refiner）和 Midjourney 5.1 的对比结果，可见 SDXL 的结果同样略胜一筹：

SDXL-Turbo

SDXL-Turbo 在模型上没有什么修改，主要是引入蒸馏技术，以便减少 LDM 的生成步数，提升生成速度。大致的流程为：

从 $T_{student}$ 中采样步长 s，对于原始图像 $x_0$ 进行 s 步的前向扩散过程，生成加噪图像 $x_s$ 。
使用学生模型 ADD-student 对 $x_s$ 进行去噪，生成去噪图像 $x_θ$ 。
基于原始图像 $x_0$ 和去噪图像 $x_θ$ 计算对抗损失（adversarial loss）。
从 $T_{teacher}$ 中采样步长 t，对去噪后的图像 $x_θ$ 进行 t 步的前向扩散过程，生成 $x_{θ,t}$ 。
使用教师模型 DM-teacher 对 $x_{θ,t}$ 进行去噪，生成去噪图像 $x_ψ$ 。
基于学生模型去噪图像 xθ 和教师模型去噪图像 $x_ψ$ 计算蒸馏损失（distillation）。
根据损失进行反向传播（注意，教师模型不更新，因此会 stop 梯度）

在这里插入图片描述

需要说明的是，通常 ADD-student 模型需要预训练过程，然后再蒸馏。此外， $T_{student}$ 的 N 比较小，作者设置为 4，而 $T_{teacher}$ 的 N 比较大，为 1000。也就是学生模型可能只加噪 1,2,3,4 步，而教师模型可能加噪 1-1000 步。
在这里插入图片描述

此外，作者在训练中还用了其他技巧，比如使用了 zero-terminal SNR；教师模型不是直接作用于原始图像 x0，而是作用于学生模型恢复出的图像 xθ，否则会出现 OOD（out of distribution）问题；作者还应用了 Score Distillation Loss，并且与最新的 noise-free score distillation 进行了对比。

Stable Diffusion V3

参考：Stable Diffusion 3 论文及源码概览

文章的核心贡献如下：

从方法设计上：
首次在大型文生图模型上使用了整流模型。
用一种新颖的 Diffusion Transformer (DiT) 神经网络来更好地融合文本信息。
使用了各种小设计来提升模型的能力。如使用二维位置编码来实现任意分辨率的图像生成。
从实验上：
开展了一场大规模、系统性的实验，以验证哪种扩散模型/整流模型的学习目标最优。
开展了扩增模型参数的实验 (scaling study)，以证明提升参数量能提升模型的效果。

流匹配原理简介
流匹配是一种定义图像生成目标的方法，它可以兼容当前扩散模型的训练目标。流匹配中一个有代表性的工作是整流 (rectified flow)，它也正是 SD3 用到的训练目标。我们会在本文中通过简单的可视化示例学习流匹配的思想。

SD3 中的 DiT
我们会从一个简单的类 ViT 架构开始，学习 SD3 中的去噪网络 DiT 模型是怎么一步一步搭起来的。读者不需要提前学过 DiT，只需要了解 Transformer 的结构，并大概知道视觉任务里的 Transformer 会做哪些通用的修改（如图块化），即可学懂 SD3 里的 DiT。

SD3 模型与训练策略改进细节
除了将去噪网络从 U-Net 改成 DiT 外，SD3 还在模型结构与训练策略上做了很多小改进：

改变训练时噪声采样方法
将一维位置编码改成二维位置编码
提升 VAE 隐空间通道数
对注意力 QK 做归一化以确保高分辨率下训练稳定

为了方便学习具体的改进，我将MM-DiT（SD3使用的）的结构图展示如下：
在这里插入图片描述

大型消融实验
对于想训练大型文生图模型的开发者，SD3 论文提供了许多极有价值的大型消融实验结果。本文会简单分析论文中的两项实验结果：各训练目标在文生图任务中的表现、SD3 的参数扩增实验结果。

自回归模型

VQ-VAE

建议阅读此文把握早期图像生成的发展脉络：Stable Diffusion 解读（一）：回顾早期工作
具体做法：深入理解 VQ-VAE

VQ-VAE的codebook虽然需要学习，但它是固定的，不是概率分布，没法随机采样。准确地说，VQ-VAE不像是VAE，更像是AE。它学到的codebook和中间特征图，不是用来做图像生成，而是做分类等其它图像任务。

如何用VQ-VAE做图像生成呢？需要再给它训练一个prior网络。在VQ-VAE原论文中，通过训练一个pixelCNN当作prior网络，来做图像生成。

pixelCNN其实是一个auto regressive的模型，而openAI的看家本领GPT也是auto regressive模型，于是他们很快想到把pixelCNN替换成GPT。这便有了DALLE-1。

DALLE

系列未开源

参考阅读：【论文阅读】DALL·E: Zero-Shot Text-to-Image Generation

自回归式的模型处理图片的时候，如果直接把像素拉成序列，当成image token来处理，如果图片分辨率过高，一方面会占用过多的内存，另一方面Likelihood的目标函数会倾向于建模短程的像素间的关系，因而会学到更多的高频细节，而不是更能被人辨认的低频结构。

与 GPT-3 一样，DALL·E是一个 Transformer 语言模型。它将文本和图像作为单个数据序列接收，通过Transformer进行自回归。

Zero-Shot：对于GPT-3来说，我们可以指示它仅根据描述和提示执行多种任务，来生成我们想要的答案，而无需任何额外的培训。例如，当提示“here is the sentence ‘a person walk his dog in the park’ translated into French: ”时，GPT-3 会回答“un homme qui promène son chien dans le parc”。这种能力称为零样本推理。研究者发现 DALL·E 将这种能力扩展到了视觉领域，在以正确的方式提示时，它能够执行多种图像到图像的转换工作。比如最直接的方法是输入一张照片，然后用文本描述“粉红色的照片”或“倒影的照片”，这也往往是最可靠的，尽管照片通常没有被完美地复制或反映，但也从一定程度上表明这个功能有望实现。实现图像版GPT-3。

四大功能

创造拟人的动物和物体
Dall-E能够很聪明地捕捉到每个事物的特性，并且合理地组织在了一起。
根据文本自动渲染真实场景图片，其仿真程度与真实照片十分接近
根据文本指令改变和转换现有图片风格

方法

训练与推理：
1.第一个阶段，训练一个dVAE（discrete variational autoencoder离散变分自动编码器），其将256256的RGB图片转换为3232的图片token。目的：降低图片的分辨率，从而解决计算量的问题。图片token的词汇量大小是8192个，即每个位置有8192种可能的取值(也就是说dVAE的encoder输出是维度为32x32x8192的logits，然后通过logits索引codebook的特征进行组合，codebook的embedding是可学习的)。第一阶段同时训练dVAE编码器和dVAE解码器。

2.第二阶段，用BPE Encoder对文本进行编码，得到最多256个文本token，token数不满256的话padding到256，然后将256个文本token与1024个图像token（32*32=1024）进行拼接，得到长度为1280的数据，用拼接的数据去训练一个自回归transformer来建模文本和图片token的联合分布。

3.最后的推理阶段，给定一张候选图片和一条文本，通过transformer可以得到融合后的token，然后用dVAE的decoder生成图片，最后通过预训练好的CLIP计算出文本和生成图片的匹配分数，得到不同采样图片的分数排序，最终找到跟文本最匹配的图片。

从以上流程可知，dVAE、Transformer和CLIP三个模型都是不同阶段独立训练的。

在这里插入图片描述

下面讲一下dVAE、Transformer和CLIP三个部分：

dVAE
由于图像特征的密集性和冗余性，不能直接提供给Transformer进行训练。目前主流的方式，例如ViT，Swin-Transformer等都是将图像的Patch作为模型的输入，然后通过一个步长等于Patch大小的大卷积核得到每个Patch的特征向量。
DALL-E提供的方案是使用离散的变分自编码器（dVAE）将大小为256x256的RGB图像压缩到大小为32x32的，通道数为8192的one-hot token的分布，变分自编码器的架构如图2所示。换句话说，阶段1的作用是将图像映射到一个大小为8192的图表中。这里通道数为8192的one-hot向量可以看做是一个词表，它的思想是通过离散VAE，实现图像特征空间想文本特征空间的映射。

通过Encoder学习出中间编码，然后通过最邻近搜索将中间编码映射为codebook中K个向量之一，然后通过Decoder对latent code进行重建。
另外由于最邻近搜索使用argmax来找codebook中的索引位置，导致不可导问题，VQVAE通过stop gradient操作来避免最邻近搜索的不可导问题，也就是latent code的梯度跳过最近邻搜索直接复制到中间编码上。
VQVAE相比于VAE最大的不同是，直接找每个属性的离散值，通过类似于查表的方式，计算codebook和中间编码的最近邻作为latent code。由于维护了一个codebook，编码范围更加可控，VQVAE相对于VAE，可以生成更大更高清的图片(这也为后续DALLE和VQGAN的出现做了铺垫)。

作用：
dVAE主要用来为图像的每个patch生成token表示，这次openAI开出的代码就是dVAE的推理代码。dVAE的encoder和decoder的机构较为简单，都是由bottleneck-style的resblock组成，相比于普通的VAE，dVAE有两点区别：
1.和VQVAE方法相似，dVAE的encoder是将图像的patch映射到8192的词表中，论文中将其分布设为在词表向量上的均匀分类分布，这是一个离散分布，由于不可导，此时不能采用重参数技巧，DALL·E使用Gumbel Softmax trick来解决这个问题。
2、在重建图像时，真实的像素值是在一个有界区间内，而VAE中使用的Gaussian分布和Laplace分布都是在整个实数集上，这造成了不匹配的问题。为了解决这个问题，论文中提出了logit-Laplace分布，如下式所示：
Transformer
Dall-E中的Transformer结构由64层attention层组成，每层的注意力头数为62，每个注意力头的维度为64，因此，每个token的向量表示维度为3968。如图2所示，attention层使用了行注意力mask、列注意力mask和卷积注意力mask三种稀疏注意力。

训练损失：
CLIP

CLIP的推理过程：

1.先将预训练好的CLIP迁移到下游任务，如图(2)所示，先将下游任务的标签构建为一批带标签的文本(例如 A photo of a {plane})，然后经过Text Encoder编码成一批相应的word embedding。

2.然后将没有见过的图片进行zero-shot预测，如图(3)所示，通过Image Encoder将一张小狗的图片编码成一个feature embedding，然后跟(2)编码的一批word embedding先归一化然后进行点积，最后得到的logits中数值最大的位置对应的标签即为最终预测结果。

在DALL·E中，CLIP的用法跟上述过程相反，提供输入文本和一系列候选图片，先通过Stage One和Stage Two生成文本和候选图片的embedding，然后通过文本和候选图片的匹配分数进行排序，最后找到跟文本最匹配的图片。

总结

总的来说，目前公开的DALL-E的实现在模型结构上并没有太多创新，而是合理利用了现有的模型结构进行组合，并采用了一些trick解决了遇到的问题，从而在大数据集上训练得到超大规模的模型，取得了令人惊艳的效果，这也符合openAI的一贯风格。但无论如何，DALL-E在深度学习能力边界探索的道路上又前进了一步，也再一次展示了大数据和超大规模模型的魅力。美中不足的是，DALL-E包含了三个模块，更像是一个pipeline，而对于普通的研究者来说，要运行这样一个复杂的大规模模型是一件很困难的事情。

DALLE-2

论文的训练数据集由图像 $x$ 及其相应标题（captions） $y$ 组成。给定图像 $x$ ，经过训练好的CLIP模型分别得到文本特征 $z_t$
和图像特征 $z_i$ 。然后训练两个组件来从标题生成图像：

prior：先验模型 $P(z_i|y)$ ，根据标题 $y$ 生成CLIP的图像特征 $z_i$
decoder：解码器 $P(x|z_i, y)$ ，生成以CLIP图像特征 $z_i$ （和可选的文本标题 $y$ ）为条件的图像 $x$

采用两阶段而不是一阶段的关键提升在于：可以利用clip的图像编码器生成的特征去监督根据文本特征生成图像特征。在训练的时候，可以根据图像获得图像特征，拿该图像特征作为ground truth来监督。

在这里插入图片描述
所以DALL·E2是一个两阶段的生成器。在先训练好CLIP之后，prior生成图片特征，最后利用decoder解码图像特征得到生成的图片。

在这里插入图片描述

DALLE-3

提出问题： 现有的文本到图像模型很难遵循详细的图像描述，并且经常忽略单词或混淆提示的含义。
作出假设： 假设这个问题是因为训练数据集中的噪声和不准确的图像标题。
解决方案： 训练一个图片标注器用于生成标注，然后使用它来重新捕获训练数据集。
主要工作： 训练了DALLE3，发现对生成的图片标注进行训练确实能提升生成模型的性能，并且提出了一套评估生成模型的方案。