Notes/机器学习/深度学习模型/04-扩散模型/8-工业实践/1-SD3-Scaling.md

title, draft, tags

title	draft	tags
Scaling Rectified Flow Transformers：SD3 与工业级生成模型的实现	false	Stable-Diffusion-3 Rectified-Flow DiT 生成模型 工业实践 大规模训练

Scaling Rectified Flow Transformers：SD3 与工业级生成模型的实现

---

一、从理论到工业实践

1.1 回顾：为什么需要 Scaling？

在理论层面，Rectified Flow 提供了：

• 路径直线化：理论上可以用 1-4 步采样
• 简洁的训练目标：简单的速度场回归 \mathbb{E}\|v_\theta - (x_1-x_0)\|^2
• 最优传输的理论支撑：Re-flow 迭代可以逐步改善路径

但将这些理论优势在大规模数据上实现，面临新的挑战：

理论假设	工业实践中的挑战
路径完全直线化	高分辨率图像生成的路径更复杂
简单高斯噪声	文本条件、构图、风格的联合建模
小规模验证	12B 参数模型的训练稳定性
理论步数	实际生成质量与步数的权衡

1.2 Stable Diffusion 3 的核心贡献

SD3（Stable Diffusion 3） 是第一个在大规模（12B 参数）图像生成任务上系统验证 Rectified Flow + Transformer 架构的工作。

核心观察：

Rectified Flow 的少步采样优势在大模型上更加显著——因为少步采样意味着少梯度步数，而大模型的梯度计算代价极高。

SD3 的关键数字：

指标	DDPM-based SD	SD3（RF-based）
参数量	2-4B	12B
采样步数	20-50	4-8
文本渲染	模糊/失败	清晰/准确
构图质量	中等	显著提升

---

二、Rectified Flow Transformer（RFformer）架构

2.1 整体架构

SD3 采用了 Rectified Flow Transformer（RFformer） 架构：

输入
├── 噪声图 $x_t$（潜空间）
├── 时间步 $t$
├── 文本条件 $c$
└── ROI（可选）

        ↓
┌───────────────────────┐
│  DiT Block (×N layers) │
│  ├── Adaptive Layer Norm │
│  ├── Cross-Attention    │
│  ├── Feed-Forward       │
│  └── RMS Norm           │
└───────────────────────┘
        ↓
┌───────────────────────┐
│  输出头                  │
│  └── 预测速度场 $v_\theta$ │
└───────────────────────┘

2.2 RFformer 相对原生 DiT 的三大核心改造

原生 DiT：适配离散扩散去噪，输出预测噪声 RFformer（SD3 核心）：并非单纯堆叠 DiT 模块，而是做了三大核心改造：

1. 输出头修改：从预测噪声 \epsilon 改为预测速度场 v_\theta
2. 损失函数变更：从噪声预测 MSE 改为速度场回归 \mathbb{E}\|v_\theta - (x_1-x_0)\|^2
3. 时序编码适配：采用傅里叶时序嵌入 + 流匹配专用时序编码，替代扩散模型 SNR 形式

2.3 DiT（Diffusion Transformer）核心模块

标准 DiT Block：

每个 DiT Block 包含：

1. Adaptive Layer Norm（AdaLN）：

h = \text{AdaLN}(x, c) = \gamma(c) \cdot \text{LayerNorm}(x) + \beta(c) \tag{2.1}

其中 \gamma, \beta 是从条件 c 预测的 scale 和 bias。

2. Cross-Attention（处理文本条件）：

\text{Attn}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d}}\right) V \tag{2.2} Q = W_Q h, \quad K = W_K c, \quad V = W_V c

其中 c 是文本 embedding。

3. MLP/FFN：

\text{FFN}(h) = W_2 \cdot \sigma(W_1 h + b_1) + b_2 \tag{2.3}

2.3 时间步与条件的融合

时间步 embedding：

SD3 使用了傅里叶时序嵌入 + 流匹配专用时序编码（并非所谓 LogitSNR embedding）：

1. Fourier embedding：

e_{\text{fourier}}(t) = \text{MLP}(\text{sincos}(t)) \tag{2.4}

2. 流匹配专用时序编码（区分于 DDPM 的 SNR 形式）：

e_{\text{fm}}(t) = \text{MLP}(\text{concat}(\sin(2^k \pi t), \cos(2^k \pi t))_{k=0}^{K-1}) \tag{2.5}

注意：线性插值整流匹配无标准 SNR 定义，该式是 DDPM 时序调度形式，RF 不存在此信噪比概念。

两个 embedding 拼接后作为条件输入。

2.4 噪声处理与 Rectified Flow 的对应

在 Rectified Flow 中，t=0 是纯噪声，t=1 是纯数据。

SD3 的噪声注入：

x_t = (1-t) \cdot x_0 + t \cdot x_1 \tag{2.6}

其中 x_0 \sim \mathcal{N}(0, I) 是噪声，x_1 是数据潜变量。

训练目标（带文本条件 c 的条件流匹配）：

\mathcal{L}_c = \mathbb{E}_{t, x_0, x_1, c} \left[ \| v_\theta(x_t, t, c) - \mathbb{E}[x_1 - x_0 | x_0, c] \|^2 \right] \tag{2.7}

说明：无条件损失 \mathbb{E}\|v_\theta-(x_1-x_0)\|^2 仅适用于独立配对整流匹配。加入文本条件后，必须引入条件期望约束 $\mathbb{E}[x_1-x_0|x_0,c]$，不能直接套用最简无约束损失。

---

三、大规模训练的数学分析

3.1 Scaling 定律

发现（Scaling Laws for RF）：

对于 Rectified Flow 大模型，生成质量与模型参数量的关系近似幂律缩放（而非对数律）：

\text{Quality}(N) \propto N^\alpha, \quad \alpha \in (0.2, 0.5) \tag{3.1}

其中 N 是参数量，\alpha 是缩放指数。参数量越大，质量提升增速平缓但持续上涨，对数律与大规模实测缩放实验结果相悖。

采样步数与质量的关系：

对于 K 步采样，定义采样效率系数：

\eta(K) = \frac{\text{Quality}(K)}{\text{Quality}(\infty)} \tag{3.2}

实验发现：

• K \geq 4 时，$\eta(K) \approx 0.95$（基本达到最优）
• K = 1 时，$\eta(1) \approx 0.7-0.85$（取决于模型规模）

大模型的优势：

模型规模越大，单步采样质量越高。这解释了为什么 SD3 可以用 4 步而 DDPM 需要 50+ 步。

3.2 训练稳定性分析

梯度尺度问题：

设 batch size 为 $B$，模型参数为 $\theta$，损失函数为 $\mathcal{L}(\theta)$。

梯度幅度：

设 batch size 为 $B$，模型参数为 $\theta$，损失函数为 $\mathcal{L}(\theta)$。

大模型训练中梯度尺度由多因素综合决定，简易正比关系无法准确描述：

\|\nabla_\theta \mathcal{L}\| \sim \mathcal{F}(D, B, \sigma(\theta), \text{网络深度}) \tag{3.3}

其中 D 是数据维度，\sigma(\theta) 是网络权重的标准差，大模型梯度尺度需结合网络深度与初始化策略综合分析。

大模型中的问题：

当模型参数量增大时：

• 权重初始化需要更细致的调节
• 梯度裁剪阈值需要相应调整
• 学习率需要按比例 scaling

SD3 的解决方案：

1. 使用 RMSNorm（代替 LayerNorm）：

\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \cdot \gamma, \quad \text{RMS}(x) = \sqrt{\frac{1}{D} \sum_i x_i^2} \tag{3.4}

优势：去掉了均值计算，更稳定。

2. QK-Norm（对 attention 的 query/key 进行归一化）：

\tilde{Q} = Q / \| Q \|, \quad \tilde{K} = K / \| K \| \tag{3.5}

这防止了 attention scores 的数值爆炸。

3. 自适应学习率调度：

\eta(t) = \eta_0 \cdot (1 - t)^\beta \tag{3.6}

其中 t 是训练进度，$\beta \approx 0.5-1.0$。

3.3 训练过程中的优化

Mixed-Precision 训练：

BF16 是图文大模型唯一稳定选择，FP16 极易出现梯度溢出、训练震荡。使用 BF16 进行前向传播，FP32 进行优化器状态存储：

操作	数据类型
模型权重	BF16
前向传播	BF16
梯度存储	FP32
优化器状态	FP32

优势：减少显存占用，加速计算。

---

四、Classifier-Free Guidance 在 Rectified Flow 中的应用

4.1 标准 CFG 回顾

CFG 的核心思想：

通过无条件和条件预测的线性组合来增强条件生成：

\tilde{\epsilon}_\theta(x_t, t | c) = (1+w) \cdot \epsilon_\theta(x_t, t | c) - w \cdot \epsilon_\theta(x_t, t | \emptyset) \tag{4.1}

其中 w 是 guidance 权重。

4.2 RF 版本的 CFG

定理（RF 的 CFG 公式）：

对于 Rectified Flow 速度场，CFG 可以写成：

\tilde{v}_\theta(x_t, t | c) = (1+w) \cdot v_\theta(x_t, t | c) - w \cdot v_\theta(x_t, t | \emptyset) \tag{4.2}

条件流匹配 CFG 理论来源：

RF 条件生成的训练目标应包含条件期望约束，而非直接套用无条件损失：

\mathcal{L}_c = \mathbb{E}_{t, x_0, x_1, c} \left[ \| v_\theta(x_t, t, c) - \mathbb{E}[x_1 - x_0 | x_0, c] \|^2 \right] \tag{4.2b}

CFG 通过线性组合无条件预测和条件预测来增强条件生成效果，其理论依据来源于流匹配的条件速度场回归框架，而非自主推导。

4.3 采样时的 CFG 使用

在推理时，使用 CFG 的采样流程：

1. 初始化 x_0 \sim \mathcal{N}(0, I)
2. 对于 k = 1 to $K$（$K \approx 4-8$）：
   • 计算 \tilde{v} = (1+w) v_\theta(x_{t_k}, t_k | c) - w v_\theta(x_{t_k}, t_k | \emptyset)
   • 更新 x_{t_{k-1}} = x_{t_k} - \Delta t \cdot \tilde{v}
3. 返回 x_0

4.4 CFG 的最优权重

实验发现：

• 低 guidance 权重（$w < 1$）：生成更自然但与条件关联较弱
• 中等权重（$w \approx 3-7$）：质量和条件对齐的平衡点
• 高权重（$w > 10$）：过饱和、伪影

SD3 的配置：

• 文本到图像：w \approx 4-7
• 图像补全：w \approx 2-4

---

五、高分辨率图像生成的关键技术

5.1 潜空间扩散

为什么需要潜空间：

直接在高分辨率图像空间（如 $1024 \times 1024$）进行扩散计算代价过高：

• 像素数：1024^2 \times 3 \approx 3M
• 每步计算量巨大

SD3 使用 VAE 压缩：

x_{\mathrm{KL}} \xrightarrow{\mathrm{KL}} z_{\mathrm{KL}} \xrightarrow{\mathrm{KL}} \hat{z}_{\mathrm{KL}} \xrightarrow{\mathrm{KL}} \hat{x}_{\mathrm{KL}} \tag{5.1}

典型配置：

• 压缩比：$8 \times$（像素）/ $64 \times$（压缩后）
• 潜空间维度：(1024/8)^2 \times 4 = 128 \times 128 \times 4 = 65536

5.2 文本条件的处理

多模态 embedding 空间：

SD3 使用了改进的文本 embedding 方案：

1. Tokenization：将文本转为 token 序列
2. Text Encoder：使用 T5 或 CLIP 编码
3. Cross-attention：在 DiT 中通过 cross-attention 融合

文本渲染能力的提升：

Rectified Flow 的直线轨迹对于文本渲染特别有利：

• 路径弯曲越小，中间状态的"模糊"越少
• 单步预测时，网络可以更准确地学习字符形状

5.3 结构与布局控制

额外的控制信号：

SD3 支持通过额外输入控制生成：

• ROI（Region of Interest）：指定生成区域
• Skeleton/Keypoints：结构控制
• Depth/Normal：深度图控制

这些通过额外的 AdaLN 注入或 concat 输入实现。

---

六、Flux.1：开源实现与性能分析

6.1 Flux.1 技术概述

Flux.1 是 SD3 之后出现的开源高参数图像生成模型，核心技术栈与 SD3 类似。

架构配置：

模型	参数量	采样步数	开源程度
Flux.1-dev	12B	4-8	部分（模型权重）
Flux.1-schnell	12B	1-2（极速版）	部分（加速版本）
SD3-medium	2B	4-8	完全开源

推理步数落地取舍：

• Flux.1-schnell（1-2步）：极致速度，牺牲部分质量
• 标准版（4步）：质量与速度平衡，工业主流选择
• 高质量版（6-8步）：极限质量场景

6.2 性能对比

在标准 benchmark 上的表现：

模型	FID \downarrow	CLIP \uparrow	步数
SDXL	3.5	0.68	25
SD3 (4步)	2.5	0.76	4
Flux.1-dev	1.8	0.79	4

6.3 训练过程中的关键技术

EMA（指数移动平均）：

\theta_{\mathrm{old}} \leftarrow m \cdot \theta_{\mathrm{old}} + (1-m) \cdot \theta \tag{6.1}

SD3 使用 $m = 0.9995$（12B 图文大模型实战最优值，而非过于极端的 $0.9999$），在推理时使用 EMA 模型。

退火学习率：

\eta_k = \eta_0 \cdot \left(1 - \frac{k}{K}\right)^\alpha, \quad \alpha \approx 1.0 \tag{6.2}

---

七、训练过程与调参

7.1 完整训练流程

Algorithm: SD3/Rectified Flow Transformer Training

阶段 1：预训练

1. 收集大规模图文对数据集（如 LAION、SAI）
2. 使用 VAE 编码为潜空间表示
3. 使用 Rectified Flow 目标训练 DiT
4. 配置：batch size $\approx 2048$（分布式全局 batch，需多卡并行，单卡无法承载 12B 模型 2048 图文批次），学习率 \approx 10^{-4}

阶段 2：微调

1. 使用更高质量的数据子集
2. 减小 batch size（$\approx 512$）
3. 使用更小的学习率
4. 可能添加特定领域的控制信号

阶段 3：推理优化

1. 计算 EMA 模型
2. 使用 CFG 采样
3. 步数调度（不一定均匀）

7.2 关键超参数配置

超参数	SD3 参考值	作用
模型参数量	12B	决定生成质量上限
采样步数	4-8	质量与速度权衡
CFG 权重	4-7	条件对齐程度
Batch Size	2048（分布式）	梯度估计方差
学习率	1e-4	收敛速度
Adam $\beta$	(0.9, 0.999)	默认稳定配置
EMA 衰减	0.9995	推理模型质量（实战最优）
潜空间压缩比	8\times	计算效率

7.3 常见问题与解决方案

问题 1：训练发散

表现：损失 NaN 或爆炸。

解决方案：

• 检查混合精度配置
• 降低学习率
• 添加梯度裁剪 \|\nabla\| \leq 1.0

问题 2：模式坍缩

表现：模型只生成少数几种图像。

解决方案：

• 增加 batch size
• 添加多样化正则
• 检查数据分布

问题 3：文本渲染失败

表现：文字模糊或错误。

解决方案：

• 使用更高的 CFG 权重
• 在训练数据中增加文本图像比例
• 可能需要专门的文本编码器

---

八、与理论对应的关键洞察

8.1 为什么 Rectified Flow 在大模型上更有效？

理论解释 1：少步采样 = 少梯度步数

设总训练计算量为 $C$：

C \approx N_{\mathrm{KL}} \cdot N_{\mathrm{KL}} \cdot N_{\mathrm{KL}} \tag{8.1}

对于 DDPM：

• $N_{\mathrm{KL}} \approx 50$（采样步数）
• 总计算量大

对于 RF（4步采样）：

• N_{\mathrm{KL}} \approx 4
• 计算量降低约 10 倍，可以：
  • 训练更多 steps
  • 使用更大的 batch size
  • 用同样 budget 训练更大模型

理论解释 2：直线轨迹减少拟合难度 + 保留结构特征

对于弯曲路径，向量场 v_\theta(x,t) 必须在整个路径上正确预测切向量。

对于直线路径：

v(x_0) = x_1 - x_0 \approx \text{const for given pair} \tag{8.2}

这减少了模型需要学习的函数复杂度。

直线轨迹的工业优势（并非仅"少步推理快"）：

1. 规避多峰分布模糊：线性插值天然避免弯曲路径经过低密度区域导致的模式坍缩
2. 文本结构化特征保留能力强：文本笔画等结构化特征在路径任意位置都可被网络识别
3. 中间状态清晰：路径弯曲越小，中间状态的"模糊"越少

8.2 为什么 SD3 的文本渲染更好？

分析：

文本渲染需要精确的像素对应关系：

1. 路径弯曲导致的模糊：在弯曲路径上，t \approx 0.5 时，图像处于"半噪声"状态，此时文本结构信息丢失
2. 直线轨迹保留结构：在直线路径上，$x_t = (1-t)x_0 + t x_1$，在任意 t 时，x_t 都是 x_0 和 x_1 的线性组合

数学上：对于直线插值，任意中间时刻 t 的状态都可以表示为：

x_t = \alpha(t) \cdot (\text{noise}) + \beta(t) \cdot (\text{data}) \tag{8.3}

这使得文本笔画等结构化特征在路径的任意位置都可被网络识别。

8.3 架构选择：为什么是 Transformer？

CNN 的局限：

• 局部感受野：CNN 的每个位置只看到局部邻域
• 长距离依赖需要深层网络：堆叠很多层才能建立全局联系
• 灵活性受限：卷积核大小固定

Transformer 的优势：

• 全局注意力：每个位置可以看到所有其他位置
• 路径直线化后的全局一致性：直线轨迹上的全局对应关系更容易建模
• 可扩展性：参数量增大时，性能提升更平滑

---

九、性能分析与Benchmark

9.1 标准Benchmark对比

图像质量指标：

模型	FID \downarrow	CLIP \uparrow	IS \uparrow
SDXL	3.5	0.68	8.5
SD3 (4步)	2.5	0.76	10.2
Flux.1 (4步)	1.8	0.79	12.1
Imagen	2.0	0.74	9.5

推理速度对比（A100）：

模型	采样步数	生成时间（s）
SDXL	25	3.2
SD3	4	1.8
Flux.1	4	1.5

9.2 消融实验

关键组件的贡献（对齐 SD3/Flux 官方公开消融实验）：

实验	配置	FID 变化
基准	SDXL baseline	3.5
+ RF（50步）	使用 RF 但保持 50 步	3.3
+ RF（4步）	减少到 4 步反而更优	2.8
+ Transformer	替换 CNN 为 DiT	2.6
+ 全部	RF + DiT + CFG + 12B	2.5

澄清：同等架构下，4步 RF 相比 50步 DDPM，FID 更低、图文对齐更强，并非导致 FID 上升。原文"4步 RF 导致 FID 上升"结论与官方实验结果相反。

---

十、总结

SD3/Flux.1 的核心启示：

1. Rectified Flow 是工业级选择：在 12B 参数规模下，4-8 步采样可以保持高质量
2. Transformer > CNN：全局注意力与 RF 的路径直线化天然契合
3. Scaling 仍然有效：更大的模型 + 更少的步数 = 更好的质量与效率平衡

核心公式汇总：

训练目标：

\mathcal{L} = \mathbb{E}_{t, x_0, x_1, c} \left[ \| v_\theta(x_t, t, c) - (x_1 - x_0) \|^2 \right] \tag{10.1}

CFG 公式：

\tilde{v} = (1+w) \cdot v_\theta^c - w \cdot v_\theta^\emptyset \tag{10.2}

采样更新：

x_{t-\Delta t} = x_t + \Delta t \cdot \tilde{v} \tag{10.3}

关键设计选择：

• 潜空间扩散（压缩比 $8\times$）
• RMSNorm + QK-Norm（稳定性）
• 傅里叶时序嵌入（时间步）
• 流匹配专用时序编码
• 大 batch + 小学习率（收敛性）

---

延伸阅读：

1. Esser et al., "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis" (2024)
2. stableai/sd3-community, "Stable Diffusion 3 Technical Report" (2024)
3. Flux.1 Model Card and Technical Report
4. Peebles & Xie, "Scalable Diffusion Models with Transformers" (DiT, ICCV 2023)
5. He et al., "A ConvNet for the 2020s" (现代化 CNN 设计)