基于轻量级 DiT (Diffusion Transformer) 在 MNIST 上训练 DDPM 和 Flow Matching，深入可视化两种生成范式的训练过程、采样机制与数学本质。
- 1. 项目概述
- 2. 模型架构：轻量级 DiT
- 3. DDPM：去噪扩散概率模型
- 4. Flow Matching：流匹配
- 5. DDPM vs Flow Matching：核心对比
- 6. 关键发现与教训
- 7. 运行方式
本项目的目标不是追求生成质量，而是深入理解两种主流生成范式的数学本质：
- DDPM（Denoising Diffusion Probabilistic Model）：通过 1000 步随机马尔科夫链，逐步将噪声转化为数据
- Flow Matching：通过确定性 ODE 沿直线路径，50 步将噪声映射到数据
使用 MNIST 手写数字（28×28, padded to 32×32）作为实验数据，共享同一个 DiT backbone 以公平对比。
dit/
├── models.py # DiT backbone + VAE/DDPM/FlowMatching 模型定义
├── train.py # 训练脚本
├── visualize_ddpm.py # DDPM 可视化（5 种视图）
├── visualize_fm.py # Flow Matching 可视化（2 种视图）
└── outputs/ # 所有生成的图片和模型权重
DiT（Diffusion Transformer）是将 Transformer 用于扩散模型的架构。本项目使用一个轻量版本：
输入图片 (1, 32, 32)
↓
PatchEmbed: 4×4 patch → 64 个 token, 每个 128 维
↓
+ Positional Embedding (可学习)
↓
DiTBlock × 3:
├── adaLN-Zero 调制（时间条件注入）
├── Multi-Head Self-Attention (4 heads)
└── MLP (128 → 512 → 128)
↓
Linear Head → unpatchify → 输出图片 (1, 32, 32)
关键设计：adaLN-Zero（自适应层归一化）将时间步 t 通过正弦编码 + MLP 转化为 6 个调制参数 (γ₁, β₁, α₁, γ₂, β₂, α₂)，控制每个 Transformer block 的行为。这使得同一个网络可以在不同噪声水平下做出不同响应。
参数量：DDPM 和 Flow Matching 各约 936K 参数，对于 MNIST 足够。
DDPM 的核心思想：定义一个前向过程（加噪）和一个反向过程（去噪），让模型学会反转噪声。
给定真实图片 x₀，逐步加入高斯噪声：
q(xₜ | xₜ₋₁) = N(xₜ; √(1-βₜ) · xₜ₋₁, βₜI)
其中 βₜ 是噪声调度表，从 β₁=0.0001 线性增长到 β_T=0.02。
由于高斯分布的叠加性，可以直接从 x₀ 跳到任意 t 步：
q(xₜ | x₀) = N(xₜ; √ᾱₜ · x₀, (1-ᾱₜ)I)
其中 αₜ = 1 - βₜ，ᾱₜ = ∏ᵢ₌₁ᵗ αᵢ（累积乘积）。
物理意义：
- 均值 √ᾱₜ · x₀：原始信号被衰减，t 越大信号越弱
- 方差 1-ᾱₜ：噪声占比，t 越大越接近 1
- 终态 t→T：ᾱ_T ≈ 0，所以 x_T ≈ N(0, I)，变成标准高斯噪声
从纯噪声 x_T ~ N(0,I) 出发，逐步去噪：
p_θ(xₜ₋₁ | xₜ) = N(xₜ₋₁; μ_θ(xₜ, t), βₜI)
其中模型预测噪声 ε_θ(xₜ, t)，然后通过公式计算均值：
μ_θ = (1/√αₜ) · (xₜ - βₜ/√(1-ᾱₜ) · ε_θ(xₜ, t))
采样时：
xₜ₋₁ = μ_θ + √βₜ · z, z ~ N(0, I) (t > 0)
xₜ₋₁ = μ_θ (t = 0)
MNIST 原始数据 x ∈ [0, 1]，但 DDPM 的前向过程假设终态分布为 N(0, I)。如果数据不归一化到 [-1, 1]：
- 加噪后终态近似 N(0.1√ᾱ_T, 1-ᾱ_T) ≈ N(0.003, 1)，有微小偏移
- 采样起点 N(0, 1) 与此不完全匹配
- 微小误差在 1000 步反向传播中被放大，导致生成完全失败（黑色/碎片图像）
解决方案：训练时 x = x * 2 - 1 映射到 [-1, 1]，采样后 (x * 0.5 + 0.5).clamp(0, 1) 映射回 [0, 1]。
DDPM 训练 400 epochs，可以从不同 epoch 的采样结果看到质量逐步提升：
| Epoch 1 | Epoch 100 | Epoch 200 | Epoch 400 |
|---|---|---|---|
| 纯噪声 | 初步轮廓 | 清晰数字 | 精细数字 |
损失曲线：DDPM（左）和 Flow Matching（右）独立子图，因为两者的 loss scale 不同。
400 epochs DDPM vs 200 epochs Flow Matching 的最终生成质量对比。
思路：将 1024 维的图片看成高维空间中的点，用 PCA 降到 2D，观察前向（加噪）和反向（去噪）轨迹的全局走向。
观察要点：
- 前向轨迹（实线）：从数据区域的各个角落出发，逐步收敛到噪声区域（原点附近），对应信号逐步消失
- 反向轨迹（虚线）：从噪声区域出发，逐步发散到数据区域的不同位置，对应模型"想象"出不同数字
- 高斯圆：在关键时间步绘制 q(xₜ|x₀) 的等效圆（圆心 = √ᾱₜ·x₀，半径 ∝ √(1-ᾱₜ)），可见随 t 增大，圆心向原点移动，半径增大
- 底部图片条：展示 xₜ 的实际外观，从清晰到模糊到纯噪声
思路：V1 只画了一张图片的高斯圆。V2 将 4 张不同图片的高斯圆全部叠加在同一坐标系中，同色代表同一张图片。
观察要点：
- 不同图片在 t=0 时高斯圆互不重叠（数据点分布在不同位置）
- 随 t 增大，所有圆的圆心都向原点移动，半径增大
- 到 t=999 时，所有圆几乎完全重合——不同图片变得"无法区分"，都退化为标准高斯噪声
- 这就是信息丢失的过程：每一步噪声都在抹去原始图片的信息
思路：将每张图片独立成一行，避免不同图片的轨迹和高斯圆互相遮挡。
观察要点：
- 每张图片独立展示前向轨迹、反向轨迹、高斯圆、图片条
- 可以清楚看到不同数字的轨迹形状差异
- 高斯圆的收缩/膨胀过程在每张图片上独立可见
思路：将 DDPM 展示为一条马尔科夫链——每一步只依赖上一步的状态。用图片节点 + 箭头 + 条件分布标注来表达链结构。
观察要点：
- 上行（前向）：q(xₜ|xₜ₋₁) 箭头，图片从清晰到噪声，ᾱₜ 值逐步减小
- 下行（反向）：p(xₜ₋₁|xₜ) 箭头，图片从噪声到清晰，模型预测去噪方向
- 节点边框颜色表示噪声程度：绿色（低噪声）→ 红色（高噪声）
- 马尔科夫性：每一步只看当前状态，不需要记住历史
思路：V4 展示了链结构，但没有展示每步内部的具体力学。V5 在 PCA 空间中画出每步的两个动作：
- 确定性均值偏移（灰色实线箭头）：xₜ₋₁ 被缩放为 μₜ = √(1-βₜ) · xₜ₋₁
- 随机噪声采样（彩色虚线箭头）：从 N(μₜ, βₜI) 中采样得到 xₜ
观察要点：
-
上图（前向过程，红色）：
- 虚线圆 = 当前步的条件高斯分布（以 μ 为圆心）
+号 = 均值 μ（确定性缩放的结果）- 圆点 = 实际采样点 xₜ（随机采样的结果）
- 早期步（t 小）：圆很小，采样点紧贴均值，噪声影响小
- 后期步（t 大）：圆很大，采样范围大，噪声主导
- 圆圈之间有明显重叠，这就是高斯分布逐步膨胀的过程
-
下图（反向过程，蓝色）：
- 模型从纯噪声出发
- 每一步预测去噪均值（
+号），然后在其附近采样 - 早期步（t 大）：圆大，探索范围大
- 后期步（t 小）：圆小，精细调整
- 这就是从混沌中逐步聚焦到数据的过程
Flow Matching 采用了完全不同的思路：用一个确定性 ODE（常微分方程）定义从噪声到数据的映射。
训练时，定义从噪声 ε ~ N(0,I) 到数据 x₀ 的线性插值路径：
xₜ = (1-t) · ε + t · x₀, t ∈ [0, 1]
对 t 求导，得到速度场：
v = dxₜ/dt = x₀ - ε
关键特性：
- 速度 v 是常数，不随 t 变化——因为路径是直线
- 每个粒子从 ε 笔直飞向 x₀，这是最优传输路径
- 模型学习的目标是：给定 (xₜ, t)，预测 v
速度场和概率密度通过连续性方程联系：
∂ρ/∂t + ∇·(ρv) = 0
这和流体力学的守恒定律一模一样：概率密度的变化完全由速度场的"搬运"决定，没有扩散项。
在时刻 t，条件分布的均值和采样点有精确关系：
μₜ = t · x₀ ← 均值：从原点 0 直线走向 x₀
xₜ = μₜ + (1-t) · ε ← 采样点：偏离均值 (1-t)·ε
- μₜ 轨迹：从 0 → x₀ 的直线
- xₜ 轨迹：从 ε → x₀ 的直线
- 两条线在 t=1 时交汇于 x₀
- 偏移 (1-t)·ε 的方向始终不变（同一个 ε），长度线性收缩
这种"所有偏移连线平行"的现象是线性插值的几何本质——每条轨迹用的是同一个 ε，不像 DDPM 每步重新采样噪声。
从 x₀ = ε ~ N(0,I) 出发，用欧拉法积分 ODE：
x_{t+Δt} = xₜ + v_θ(xₜ, t) · Δt
50 步就能得到高质量结果，因为路径接近直线，大步不会偏离。
Flow Matching 训练 200 epochs（DDPM 需要 400）。训练目标是回归速度：
L = ||v_θ(xₜ, t) - (x₀ - ε)||²
同样需要 [-1, 1] 归一化，原因和 DDPM 一致——需要 t=0 时分布以 0 为中心。
思路：在 PCA 空间中同时展示三层几何结构——样本轨迹、均值轨迹、收缩高斯分布，以及模型学到的速度场。
上图——几何关系：
- 实线 = 样本轨迹 xₜ = (1-t)ε + t·x₀，从噪声到图片
- 虚线 = 均值轨迹 μₜ = t·x₀，从原点到图片
- 点线 = 偏移连线 μₜ → xₜ = (1-t)·ε
- 关键观察：所有点线互相平行（因为方向都是 ε），长度随 t 线性缩短
- 这直接体现了 Flow Matching 用同一个 ε 贯穿整条轨迹的本质
- 虚线圆 = 以 μₜ 为圆心的高斯分布，半径 ∝ (1-t)
- t=0 时最大：标准高斯噪声分布
- t=1 时收缩为点：确定性地到达 x₀
下图——速度场：
- 每个采样点处的彩色箭头 = 模型预测的 v_θ(xₜ, t)
- 黑色箭头 = 真实速度 v = x₀ - ε
- 关键观察：所有箭头方向高度一致——验证了 v 是常数，路径是直线
- 这解释了为什么 50 步就够：沿直线走，不需要小步修正方向
思路：与 DDPM V3 类似，每张图片独立一行，避免互相遮挡。左列为几何关系，右列为速度场 + 图片变化。
左列（几何关系）：
- 每张图片独立展示轨迹 + 均值线 + 偏移连线 + 高斯圆
- 可以清楚看到每张图片的高斯圆从大到小的收缩过程
- 偏移连线从长到短，但方向始终一致
右列（速度场 + 图片）：
- 模型预测的 v_θ 箭头沿轨迹分布
- 底部小图展示从噪声到清晰数字的渐变过程
- 不同数字的速度方向不同，但每张图片内的箭头方向一致
| 维度 | DDPM | Flow Matching |
|---|---|---|
| 演化方程 | Fokker-Planck（带扩散项） | 连续性方程（无扩散项） |
| 采样过程 | 随机 SDE：drift + 噪声 | 确定性 ODE：dx/dt = v |
| 单步操作 | 均值偏移 + 随机采样（两个动作） | 沿速度场走一步（一个动作） |
| 每步噪声 | 重新采样 εₜ（每步不同） | 同一个 ε 贯穿始终 |
| 轨迹形状 | 锯齿形随机游走 | 笔直的直线 |
| 偏移连线 | 方向随机跳变 | 始终平行（方向 = ε） |
| 采样步数 | 1000 步 | 50 步 |
| 训练目标 | 预测噪声 ε_θ(xₜ, t) | 预测速度 v_θ(xₜ, t) |
| 训练 epochs | 400 | 200 |
| 密度变化 | 漂移 + 扩散两个力 | 纯速度场搬运 |
一句话总结：DDPM 像醉汉回家——知道大方向但每步都在晃；Flow Matching 像导航——每个粒子拿到一条直线航路，直飞目的地。
DDPM 和 Flow Matching 都需要将 [0, 1] 数据映射到 [-1, 1]。不做这一步，训练 loss 看起来正常（因为训练时 xₜ 的分布是对的），但采样从 N(0,1) 开始时会产生分布偏移，1000 步放大后完全失败。
这是一个隐蔽的 bug：训练 loss 无法检测到这个问题，只有看生成结果才能发现。
1024 维空间降到 2 维必然丢失大量信息。但 PCA 保留了方差最大的方向，足以展示：
- 轨迹的全局走向（数据区 ↔ 噪声区）
- 高斯分布的收缩/膨胀趋势
- DDPM 的随机游走 vs FM 的直线路径
936K 参数的 DiT 对 MNIST 来说绰绰有余。验证方式：同架构的 VAE（1.97M 参数）可以正常工作，说明 backbone 的表达能力不是瓶颈。
- Python 3.10+
- PyTorch 2.0+（支持 MPS/CUDA/CPU）
- torchvision, scikit-learn, matplotlib, tqdm
python train.py训练 DDPM (400 epochs) + Flow Matching (200 epochs)，生成过程中每 20 个 epoch 保存采样图片。
python visualize_ddpm.py生成 5 种可视化（如无模型权重会自动训练 100 epoch 的快速版本）：
ddpm_trajectory.png— V1: PCA 轨迹总览ddpm_trajectory_all_gauss.png— V2: 多图片高斯叠加ddpm_trajectory_individual.png— V3: 逐图片独立展示ddpm_markov_chain.png— V4: 马尔科夫链结构ddpm_markov_steps.png— V5: 逐步马尔科夫过程
python visualize_fm.py生成 2 种可视化：
fm_geometry.png— 几何关系总览（轨迹 + 高斯 + 速度场）fm_individual.png— 逐图片几何与速度场