Flow Matching

0. 本讲导读

上一讲我们造好了「发动机」：给定一个向量场（vector field）$\vf_t^\theta$，我们知道怎么从噪声出发、沿着 ODE 一路积分到 $t=1$，得到一个样本。但我们从来没说过 $\theta$ 该取什么值。一个随机初始化的神经网络当然给不出猫的照片——它给出的是一团随机涂鸦。这一讲要解决的，就是「训练」这两个字。

问题的形状很清楚。我们希望

这是一个关于轨迹终点分布的要求。困难在于：我们手上只有一堆数据样本 $z^{(1)},\dots,z^{(N)} \sim \data$，既不知道 $\data$ 的密度，更不知道「能把 $\simple$ 搬运到 $\data$ 的那个向量场」长什么样。没有标签，怎么做回归？

Flow matching 的回答分三步走，每一步都会在本讲里被完整推导：

1. 先规定路线，再反推速度。 我们不去猜向量场，而是先人为指定一条从 $\simple$ 到 $\data$ 的分布插值路径 $(p_t)_{0\le t\le 1}$，称为概率路径（probability path）。路径一旦定死，能生成它的向量场就被连续性方程钉住了。
2. 把「对整个数据分布」的难题拆成「对单个数据点」的易题。 对每个固定的数据点 $z$，条件路径 $p_t(\cdot|z)$ 和它的条件向量场 $\vf_t^{\text{target}}(\cdot|z)$ 都有闭式解——我们能用纸笔算出来。然后由边际化技巧（marginalization trick），把这些条件向量场按后验概率加权平均，就得到了真正想要的边际向量场。
3. 把「不可算的回归目标」换成「可算的回归目标」。 边际向量场虽然存在，却是一个高维积分，训练时根本算不出来。关键定理告诉我们：直接对条件向量场做回归，其损失与对边际向量场做回归的损失只差一个与 $\theta$ 无关的常数，因此两者梯度完全相同。于是我们可以放心地用可算的目标训练。

整个方法的骨架可以用一张「Flow Matching 矩阵」概括：横向是「概率路径 → 向量场 → 损失函数」三级递进，纵向是「条件（单个数据点）／边际（整个数据分布）」两层。上面一行样样可算，下面一行样样不可算，而这一讲的全部技术含量，就是证明「优化上面一行 = 优化下面一行」。

1. 训练问题：我们到底要拟合什么

先把对象和维度钉死，后面所有推导都以此为准。

本讲限定在 flow model

上一讲同时讲了 flow model（ODE）和 diffusion model（SDE）。本讲只处理 flow model：生成过程就是

模拟到 $t=1$，返回 $X_1$ 作为样本。SDE 版本要等到下一讲引入 score function 之后才能处理。

于是「怎么训练」这个含糊的问题，被翻译成一个非常具体的数学问题：

为什么这个问题不是普通的监督学习

把它和一个标准回归任务对比一下，困难就浮出来了。监督学习里我们有成对的 $(\text{输入},\text{标签})$；这里我们只有 $\data$ 的样本，没有任何一条数据告诉我们「在时刻 $t$、位置 $x$ 处速度应该是多少」。更糟的是：

• 损失定义在终点分布上，而终点是长长一串 ODE 积分的结果。如果直接对「终点分布与 $\data$ 的某种距离」求梯度，就必须反向传播穿过整条模拟轨迹（这正是 neural ODE / continuous normalizing flow 的做法）。代价极高，而且要算 Jacobian 的迹或行列式。
• 中间时刻 $0 \lt t \lt 1$ 发生什么，问题本身没有规定。也就是说约束只在两端，中间是完全自由的。

第二点看起来是坏消息，实际上是本讲全部智慧的来源：既然中间没被规定，那就由我们自己规定。一旦把中间过程也钉死，「向量场该是什么」就从一个开放问题变成了一个有唯一答案（准确说是一族答案）的计算题，而且这个答案可以写成解析式。这就是「概率路径」这个概念要干的事。

2. 条件概率路径与边际概率路径

2.1 定义：条件概率路径

先引入一个记号。对数据点 $z \in \R^d$，用 $\delt{z}$ 表示 Dirac delta（狄拉克 δ）「分布」：从 $\delt{z}$ 采样永远返回 $z$，它是最简单的确定性分布。严格来说它不是密度函数而是测度，但在本讲的所有计算里，只要记住 $\int f(x)\delt{z}(x)\dd{x} = f(z)$ 就够用了。

用白话说：固定一个数据点 $z$，条件概率路径描述「整个初始分布逐渐塌缩到这一个点 $z$ 上」的过程。$t=0$ 时它是整个 $\simple$，$t=1$ 时它已经完全集中在 $z$ 这一个点上。可以把概率路径理解为分布空间里的一条轨迹：每个时刻 $t$ 对应一个分布，$t$ 从 $0$ 走到 $1$，分布从 $\simple$ 走到 $\delt{z}$。

2.2 定义：边际概率路径

单个数据点显然不够。把 $z$ 按数据分布随机抽取，就得到边际版本。

这两行的地位极不对称，这个不对称贯穿整门课，务必看清楚：

2.3 验算：边际路径确实连接了噪声与数据

讲义在这里写了一句「Check for yourself」。我们把这个检验真的做出来——它只有两行，但把「为什么条件端点要取 $\delt{z}$ 而不是 $\data$」这个疑问彻底解决了。

这段验算解释了一个初学者常见的困惑：为什么条件路径的终点是单个数据点 $\delt{z}$，而不是数据分布 $\data$？ 因为「对 $z\sim\data$ 取平均」这件事已经由边际化那一步完成了。每个条件路径只负责瞄准一个靶心，所有靶心按 $\data$ 的权重叠加起来，正好还原出整个数据分布。让条件路径直接瞄准 $\data$ 反而是错的——那样边际路径的终点会变成 $\data$ 与自身的某种混合，而且条件路径也就不再有闭式解了。

2.4 一个致命的细节：概率路径不含动力学

这是最容易被跳过、却在后面反复咬人的一点。

3. 高斯概率路径：把一切都算出来

概率路径的定义只给了两端的约束，中间怎么走完全自由。高斯概率路径（Gaussian probability path）是迄今为止最重要、也是几乎所有 state-of-the-art 模型实际使用的选择。这一节把它彻底算穿。

3.1 定义与边界条件的验证

先验证它确实是一条合法的条件插值路径。

3.2 采样：重参数化

高斯路径最实用的性质是采样极其简单。由高斯分布的线性变换性质（若 $\epsilon\sim\N(0,I_d)$，则 $A\epsilon + b \sim \N(b, AA^\top)$，此处 $A = \beta_t I_d$）：

而由边际路径的采样定义，同一个 $x$ 也就是边际路径的样本：$x \sim p_t$。这一行是整个训练算法的数据生成器：从数据集里取一张图 $z$，抽一个标准高斯噪声 $\epsilon$，抽一个时间 $t$，做一次线性混合，就得到了一个训练输入。没有任何模拟、没有任何积分。

3.3 具体算例一：两点数据集

抽象公式必须落到具体例子上才算真懂。取 $d=1$，数据分布是两个点的均匀混合 $\data = \frac12\left(\delt{z_1} + \delt{z_2}\right)$。此时边际路径是

也就是一个两分量高斯混合：两个分量的中心分别以速度 $\dot\alpha_t z_1$ 和 $\dot\alpha_t z_2$ 从原点向外滑动，同时各自的标准差 $\beta_t$ 从 $1$ 收缩到 $0$。$t=0$ 时两个分量完全重合成一个 $\N(0,1)$（因为 $\alpha_0=0$），$t=1$ 时收缩成两根 delta 尖峰。整条路径就是「一个高斯团裂变成两个尖峰」的过程。推广到 $K$ 个数据点是显然的：$\data = \sum_k w_k \delt{z_k}$ 给出 $p_t = \sum_k w_k \N(\alpha_t z_k, \beta_t^2 I_d)$。

注意即使在这个最简单的例子里，$p_t(x)$ 也已经是一个求和；数据集有 $N$ 个样本时就是 $N$ 项求和，$N$ 是百万量级时求值一次的代价就已经不可接受了。这个玩具例子恰恰说明了为什么真实场景里 $p_t(x)$ 不可算。

3.4 具体算例二：高斯数据分布（全解析）

再看一个能一路算到底的例子：设数据分布本身就是高斯，$\data = \N(\mu, \sigma^2 I_d)$。这时边际路径也有闭式解。

值得玩味的是：边际路径 $p_t = \N(\alpha_t\mu, \gamma_t^2 I_d)$ 自己也是一条「同样形状」的高斯路径，只不过把「数据点 $z$」换成了「数据均值 $\mu$」、把「噪声尺度 $\beta_t$」换成了 $\gamma_t$。这个观察在 §7.4 会被用来验证边际化技巧——我们将用两种完全不同的方法算出同一个边际向量场，从而给公式一次独立的正确性检验。

3.5 常用调度对照

满足边界条件的 $(\alpha_t,\beta_t)$ 有无穷多种。下表列出两个典型选择，两者都可以直接代入后面推出的所有公式。

两者的边界条件都易验：$\sin 0 = 0$、$\sin\frac{\pi}{2}=1$、$\cos 0 = 1$、$\cos\frac{\pi}{2}=0$，且在 $[0,1]$ 上分别单调增／单调减。本讲的所有定理对任意合法调度都成立；调度的选择只影响常数因子与数值稳定性，不影响正确性。

4. 条件向量场与边际向量场

概率路径规定了「每一帧应该长什么样」，但没规定粒子怎么动（§2.4）。现在补上动力学。

4.1 定义：条件向量场

三点必须说清楚：

• 这是一个要求，不是一个构造。任何满足这个性质的向量场都算数（回忆 §2.4：这样的向量场不唯一）。
• 它对每个 $z$ 是一个不同的向量场——「条件」二字就是「以 $z$ 为条件」。$z$ 在这里是参数，不是自变量。
• 关键优势：因为条件路径是我们自己挑的简单分布（比如高斯），$\vf_t^{\text{target}}(\cdot|z)$ 常常可以用手推出来。§5 就做这件事。

4.2 条件向量场单独看毫无用处

先泼一盆冷水，因为这盆冷水正是理解边际化技巧的入口。

4.3 定理：边际化技巧

现在是本讲的第一根支柱。

4.4 先读懂公式再证明它

在动手证明之前，先把这个积分读明白。注意到由贝叶斯公式，

正是「在时刻 $t$ 观测到带噪点 $x$ 时，它来自哪个数据点 $z$」的后验分布：分子是「先验 $\data(z)$ × 似然 $p_t(x|z)$」，分母 $p_t(x) = \int p_t(x|z)\data(z)\dd{z}$ 是归一化常数（这也再次说明了为什么 $p_t(x)$ 是配分函数式的、不可算的量）。于是边际化技巧可以写成一句话：

其中期望是关于联合分布 $Z \sim \data,\ X_t \sim p_t(\cdot|Z)$ 取的。边际向量场 = 条件向量场关于后验的条件期望。 这个形式在 §7 证明梯度等价定理时会立刻派上用场，请务必记住。

4.5 混合数据分布下的显式形式

把 §3.3 的有限数据集代进去，可以把这个「加权平均」写得完全具体。设 $\data = \sum_{k=1}^{K} w_k \delt{z_k}$（$w_k \ge 0$，$\sum_k w_k = 1$），高斯路径。则积分退化成求和：

这里高斯密度前的常数因子 $(2\pi\beta_t^2)^{-d/2}$ 在分子分母中约掉了。可以看到 $w_k(t,x)$ 就是一个以 $-\norm{x-\alpha_t z_k}^2/(2\beta_t^2)$ 为 logit 的 softmax，「温度」是 $\beta_t^2$。当 $\beta_t \to 0$（即 $t\to1$）时温度趋于零，softmax 退化成 argmax——边际向量场直接指向最近的那个训练样本。

定理 1 的完整证明需要连续性方程这个工具，我们在下一节先把工具建好，再回来收网。

5. 高斯路径的条件向量场：一步步解出来

定理 1 说「只要有条件向量场，就能造出边际向量场」。那条件向量场从哪来？对高斯路径，答案可以完全用手推出来。这一节把讲义 Example 10 的推导拆成每一步。

5.1 先造 flow，再从 flow 里提取向量场

回忆上一讲：向量场 $\vf_t$ 的流（flow）$\psi_t : \R^d \to \R^d$ 定义为 ODE 的解映射，满足

直接猜向量场很难，但猜流很容易——因为高斯路径的采样公式 $x = \alpha_t z + \beta_t \epsilon$ 已经把答案摆在脸上了：它就是「把初始噪声 $\epsilon$ 送到时刻 $t$ 的位置」的映射。于是定义

5.2 合理性检查

推完公式一定要做几个廉价的检查，这是抓错最有效的手段。

• 维度。 $\left(\dot\alpha_t - \frac{\dot\beta_t}{\beta_t}\alpha_t\right)$ 与 $\frac{\dot\beta_t}{\beta_t}$ 都是标量，$z, x \in \R^d$，所以 $\vf_t^{\text{target}}(\cdot|z) : \R^d \to \R^d$ ✓。它对 $x$ 是仿射的（一个标量乘 $x$ 加一个常向量），这也是它如此好算的原因。
• 分布中心的速度。 把分布的中心 $x = \alpha_t z$ 代进去： $$ \vf_t^{\text{target}}(\alpha_t z|z) = \left(\dot\alpha_t - \frac{\dot\beta_t \alpha_t}{\beta_t}\right)z + \frac{\dot\beta_t}{\beta_t}\alpha_t z = \dot\alpha_t z . $$ 而条件分布的均值是 $\alpha_t z$，其时间导数正是 $\dot\alpha_t z$。中心点的速度与中心的移动速度一致 ✓。
• 坐标形式。 令 $r := x - \alpha_t z$（相对于中心的偏移），则 $$ \vf_t^{\text{target}}(x|z) = \dot\alpha_t z + \frac{\dot\beta_t}{\beta_t}\, r . $$ （验证：$\left(\dot\alpha_t - \frac{\dot\beta_t\alpha_t}{\beta_t}\right)z + \frac{\dot\beta_t}{\beta_t}(r + \alpha_t z) = \dot\alpha_t z + \frac{\dot\beta_t}{\beta_t}r$ ✓。） 读法非常干净：「跟着中心平移」＋「以相对速率 $\dot\beta_t/\beta_t$ 向中心收缩」。因为 $\beta_t$ 递减，$\dot\beta_t \lt 0$，所以第二项确实是指向中心的收缩。这个形式在 §6.3 验算连续性方程时会直接用到。

5.3 沿路径代入：目标就是插值曲线的速度

训练时我们不会在任意 $x$ 处求值，而只在路径上的点 $x = \alpha_t z + \beta_t \epsilon$ 处求值。代进去看看会发生什么：

这个结果值得裱起来。记 $x_t := \alpha_t z + \beta_t \epsilon$ 是固定 $(z,\epsilon)$ 时的插值曲线，则

5.4 CondOT 特例

取 $\alpha_t = t$、$\beta_t = 1-t$，于是 $\dot\alpha_t = 1$、$\dot\beta_t = -1$。代入闭式解：

形式漂亮得不像话：从当前位置 $x$ 指向目标 $z$ 的方向，除以剩余时间 $1-t$。这正是「在剩余 $1-t$ 的时间里匀速走完剩下的 $z-x$ 这段路」所需的速度——CondOT 之所以叫「直线调度」，原因就在这里。

再用 §5.3 的结论核对一遍：路径点是 $x = tz + (1-t)\epsilon$，则

而 §5.3 的通式给出 $\dot\alpha_t z + \dot\beta_t \epsilon = 1\cdot z + (-1)\cdot\epsilon = z - \epsilon$ ✓。两条路殊途同归。

6. 连续性方程：从「分布怎么变」到「粒子怎么动」

要证明定理 1，需要一件工具：一个把「概率路径」和「向量场」严格挂钩的等价条件。这就是连续性方程（continuity equation）。

6.1 散度

先定义散度（divergence）算子。对向量场 $v_t : \R^d \to \R^d$，

其中 $v_t^i$ 是 $v_t$ 的第 $i$ 个分量。注意 $\divg(v_t)(x) \in \R$ 是标量（把向量场「压」成一个数）。物理含义：$\divg v_t(x) \gt 0$ 表示 $x$ 处是「源」（流出多于流入），$\lt 0$ 表示是「汇」。

6.2 定理与直觉

注意这是一个充要条件，而且前提里已经包含了初始条件 $X_0 \sim p_0$。这一点在用它证明定理 1 时至关重要：验证了 PDE，再加上初值对上，就足以断定整条轨迹的边际分布是 $p_t$。

6.3 验算：高斯路径确实满足连续性方程

讲义在这里留了一个脚注：「One can also double check this by plugging it into the continuity equation」。我们把它做出来——这既是对 §5 结果的独立验证，也是熟悉连续性方程的最好练习。

三项的物理意义各自清晰：$-d\,\dot\beta_t/\beta_t$ 来自分布整体的收缩（体积项），$\dot\alpha_t\inner{r}{z}/\beta_t^2$ 来自中心的平移，$\dot\beta_t\norm{r}^2/\beta_t^3$ 来自尺度变化对形状的影响。这次验算完全独立于 §5 的推导路线（那里用的是 flow，这里用的是 PDE），两者一致大大提高了对结果的信心。

7. 边际化技巧的完整证明

7.1 证明策略

要证的是定理 1：由

定义的边际向量场，其 ODE 轨迹满足 $X_t \sim p_t$。

直接从「轨迹的分布」下手是没法做的——那要求解一个非线性 ODE 再算像测度。连续性方程把这件事变成了纯粹的微积分：由定理 2（充要条件），只要验证 $\partial_t p_t = -\divg(p_t \vf_t^{\text{target}})$，再配上初值 $X_0 \sim \simple = p_0$（这是 ODE 的初始条件，本来就成立），结论就跟着来了。

我们手上有什么？只有一条信息：每个条件向量场都生成对应的条件路径，也就是（再次由定理 2，这次用在条件对象上）

所以整个证明就是：把 $\partial_t p_t$ 展开成对 $z$ 的积分，对每个 $z$ 用 (C)，再把散度提到积分号外，最后凑回定义式。

7.2 证明

7.3 每一步用了什么

7.4 独立验算：高斯数据分布

定理 1 的公式含有一个不可算的积分，看上去无法检验。但在 §3.4 的高斯数据例子里，我们可以把积分真的算出来，再和另一条完全独立的推导对照。这是对定理 1 的一次实打实的检验。

两条路线一条走的是「后验加权平均」（定理 1 的定义式），另一条走的是「把边际路径本身当成高斯路径直接套公式」，中间没有任何共享的中间结果，却给出同一个答案。这就是对边际化技巧的一次硬检验。

8. 两个损失函数：可算的与不可算的

8.1 边际 Flow Matching 损失

现在我们知道了训练目标：让神经网络等于边际向量场，$\vf_t^\theta \approx \vf_t^{\text{target}}$。最自然的做法是均方误差回归。

读法：随机抽一个时间 $t$；随机抽一个数据点 $z$ 并给它加噪得到 $x$（这就是从 $p_t$ 采样）；比较网络输出与真值向量场；取平方误差的期望。若 $\mathcal{L}_{\text{FM}}(\theta) = 0$，则 $\vf_t^\theta = \vf_t^{\text{target}}$（在 $p_t$ 的支撑上几乎处处成立），由定理 1 模拟 ODE 就能得到 $X_1\sim\data$。逻辑完全闭合。

8.2 条件 Flow Matching 损失

解决方案出人意料地朴素：既然 $\vf_t^{\text{target}}(x)$ 算不出来，而 $\vf_t^{\text{target}}(x|z)$ 有闭式解，那就直接回归后者。反正采样 $x$ 的时候我们本来就抽了一个 $z$，它就在手边。

这个损失每一项都可算：$t$ 是 $\text{Unif}[0,1]$ 采样，$z$ 从数据集里取，$x = \alpha_t z + \beta_t\epsilon$ 一行代码，目标 $\vf_t^{\text{target}}(x|z) = \dot\alpha_t z + \dot\beta_t\epsilon$（§5.3）也是一行代码。这是一个彻头彻尾的普通监督回归。

8.3 两者的对照

9. 核心定理：两个损失只差一个常数

下面给两个证明。证明一是讲义的直接展开法，机械但完全自足；证明二用条件期望的塔性质，短得多，而且解释了「为什么」。两个都值得看。

9.1 证明一：直接展开

记号约定：下文所有期望里 $t\sim\text{Unif}[0,1]$，$z\sim\data$，$x\sim p_t(\cdot|z)$（等价地 $x\sim p_t$）。「常数」一律指与 $\theta$ 无关的量。

9.2 证明二：条件期望视角

直接展开法虽然完备，却像变戏法——加加减减之后答案就出来了，但没解释「为什么」。下面这个证明短得多，并且揭示了本质。

9.3 三条必须搞清楚的推论

10. 训练算法：CondOT 路径与最小实现

10.1 把损失落到高斯路径上

定理 3 说「优化 $\mathcal{L}_{\text{CFM}}$ 就行」。现在把高斯路径的具体公式代进去，看看训练目标最终长什么样。

最终形式干净到令人怀疑：输入是噪声与数据的加权混合，标签是两个权重的导数做同样的混合。 没有 ODE 模拟、没有积分、没有除法。

10.2 CondOT：直线调度

取 $\alpha_t = t$、$\beta_t = 1-t$（$\dot\alpha_t = 1$、$\dot\beta_t=-1$），得到本讲最著名的公式：

对应的条件概率路径写作 $p_t(x|z) = \N\!\left(tz,\ (1-t)^2 I_d\right)$，称为 CondOT 概率路径（conditional optimal transport path，也叫 rectified flow / 线性插值路径）。

10.3 训练算法

10.4 最小 PyTorch 实现

下面是一份可直接运行的完整实现，每一块都与前面的公式一一对应。

(1) 网络：$\vf_t^\theta : \R^d\times[0,1]\to\R^d$。 时间作为额外输入拼接进去（真实模型会用正弦位置编码，这里从简）。

import torch
import torch.nn as nn

class VectorField(nn.Module):
    """u_t^theta(x)：输入 x 与 t，输出 R^d 中的速度向量。"""
    def __init__(self, dim=2, hidden=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim + 1, hidden), nn.SiLU(),
            nn.Linear(hidden, hidden),  nn.SiLU(),
            nn.Linear(hidden, hidden),  nn.SiLU(),
            nn.Linear(hidden, dim),
        )

    def forward(self, x, t):            # x: (bs, dim), t: (bs, 1)
        return self.net(torch.cat([x, t], dim=-1))   # (bs, dim)

(2) CFM 损失（CondOT）。 直接抄 §10.2 的公式：三行采样、一行标签、一行 MSE。

def cfm_loss(model, z):                 # z: (bs, dim)，来自 p_data
    bs = z.shape[0]
    t   = torch.rand(bs, 1, device=z.device)     # (bs, 1)    t ~ Unif[0,1]
    eps = torch.randn_like(z)                    # (bs, dim)  eps ~ N(0, I_d)

    x      = t * z + (1.0 - t) * eps             # (bs, dim)  x = alpha_t z + beta_t eps
    target = z - eps                             # (bs, dim)  u_t^target(x|z) = da z + db eps

    pred = model(x, t)                           # (bs, dim)
    return ((pred - target) ** 2).sum(dim=-1).mean()   # E ||.||^2

注意 sum(dim=-1) 对应 $\norm{\cdot}^2$（在 $d$ 个坐标上求和），mean() 对应期望 $\E$。两者顺序不能颠倒成 mean() 全平均——那只会把损失缩放 $1/d$，虽然不改变 $\arg\min$，但会影响学习率的语义。

(3) 任意调度的通用版本。 直接实现 §10.1 的一般公式 $\norm{\vf_t^\theta(\alpha_t z+\beta_t\epsilon) - (\dot\alpha_t z + \dot\beta_t\epsilon)}^2$：

def cfm_loss_general(model, z, sched):
    """sched(t) 返回 (alpha, beta, dalpha, dbeta)，形状均为 (bs, 1)。"""
    t = torch.rand(z.shape[0], 1, device=z.device)
    eps = torch.randn_like(z)
    a, b, da, db = sched(t)

    x      = a * z + b * eps                     # (bs, dim)  x ~ p_t(.|z)
    target = da * z + db * eps                   # (bs, dim)  沿路径的速度
    return ((model(x, t) - target) ** 2).sum(dim=-1).mean()

def condot(t):        # alpha_t = t, beta_t = 1 - t
    return t, 1.0 - t, torch.ones_like(t), -torch.ones_like(t)

def trig(t):          # alpha_t = sin(pi t / 2), beta_t = cos(pi t / 2)
    h = torch.pi / 2
    return (torch.sin(h * t), torch.cos(h * t),
            h * torch.cos(h * t), -h * torch.sin(h * t))

(4) 训练循环与采样器。 采样就是上一讲的 Euler 法。

def sample_checkerboard(n, device='cpu'):
    """二维棋盘分布 p_data，用拒绝采样生成。返回 (n, 2)。"""
    out = []
    total = 0
    while total < n:
        y = torch.rand(4 * n, 2, device=device) * 4.0 - 2.0     # 铺满 [-2, 2]^2
        keep = (torch.floor(y[:, 0]) + torch.floor(y[:, 1])) % 2 == 0
        y = y[keep]
        out.append(y); total += y.shape[0]
    return torch.cat(out, dim=0)[:n]

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = VectorField(dim=2).to(device)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for step in range(20000):
    z = sample_checkerboard(2048, device)        # (bs, 2)  z ~ p_data
    loss = cfm_loss(model, z)                    # 标量
    opt.zero_grad(); loss.backward(); opt.step()

@torch.no_grad()
def sample(model, n, dim=2, steps=200, device='cpu'):
    x = torch.randn(n, dim, device=device)       # (n, dim)  X_0 ~ p_init
    dt = 1.0 / steps
    for i in range(steps):                       # Euler 法模拟 ODE
        t = torch.full((n, 1), i * dt, device=device)
        x = x + model(x, t) * dt                 # X_{t+dt} = X_t + u_t(X_t) dt
    return x                                     # (n, dim)  近似 X_1 ~ p_data

10.5 用有限数据集验证边际化技巧

最后给一段验证代码：当 $\data$ 是有限个点的均匀混合时，§4.5 的显式公式可以精确计算边际向量场。把它与训练好的网络对比，就能直接检验定理 1 与定理 3。

def marginal_vf_exact(x, t, Z):
    """把 §4.5 的求和写出来（CondOT，p_data = 均匀混合 of Z）。
       x: (bs, dim)   t: (bs, 1)   Z: (K, dim)   返回 (bs, dim)"""
    a = t                                    # (bs, 1)   alpha_t = t
    b = (1.0 - t).clamp_min(1e-4)            # (bs, 1)   beta_t = 1 - t，末端 clamp

    # 后验 logit：-||x - alpha_t z_k||^2 / (2 beta_t^2)
    diff = x[:, None, :] - a[:, None] * Z[None]        # (bs, K, dim)
    logit = -(diff ** 2).sum(dim=-1) / (2 * b ** 2)    # (bs, K)
    w = torch.softmax(logit, dim=-1)                   # (bs, K)  后验 p_t(z_k | x)

    u_cond = (Z[None] - x[:, None, :]) / b[:, None]    # (bs, K, dim)  (z - x)/(1 - t)
    return (w[..., None] * u_cond).sum(dim=1)          # (bs, dim)  加权平均

代码与公式的对应关系一目了然：w 是后验 $\frac{p_t(x|z_k)\data(z_k)}{p_t(x)}$（softmax 里已自动完成归一化，即除以 $p_t(x)$），u_cond 是条件向量场 $\frac{z-x}{1-t}$（§5.4），最后一行的加权求和就是定理 1 的积分。把训练好的 model(x, t) 与 marginal_vf_exact(x, t, Z) 在同一批 $(x,t)$ 上比较，应当高度吻合——这就是定理 3 在说的事情：用逐样本的、互相矛盾的条件标签做回归，收敛到的正是这个后验加权平均。

本讲小结

把所有关键公式集中列一遍。全部以 $x, z, \epsilon \in \R^d$、$t\in[0,1]$、$\alpha_t,\beta_t\in\R$ 为准。

三个母题在本讲的进度

延伸阅读

本讲的原始文献

• Flow Matching for Generative Modeling (Lipman et al., 2022) — 本讲全部内容的出处。条件流匹配、边际化技巧、CondOT 路径都出自这篇。先读它的第 3、4 节，与本讲逐节对应。
• Flow Straight and Fast: Rectified Flow (Liu et al., 2022) — 独立提出线性插值路径，并给出「反复重配对以拉直轨迹」的 reflow 过程，解释了 §10.2 里「轨迹为什么是弯的、怎么把它弄直」。
• Stochastic Interpolants (Albergo, Boffi & Vanden-Eijnden, 2023) — 把 $\alpha_t z+\beta_t\epsilon$ 推广成任意插值 $I_t(z,\epsilon)$ 的统一框架，同时涵盖 ODE 与 SDE，是理解「调度到底能有多自由」的最佳参考。
• Flow Matching Guide and Code (Lipman et al., 2024) — 综述加代码，覆盖离散、黎曼流形等推广。想动手时从这里抄实现最省事。

把本讲的算法用到工业规模

• Scaling Rectified Flow Transformers for High-Resolution Image Synthesis (Esser et al., 2024) — Stable Diffusion 3。用的正是 §10.2 那个 CondOT 目标，重点看它对时间 $t$ 采样分布的改造（不用均匀分布而是偏向中间时刻）。
• Movie Gen: A Cast of Media Foundation Models (Polyak et al., 2024) — Meta 的视频生成模型，同样是 $\alpha_t=t,\ \beta_t=1-t$ 的条件流匹配目标，说明本讲的公式在 $d$ 极大时依然照搬即可。
• SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers (Ma et al., 2024) — 在同一网络架构下系统对比不同插值路径与训练目标，是判断「调度选择到底有多重要」的实证参考。

推广方向

• Flow Matching on General Geometries (Chen & Lipman, 2023) — 把 $\R^d$ 换成黎曼流形。$x=\alpha_t z+\beta_t\epsilon$ 这种线性插值在流形上没有意义，需要用测地线替代，是检验自己是否真正理解「路径—向量场」对应关系的好材料。
• Scalable Diffusion Models with Transformers (Peebles & Xie, 2023) — DiT。本讲把 $\vf_t^\theta$ 当黑箱，这篇讲这个黑箱在图像上应该怎么造。