Flow Matching and Diffusion Models
Generative AI with Stochastic Differential Equations。从「什么是 ODE」一路推到 Stable Diffusion 3 与离散扩散语言模型,全部公式逐步推导,全部 lab 亲手实现并实测。
0. 这门课在讲什么
今天几乎所有最好的生成模型——图像的 Nano Banana 与 FLUX、视频的 VEO-3、蛋白质结构的 AlphaFold3——底层都是同一套东西:把噪声一步步变成数据。这个「一步步」在数学上就是模拟一条常微分方程(ODE)或随机微分方程(SDE)的轨迹。
这门课的全部内容,就是回答三个问题:
- (a)怎样构造这条微分方程?——我们想要的向量场 $u_t(x)$ 是一个含有对全部数据求积分的量,根本算不出来。答案是边际化技巧:用一个可以算的条件版本去替代它。
- (b)为什么可以用条件目标训练边际模型?——因为两个损失只差一个与参数无关的常数,梯度完全相同。这一条定理撑起了 flow matching 和 score matching 两套算法。
- (c)为什么 ODE 和 SDE 会给出同一个分布?——因为往漂移项里加上 $\frac{\sigma_t^2}{2}\nabla\log p_t$ 之后,Fokker–Planck 方程里多出来的扩散项恰好被抵消。
把这三条想透,剩下的(U-Net 还是 DiT、要不要在隐空间里做、guidance 怎么加、文本这种离散数据怎么办)都只是工程选择与结构平移。
这份文档做了什么
官方讲义写得很好,但和所有数学讲义一样,充满了 "one can show that…"、"a straightforward calculation gives…"。这份文档的存在意义就是:把这些被省略的计算全部做出来。
- 连续性方程、Fokker–Planck 方程、边际化技巧、CFM 与 FM 的梯度等价性、SDE 的边际不变性、Tweedie 公式、CFG 的贝叶斯分解、ELBO、Kolmogorov 前向方程——每一个都给完整证明,每一步标注依据。
- 抽象公式一律配上一维/二维高斯的可手算例子,把 $\alpha_t,\beta_t$ 代进去把闭式解真的算出来。
- 三个 lab 的全部 26 处 TODO 都独立实现、真实跑通,并配有 47 项数值测试与真实结果图。
预备知识
- 多元微积分(散度、梯度、分部积分、链式法则)与线性代数
- 概率论基础(条件期望、高斯分布、塔性质)
- PyTorch(会写
nn.Module和训练循环) - 不需要:随机微积分、测度论、扩散模型的先验知识——这些课程内从零讲起
1. 六讲笔记
把「生成一张图片」翻译成「从概率分布采样」,再翻译成「模拟微分方程」。向量场与流、Picard–Lindelöf 存在唯一性、Euler 与 Heun 的误差阶、布朗运动的三条定义性质与 √t 尺度、Euler–Maruyama 为什么是 √h、OU 过程的两条闭式解路线,最后落到连续性方程与 Fokker–Planck 方程的完整证明。
全课的心脏。条件与边际概率路径、高斯路径 $p_t(\cdot|z)=\N(\alpha_t z,\beta_t^2 I)$ 的全套闭式解、边际化技巧的完整证明、条件向量场的逐步推导,以及「CFM 损失与 FM 损失只差一个常数」这一核心定理的两种独立证明。最后落到 CondOT 路径与十行可运行的训练循环。
把 flow matching 的结论整体翻译成「对数似然梯度」的语言。得分函数为什么与归一化常数无关、高斯路径下的条件得分闭式解、「任意 $\sigma_t$ 下 SDE 与 ODE 边际相同」的 Fokker–Planck 验证、条件 score matching 的等价性证明、Tweedie 公式与 $x_0$/$\epsilon$/$v$ 三种参数化的换算,以及与 DDPM 的接口。
把 $y$ 塞进网络很容易,让模型「真的听话」才难。从贝叶斯分解 $\nabla\log p_t(x|y)=\nabla\log p_t(x)+\nabla\log p_t(y|x)$ 出发,证明引导强度 $w$ 对应在采样一个被锐化的分布 $\propto p_t(x)p_t(y|x)^w$,再一路推到 classifier-free guidance 与 label dropout 的训练技巧。含两个可手算的高斯算例。
前三讲告诉我们拟合什么,这一讲回答 $u_t^\theta$ 长什么样、在哪个空间里跑。时间条件化(Fourier 特征、AdaLN)、U-Net 的跳连为什么必需、Diffusion Transformer 与 adaLN-Zero 的零初始化技巧、VAE 的 ELBO 完整推导与重参数化技巧、KL 项的闭式解,最后是 Stable Diffusion 3 与 Meta Movie Gen 的整体 pipeline。
文本是离散的,$\R^d$ 上的 ODE/SDE 不存在了。整套配方被平移到有限状态空间:速率矩阵 $Q_t$ 取代向量场、Kolmogorov 前向方程取代连续性方程、离散版边际化技巧取代原版。掩码(masking)构造下的条件速率矩阵推导,以及与自回归语言模型的取舍对比。
2. 三个 Lab:全部实现并实测
三个 lab 共 26 处代码 TODO,全部独立实现、在 RTX 5080 上真实跑通(0 错误),并配有 47 项数值正确性测试与真实结果图。每份解析按「题目在问什么 → 数学依据 → 参考实现 → 为什么这样写 → 常见错误」五段式展开。
手写 Euler 与 Euler–Maruyama 求解器,实现布朗运动、OU 过程与 Langevin 动力学,并亲眼看到 Langevin 把一个宽高斯推成五峰混合分布。核心是想透为什么噪声项是 $\sqrt{h}$ 而不是 $h$。
从零搭出完整的 flow matching 训练管线:$\alpha_t,\beta_t$ 调度、条件路径采样、条件向量场与条件得分的闭式解、两个训练器,以及由向量场反解 score 的转换公式。附带两个真实的坑:学生版 notebook 里的 doubly-stochastic bug,以及转换公式在 $t=0$ 处的固有奇点。
在 MNIST 上造一个完整的条件生成模型:带 label dropout 的 CFG 训练、从零手写多头注意力、Fourier 时间编码、patchify/depatchify、adaLN-Zero 的 DiT block,再训一个带注意力的 VAE,最后把两者拼成 latent diffusion。
原课程明确建议不要让大模型替你写 lab 代码——那样你什么也学不到。这三份解析的定位是写完之后的对照与查错,以及卡住时的分步提示。真正的收获来自你自己把张量形状对上、把损失调收敛的那个过程。官方参考答案本身也是公开的(见 labs 仓库 的 solutions/)。
3. 全课的一条逻辑链
六讲看似分散,其实是同一个模板反复套用。把它写成一张表,整门课就闭合了:
| 环节 | 连续空间 $\R^d$(Lecture 01–04) | 离散空间 $\Vocab^N$(Lecture 05) |
|---|---|---|
| 演化的生成元 | 向量场 $u_t:\R^d\times[0,1]\to\R^d$ | 速率矩阵 $Q_t\in\R^{|\Vocab|\times|\Vocab|}$ |
| 密度演化方程 | 连续性方程 $\partial_t p_t=-\divg(p_t u_t)$ | Kolmogorov 前向方程 $\partial_t p_t(x)=\sum_y Q_t(y,x)p_t(y)$ |
| 条件路径 | $p_t(\cdot|z)=\N(\alpha_t z,\beta_t^2 I)$ | 逐位置以概率 $1-\alpha_t$ 变成 MASK |
| 边际化技巧 | $u_t(x)=\int u_t(x|z)\frac{p_t(x|z)\data(z)}{p_t(x)}\dd{z}$ | $Q_t(x,y)=\sum_z Q_t(x,y|z)\frac{p_t(x|z)\data(z)}{p_t(x)}$ |
| 可算的训练目标 | $\E\norm{u_t^\theta(x)-u_t(x|z)}^2$ | 预测原始 token 的交叉熵 |
| 为什么成立 | 条件损失与边际损失只差一个与 $\theta$ 无关的常数 ⇒ 梯度相同 | |
而 score matching(03-A)是把左列整体换成 $\nabla\log p_t$ 的语言,guidance(03-B)是把 $\data$ 换成 $\data(\cdot|y)$,隐空间扩散(04)是把 $x$ 换成 VAE 编码后的 $z$。模板从未改变。
4. 环境与复现
全部 lab 在下述环境实测通过:
| 项 | 版本 / 配置 |
|---|---|
| Python | 3.11(conda 环境 diff) |
| PyTorch | 2.11.0+cu128 |
| GPU | NVIDIA RTX 5080(Blackwell,sm_120,16 GB) |
| 其他 | numpy · matplotlib · seaborn · scikit-learn · einops · torchvision |
# 建环境
conda create -y -n diff python=3.11 && conda activate diff
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
pip install numpy matplotlib scipy tqdm jupyter seaborn einops scikit-learn
# 跑数值测试(47 项,纯 CPU,约 1 分钟)
python labs/tests/test_labs.py
# 跑完整 lab(已填空并附执行结果)
jupyter lab labs/lab1/lab_one_solved.ipynb
目录结构:
diffusion-lesson/
├── index.html 本页
├── lecture-01 … lecture-05.html 六讲笔记
├── lab-1 … lab-3.html 三份 lab 解析
├── assets/ 样式、MathJax 配置、目录脚本
├── images/ lab 真实运行结果图
├── slides/lecture_01…05/ 官方幻灯片逐页渲染图
├── labs/
│ ├── lab1|lab2|lab3/*_solved.ipynb 已填空并执行过的 notebook
│ └── tests/test_labs.py 47 项数值正确性测试
└── _src/ 原始 PDF、抽取文本、写作规范、校验脚本
5. 原始材料与延伸阅读
课程材料
- 课程主页——幻灯片、录像、lab 链接
- An Introduction to Flow Matching and Diffusion Models(84 页)——Holderrieth & Erives 的官方讲义,本文档的主要依据
- iap-diffusion-labs——lab 与官方参考答案
奠基论文
- DDPM: Denoising Diffusion Probabilistic Models (2020)——把扩散模型带进实用阶段的那篇
- Score-Based Generative Modeling through SDEs (2021)——连续时间视角的统一框架,本课 Lecture 03-A 的来源
- Flow Matching for Generative Modeling (2022)——flow matching 原始论文
- Rectified Flow (2022)——线性路径的独立发现
- Classifier-Free Diffusion Guidance (2022)——CFG
架构与系统
- Scalable Diffusion Models with Transformers (DiT, 2022)——adaLN-Zero 出处,Lab 3 的实现对象
- High-Resolution Image Synthesis with Latent Diffusion Models (2021)——latent diffusion / Stable Diffusion
- Scaling Rectified Flow Transformers (SD3, 2024)——把 flow matching 与 DiT 组合到最前沿
- Auto-Encoding Variational Bayes (2013)——VAE 与重参数化技巧