DiT- Scalable Diffusion Models with Transformers论文笔记

• 论文：<https://arxiv.org/abs/2212.09748>

核心思想：把扩散模型的骨干网络从 U-Net 换成 Transformer，在 latent patches 上做 self-attention。

传统扩散模型的骨干：U-Net

flowchart TD
    subgraph Encoder["编码器 (下采样)"]
        E1["Conv Block 1<br>256×256, 64ch"]
        E2["Conv Block 2<br>128×128, 128ch"]
        E3["Conv Block 3<br>64×64, 256ch"]
        E4["Conv Block 4<br>32×32, 512ch"]
    end

    B["Bottleneck<br>16×16, 1024ch"]

    subgraph Decoder["解码器 (上采样)"]
        D4["Conv Block 5<br>32×32, 512ch"]
        D3["Conv Block 6<br>64×64, 256ch"]
        D2["Conv Block 7<br>128×128, 128ch"]
        D1["Conv Block 8<br>256×256, 64ch"]
    end

    Input["输入图像<br>256×256"] --> E1
    E1 -->|"MaxPool ↓2"| E2
    E2 -->|"MaxPool ↓2"| E3
    E3 -->|"MaxPool ↓2"| E4
    E4 -->|"MaxPool ↓2"| B

    B -->|"UpConv ↑2"| D4
    D4 -->|"UpConv ↑2"| D3
    D3 -->|"UpConv ↑2"| D2
    D2 -->|"UpConv ↑2"| D1
    D1 --> Output["输出<br>256×256"]

    E4 -.->|"Skip Connection"| D4
    E3 -.->|"Skip Connection"| D3
    E2 -.->|"Skip Connection"| D2
    E1 -.->|"Skip Connection"| D1

• 编码器：一路下采样，提取越来越抽象的特征

• 解码器：一路上采样，恢复分辨率

• Skip Connection：跳跃连接，保留细节信息

• 整体形状像字母 U，所以叫 U-Net

DiT 的做法：用 Transformer 替代 U-Net

flowchart TD
    A["带噪 Latent<br>(从 VAE 编码器得到)"] --> B["切成 Patches<br>(类似 ViT)"] --> C["Patch Embedding<br>+ 位置编码"]
    C --> D["Transformer Block ×N<br>(Self-Attention + FFN)"]
    T["条件输入<br>(timestep, class label)"] --> D
    D --> E["Linear Decoder<br>还原 patch → latent"]
    E --> F["去噪后的 Latent"]
    F --> G["VAE 解码器 → 图像"]

DiT vs U-Net 对比

	U-Net（CNN）	DiT（Transformer）
感受野	局部，靠堆层数扩大	全局，self-attention 天然看全图
可扩展性	较差，加大模型收益递减	好，scaling law 明显
评估指标	FID 2~3	FID 2.27（ImageNet 256×256 SOTA）
核心优势	成熟稳定	模型越大效果越好，scalability 强

Diffusion 数学公式拆解

式子 1：前向加噪过程

$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)\mathbf{I})$

$\mathcal{N}(x_t;\mu ,{\sigma }^2)$ = ” $x_t$ 服从均值为$\mu$、方差为 ${\sigma }^2$ 的高斯分布” ✅

给定原始干净图像 $x_0$，第 $t$ 步的带噪图像 $x_t$ 服从一个高斯分布：

• 均值 = $\sqrt{{\bar{\alpha }}_t}\cdot x_0$（原图被缩小了一点）

• 方差 = $(1-{\bar{\alpha }}_t)\cdot \mathbf{I}$（加了一些噪声）

${\bar{\alpha }}_t$ 是一个从 1 逐渐减小到 ~0 的系数：

• t 很小时： ${\bar{\alpha }}_t\approx 1$ → 均值 ≈ $x_0$，方差 ≈ 0 → 图像几乎没变

• t 很大时： ${\bar{\alpha }}_t\approx 0$ → 均值 ≈ 0，方差 ≈ 1 → 几乎是纯噪声

直觉： ${\bar{\alpha }}_t$ 控制”原图还剩多少”， $(1-{\bar{\alpha }}_t)$ 控制”噪声加了多少”，两者加起来恒为 1。

# 式子 1：采样噪声
epsilon = torch.randn_like(x0)  # 采样 ε ~ N(0, I)，shape 跟 x0 一样

• 📌 ${\bar{\alpha }}_t$ 和 ${\beta }_t$ 是怎么来的？

  先定义一组 ${\beta }_t$（每一步加多少噪声），然后递推：

  ${\alpha }_t=1-{\beta }_t$

  ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s={\alpha }_1\cdot {\alpha }_2\cdots {\alpha }_t$

  ${\bar{\alpha }}_t$ 是所有 $\alpha$ 的累积乘积。因为每个 ${\alpha }_t<1$，越乘越小。

  举个具体例子（线性 schedule，总步数 $T=1000$， $\beta$ 从 0.0001 线性增长到 0.02）：

  $t$	${\beta }_t$	${\alpha }_t=1-{\beta }_t$	${\bar{\alpha }}_t$ （累乘）
  1	0.0001	0.9999	≈ 0.9999
  100	0.003	0.997	≈ 0.74
  500	0.01	0.99	≈ 0.05
  1000	0.02	0.98	≈ 0.00003

  常见的 noise schedule：

  • Linear： ${\beta }_t$ 线性增长（DDPM 原版用的）

  • Cosine： ${\bar{\alpha }}_t$ 按余弦曲线衰减，前期衰减慢、后期快，效果通常更好（Improved DDPM 用的）

  ${\bar{\alpha }}_t$ 是超参数，不是模型学出来的。训练前把整条 schedule 算好，存成一个长度为 $T$ 的数组，训练时直接查表用 ✅

---

式子 2：重参数化采样

$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t,{ϵ}_t\sim \mathcal{N}(0,\mathbf{I})$

式子 1 的等价操作形式 —— 直接告诉你怎么算：

1. 先采一个跟图像同样大小的随机噪声 ${ϵ}_t$

2. 用 $\sqrt{{\bar{\alpha }}_t}$ 缩放原图（信号变弱）

3. 用 $\sqrt{1-{\bar{\alpha }}_t}$ 缩放噪声（噪声变强）

4. 两个加起来 = 第 $t$ 步的带噪图像

这就是重参数化技巧（reparameterization trick）：把”从分布采样”变成”确定性运算 + 一个固定的随机变量”，这样就能对 $x_0$ 求梯度了（反向传播需要）。

# 式子 2：加噪 = 缩放原图 + 缩放噪声
x_t = sqrt_alpha_bar_t * x0 + sqrt_one_minus_alpha_bar_t * epsilon

---

详细训练过程（原文拆解）

Diffusion 模型训练的是一个反转前向加噪过程的反向过程：

$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}({\mu }_{\theta }(x_t),{\Sigma }_{\theta }(x_t))$

神经网络用来预测 $p_{\theta }$ 的统计量（均值 ${\mu }_{\theta }$ 和方差 ${\Sigma }_{\theta }$）。

训练目标的推导：

反向过程模型通过 $x_0$ 的对数似然的 变分下界（variational lower bound, VLB） 来训练，简化后为：

$\mathcal{L}(\theta )=-p(x_0|x_1)+{\sum }_t{\mathcal{D}}_{KL}(q^{\ast }(x_{t-1}|x_t,x_0)\Vert p_{\theta }(x_{t-1}|x_t))$

其中 $q^{\ast }(x_{t-1}|x_t,x_0)$ 是前向过程的后验分布（已知 $x_0$ 和 $x_t$ 时， $x_{t-1}$ 的真实分布）。因为 $q^{\ast }$ 和 $p_{\theta }$ 都是高斯分布，它们的 KL 散度可以用均值和方差直接计算。

• 📝 VLB 是怎么推导出来的？

  目标：最大化数据的对数似然 $\log p_{\theta }(x_0)$，但它没法直接算。

  Step 1：引入边缘化

  把所有中间步骤 $x_{1:T}$ 引入：

  $log{p}_{\theta }(x_0)=\log \int p_{\theta }(x_{0:T})d{x}_{1:T}$

  这个积分维度太高，算不了。

  Step 2：用 Jensen 不等式取下界

  引入前向过程 $q(x_{1:T}|x_0)$ 作为辅助分布：

  $\log p_{\theta }(x_0)=\log \int \frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}\cdot q(x_{1:T}|x_0)d{x}_{1:T}$

  由 Jensen 不等式（ $\log$ 是凹函数）：

  $\log p_{\theta }(x_0)\ge {\mathbb{E}}_q\left[\log \frac{p_{\theta }(x_{0:T})}{q(x_{1:T}|x_0)}\right]=\text{VLB}$

  这就是变分下界！因为 Jensen 不等式“压低”了真实值，所以是 lower bound。

  Step 3：展开并分解

  把 $p_{\theta }(x_{0:T})$ 和 $q(x_{1:T}|x_0)$ 分别展开为每一步的乘积：

  $p_{\theta }(x_{0:T})=p(x_T){\prod }_{t=1}^T{p}_{\theta }(x_{t-1}|x_t)$

  $q(x_{1:T}|x_0)={\prod }_{t=1}^{T}q(x_t|x_{t-1})$

  代入 VLB 并整理，最终得到：

  $\text{VLB}=\underset{\text{重建项}}{\underbrace{{\mathbb{E}}_q[\log p_{\theta }(x_0|x_1)]}}-\underset{\text{每步的 KL 散度}}{\underbrace{\sum _{t=2}^T{\mathcal{D}}_{KL}(q(x_{t-1}|x_t,x_0)\Vert p_{\theta }(x_{t-1}|x_t))}}-\underset{\text{常数，可忽略}}{\underbrace{{\mathcal{D}}_{KL}(q(x_T|x_0)\Vert p(x_T))}}$

  每项的含义：

  • 重建项 $\log p_{\theta }(x_0|x_1)$：模型从 $x_1$ 还原 $x_0$ 的能力

  • KL 散度项 ${\mathcal{D}}_{KL}(q\Vert p_{\theta })$：模型预测的 $p_{\theta }(x_{t-1}|x_t)$ 和真实后验 $q(x_{t-1}|x_t,x_0)$ 的差距

  • 最后一项： $q(x_T|x_0)$ 和 $p(x_T)=\mathcal{N}(0,\mathbf{I})$ 的差距， $T$ 足够大时这项趋近 0

  Step 4：从 KL 到 MSE

  因为 $q(x_{t-1}|x_t,x_0)$ 和 $p_{\theta }(x_{t-1}|x_t)$ 都是高斯分布，它们的 KL 散度可以用均值差的平方表示。再把均值 ${\mu }_{\theta }$ 重参数化为噪声预测 ${ϵ}_{\theta }$，就得到了：

  ${\mathcal{L}}_{simple}=||{ϵ}_{\theta }(x_t)-{ϵ}_t|{|}_2^2$

  从对数似然 → Jensen 不等式取下界 → 展开为 KL 散度之和 → 高斯 KL 简化为均值差 → 重参数化为噪声预测 MSE。这就是从 VLB 到 ${\mathcal{L}}_{simple}$ 的完整链路 ✅

从 KL 散度到 MSE：

通过把 ${\mu }_{\theta }$ 重参数化为噪声预测网络 ${ϵ}_{\theta }$，模型可以用简单的 MSE 训练：

${\mathcal{L}}_{simple}(\theta )=||{ϵ}_{\theta }(x_t)-{ϵ}_t|{|}_2^2$

两个输出头的训练策略（沿用 Improved DDPM [36] 的做法）：

• Noise 头 ${ϵ}_{\theta }$：用 ${\mathcal{L}}_{simple}$（MSE）训练 → 更稳定

• 方差头 ${\Sigma }_{\theta }$：用完整的 $\mathcal{L}$（包含 ${\mathcal{D}}_{KL}$ 项）训练 → 学习每步的不确定性

推理（采样生成）：

训练完成后，从纯噪声开始采样：

1. 初始化 $x_{t_{max}}\sim \mathcal{N}(0,\mathbf{I})$

2. 逐步采样 $x_{t-1}\sim p_{\theta }(x_{t-1}|x_t)$（同样用重参数化技巧）

3. 重复直到 $x_0$

---

式子 3：反向去噪过程

$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}({\mu }_{\theta }(x_t),{\Sigma }_{\theta }(x_t))$

模型要学的就是这个 —— 给你一个噪声图 $x_t$，预测”退回一步”后 $x_{t-1}$ 的分布：

• ${\mu }_{\theta }$：神经网络预测的均值（“我觉得去掉噪声之后应该长这样”）

• ${\Sigma }_{\theta }$：神经网络预测的方差（“我有多确定”）

这就是 DiT 输出的 Noise + Σ 的来源！

---

式子 4：训练损失（简化版）

${\mathcal{L}}_{simple}(\theta )=||{ϵ}_{\theta }(x_t)-{ϵ}_t|{|}_2^2$

训练目标超级简单：

1. 取一张干净图 $x_0$，随机采噪声 ${ϵ}_t$，用式子 2 算出 $x_t$

2. 把 $x_t$ 喂给模型，让它猜”加了什么噪声” → 输出 ${ϵ}_{\theta }(x_t)$

3. 跟真实噪声 ${ϵ}_t$ 算 MSE（均方误差）

4. 反向传播，更新参数

本质上就是一个预测噪声 → 跟真实噪声对比 → 缩小差距的回归任务。

---

串起来看

训练时：x₀ → 加噪得到 xₜ → 模型预测噪声 εθ → 跟真实 εₜ 算 loss → 更新参数
推理时：从纯噪声 x_T 开始 → 模型预测噪声 → 减掉 → 得到 x_{T-1} → 重复 → x₀

概率论的部分说白了就是：用高斯分布建模每一步的”不确定性”，让加噪和去噪都有数学保证。但实际训练就是 MSE 回归，没那么玄 ✅