【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 guidancelatent_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 MLtimestep = 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-1latents_dtype = latents.dtypelatents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
- 最后,将latents经过vae的解码器,即可生成最终的图片。
latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factorimage = 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 时刻数据 zt 由真实图像 x0 变换成纯噪声 ϵ \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)。
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) |