Stable Diffusion XL总结

Stable Diffusion XL是一个二阶段的级联扩散模型，包括Base模型和Refiner模型。其中Base模型的主要工作和Stable Diffusion一致，具备文生图，图生图，图像inpainting等能力。在Base模型之后，级联了Refiner模型，对Base模型生成的图像Latent特征进行精细化，其本质上是在做图生图的工作。

Base模型由U-Net，VAE，CLIP Text Encoder（两个）三个模块组成，在FP16精度下Base模型大小6.94G（FP32：13.88G），其中U-Net大小5.14G，VAE模型大小167M以及两个CLIP Text Encoder一大一小分别是1.39G和246M。

Refiner模型同样由U-Net，VAE，CLIP Text Encoder（一个）三个模块组成，在FP16精度下Refiner模型大小6.08G，其中U-Net大小4.52G，VAE模型大小167M（与Base模型共用）以及CLIP Text Encoder模型大小1.39G（与Base模型共用）。

比起Stable Diffusion，Stable Diffusion XL的参数量增加到了66亿（Base模型35亿+Refiner模型31亿），并且先后发布了模型结构完全相同的0.9和1.0两个版本。Stable Diffusion XL 1.0使用更多训练集+RLHF来优化生成图像的色彩，对比度，光线以及阴影方面，使得生成图像的构图比0.9版本更加鲜明准确。
       Stable DiffusionV1.4/1.5的U-Net参数量只有860M，就算是Stable DiffusionV2.0/2.1，其参数量也不过865M。但等到Stable Diffusion XL，U-Net模型（Base部分）参数量就增加到2.6B，参数量增加幅度达到了3倍左右。

可以看到，其中增加的SDXL_Spatial Transformer_X模块（主要包含Self Attention + Cross Attention + FeedForward）数量占新增参数量的主要部分U-Net的Encoder和Decoder结构也从原来的4stage改成3stage（[1,1,1,1] -> [0,2,10]），说明SDXL只使用两次下采样和上采样，而之前的SD系列模型都是三次下采样和上采样。并且比起Stable DiffusionV1/2，Stable Diffusion XL在第一个stage中不再使用Spatial Transformer Blocks，而在第二和第三个stage中大量增加了Spatial Transformer Blocks（分别是2和10）。首先，在第一个stage中不使用SDXL_Spatial Transformer_X模块，可以明显减少显存占用和计算量。然后在第二和第三个stage这两个维度较小的feature map上使用数量较多的SDXL_Spatial Transformer_X模块，能在大幅提升模型整体性能（学习能力和表达能力）的同时，优化了计算成本。整个新的SDXL Base U-Net设计思想也让SDXL的Base出图分辨率提升至1024x1024。在参数保持一致的情况下，Stable Diffusion XL生成图片的耗时只比Stable Diffusion多了20%-30%之间

在SDXL U-Net的Encoder结构中，包含了两个CrossAttnDownBlock结构和一个SDXL_DownBlock结构；在Decoder结构中，包含了两个CrossAttnUpBlock结构和一个SDXL_UpBlock结构；与此同时，Encoder和Decoder中间存在Skip Connection，进行信息的传递与融合。

BasicTransformer Block模块是整个框架的基石，由SelfAttention，CrossAttention和FeedForward三个组件构成，并且使用了循环残差模式，让SDXL Base U-Net不仅可以设计的更深，同时也具备更强的文本特征和图像体征的学习能力。

Stable Diffusion XL中的Text Condition信息由两个Text Encoder提供（OpenCLIP ViT-bigG和OpenAI CLIP ViT-L），通过Cross Attention组件嵌入，作为K Matrix和V Matrix。与此同时，图片的Latent Feature作为Q Matrix。

当输入是图片时，Stable Diffusion XL和Stable Diffusion一样，首先会使用VAE的Encoder结构将输入图像转换为Latent特征，然后U-Net不断对Latent特征进行优化，最后使用VAE的Decoder结构将Latent特征重建出像素级图像。除了提取Latent特征和图像的像素级重建外，VAE还可以改进生成图像中的高频细节，小物体特征和整体图像色彩。

当Stable Diffusion XL的输入是文字时，这时我们不需要VAE的Encoder结构，只需要Decoder进行图像重建。VAE的灵活运用，让Stable Diffusion系列增添了几分优雅。

Stable Diffusion XL使用了和之前Stable Diffusion系列一样的VAE结构，但在训练中选择了更大的Batch-Size（256 vs 9），并且对模型进行指数滑动平均操作（EMA，exponential moving average），EMA对模型的参数做平均，从而提高性能并增加模型鲁棒性。在损失函数方面，使用了久经考验的生成领域“交叉熵”—感知损失（perceptual loss）以及回归损失来约束VAE的训练过程。

与此同时，VAE的缩放系数也产生了变化。VAE在将Latent特征送入U-Net之前，需要对Latent特征进行缩放让其标准差尽量为1，之前的Stable Diffusion系列采用的缩放系数为0.18215，由于Stable Diffusion XL的VAE进行了全面的重训练，所以缩放系数重新设置为0.13025。

CLIP模型主要包含Text Encoder和Image Encoder两个模块，在Stable Diffusion XL中，和之前的Stable Diffusion系列一样，只使用Text Encoder模块从文本信息中提取Text Embeddings。

不过Stable Diffusion XL与之前的系列相比，使用了两个CLIP Text Encoder，分别是OpenCLIP ViT-bigG（1.39G）和OpenAI CLIP ViT-L（246M），从而大大增强了Stable Diffusion XL对文本的提取和理解能力。

与传统深度学习中的模型融合类似，Stable Diffusion XL分别提取两个Text Encoder的倒数第二层特征，并进行concat操作作为文本条件（Text Conditioning）。其中OpenCLIP ViT-bigG的特征维度为77x1280，而CLIP ViT-L的特征维度是77x768，所以输入总的特征维度是77x2048（77是最大的token数），再通过Cross Attention模块将文本信息传入Stable Diffusion XL的训练过程与推理过程中。和Stable Diffusion一致的是，Stable Diffusion XL输入的最大Token数依旧是77，当输入text的Token数量超过77后，将通过Clip操作拉回77；如果Token数不足77则会padding操作得到77x2048

Refiner模型和Base模型一样是基于Latent的扩散模型，也采用了Encoder-Decoder结构，和U-Net兼容同一个VAE模型，不过Refiner模型的Text Encoder只使用了OpenCLIP ViT-bigG。在Stable Diffusion XL推理阶段，输入一个prompt，通过VAE和U-Net（Base）模型生成Latent特征，接着给这个Latent特征加一定的噪音，在此基础上，再使用Refiner模型进行去噪，以提升图像的整体质量与局部细节。

【一】图像尺寸条件化

之前在Stable Diffusion的训练过程中，主要分成两个阶段，一个是在256x256的图像尺寸上进行预训练，然后在512x512的图像尺寸上继续训练。

而这两个阶段的训练过程都要对最小图像尺寸进行约束。第一阶段中，会将尺寸小于256x256的图像舍弃；同样的，在第二阶段，会将尺寸小于512x512的图像筛除。这样的约束会导致训练数据中的大量数据被丢弃，从而很可能导致模型性能和泛化性的降低。

Stable Diffusion XL为了在解决数据集利用率问题的同时不引入噪声伪影，将U-Net（Base）模型与原始图像分辨率相关联，核心思想是将输入图像的原始高度和宽度作为额外的条件嵌入U-Net模型中 。height和width都使用傅里叶特征编码进行独立嵌入，然后将特征concat后加在Time Embedding上，将图像尺寸引入训练过程。这样以来，模型在训练过程中能够学习到图像的原始分辨率信息，从而在推理生成阶段更好地适应不同尺寸的图像生成，而不会产生噪声伪影的问题。

【二】图像裁剪参数条件化

之前的Stable Diffusion系列模型，由于需要输入固定的图像尺寸用作训练，很多数据在预处理阶段会被裁剪。生成式模型中典型的预处理方式是先调整图像尺寸，使得最短边与目标尺寸匹配，然后再沿较长边对图像进行随机裁剪或者中心裁剪。虽然裁剪是一种数据增强方法，但是训练中对图像裁剪导致的图像特征丢失，可能会导致模型在图像生成阶段出现不符合训练数据分布的特征。

Stable Diffusion XL使用了一种简单而有效的条件化方法，即图像裁剪参数条件化策略。其主要思想是在加载数据时，将左上角的裁剪坐标通过傅里叶编码并嵌入U-Net（Base）模型中，并与原始图像尺寸一起作为额外的条件嵌入U-Net模型，从而在训练过程中让模型学习到对“图像裁剪”的认识。

【三】多尺度训练

Stable Diffusion XL采用了多尺度训练策略，这个在传统深度学习时代头牌模型YOLO系列中常用的增强模型鲁棒性与泛化性策略，终于在AIGC领域应用并固化了，并且Stable Diffusion XL在多尺度的技术上，增加了分桶策略。

Stable Diffusion XL首先在256x256和512x512的图像尺寸上分别预训练600000步和200000步（batch size = 2048），总的数据量约等于 （600000 + 200000） x 2048 = 16.384亿。

接着Stable Diffusion XL在1024x1024的图像尺寸上采用多尺度方案来进行微调，并将数据分成不同纵横比的桶（bucket），并且尽可能保持每个桶的像素数接近1024×1024，同时相邻的bucket之间height或者width一般相差64像素左右。在训练过程中，一个Batch从一个桶里的图像采样，并且我们在每个训练步骤中在不同的桶大小之间交替切换。除此之外，aspect ratio也会作为条件嵌入到U-Net（Base）模型中，让模型能够更好地学习到“多尺度特征”。完成了多尺度微调后，Stable Diffusion XL就可以进行不同aspect ratio= Height / Width的AI绘画了。

【四】使用Offset Noise

在SDXL进行微调时，使用了Offset Noise操作，能够让SDXL生成的图像有更高的色彩自由度（纯黑或者纯白背景的图像）。SD v1和SD v2一般只能生成中等亮度的图片，即生成平均值相对接近 0.5 的图像（全黑图像为 0，全白图像为 1），之所以会出现这个问题，是因为SD模型训练和推理过程的不一致造成的。

SD模型在训练中进行noise scheduler流程并不能将图像完全变成随机高斯噪声，但是推理过程中，SD模型是从一个随机高斯噪声开始生成的，因此就会存在训练与推理的噪声处理过程不一致。

Offset Noise操作是解决这个问题的一种直观并且有效的方法，我们只需要在SD模型的微调训练时，把额外在从高斯分布中采样的偏置噪声引入图片添加噪声的过程中，这样就对图像的色彩均值造成了破坏，从而提高了SDXL生成图像的“泛化性能”。

SDXL Turbo

SDXL Turbo模型是在SDXL 1.0模型的基础上设计了全新的蒸馏训练方案（Adversarial Diffusion Distillation，ADD），让模型只需要1-4步就能够生成高质量图像。StabilityAI官方发布的报告中表示SDXL Turbo在推理速度上提供了重大改进。在A100上，SDXL Turbo以207ms的速度生成一张512x512的图像（prompt encoding + a single denoising step + decoding, fp16），其中U-Net部分耗时占用了67ms。

SDXL Turbo有上面的效果，其主要还是得益于最新的Adversarial Diffusion Distillation（ADD）的蒸馏方案，其整体架构如下图所示：

ADD蒸馏算法中主要通过优化两个loss来对SDXL Turbo进行训练：

1. adversarial loss：借鉴了GAN的思想，设计了SDXL的adversarial loss，通过一个Discriminator来辨别生成的图像和真实的图像，以确保即使在一个或两个采样步骤的低步骤状态下也能确保高图像保真度，同时避免了其他蒸馏方法中常见的失真或模糊问题。
2. distillation loss：使用分数蒸馏来让SDXL 1.0模型作为教师模型，让student的输出和teacher的输出一致。