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进行训练：