【Dive Into Stable Diffusion v3.5】2：Stable Diffusion v3.5原理介绍

【Dive Into Stable Diffusion v3.5】系列博文：

第1篇：开源项目正式发布——深入探索SDv3.5模型全参/LoRA/RLHF训练
第2篇：Stable Diffusion v3.5原理介绍

1 前言
- 1.1 扩散模型的原理
- 1.2 损失函数
- 1.3 加噪流程
- 1.4 推理流程
- 1.5 negative prompts
2 SD3原理介绍
- 2.1 整流模型介绍
- 2.2 非均匀训练噪声采样
- 2.3 网络整体架构
- - 2.3.1 多模态 DiT (MM-DiT)
  - 2.3.2 比例可变的位置编码
  - 2.3.3 训练数据预处理
3 SD3.5原理介绍
- 3.1 QK Normalization（查询-键归一化）
- 3.2 双注意力层（Dual Attention Layers）
- 3.3 模型变体与性能对比

1 前言

开源项目地址：https://github.com/Donvink/dive-into-stable-diffusion-v3-5
如果觉得有用，别忘了点个 ⭐️ 支持开源哦！

1.1 扩散模型的原理

首先，我们用最通俗易懂的语言介绍一下扩散模型的原理。

如下图所示，有一张图片，假如我们每次都往图片上增加一点噪声，例如图中的蓝色点，我们累计加了1000步噪声之后，整张图片就接近0-1分布的白噪声了。

加噪声

对于每一次增加噪声的时候，如果我们把增加噪声后的图片当做输入，把增加前的图片当做输出，这样就可以试着让模型去学习如何复原图片了。但是学习复原图片是一件非常难的事情，那么我们是否可以换一个角度，去预测噪声呢？

答案是肯定的，如下图所示，同样，我们把增加噪声后的图片当做输入，把增加前的图片当做输出，每次让模型去预测噪声，类似于ResNet，用模型去预测残差。

预测噪声

1.2 损失函数

Loss计算相对比较简单，只要计算加权后的L2-Loss即可。

# Predict the noise residual
model_pred = transformer(
    hidden_states=noisy_model_input,
    timestep=timesteps,
    encoder_hidden_states=prompt_embeds,
    pooled_projections=pooled_prompt_embeds,
    return_dict=False,
)[0]

# Preconditioning of the model outputs.
if args.precondition_outputs:
    model_pred = model_pred * (-sigmas) + noisy_model_input

# these weighting schemes use a uniform timestep sampling
# and instead post-weight the loss
weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas)

# flow matching loss
if args.precondition_outputs:
    target = model_input
else:
    target = noise - model_input

# Compute regular loss.
loss = torch.mean(
    (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
    1,
)
loss = loss.mean()

1.3 加噪流程

添加噪声的方法有很多，我们可以均匀地添加噪声，每一步添加的噪声强度都是固定大小sigma。也可以非均匀地添加噪声，例如刚开始的时候噪声弱一点，让模型先学一些细节，越往后噪声越强，让模型学一些轮廓。

上一节我们提到，每张图片需要迭代1000步之后，最终才会变成均值为0方差为1的白噪声，那我们在训练的时候是不是需要每个图片文本对都需要迭代1000步呢，这训练的耗时会有变多大？

实际上，在初始化的时候，我们就构建好从0 ~ 999步的噪声采样schedule，对于每个图片文本对，我们都随机地从schedule中取出一个时间步，获取该时间步对应的噪声强度，添加到模型输入中，forward和backward的流程和第一节中提到的一样。代码如下：

# Sample noise that we'll add to the latents
noise = torch.randn_like(model_input)
bsz = model_input.shape[0]

# Sample a random timestep for each image
# for weighting schemes where we sample timesteps non-uniformly
u = compute_density_for_timestep_sampling(
    weighting_scheme=args.weighting_scheme,
    batch_size=bsz,
    logit_mean=args.logit_mean,
    logit_std=args.logit_std,
    mode_scale=args.mode_scale,
)
indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device)

# Add noise according to flow matching.
# zt = (1 - texp) * x + texp * z1
sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise

# Predict the noise residual
model_pred = transformer(
    hidden_states=noisy_model_input,
    timestep=timesteps,
    encoder_hidden_states=prompt_embeds,
    pooled_projections=pooled_prompt_embeds,
    return_dict=False,
)[0]

只要我们的训练集够大，迭代步数够多，那么我们就可以遍历了完整的0~999步加噪过程，同时，训练耗时也会大大降低。

1.4 推理流程

在推理的时候，我们只需要输入prompt，迭代一定步数，就能生成一张很好的图片。流程如下：

首先将prompt输入到text encoder中获取对应的文本编码Embedding；
然后生成一个随机的潜空间张量，latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)；
再将文本Embedding和latents一起输入到unet / ViT block中，迭代n步，上一步迭代生成的latents去噪后是下一次迭代的输入；

for i, t in enumerate(timesteps):
	if self.interrupt:
	    continue
	
	# expand the latents if we are doing classifier free guidance
	latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
	# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
	timestep = t.expand(latent_model_input.shape[0])
	
	noise_pred = self.transformer(
	    hidden_states=latent_model_input,
	    timestep=timestep,
	    encoder_hidden_states=prompt_embeds,
	    pooled_projections=pooled_prompt_embeds,
	    joint_attention_kwargs=self.joint_attention_kwargs,
	    return_dict=False,
	)[0]
	
	# compute the previous noisy sample x_t -> x_t-1
	latents_dtype = latents.dtype
	latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]

最后，将latents经过vae的解码器，即可生成最终的图片。

latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor

image = self.vae.decode(latents, return_dict=False)[0]
image = self.image_processor.postprocess(image, output_type=output_type)

1.5 negative prompts

有时候我们不仅会输入positive prompts，也会输入negative prompts，目的是不让生成的图片中包含negative prompts的相关信息。如何实现呢？

假设我们不想让图片中的人物有胡子，那么我们将【胡子】作为negative prompts输入到模型中，【胡子】在潜空间的特征分布是固定的，我们只需要将positive prompts生成的Latents减去negative prompts生成的Latents，就能去掉【胡子】的特征了，最终生成的图片中人物是没有胡子的。

# perform guidance
if self.do_classifier_free_guidance:
    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)

严格上来讲，Stable Diffusion通过条件嵌入与非条件嵌入的差值实现引导，而非直接对潜空间特征做减法。正向提示词（positive prompts）和反向提示词（negative prompts）的文本会被分别编码为条件嵌入向量（text_embeddings）和无条件嵌入向量（uncond_embeddings）。

反向提示词并非直接修改潜空间（Latents），而是通过调整条件概率分布影响去噪过程。例如，输入"胡子"作为反向提示词，模型会降低生成图像中出现胡子的概率分布权重。
动态调整而非固定减法：潜空间的特征分布是动态变化的，反向提示词的作用是通过梯度更新方向抑制特定特征，而非固定数值的减法。

2 SD3原理介绍

Stable Diffusion 3 (SD3) 的文章标题为 Scaling Rectified Flow Transformers for High-Resolution Image Synthesis。

这篇文章的内容很简明，就是：用整流 (rectified flow) 生成模型、Transformer 神经网络做了模型参数扩增实验，以实现高质量文生图大模型。 文章的核心贡献如下：

从方法设计上

首次在大型文生图模型上使用了整流模型。
用一种新颖的 Diffusion Transformer (DiT) 神经网络来更好地融合文本信息。
使用了各种小设计来提升模型的能力。如使用二维位置编码来实现任意分辨率的图像生成。

从实验上

开展了一场大规模、系统性的实验，以验证哪种扩散模型/整流模型的学习目标最优。
开展了扩增模型参数的实验 (scaling study)，以证明提升参数量能提升模型的效果。

2.1 整流模型介绍

所谓图像生成，其实就是让神经网络模型学习一个图像数据集所表示的分布，之后从分布里随机采样。为了直观理解，我们可以用二维点来表示一张图像的数据。比如在下图中我们希望学习红点表示的分布，即我们希望随机生成点，生成的点都落在红点处，而不是落在灰点处。

我们可以用一种巧妙的方法间接学习图像生成路线，知道了预定义的数据到噪声的路线后，我们其实就知道了数据在路线上每一位置的速度（红箭头）。那么，我们可以以每一位置的反向速度（蓝箭头）为真值，学习噪声到真实数据的速度场，这样的学习目标被称为流匹配。

流匹配

对于不同的扩散模型及流匹配模型，其本质区别在于图像到噪声的路线的定义方式。在扩散模型中，图像到噪声的路线是由一个复杂的公式表示的。而整流模型将图像到噪声的路线定义为了直线。根据论文的介绍，整流中 t 时刻数据 z_t 由真实图像 x₀ 变换成纯噪声 $\mathbf{\epsilon}$ 的位置为：

流匹配

由于整流最后学习出来的生成路线近乎是直线，这种模型在设计上就支持少步数生成。

2.2 非均匀训练噪声采样

在学习这样一种生成模型时，会先随机采样一个时刻 t ∈ [0, 1]，根据公式获取此时刻对应位置在生成路线上的速度，再让神经网络学习这个速度。刚开始和快到终点的路线很好学，而路线的中间处比较难学，因此在采样时刻 t 时，SD3 使用了一种非均匀采样分布。

如下图所示，SD3 主要考虑了两种公式: mode（左）和 logit-norm （右）。二者的共同点是中间多，两边少。mode 相比 logit-norm，在开始和结束时概率不会过分接近 0。

噪声采样

2.3 网络整体架构

和之前版本的 SD 一样，SD3 主要基于隐扩散模型（latent diffusion model, LDM），是一个两阶段的生成方法：

先用一个 LDM 生成隐空间低分辨率的图像，
再用一个自编码器把图像解码回真实图像。

LDM 会使用一个神经网络模型来对噪声图像去噪，SD3 的主要改进是把去噪模型的结构从 U-Net 变成了 DiT。

2.3.1 多模态 DiT (MM-DiT)

SD3 的去噪模型是一个 Diffusion Transformer (DiT)，如果去噪模型只有带噪图像这一种输入的话，DiT 则会是一个结构非常简单的模型，和标准 ViT 一样，但是，扩散模型中的去噪网络一定得支持带约束生成，因为扩散模型约束于去噪时刻 t 。此外，作为文生图模型，SD3 还得支持文本约束。DiT 及本文的 MM-DiT 把模型设计的重点都放在了处理额外约束上。

模型总览图

时间约束

如下图所示，SD3 的模块保留了 DiT 的设计，用自适应 LayerNorm (Adaptive LayerNorm, AdaLN) 来引入额外约束。具体来说，过了 LayerNorm 后，数据的均值、方差会根据时刻约束做调整。另外，过完 Attention 层或 FF 层后，数据也会乘上一个和约束相关的系数。

时间约束

文本约束

文本约束以两种方式输入进模型：与时刻编码拼接、在注意力层中融合。具体数据关联细节可参见下图。如图所示，为了提高 SD3 的文本理解能力，描述文本 (“Caption”) 经由三种编码器编码，得到两组数据。一组较短的数据会经由 MLP 与文本编码加到一起；另一组数据会经过线性层，输入进 Transformer 的主模块中。

文本约束

DiT 的子模块

SD3 的 DiT 的子模块结构图如下所示，和标准 DiT 子模块一样，时刻编码 y 通过修改 LayerNorm 后数据的均值、方差及部分层后的数据大小来实现约束。输入的图像编码 x 和文本编码 c 以相同的方式做了 DiT 里的 LayerNorm, FFN 等操作，但此模块用了一种特殊的融合注意力层。具体来说，在过注意力层之前，x 和 c 对应的 Q, K, V 会分别拼接到一起，而不是像之前的模型一样，Q 来自图像，K 和 V 来自文本。过完注意力层，输出的数据会再次拆开，回到原本的独立分支里。由于 Transformer 同时处理了文本、图像的多模态信息，所以作者将模型取名为 MM-DiT (Multimodal DiT)。

dit

2.3.2 比例可变的位置编码

大部分论文在使用ViT 架构时，都会把图像的图块从左上到右下编号，把二维图块拆成一维序列，再用这种一维位置编码来对待图块，这样做有一个很大的缺点：生成的图像的分辨率是无法修改的。

解决此问题的方法很简单，只需要将一维的编码改为二维编码，这样 Transformer 就不会搞混二维图块间的关系了。

位置编码

SD3 的 MM-DiT 一开始是在 256 × 256 固定分辨率上训练的，之后在高分辨率图像上训练时，开发者用了一些巧妙的位置编码设置技巧，让不同比例的高分辨率图像也能共享之前学到的这套位置编码。

2.3.3 训练数据预处理

在大规模训练前，作者用三个方式过滤了数据：

用了一个 NSFW 过滤器过滤图片，似乎主要是为了过滤色情内容；
用美学打分器过滤了美学分数太低的图片；
移除了看上去语义差不多的图片。

在训练 LDM 时，自编码器和文本编码器是不变的，因此可以提前处理好所有训练数据的图像编码和文本编码。

参考资料：https://zhouyifan.net/2024/07/14/20240703-SD3/

3 SD3.5原理介绍

对于 SD3.5-large 使用的 transformer 模型，其结构基本和 SD3-medium 里的相同，但有以下更改:

QK normalization: 对于训练大型的 Transformer 模型，使用 QK normalization 已经成为标准做法，所以 SD3.5-large 也不例外。
双注意力层: 在 MMDiT 结构中，文本和图像两个模态都在使用同一个注意力层；而 SD3.5-large 则使用了两个注意力层。

除此之外，文本编码器 (text encoder)、图像的变分自编码器 (VAE) 以及噪声调度器 (noise scheduler) 均和 SD3-medium 保持一致。

3.1 QK Normalization（查询-键归一化）

在Transformer的自注意力机制中，查询（Query, Q）和键（Key, K）矩阵的点积计算决定了不同位置之间的相关性权重。原始的缩放点积注意力通过除以来缓解点积值过大的问题，但这一操作在大型模型中可能仍不足以保证数值稳定性。

常见方法包括：

L2归一化：对Q和K的每个向量进行L2归一化，使点积结果转化为余弦相似度（范围在[-1, 1]），避免数值过大。
层归一化（LayerNorm）：在计算点积前，对Q和K矩阵分别应用层归一化，稳定分布。
缩放因子调整：结合归一化与可学习的缩放参数，增强模型灵活性。

优势：

训练稳定性：防止梯度爆炸/消失，尤其对深层或参数量大的模型（如SD3.5-large）至关重要。
收敛速度：归一化后的注意力分数分布更均匀，加速模型收敛。
鲁棒性：减少输入尺度变化对注意力机制的影响，提升泛化能力。

3.2 双注意力层（Dual Attention Layers）

在原始MMDiT（多模态扩散Transformer）中，文本和图像特征共享同一个注意力层。这种设计可能导致以下问题：

模态特性冲突：文本（序列数据）与图像（空间数据）的结构差异大，共享参数可能抑制模态特异性特征的学习。
容量限制：单一注意力层需同时适配两种模态，可能限制模型表达能力。

SD3.5-large的改进在于为文本和图像分别提供独立的注意力层：

文本注意力层：专注于序列内词与词的关系（如语法、语义依赖），采用自注意力机制。
图像注意力层：针对空间特征（如局部纹理、全局构图），可能结合空间注意力或窗口注意力机制。

实现方式：

并行处理：文本和图像特征分别输入各自的注意力层，输出后融合（如拼接或相加）。
交叉交互：在独立处理模态内特征后，通过交叉注意力层实现跨模态交互（如文本引导图像生成）。

优势：

模态特异性优化：独立参数允许模型捕捉文本和图像的独特模式。
增强表达能力：双分支结构提升多模态融合的精细度，改善生成质量（如更准确的图文对齐）。
训练效率：减少模态间干扰，加速收敛。

3.3 模型变体与性能对比

SD3.5提供三种版本，技术参数对比如下：

版本	参数量	分辨率范围	生成步数	硬件需求（显存）
SD3.5 Large	8.1B	256px-2Mpx	25步	≥16GB（A100）
SD3.5 Large Turbo	蒸馏优化	512px-1Mpx	4步	12GB（RTX 4090）
SD3.5 Medium	2.5B	256px-0.5Mpx	20步	8GB（RTX 3060）