Stability AI解释即将推出的Stable Diffusion 3技术细节

在Stability AI即将公开预览的Stable Diffusion 3（SD3）之际，官方先解释了SD3的技术细节。SD3采用专为多模态设计的MMDiT架构，能够良好地处理图像和文本Token资讯，使得模型整体理解和排版能力都获得提升。SD3在基准测试上，包括提示词遵循、排版和视觉美感，皆优于目前最先进的文本转图像模型。

文本转图像生成过程，模型需要同时考虑文本和图像两种模态，而SD3所采用的MMDiT新架构（下图），便是针对多模态数据设计。与过去Stable Diffusion的版本类似，开发团队使用预训练模型得出适当的文本和图样表示，更具体的说，SD3使用两个CLIP模型和一个T5模型来编码文本表示，并使用一个自动编码模型来编码图像Token。

SD3所使用的MMDiT架构创建在Diffusion Transformer（DiT）的基础之上。值得一提的是，DiT是OpenAI视频生成模型Sora的共同研究领导人William Peebles，在2023年于ICCV所发布的论文内容，同时也是Sora之所以能够在视频生成上有重大突破的原因，DiT结合了扩散模型和Transformer的优点。

Stability AI指出，文本和图像嵌入在概念上完全不同，因此他们对两种模态使用两组独立的权重，也就是说每种模态都有两个独立的Transformer，只不过在运行上，是将两种模态的串行连接起来执行注意力操作，这样两种表示便能在自己的空间中运行，而又同时可以考虑另一种表示。

借由这种方法，资讯可以在图像和文本Token间流动，以提高图像输出的品质，开发团队还提到，这种架构也可以扩展至视频等多种模态。由于SD3改进了提示词遵循能力，因此模型能够根据用户的具体要求，创建各种不同主题和特征的图像，同时还可保持图像风格的灵活度。

SD3使用了一种称为校正流（Rectified Flow）的公式来改善图片品质，校正流公式会在模型学习时，把数据和噪声以线性轨迹连接，在较直的推理路径中生成更加清晰的图片，并且还可减少完成的步骤。也就是让模型在训练的时候，通过逐渐将噪声添加到资讯中，模型便能够学习从具有噪声的图像中，恢复出原始干净的图像，这让模型学会更好的处理图像，去除不需要的噪声，最终生成更真实的图像。

另外，开发团队还在SD3训练过程，采用了新的轨迹采样调度，在轨迹中段给予更多的加权，使得SD3能够更有效地生成图像，特别是在较少采样步骤的时候，这种改进可以协助模型在处理复杂图像生成任务时，表现得更好。

针对SD3训练所采用的新技术，包括使用重新分配加权的校正流公式，以及MMDiT架构主干，Stability AI研究其对模型规模的影响，研究发现，无论是模型的大小还是训练的时间增加，错误率都会逐渐减少，并使得模型性能持续上升。这代表Stability AI所采用的这些技术具有发展性，未来仍有改善模型的空间。

Stability AI在消费级硬件具有24GB VRAM的RTX 4090显卡上，运行80亿参数的SD3模型，执行50个采样步骤，总共花费34秒生成1024x1024的图像。不过，Stability AI的SD3模型具有多种大小，从8亿参数到80亿参数都有，能够符合不同硬件需求。

另外，SD3在处理文本采用了一个稍具规模的文本编码器T5，拥有47亿参数，这增加了模型运行需要的内存，在不使用T5编码器的情况下可大幅减少内存，且造成很少的性能损失。虽然移除对图像美感没有影响，但是会降低文本遵循度，也就是说，不使用T5时，生成的图像和文本描述可能会有一些落差，部署模型的用户可以权衡提示词遵循度和内存的限制，灵活选择是否要移除T5编码器。