Flow Matching Guide and Code 第5章解读：FlatTorus Riemannian Flow Matching 训练逻辑技术文档

本文档详细说明 examples/2d_riemannian_flow_matching_flat_torus.ipynb 中训练块（超参数设置 + 单步训练循环）的代码逻辑、数学原理与函数调用链。

1. 总览

训练目标：在平坦环面 $M = [0, 2\pi)^2$ 上学习一个速度场 $v_\theta(x, t)$，使得从先验 $p_0$ 沿 ODE $\dot{x} = v_\theta(x, t)$ 积分到 $t=1$ 时，得到数据分布 $p_1$。

• 概率路径：测地线路径，$X_t = \exp_{X_0}(\alpha_t \cdot \log_{X_0}(X_1))$，$\alpha_t$ 由调度器给出（本示例中 CondOTScheduler 取 $\alpha_t = t$）。
• 损失：Flow matching 的 L2 损失，拟合条件速度 $\dot{X}_t$。

2. 代码结构概览

2.0 模块与对象层次

训练循环入口：单步顺序为
采样端点 → wrap 上流形 → 采样 $t$ → path.sample 得 $(x_t, \dot{x}_t)$ → vf(x_t, t) 与 dx_t 算 L2 损失 → backward / step。

数据流：inf_train_gen / randn_like → wrap(manifold, ·) → path.sample → PathSample(x_t, dx_t, ...) → vf(x_t, t) 与 dx_t 求差、均方 → loss。

3. 超参数与对象初始化

3.1 超参数

3.2 速度场模型 vf

vf = ProjectToTangent(MLP(...), manifold=manifold)
vf.to(device)

• MLP：输入 $(x, t)$（流形坐标 + 时间），输出 $\mathbb{R}^{\mathrm{dim}}$ 的向量（欧氏速度）。
• ProjectToTangent：对输入 $x$ 做 manifold.projx(x)，再对 MLP 输出做 manifold.proju(x, v)，保证输出是 $x$ 处切空间中的向量。
• 调用链（推理/训练时）：
  • vf(x, t) → ProjectToTangent.forward(x, t)
  • → x' = manifold.projx(x)（环面：$x' = x \bmod 2\pi$）
  • → v = MLP(x', t)
  • → v' = manifold.proju(x', v)（FlatTorus 上恒等）
  • → 返回 v'

3.3 概率路径 path

path = GeodesicProbPath(scheduler=CondOTScheduler(), manifold=manifold)

• GeodesicProbPath：基于流形测地线的概率路径；需要 ConvexScheduler 提供 $\alpha_t$。
• CondOTScheduler：$\alpha_t = t$，$\sigma_t = 1 - t$，$\dot{\alpha}_t = 1$，$\dot{\sigma}_t = -1$。
• path.sample(t, x_0, x_1) 的数学与实现见下文 §5。

3.4 优化器

optim = torch.optim.Adam(vf.parameters(), lr=lr)

仅优化 vf（即 ProjectToTangent 内的 MLP）的参数。

4. 单步训练循环：逻辑与数据流

每一步迭代完成以下流程（顺序与代码一致）。

4.1 清空梯度

optim.zero_grad()

为当前步的 loss.backward() 做准备。

4.2 从耦合 $\pi(X_0, X_1)$ 采样

• x_1：数据端

  • 调用 inf_train_gen(batch_size=batch_size, device=device)
  • 在平面区域上生成棋盘格状样本，形状 (batch_size, 2)，数值约在 $[-2,2]^2$ 或类似范围。

• x_0：先验端

  • x_0 = torch.randn_like(x_1).to(device)
  • 即 $X_0 \sim \mathcal{N}(0, I)$（与 x_1 同 shape）。

4.3 将端点投影到流形 $[0, 2\pi)^2$

x_1 = wrap(manifold, x_1)
x_0 = wrap(manifold, x_0)

• wrap(manifold, samples) 实现：
  • center = zeros_like(samples)
  • return manifold.expmap(center, samples)
• 对 FlatTorus：expmap(0, u) = u % (2π)，因此
  • x_1、x_0 被映射到 $[0, 2\pi)^2$，保证路径两端都在流形上。

调用关系：
wrap → manifold.expmap(0, samples) → FlatTorus 上等价于 samples % (2π)。

4.4 采样时间 $t$

t = torch.rand(x_1.shape[0]).to(device)

• 每个样本一个 $t \in [0,1]$，shape (batch_size,)。

4.5 沿概率路径采样 $(X_t, \dot{X}_t)$

path_sample = path.sample(t=t, x_0=x_0, x_1=x_1)

• path：GeodesicProbPath(scheduler=CondOTScheduler(), manifold=manifold)
• path.sample(x_0, x_1, t)（即 GeodesicProbPath.sample(x_0, x_1, t)）返回 PathSample，包含：
  • x_t：路径在时间 $t$ 上的点，流形坐标 $[0,2\pi)^2$
  • dx_t：路径在 $t$ 处对时间的导数（条件目标速度）
  • x_0, x_1, t（同输入）

内部函数调用链（见 §5）：

1. expand_tensor_like(t, x_1[..., 0:1])，使 t 与 batch 维度对齐。
2. 对每个 $(x_0^{(i)}, x_1^{(i)}, t^{(i)})$（通过 vmap）：
   • geodesic(manifold, x_0, x_1) → 得到可调用对象 path(τ)，满足
     $path(\tau) = \exp_{x_0}(\tau \cdot \log_{x_0}(x_1))$，$\tau \in [0,1]$。
   • scheduler(t) → SchedulerOutput(alpha_t=t, sigma_t=1-t, d_alpha_t=1, d_sigma_t=-1)（CondOT 下 $\alpha_t = t$）。
   • path(alpha_t) = path(t) = 测地线在 $t$ 处的点 $X_t$。
   • jvp：对 $\tau \mapsto path(\mathrm{scheduler}(\tau).\mathrm{alpha\_t})$ 在 $\tau=t$ 处求导（CondOT 下即对 $path(\tau)$ 在 $\tau=t$ 求导），得到 $X_t$ 与 $\frac{\partial}{\partial t}X_t = \dot{X}_t$（即 dx_t）。
3. 将 vmap 结果 reshape 成与 x_1 相同 shape，构造 PathSample(x_t, dx_t, x_1, x_0, t) 并返回。

4.6 计算 Flow Matching L2 损失

loss = torch.pow(vf(path_sample.x_t, path_sample.t) - path_sample.dx_t, 2).mean()

• path_sample.x_t：流形上点，坐标在 $[0, 2\pi)^2$（与 FlatTorus 定义一致）。
• path_sample.t：与 x_t 一一对应的时间，shape 与 batch 兼容。
• vf(path_sample.x_t, path_sample.t)：
  • 调用 ProjectToTangent.forward(x_t, t)
  • 内部：projx(x_t) → MLP → proju(x_t, v)，输出切空间中的预测速度，shape 与 x_t 一致。
• path_sample.dx_t：测地线在 $t$ 处的真实条件速度 $\dot{X}_t$（目标）。
• 损失：$\mathcal{L} = \mathbb{E}\bigl[\|v_\theta(X_t, t) - \dot{X}_t\|^2\bigr]$，即 L2 拟合速度场。

4.7 反向传播与参数更新

loss.backward()
optim.step()

• 梯度只流入 vf（MLP + ProjectToTangent 的 projx/proju 若可微则参与，FlatTorus 的 projx/proju 可微）。

4.8 日志

每隔 print_every 步打印当前迭代数、平均每步耗时和当前 loss.item()。

5. 关键函数调用说明

5.1 函数签名与角色速查

5.2 geodesic(manifold, x_0, x_1)（flow_matching.utils.manifolds.utils）

• 作用：构造从 $x_0$ 到 $x_1$ 的测地线，参数化在 $\tau \in [0,1]$。
• 数学：
  • $u = \mathrm{logmap}(x_0, x_1)$（从 $x_0$ 指向 $x_1$ 的切向量）
  • $path(\tau) = \mathrm{expmap}(x_0, \tau \cdot u)$
• FlatTorus：
  • logmap(x, y) = atan2(sin(y-x), cos(y-x))（最短方向，考虑周期）
  • expmap(x, u) = (x + u) % (2π)
• 返回：可调用对象 path(τ)，输入 τ 的 shape 与 batch 兼容，输出流形上点的 shape 与 x_0/x_1 一致。

5.3 CondOTScheduler.__call__(t)

• 输入：t，shape (batch_size,) 或可广播。
• 输出：SchedulerOutput：
  • alpha_t = t
  • sigma_t = 1 - t
  • d_alpha_t = 1
  • d_sigma_t = -1
• 本示例中只用到 alpha_t，用于测地线参数：$X_t = path(\alpha_t) = path(t)$。

5.4 GeodesicProbPath.sample(x_0, x_1, t) 内部

• assert_sample_shape：检查 x_0, x_1, t 的 batch 等维度兼容性。
• t 经 expand_tensor_like 与 x_1[..., 0:1] 对齐，便于与 x_0, x_1 一起做 vmap。
• cond_u(x_0, x_1, t)（单样本逻辑，再被 vmap）：
  1. path = geodesic(manifold, x_0, x_1)
  2. alpha_t = scheduler(t).alpha_t（CondOT 下即为 t）
  3. x_t, dx_t = jvp(lambda t: path(scheduler(t).alpha_t), (t,), (1,))
     • 对 $\tau \mapsto path(\tau)$ 在 $\tau = t$ 处求导，得到 $path(t)$ 和 $\frac{d}{d\tau}path(\tau)\big|_{\tau=t}$；由于 $\tau=t$ 时 $\frac{d}{dt}\alpha_t = 1$，这里得到的就是 $\dot{X}_t$。
• vmap(cond_u)(x_0, x_1, t)：对整批做上述计算。
• reshape_as(x_1)：保证 x_t、dx_t 的 shape 与 x_1 一致。
• 返回：PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)。

5.5 vf(x, t) = ProjectToTangent.forward(x, t)

• x = manifold.projx(x)：FlatTorus 上为 x % (2π)，保证 $x$ 在 $[0, 2\pi)^2$。
• v = vecfield(x, t)：MLP 在欧氏空间预测速度向量。
• v = manifold.proju(x, v)：FlatTorus 上切空间即 $\mathbb{R}^2$，故恒等。
• 返回的 v 即模型给出的切空间速度，与 path_sample.dx_t 的坐标与物理意义一致。

5.6 函数调用关系总览

训练循环 (每步)
├── inf_train_gen(batch_size, device)     → x_1
├── torch.randn_like(x_1)                 → x_0
├── wrap(manifold, x_1), wrap(manifold, x_0)
│   └── manifold.expmap(0, samples)
├── torch.rand(...)                       → t
├── path.sample(t, x_0, x_1)
│   ├── expand_tensor_like(t, x_1)
│   └── vmap(cond_u)(x_0, x_1, t)
│       ├── geodesic(manifold, x_0, x_1)  → path(τ)，内部用 logmap/expmap
│       ├── scheduler(t)                  → alpha_t, sigma_t, ...
│       └── jvp(path(scheduler(·).alpha_t), t, 1) → x_t, dx_t
├── vf(path_sample.x_t, path_sample.t)
│   ├── manifold.projx(x)
│   ├── MLP(x, t)
│   └── manifold.proju(x, v)
└── loss = mean((vf(...) - dx_t)²); loss.backward(); optim.step()

6. 数据与坐标约定

• 流形坐标：FlatTorus 全程使用 $[0, 2\pi)^2$。

  • inf_train_gen 生成的数据经 wrap 后落入该域；
  • x_0 经 wrap 后也在该域；
  • path.sample 得到的 x_t、dx_t 以及 vf 的输入/输出均在该坐标下。

• 时间：$t \in [0,1]$，$t=0$ 对应先验端，$t=1$ 对应数据端；CondOT 下测地线参数 $\alpha_t = t$，故 $X_t$ 即为从 $x_0$ 到 $x_1$ 的测地线在 $t$ 处的点。

7. 小结：单步训练的数据与调用顺序

inf_train_gen(batch_size, device)  →  x_1 (平面棋盘格)
torch.randn_like(x_1)              →  x_0 (高斯)
wrap(manifold, x_1), wrap(manifold, x_0)  →  x_1, x_0 ∈ [0,2π)²
torch.rand(...)                    →  t ∈ [0,1]

path.sample(t, x_0, x_1)
  → geodesic(manifold, x_0, x_1)   →  path(τ)
  → scheduler(t).alpha_t           →  α_t = t
  → jvp(path(α_t), t, 1)           →  x_t, dx_t
  → PathSample(x_t, dx_t, x_1, x_0, t)

vf(path_sample.x_t, path_sample.t)
  → projx(x_t) → MLP(x_t, t) → proju(x_t, v)  →  v_θ(x_t, t)

loss = mean(|v_θ(x_t,t) - dx_t|²)
loss.backward(); optim.step()

以上即平坦环面示例中训练块的完整逻辑与函数调用说明。