一文解读：Stable Diffusion 3究竟厉害在哪里？_stable diffusion 3 rectified flows

知乎原文：叫我Alonzo就好了

前言

最近一段时间，正当所有人都在为OpenAI发布Sora狂欢时，Stability AI更是推出了Stable Diffusion 3的技术报告。**这两项技术不约而同都采用了Diffusion Transformer的架构设计，之前我也在我的文章中进行了解读：Diffusion Transformer究竟好在哪里？感兴趣的朋友可以移步这篇文章，里面也对Stable Diffusion、Transformer、VQGAN等一些前置知识作了介绍。

值得注意的是，Stable Diffusion 3的强大性能其实并不仅限于Diffusion Transformer在架构上所带来的增益，其在提示词、图像质量、文字拼写方面的能力都得到了极大的提升。那么究竟是什么让Stable Diffusion 3如此强大？今天我们就从Stable Diffusion 3的技术报告中解读stable diffusion 3强大背后的技术原理。

研究动机

Stable Diffusion 3的研究动机其实是从数据加噪的forward路径上出发的。 作者指出，如果forward路径并非最优路径，那么就会导致扩散模型在训练过程中没有办法从输入加噪数据中完全去除噪声，那么在测试的过程中，模型就会生成一些带有artifacts的结果，例如灰度图像。
特别地，forward的选择也会影响backward路径的计算复杂度，也就是采样效率（sampling efficiency），进一步影响到图片生成的速度。
尽管我们可以整合额外采样步数，通过曲线路径来尽可能地仿真最优路径，但是曲线路径仍然存在错误累计的问题。对此，直线路径（也就是通过一步直接进行仿真）会减少错误累积的发生。
对于这个问题，有一个特别的解决方案，也就是rectified flow。recitified flow已经在小中型的实验上证明了其有效性，而在Stable Diffusion 3的设定中，模型规模从800M进一步扩大至了8B，这项技术在大规模实验上的有效性亟待证明，这也是Stable Diffusion 3这篇工作的创新点之一。
Stable Diffusion 3从模型forward的过程出发展开研究，读完intro之后的第一感觉就是不得不佩服这群人的理论功底和调研能力。forward过程的研究，属于是大家都知道可以做，但是非常难做的类型，不是仅仅做好工程问题就能搞定的。 forward过程跟backward过程往往紧密相连，一旦改动了diffusion model的forward过程，那就意味着backward过程的目标也要进行调整，只有对于diffusion model的全生命周期熟稔于心的团队才能handle这样的创新。

技术路线

技术路线的解读分为两个方面进行，一是关于flow matching的改进；另一方面是在Diffusion Transformer上的具体设计。

Simulation-Free Training of Flows

这里直接沿用了原文的章节标题，也非常言简意骇。「Simulation-Free」指的是扩散模型的模型参数不再用来直接拟合带噪数据和干净数据分布之间的映射关系，而是通过拟合一个向量场，通过这个向量场将加噪/去噪过程给联系起来。「Flows」指代的就是通过flow matching的技术来实现这件事。本质上，retified flow对diffusion model原本的运作方式进行了降维，进而提升了forward、backward过程的一致性以及采样过程的效率，与Stable Diffusion 3的研究动机相契合。

关于rectified flow会怎样影响diffusion model的损失函数，原文中给了详尽的推导过程，补充材料中也给了严谨的数学证明，这里就不再一一赘述。损失函数的最终推导形式如下：

除了rectified flow之外，文章还简单阐述了几种其他的变式的forward过程，这里也不再具体展开。

前面说到，在对forward过程进行修改过后，也要同时对backward过程进行调整。Stable Diffusion 3的做法则是提出了定制化的采样策略，希望模型在采样的过程中，在中间的步数拥有更高的权重。这里，Stable Diffusion 3提出了一个密度函数，通过对每一步的scale和bias进行控制来重新调整采样的分布，具体实现的密度函数表达式如下

对于这类型看似非常晦涩难懂的公式，有一部分工作可能是依赖大量的工程经验得来的，也就是说经过大量的实验试错，因此往往很难得到有效的定性结论。同样地，由于笔者本人水平有限，在此的相关分析及讨论也仅作参考。

值得注意的是，这里说到的中间采样部署权重过低的问题。其实在T2I-Adapter[4]这篇工作中也有明确指出：Stable diffusion在生成图片的时候，重要的采样步数主要集中在采样过程的前半段，因此在前半段加入控制才是最合适的。而Stable Diffusion 3从rectified flow的角度重新address了这个问题。

针对于cosine的噪声策略，Stable Diffusion 3也提出了相应的密度函数公式。这里值得注意的一个点是对于扩散模型关于timestep采样策略的改进，原本的扩散模型timestep是从[0,T] 的均匀分布中采样得到，而Stable Diffusion 3中不再对timestep进行均匀采样，而是从非线性分布中进行采样，同理也可以获得cosine噪声策略下的密度函数。

MM-DiT

另一方面则是关于Diffusion Transformer上的具体设计。Stable Diffusion 3同样沿用了前面两代模型的VQGAN设计，同时也沿用了CLIP作为文本编码器之一。除了上面Stable Diffusion 3与DiT的相似之处，还有一些Stable Diffusion 3中独有的创新点：

• 多个文本编码器对于「文本」和「图像」语义的解耦。 这里原文中写的是text and image embeddings，结合原文中上下文的描述后可以确定指代的是图像中的「文字」信息，以及提示词关于生成图像的语义描述信息。Stable Diffusion 3使用了多个文本编码器完成这件事，具体地，CLIP用于编码一个粗粒度的文本表征作为训练、测试过程的条件；而T5 XXL以及OpenCLIP则用于获取更细粒度的表征，以增强模型在生成具体文字时的表现，具体的做法是将各自模型提取到的表征padding对齐后，再通过Transformer中self-attention的运算融入模型训练中去。
  
• RMS-Norm。 Stable Diffusion 3的另外一个创新点就是在每一个attention运算之前，对Q和K进行了RMS-Norm归一化，用于增强模型训练的稳定性。
  
• 扩充特征维度后的压缩模型。思路其实很简单直接，就是VQGAN压缩得到的特征维度从原来的4维提升到16维，通过扩充特征维度来让压缩后的特征能包含更多信息量，进而生成模型也能学得更好，本质上也就是一个工程问题，将维度扩充之后的VQGAN训练好即可。具体来说，通道数增加之后的压缩模型的重构性能得到了显著的提升：
  
• 使用合成的caption。这点跟DALL-E 3[5]中提出训练一个image captioner，通过合成caption来进行数据增强是一致的。
• Timestep Shifting。这里的本质是指在高分辨率图像上进行训练时，采用特殊的timestep采样策略，让均匀分布采样得到的timestep能够得到一个更高数值。这一点在论文《Late-Constraint Diffusion Guidance for Controllable Image Synthesis》[6]中的timestep resampling策略中也有提到，在均匀采样得到一个timestep后，通过重采样的策略，能够获得更大的timestep值，也就是说对于干净数据输入的平均污染程度会比原来加剧，从而让模型在生成结果时去噪的效果更加鲁棒、稳定。在Stable Diffusion 3的实验设定下其实也不难理解，在对高分辨率图像进行生成时，模型需要对更多的像素点进行去噪，也就需要在训练过程中加剧噪声污染程度。
  
  上图为Stable Diffusion 3的timestep shifting：左上角为采样前timestep（横轴）和采样后timestep（纵轴）的曲线分布，右上角为不同比率下人类偏好程度，结果图中，上半部分为使用前的效果，下半部分为使用后的效果。

补充材料

补充材料中主要包括更多的实验结果、必要的数学推导、一些补充实验等等。除此之外，还有一些比较有意思的实验。

Instruction-based Image Editing

Stable Diffusion 3还提供了在InstructPix2Pix上微调后的2B版本模型，从而能够让模型基于指令对图像进行编辑。其结果展示如下：

通过Image Copy Detection去除数据冗余

另一个比较有意思的发现是扩散模型倾向于记住训练中出现的样本，因此对数据中冗余的、重复出现的数据样本进行清洗很有必要，Stable Diffusion 3也设计了具体的数据清洗方案。具体来说，Stable Diffusion 3借助SSCD的image copy generation方法来实现这一点，具体流程在原文中提供了详尽的分析及伪代码，需要的朋友可以自取。其中，检测聚类中重复图片流程的伪代码为：

检测生成图片中被记住样本的伪代码：

结语

Stable Diffusion 3除去Diffusion Transformer的架构设计之外，在forward过程的改进、sampling过程的设计、压缩模型、captions、数据清洗等多个角度展现了其独到之处，无论是从工程项目还是从学术研究的角度来看，都有许多的可取之处。Stable Diffusion 3能否在CVPR 2022年的Latent Diffusion Models那样给CV社区重新注入活力，乃至给文生视频领域更多的启发？让我们拭目以待。

---

最后，欢迎大家关注我的个人GitHub账号，我总结了一些课题的顶会paper list，有需要的小伙伴可以参考，paper list会持续更新最新发布的顶会文章。你们的一键三连、GitHub Star就是对我最大的支持！

Video Generation Paper List

Text-to-Image Generation Paper List

Image Inpainting Paper List