Lab 2 解析：Flow Matching 与 Score Matching

0. 本 lab 导读

Lab 1 教会我们的是「怎么把一个已知的 ODE / SDE 模拟出来」——那时候漂移项（drift）和扩散项（diffusion）都是白送的。Lab 2 要回答的是上一个层级的问题：这些漂移项本身是从哪来的？答案就是 flow matching 与 score matching：先人为设计一条条件概率路径（conditional probability path）$p_t(x\mid z)$，把它的条件向量场和条件得分解析地写出来，然后用回归的方式把神经网络训练成对应的边际（marginal）量。

整个 lab 的逻辑主线只有一句话：边际的东西不可解，但条件的东西可解；而对条件量做回归，最优解恰好就是边际量。八个 TODO 全部围绕这一句展开——前四个负责把「条件的东西」写出来，后四个负责把「回归」做出来。

1. 先看骨架：三层抽象与张量形状约定

动手之前花五分钟看懂 Part 0 / Part 1 给的类层次，后面所有 TODO 都会顺很多。这个 lab 的代码是三层结构：

路径层里有一个已经写好、但极其重要的方法，值得单独拎出来读：

def sample_marginal_path(self, t: torch.Tensor) -> torch.Tensor:
    num_samples = t.shape[0]
    z = self.sample_conditioning_variable(num_samples)  # (num_samples, dim)
    x = self.sample_conditional_path(z, t)              # (num_samples, dim)
    return x

这三行就是整个 flow matching 的哲学缩影：边际分布 $p_t(x)=\int p_t(x\mid z)p_{\text{data}}(z)\,\dd{z}$ 我们写不出闭式，但我们能采样它——先采 $z$，再采 $x\mid z$，扔掉 $z$，剩下的 $x$ 就服从边际。后面训练时用的正是同一个双层采样结构，只不过 $z$ 不扔掉，而是留下来算回归目标。

另外注意 Alpha / Beta 基类的 __init__ 里塞了两个 assert：

class Alpha(ABC):
    def __init__(self):
        assert torch.allclose(self(torch.zeros(1,1)), torch.zeros(1,1))  # alpha_0 = 0
        assert torch.allclose(self(torch.ones(1,1)),  torch.ones(1,1))   # alpha_1 = 1

也就是说，只要你 LinearAlpha() 一实例化不报错，边界条件就已经自动被检查过了。这是个很贴心的设计，但也意味着：如果你在 __init__ 里报 AssertionError，问题一定出在 __call__ 的边界值上，而不是别的地方。

2. Part 2：高斯条件概率路径

本节要实现的对象是

直觉上，$z$ 是一个「目标点」（在高斯路径里 $z\sim p_{\text{data}}$，也就是一个真实数据点），$\alpha_t$ 控制这个点被保留了多少，$\beta_t$ 控制还剩多少噪声。$t=0$ 时全是噪声（标准高斯），$t=1$ 时全是信号（退化成 $\delta_z$）。反过来看，从 $t=1$ 走到 $t=0$ 就是一个把数据点逐步加噪、最终抹成纯噪声的腐蚀过程（corruption process）——生成就是把它反过来走。

问题 2.1：实现 $\alpha_t$ 与 $\beta_t$

题目在问什么。实现 LinearAlpha.__call__ 返回 $\alpha_t=t$，SquareRootBeta.__call__ 返回 $\beta_t=\sqrt{1-t}$。基类已经用 vmap(jacrev(...)) 给出了通用的自动微分版 dt，但解答里把它们替换成了解析导数（更快也更可控）。

数学依据。四个边界条件不是随便定的，每一条都有明确含义：

换句话说，这四个条件就是「这条路径确实连接 $p_{\text{simple}}$ 和 $p_{\text{data}}$」的全部要求。中间怎么走（$\alpha_t,\beta_t$ 具体取什么形状）是完全自由的设计空间——这就是为什么后来会有 rectified flow、EDM、cosine schedule 等等一堆变体。

本节选的 $\beta_t=\sqrt{1-t}$ 有一个特别的性质：$\alpha_t^2+\beta_t^2=t^2+1-t$，在 $t\in(0,1)$ 内小于 1，所以它不是方差保持（variance-preserving）路径；但它让后面 $\beta_t^2\dot\alpha_t-\alpha_t\dot\beta_t\beta_t=1-\frac t2$ 变得极其干净（问题 3.3 会用到）。

参考实现。

class LinearAlpha(Alpha):
    """alpha_t = t"""
    def __call__(self, t: torch.Tensor) -> torch.Tensor:
        # t: (num_samples, 1) -> alpha_t: (num_samples, 1)
        # alpha_0 = 0（纯噪声），alpha_1 = 1（纯数据）
        return t

    def dt(self, t: torch.Tensor) -> torch.Tensor:
        # d/dt t = 1，形状必须和 t 一样，不能返回标量
        return torch.ones_like(t)          # (num_samples, 1)


class SquareRootBeta(Beta):
    """beta_t = sqrt(1 - t)"""
    def __call__(self, t: torch.Tensor) -> torch.Tensor:
        # beta_0 = 1（噪声标准差为 1），beta_1 = 0（无噪声）
        return torch.sqrt(1 - t)           # (num_samples, 1)

    def dt(self, t: torch.Tensor) -> torch.Tensor:
        # d/dt sqrt(1-t) = -1/(2 sqrt(1-t))，t->1 时发散，加 1e-4 兜底
        return -0.5 / (torch.sqrt(1 - t) + 1e-4)   # (num_samples, 1)

为什么这样写。

• return t 而不是 return t.clone()：这里没有原地修改，直接返回引用是安全的，而且能省一次拷贝。但如果你之后要在 __call__ 结果上做原地操作，就必须 clone()。
• torch.ones_like(t) 而不是 1.0：返回值必须是 (bs,1) 的张量，且 dtype / device 要跟着 t 走。返回 Python 标量在后面 (dt_alpha - ...) * z 的广播里碰巧也能work，但一旦有人写 alpha.dt(t).shape 就会炸，而且丢了 device 信息。
• torch.sqrt(1 - t) 而不是 (1-t) ** 0.5：两者数值等价，但前者在 1-t 为负（由于浮点误差 $t$ 略大于 1）时返回 nan 而不是复数报错，行为更可预测。
• 把 dt 从自动微分版换成解析版：基类的 vmap(jacrev(self)) 每次都要建计算图，在 5000 个 epoch × 每 epoch 一次的调用量下是纯浪费。数值测试确认 alpha.dt 与自动微分误差为 0.00e+00（因为常数 1 是精确的）。

问题 2.2：从条件路径采样

题目在问什么。给定 $z$（形状 (bs,dim)）与 $t$（形状 (bs,1)），返回一个服从 $\N(\alpha_t z,\beta_t^2 I_d)$ 的样本，形状 (bs,dim)。

数学依据。就是重参数化技巧（reparameterization trick）：若 $\epsilon\sim\N(0,I_d)$，则

代入 $\mu=\alpha_t z$、$\sigma=\beta_t$ 得

这条式子有一个远超「采样技巧」的身份：它同时定义了条件流（conditional flow）

即把 $t=0$ 时的一个具体噪声点 $x_0$ 搬运到 $t$ 时刻位置的那个确定性映射。问题 2.3 的整个推导就建立在 $\psi_t$ 之上，所以这里写下的这一行，其实是后面所有东西的源头。

参考实现。

def sample_conditional_path(self, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """
    z: (num_samples, dim)   t: (num_samples, 1)
    return: (num_samples, dim)
    """
    # alpha(t), beta(t) 均为 (num_samples, 1)，与 z 的 (num_samples, dim) 广播
    return self.alpha(t) * z + self.beta(t) * torch.randn_like(z)

为什么这样写。

• torch.randn_like(z) 而不是 torch.randn(z.shape)：前者自动继承 z 的 dtype 和 device。本 lab 全程在 GPU 上跑，后者会在 CPU 上造张量然后抛 RuntimeError: Expected all tensors to be on the same device。
• 广播方向：(bs,1) * (bs,dim) -> (bs,dim)。这里的关键是同一个样本的两个坐标共享同一个 $\alpha_t$，因为高斯是各向同性的。如果你不小心把 t 传成了 (bs,)，广播会变成 (bs,) * (bs,dim) —— PyTorch 会把 (bs,) 当成最后一维去对齐 dim，在 bs != dim 时直接报错，在 bs == dim 时静默算错。这是本 lab 最阴险的形状陷阱。
• 每次调用都重新采一次 $\epsilon$：这是必须的。训练时同一个 $(z,t)$ 对应的 $x$ 必须是随机的，否则回归目标就退化成了确定性映射，学不到条件期望。

问题 2.3：条件向量场

题目在问什么。实现

题面直接把答案给了，所以代码本身是送分的。但如果只抄公式，你会错过这个 lab 最核心的一次推导。下面完整重做一遍。

参考实现。

def conditional_vector_field(self, x, z, t):
    """
    x: (num_samples, dim)   z: (num_samples, dim)   t: (num_samples, 1)
    return: (num_samples, dim)。注意只在 t in [0,1) 上有定义
    """
    alpha_t    = self.alpha(t)       # (num_samples, 1)
    beta_t     = self.beta(t)        # (num_samples, 1)
    dt_alpha_t = self.alpha.dt(t)    # (num_samples, 1)
    dt_beta_t  = self.beta.dt(t)     # (num_samples, 1)
    return (dt_alpha_t - dt_beta_t / beta_t * alpha_t) * z + dt_beta_t / beta_t * x

为什么这样写。

• 四个标量系数全是 (bs,1)，与 (bs,dim) 的 z、x 广播，结果 (bs,dim)。没有任何 reshape/unsqueeze 的必要——这正是前面坚持让 t 是 (bs,1) 的回报。
• 把 alpha(t) 等先存成局部变量，而不是在表达式里调四次：self.alpha.dt(t) 在基类实现下是一次 vmap(jacrev(...))，重复调用非常贵。即使换成解析版，可读性也更好。
• dt_beta_t / beta_t 在 $t\to1$ 时是 $-\frac{1}{2(1-t)}\to-\infty$。这不需要（也不应该）加保护：题面明确说了这个函数只在 $t\in[0,1)$ 上有定义，训练时 torch.rand 的取值区间是 $[0,1)$，模拟时最后一步也停在 $t<1$。人为把它 clamp 住反而会引入错误的动力学。

问题 2.4：条件得分

题目在问什么。实现 $\nabla_x\log p_t(x\mid z)$，即条件分布对数密度关于位置 $x$（不是关于参数、不是关于 $t$）的梯度。

数学依据。这是本 lab 唯一一处需要真正「求导一个密度」的地方，但因为分布是高斯，两行就能算完。

参考实现。

def conditional_score(self, x, z, t):
    """
    x: (num_samples, dim)   z: (num_samples, dim)   t: (num_samples, 1)
    return: (num_samples, dim)。同样只在 t in [0,1) 上有定义
    """
    alpha_t = self.alpha(t)   # (num_samples, 1)
    beta_t  = self.beta(t)    # (num_samples, 1)
    return (z * alpha_t - x) / beta_t ** 2

为什么这样写。只有一行，但有三处值得说：

• 符号：分子是 $\alpha_t z-x$（指向均值），不是 $x-\alpha_t z$。写反了会让后面的 Langevin 项变成「把样本推离模式」，SDE 采样直接发散。
• 分母是 $\beta_t^2$ 不是 $\beta_t$：score 的量纲是 $1/\text{长度}$，来自 $\frac{\text{长度}}{\text{方差}}$。这是最高频的错误之一，而且在 $\beta_t$ 接近 1 的地方（$t$ 小）几乎看不出差别，只有 $t$ 大时才会暴露。
• 不需要 torch.log、不需要 autograd：闭式解是解析的。用 torch.autograd.grad 去数值求这个梯度虽然也对，但慢几十倍，而且在 $t\to1$ 处更不稳定。

3. 插曲：学生版 notebook 里的一个真实 bug

这一节不对应任何 TODO，但如果你严格照着学生版 notebook 跑，你模拟的不是题面写的那个 SDE。这是我在本地完整执行时实测发现的，值得单独讲。

先看 Lab 1 就已经给好、Lab 2 原封不动搬过来的 Euler–Maruyama 求解器：

class EulerMaruyamaSimulator(Simulator):
    def step(self, xt, t, h):
        return (xt
                + self.sde.drift_coefficient(xt, t) * h
                + self.sde.diffusion_coefficient(xt, t) * torch.sqrt(h) * torch.randn_like(xt))

注意最后一项：求解器自己乘了 torch.sqrt(h) * torch.randn_like(xt)。也就是说，diffusion_coefficient 应该只返回确定性的 $\sigma_t$，随机性由求解器负责。

但学生版 notebook 里 ConditionalVectorFieldSDE 和 LangevinFlowSDE 的 diffusion_coefficient 都写着：

def diffusion_coefficient(self, x, t):
    return self.sigma * torch.randn_like(x)   # <-- 这里多采了一次噪声

于是每一步的随机增量变成了

即两个独立高斯的逐元素乘积，而不是题面要求的 $\sigma\sqrt h\,\epsilon$。

官方 solutions 已经把它修成 return self.sigma（课程仓库 2026-01-11 的 changelog 提到修了这个 bug），但修复只进了 solutions 分支，学生版 notebook 至今仍是错的。本页解答里写成：

def diffusion_coefficient(self, x, t):
    # 扩散系数是确定性的；随机性由 EulerMaruyamaSimulator.step 里的
    # sqrt(h) * randn_like(xt) 提供，这里再采一次噪声就是采了两次
    return self.sigma * torch.ones_like(x)    # (bs, dim)

用 torch.ones_like(x) 而不是裸的 self.sigma，是为了保持返回值形状为 (bs, dim)、与基类文档字符串一致；数值上与 return self.sigma 等价（后者靠标量广播）。两种写法都对。

4. Part 3：训练——Flow Matching 与 Score Matching

前面四个 TODO 把「条件的东西」全写完了。现在进入 lab 的第二幕：把条件量回归成边际量。

要学的对象是边际向量场

其中 $p_t(z\mid x)\propto p_t(x\mid z)p_{\text{data}}(z)$ 是后验。这个期望我们既算不出来也采不了样（后验里含未知的 $p_t(x)$）。Flow matching 的全部技巧就在于：条件期望是 $L^2$ 意义下的最优回归量，所以

把它对 $x\sim p_t$ 再取一次期望、换成联合采样，就得到可直接蒙特卡洛估计的条件 flow matching 损失

问题 3.1：条件 Flow Matching 的训练损失

题目在问什么。用一个 batch 的蒙特卡洛样本估计上面那个期望，返回一个标量张量。

数学依据。照抄公式里的下标即可，采样顺序严格按依赖关系来：

注意 $x_i$ 的分布依赖 $z_i$ 和 $t_i$，所以必须最后采；$z_i$ 与 $t_i$ 相互独立，先后无所谓。

参考实现。

class ConditionalFlowMatchingTrainer(Trainer):
    def __init__(self, path, model: MLPVectorField, **kwargs):
        super().__init__(model, **kwargs)
        self.path = path

    def get_train_loss(self, batch_size: int) -> torch.Tensor:
        # 第一步：采条件变量 z ~ p_data
        z = self.path.sample_conditioning_variable(batch_size)          # (bs, dim)
        # 第二步：采时间 t ~ U[0,1)。.to(z) 同时对齐 dtype 和 device
        t = torch.rand(batch_size, 1).to(z)                             # (bs, 1)
        # 第三步：采位置 x ~ p_t(. | z)
        x = self.path.sample_conditional_path(z, t)                     # (bs, dim)

        # 第四步：回归条件向量场
        ut_theta = self.model(x, t)                                     # (bs, dim)
        ut_ref   = self.path.conditional_vector_field(x, z, t)          # (bs, dim)
        error = torch.einsum('bd,bd -> b', ut_theta - ut_ref, ut_theta - ut_ref)  # (bs,)
        return torch.mean(error)                                        # 标量

为什么这样写。

• torch.rand(batch_size, 1) 而不是 torch.rand(batch_size)：形状必须是 (bs,1)，理由见第 1 节的形状约定。
• .to(z) 是个惯用法：把 t 的 dtype 和 device 一次性对齐到 z。写 .to(device) 也行，但依赖外部的全局变量；.to(z) 自包含。
• torch.rand 的取值范围是 $[0,1)$，右端开区间正好避开了 $t=1$ 处 $\beta_t=0$ 的奇点。这不是巧合，是题面刻意选的。如果你自作主张写成 torch.rand(...)*1.0001 或者用 linspace(0,1,bs)，就会踩到除零。
• torch.einsum('bd,bd->b', e, e) 计算的是每个样本的平方范数 $\norm{e_i}^2$，等价于 (e**2).sum(dim=-1)。用哪种都行；einsum 版本在高维时更明确地表达了「只对特征维求和、保留 batch 维」的意图。不要写成 torch.nn.functional.mse_loss 的默认行为——那会额外除以 dim，损失数值差一个常数因子 2（本 lab 中 dim=2），虽然不影响优化方向，但会让你无法和别人的损失曲线对比。
• 不需要 detach 任何东西：ut_ref 全部由不含参数的解析式算出，本来就没有梯度路径。（但如果你的 path 里不小心用了带参数的模块，就要小心了。）

问题 3.2：条件 Score Matching 的训练损失

题目在问什么。把 3.1 的回归目标从条件向量场换成条件得分，训练 $s_t^\theta(x)\approx\nabla\log p_t(x)$。

数学依据。与 3.1 完全同构。边际化技巧对 score 同样成立：

于是

这就是文献里的去噪得分匹配（denoising score matching, DSM）。它比原始 score matching（需要计算 $\tr\nabla^2\log p$，代价 $O(d^2)$）便宜得离谱，是整个扩散模型工业化的前提。

参考实现。

class ConditionalScoreMatchingTrainer(Trainer):
    def __init__(self, path, model: MLPScore, **kwargs):
        super().__init__(model, **kwargs)
        self.path = path

    def get_train_loss(self, batch_size: int) -> torch.Tensor:
        # 前三步与 3.1 逐字相同
        z = self.path.sample_conditioning_variable(batch_size)   # (bs, dim)
        t = torch.rand(batch_size, 1).to(z)                      # (bs, 1)
        x = self.path.sample_conditional_path(z, t)              # (bs, dim)

        # 唯一的区别：回归目标换成条件得分
        s_theta = self.model(x, t)                               # (bs, dim)
        s_ref   = self.path.conditional_score(x, z, t)           # (bs, dim)
        mse = torch.einsum('bd,bd -> b', s_theta - s_ref, s_theta - s_ref)  # (bs,)
        return torch.mean(mse)

为什么这样写。结构上没有新东西，但有一个数值特性必须知道：

问题 3.3：由向量场反解得分

题目在问什么。不再单独训练一个 score 网络，而是直接从已训好的 $u_t^\theta$ 算出 $\tilde s_t^\theta$，验证「高斯路径下向量场与得分是等价信息」。

题面把结论给了：

但这条关系是整个 Lecture 3-A 的技术核心（讲义 Proposition 1），值得完整推一遍再反解。

参考实现。（下面写的是与题面字面一致的版本，也就是 ScoreFromVectorField 的标准解答；它在 $t=0$ 处会给出 nan，原因见后。）

class ScoreFromVectorField(torch.nn.Module):
    """由学到的向量场参数化 score（只对高斯条件概率路径成立）"""
    def __init__(self, vector_field: MLPVectorField, alpha: Alpha, beta: Beta):
        super().__init__()
        self.vector_field = vector_field
        self.alpha = alpha
        self.beta = beta

    def forward(self, x: torch.Tensor, t: torch.Tensor):
        """x: (bs, dim)   t: (bs, 1)   return: (bs, dim)"""
        alpha_t    = self.alpha(t)       # (bs, 1)
        beta_t     = self.beta(t)        # (bs, 1)
        dt_alpha_t = self.alpha.dt(t)    # (bs, 1)
        dt_beta_t  = self.beta.dt(t)     # (bs, 1)

        # s_t(x) = [u_t(x) - a_t x] / b_t
        num = self.vector_field(x, t) - (dt_alpha_t / alpha_t) * x        # (bs, dim)
        den = beta_t ** 2 * dt_alpha_t / alpha_t - dt_beta_t * beta_t     # (bs, 1)
        return num / den

为什么这样写。

• num 是 (bs,dim)，den 是 (bs,1)，相除自动广播——又一次靠形状约定省掉了 unsqueeze。
• 这个类没有自己的参数：vector_field 是外部传进来的已训模型，alpha/beta 是纯函数。它只是一个「换坐标系」的包装器。这也印证了核心论断——学会了 $u_t^\theta$ 就等于免费得到了 $s_t^\theta$，不需要第二次训练。
• 题面在可视化时用 ts = torch.linspace(0.0, 0.9999, 4)，右端取 $1-\varepsilon$ 而不是 1，就是为了避开 $b_t\big|_{t=1}=0$（此时 $1-\frac t2=\frac12\ne0$，所以其实是 $\dot\beta_t$ 在 $t=1$ 发散；有那个 $10^{-4}$ 兜底也不会真的炸，但取 $0.9999$ 更干净）。

5. Part 4：线性概率路径与任意分布之间的桥

前三部分的高斯路径有一个隐含限制：$p_{\text{simple}}$ 必须是标准高斯（因为 $p_0(x\mid z)=\N(0,\beta_0^2 I)=\N(0,I)$ 与 $z$ 无关这件事，是靠「高斯 + $\alpha_0=0$」才成立的）。Part 4 换一条路径，把这个限制拿掉。

线性条件概率路径（linear conditional probability path）定义为插值随机变量

$p_t(\cdot\mid z)$ 就是 $X_t$ 的分布。边界立刻成立：$t=0$ 时 $X_0\sim p_{\text{simple}}$，$t=1$ 时 $X_1\equiv z$ 即 $\delta_z$。

问题 4.1：实现线性条件概率路径

题目在问什么。实现 sample_conditional_path 与 conditional_vector_field 两个方法。

数学依据。采样直接照抄定义。向量场用与 Q2.3 完全相同的三步法（写流 → 求导 → 反解代回）：

参考实现。

class LinearConditionalProbabilityPath(ConditionalProbabilityPath):
    def __init__(self, p_simple: Sampleable, p_data: Sampleable):
        super().__init__(p_simple, p_data)

    def sample_conditioning_variable(self, num_samples: int) -> torch.Tensor:
        return self.p_data.sample(num_samples)                 # (num_samples, dim)

    def sample_conditional_path(self, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """X_t = (1-t) X_0 + t z，X_0 ~ p_simple"""
        # 关键：p_simple 不再限定为高斯，必须真的从它采样
        x0 = self.p_simple.sample(z.shape[0])                  # (num_samples, dim)
        return (1 - t) * x0 + t * z                            # (num_samples, dim)

    def conditional_vector_field(self, x, z, t) -> torch.Tensor:
        """u_t(x|z) = (z - x) / (1 - t)，只在 t in [0,1) 上有定义"""
        return (z - x) / (1 - t)                               # (num_samples, dim)

    def conditional_score(self, x, z, t) -> torch.Tensor:
        raise Exception("You should not be calling this function!")

为什么这样写。

• 必须写 self.p_simple.sample(...)，不能写 torch.randn_like(z)。这是本题唯一真正的考点。在 Part 2 里 randn_like 是对的（因为那里 $p_{\text{simple}}$ 就是标准高斯），惯性一上来就会写错。写成 randn_like 后果是：问题 4.3 里明明把 p_simple 设成了圆环分布，训出来的模型起点却是高斯——真值行和学习行在 $t=0$ 那一列会一个是圆环、一个是高斯团，一眼可见。
• z.shape[0] 而不是硬编码的 batch_size：这个方法会被 sample_marginal_path 以各种 batch 大小调用（可视化时用到 100000），必须跟着 z 走。
• (1 - t) * x0：t 是 (bs,1)，x0 是 (bs,dim)，广播照旧。
• conditional_score 主动抛异常而不是返回 None 或 0：这是好的防御式设计。如果有人不小心把 ConditionalScoreMatchingTrainer 套在线性路径上，会立刻得到明确的报错，而不是训练出一个悄悄错误的模型。

问题 4.2：用线性路径做 Flow Matching

这一小节没有新代码：直接把 Q3.1 写好的 ConditionalFlowMatchingTrainer 套到 LinearConditionalProbabilityPath 上即可。这本身就是一个设计上的验收点——如果你的 trainer 里写死了任何高斯路径特有的东西（比如直接用 path.alpha），这里就会报错。正确的 trainer 只通过 ConditionalProbabilityPath 的抽象接口说话。

path = LinearConditionalProbabilityPath(
    p_simple = Gaussian.isotropic(dim=2, std=1.0),
    p_data   = CheckerboardSampleable(device, grid_size=4),
).to(device)

linear_flow_model = MLPVectorField(dim=2, hiddens=[64, 64, 64, 64])
trainer = ConditionalFlowMatchingTrainer(path, linear_flow_model)
losses = trainer.train(num_epochs=10000, device=device, lr=1e-3, batch_size=2000)

实测：损失从 19.1 降到 8 左右后进入平台（末尾几个 epoch 在 7.9–8.3 之间）。同样地，它不会到 0，理由与 Q3.1 完全相同。

问题 4.3：在任意源与目标之间架桥

题目在问什么。把 $p_{\text{simple}}$ 从高斯换成另一个「有形状」的分布，观察会发生什么。参考配置是同心圆环 → 棋盘：

path = LinearConditionalProbabilityPath(
    p_simple = CirclesSampleable(device),                  # 两个同心圆环
    p_data   = CheckerboardSampleable(device, grid_size=4),
).to(device)

bridging_flow_model = MLPVectorField(dim=2, hiddens=[100, 100, 100, 100])
trainer = ConditionalFlowMatchingTrainer(path, bridging_flow_model)
losses = trainer.train(num_epochs=20000, device=device, lr=1e-3, batch_size=2000)

注意配置的变化：隐藏层从 64 加宽到 100，epoch 从 10000 翻到 20000。这不是随手调的——源分布本身有结构之后，需要学的映射复杂得多。

6. 如何验证自己的实现

肉眼比对散点图能抓住大部分错误，但对「系数差一点」「符号在某个分支下写反」这类问题非常不敏感。下面这套数值自查覆盖了 Lab 2 的全部 8 个 TODO，本仓库把它们写成了 labs/tests/test_labs.py，Lab 2 部分共 12 项，在本地环境全部通过。你可以照着这个表自己搭一遍——每一项都只要几行。

几条使用建议：

• 第 7 项和第 11 项性价比最高。它们不需要训练、不需要画图，几行代码，却能以 $10^{-7}$ 的分辨率钉死两个最容易出错的公式。任何时候改动了向量场，先跑这两项。
• 测试里 $t$ 一律采在 $[0.02,0.92]$ 之类的内部区间，主动避开 $t=0$ 和 $t=1$。这不是作弊——那两个端点的行为是参数化的固有性质，不属于「实现是否正确」的范畴（见第 4 节末尾的讨论）。
• 第 3 项误差恰好为 0，是因为 $\dot\alpha_t\equiv1$ 在浮点下是精确表示的；第 4 项误差 $2.17\times10^{-4}$ 恰好是 $10^{-4}$ 稳定项的量级，说明除了那个刻意引入的偏差之外没有别的误差。看误差的量级比看「是否小于阈值」信息量大得多。
• 训练相关的部分（Q3.1 / Q3.2 / Q4.2）没法做严格的单元测试，只能看两件事：损失是否收敛到一个正的平台（不是 0，也不是发散），以及真值/学习两行图是否吻合。

7. 本 lab 小结

把八个 TODO 串起来，Lab 2 其实只讲了一件事的两个版本（向量场版与得分版），外加一次「换路径」的泛化演练。

四条值得带走的结论：

做完这个 lab，你手上就有了一套完整的、能跑通的二维生成模型。Lab 3 会把同样的骨架搬到真实数据（MNIST）上，把 MLP 换成 Diffusion Transformer，并加上 VAE 潜空间与 classifier-free guidance——但训练循环的那四行，和你在 Q3.1 写下的一模一样。

延伸阅读

• Lecture 2 · Flow Matching — 本 lab Q2.1–Q2.3、Q3.1、Q4.1 的理论出处：条件路径、条件向量场、边际化技巧与 flow matching 定理。
• Lecture 3-A · 得分函数与 Score Matching — Q2.4、Q3.2、Q3.3 的理论出处：高斯 score 闭式解、SDE extension trick、$u_t$ 与 $s_t$ 的转换公式（Proposition 1）。
• Flow Matching for Generative Modeling (Lipman et al., 2022) — 本 lab 全部公式的原始论文，§3 与 §4 与 Q2.3 的推导逐行对应。
• Flow Straight and Fast: Rectified Flow (Liu et al., 2022) — Part 4 线性路径的来源，以及「多轮 reflow 把轨迹拉直」的完整方案，正好回答 Q4.3 观察里的第 4 条。
• Generative Modeling by Estimating Gradients of the Data Distribution (Song & Ermon, 2019) — 去噪得分匹配用于生成的开山之作，Q3.2 损失的直接来源。
• Score-Based Generative Modeling through SDEs (Song et al., 2020) — 把 score matching 与 SDE / 概率流 ODE 统一起来，第 3 节讲的 Euler–Maruyama 细节在这篇里有严谨表述。
• Elucidating the Design Space of Diffusion Models (Karras et al., 2022) — 关于调度函数 $\alpha_t,\beta_t$ 的设计空间与端点截断的工程实践，可以直接对照 Q2.1 里那个 $10^{-4}$ 的讨论。