""" fewstep_lib.py —— 第 7 讲的工具箱：少步生成（流映射、reflow、蒸馏、一致性模型、MeanFlow） ===================================================================================== 【这个模块提供什么】 全讲共用一个"考场"：第 1 讲的五峰分布 + CondOT 路径。它的边缘向量场有闭式解（第 2 讲 gmm_oracle）， 所以可以算出**精确的流映射** Φ(x_0) = ψ_{0→1}(x_0)——"无穷多步的 ODE 会把噪声 x_0 送到哪里"。 这是所有少步方法想要逼近的东西；有了它，就能把"方法本身的误差"和"网络没学好"分开。 1. 考场：make_problem、flow_map（精确流映射）、各种 ODE 求解器（欧拉 / 中点 / Heun）、evaluate（打分） 2. 网络：TimeMLP —— 输入 (x, 若干个时间)，输出 2 维向量（速度、平均速度或数据点） 3. 训练方法（每个都写清楚回归的是什么）： train_fm 条件流匹配（独立耦合：x_0 和 z 独立抽）——基线 train_reflow reflow：在 (x_0, Φ(x_0)) 配对上做流匹配（Liu et al. 2023） train_direct 直接蒸馏：一步回归 x_0 → Φ(x_0)（Luhman & Luhman 2021） train_consistency 一致性蒸馏（Song et al. 2023），老师是精确的边缘向量场 train_meanflow MeanFlow：学平均速度，用 MeanFlow 恒等式构造目标（Geng et al. 2025） 【约定】时间 t = 0 是噪声、t = 1 是数据（全课统一）。论文里一致性模型和 MeanFlow 用的是相反的方向， 这里的公式都已经翻译成本课的方向，翻译方法在讲义里写明。 """ import copy import os import sys import time import torch import torch.nn as nn HERE = os.path.dirname(os.path.abspath(__file__)) for sub in ("lecture01", "lecture02"): sys.path.insert(0, os.path.join(HERE, "..", sub)) from distributions import energy_distance, five_modes, mode_report # noqa: E402 from gmm_oracle import GaussianPathGMMOracle # noqa: E402 from paths import CondOT # noqa: E402 DT = torch.float64 # ===================================================================== # 1. 考场 # ===================================================================== def make_problem(n_ref=4000, seed=123): """ 五峰分布 + CondOT 路径的精确边缘向量场。 【返回】(gmm, oracle, ref)：ref 是 n_ref 个真样本，用来算能量距离 """ gmm = five_modes(DT) oracle = GaussianPathGMMOracle(gmm, CondOT()) ref = gmm.sample(n_ref, torch.Generator().manual_seed(seed)) return gmm, oracle, ref def solve(field, x0, n, t0=0.0, t1=1.0, method="euler"): """ 用 n 步求解 dx/dt = field(x, t)，从 t0 到 t1。 【参数】 field : 函数 (x, t) → (n, d)，t 是 Python float method : "euler"（每步 1 次 field，一阶）、"midpoint" / "heun"（每步 2 次，二阶） 【返回】(x_{t1}, nfe)：nfe = 调用 field 的次数（"网络前向次数"） """ x = x0.clone() h = (t1 - t0) / n for k in range(n): t = t0 + k * h if method == "euler": x = x + h * field(x, t) elif method == "midpoint": x = x + h * field(x + 0.5 * h * field(x, t), t + 0.5 * h) elif method == "heun": k1 = field(x, t) k2 = field(x + h * k1, t + h) x = x + 0.5 * h * (k1 + k2) else: raise ValueError(method) return x, n * (1 if method == "euler" else 2) def flow_map(oracle, x0, t0=0.0, t1=1.0, n=200): """ "精确"流映射 ψ_{t0→t1}(x0)：用精确的边缘向量场、Heun 法 n 步（默认 200 步，400 次求值）。 【边界】01 号脚本以 Heun 800 步为参照检验：200 步的误差中位数 2e-5、99% 分位 5e-4， 但最大可达 7e-3——在两个峰的"分界线"附近，轨迹对起点极其敏感，数值误差也被放大。 这比任何少步方法的误差小几个数量级，所以当作"真值"。10 万个点约 5 秒。 """ return solve(lambda x, t: oracle.vector_field(x, t), x0.to(DT), n, t0, t1, "heun")[0] def evaluate(x, gmm, ref): """ 给一批样本打分：能量距离（越小越好）、峰权重的 TV、落在峰之间的比例、峰内散度（理想值 1）。 【返回】dict(ed, tv, far, spread) """ m = mode_report(gmm, x.to(DT)) return {"ed": energy_distance(x.to(DT), ref), "tv": m["tv"], "far": m["far_rate"], "spread": m["spread"]} def fmt(m): """把 evaluate 的结果格式化成 "能量距离 / TV / 峰之间" 三个字符串。""" return [f"{m['ed']:.4f}", f"{m['tv']:.3f}", f"{m['far']:.1%}"] def straightness(field, x0, n=200): """ 轨迹的"不直"程度（Liu et al. 2023 的 straightness）： S = E ∫_0^1 ‖u_t(x_t) − (x_1 − x_0)‖² dt S = 0 当且仅当每条轨迹都是匀速直线——这时欧拉法一步就是精确的。 【参数】field：(x, t) → 速度；x0：(n, d) 起点；n：积分步数（Heun） 【返回】(S, x1)：S 是 float；x1 是终点 """ h = 1.0 / n xs, vs = [x0.clone()], [] x = x0.clone() for k in range(n): t = k * h k1 = field(x, t) k2 = field(x + h * k1, t + h) vs.append(k1) x = x + 0.5 * h * (k1 + k2) xs.append(x) disp = x - x0 s = torch.stack([((v - disp) ** 2).sum(dim=1) for v in vs]).mean().item() return s, x # ===================================================================== # 2. 网络 # ===================================================================== class TimeMLP(nn.Module): """ 输入 (x, t_1, ..., t_k)（x 是 2 维，后面是 k 个时间标量），输出 2 维向量。 【参数】n_times：时间输入的个数（流匹配 1 个：t；MeanFlow 2 个：t 和 s；直接蒸馏 0 个） hidden、depth：MLP 的宽度和隐藏层数 【返回】forward(x, *times)：times 里每个是 (n, 1) 张量或 Python float 【Python 小知识】*times 收集任意个位置参数成一个元组。 copy.deepcopy(model)（一致性蒸馏里用）复制出一个参数完全相同、但互不影响的新网络。 """ def __init__(self, n_times=1, hidden=128, depth=4, dim=2): super().__init__() layers, width = [], dim + n_times for _ in range(depth): layers += [nn.Linear(width, hidden), nn.SiLU()] width = hidden layers.append(nn.Linear(width, dim)) self.net = nn.Sequential(*layers) def forward(self, x, *times): cols = [x] for t in times: cols.append(t if torch.is_tensor(t) else torch.full((x.shape[0], 1), float(t), dtype=x.dtype)) return self.net(torch.cat(cols, dim=1)) def as_field(model): """把一个吃 (x, t) 的 float32 网络包装成 solve 需要的 field（输入输出 float64）。""" @torch.no_grad() def field(x, t): return model(x.float(), t).double() return field def _loop(model, loss_fn, steps, lr, seed, log_every=None): """ 通用的训练循环：Adam，学习率在最后 20% 的步数里线性降到 0（让结果更稳定）。 【参数】loss_fn(g) → 标量损失，g 是随机数生成器 【返回】(logs, elapsed) """ g = torch.Generator().manual_seed(seed) opt = torch.optim.Adam(model.parameters(), lr=lr) sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda k: min(1.0, (steps - k) / (0.2 * steps))) logs, running = [], [] t0 = time.perf_counter() model.train() for step in range(1, steps + 1): loss = loss_fn(g) opt.zero_grad() loss.backward() opt.step() sched.step() running.append(loss.item()) if log_every and step % log_every == 0: logs.append({"step": step, "loss": sum(running) / len(running)}) running = [] model.eval() return logs, time.perf_counter() - t0 # ===================================================================== # 3. 训练方法 # ===================================================================== def train_fm(model, gmm, steps=6000, batch=512, lr=1e-3, seed=0): """ 条件流匹配（第 2 讲 Algorithm 3，CondOT）：x_0 ~ N(0, I) 和 z ~ p_data **独立**抽（"独立耦合"）， x_t = (1 − t) x_0 + t z，回归目标 z − x_0。 【返回】(logs, elapsed) """ def loss_fn(g): z = gmm.sample(batch, g).float() x0 = torch.randn(batch, 2, generator=g) t = torch.rand(batch, 1, generator=g) xt = (1 - t) * x0 + t * z return ((model(xt, t) - (z - x0)) ** 2).sum(dim=1).mean() return _loop(model, loss_fn, steps, lr, seed) def train_reflow(model, pairs, steps=6000, batch=512, lr=1e-3, seed=0): """ Reflow（Rectified Flow 的第二轮）：在**配对**好的 (x_0, x_1) 上做同样的流匹配。 配对来自上一轮模型的 ODE：x_1 = Φ(x_0)。损失和 train_fm 完全一样，只是 z 换成了"x_0 的搭档"。 【参数】pairs：(x0s, x1s)，各 (N, 2) 【直觉】独立耦合下，从不同 x_0 出发的直线会交叉，边缘场只能"平均"它们，于是轨迹弯曲； ODE 的配对不交叉（ODE 的轨迹不相交），在它上面学到的边缘场 ≈ 条件场 = 直线。 【边界】配对数据是有限的（N 个）；网络在有限配对上训练，可能记住配对（这里 N = 10 万，问题不大）。 """ x0s, x1s = pairs[0].float(), pairs[1].float() def loss_fn(g): idx = torch.randint(0, x0s.shape[0], (batch,), generator=g) x0, x1 = x0s[idx], x1s[idx] t = torch.rand(batch, 1, generator=g) xt = (1 - t) * x0 + t * x1 return ((model(xt, t) - (x1 - x0)) ** 2).sum(dim=1).mean() return _loop(model, loss_fn, steps, lr, seed) def train_direct(model, pairs, steps=6000, batch=512, lr=1e-3, seed=0): """ 直接蒸馏：一个没有时间输入的网络 G(x_0)，直接回归 Φ(x_0)（均方误差）。一步生成：x_1 = G(x_0)。 【边界】Φ 在"两个峰的分界线"附近变化极快（01 号脚本），MLP 学不出这么陡的跳变， 会把分界线附近的 x_0 送到两个峰之间（"模糊"）。 """ x0s, x1s = pairs[0].float(), pairs[1].float() def loss_fn(g): idx = torch.randint(0, x0s.shape[0], (batch,), generator=g) return ((model(x0s[idx]) - x1s[idx]) ** 2).sum(dim=1).mean() return _loop(model, loss_fn, steps, lr, seed) def consistency_fn(model, x, t): """ 一致性函数 f(x, t) = x + (1 − t) F_θ(x, t)：从时刻 t 的 x 直接预测轨迹在 t = 1 的终点。 【为什么这样参数化】边界条件 f(x, 1) = x（已经在数据端了，原地不动）自动满足， 不需要网络去学。F_θ 的含义是"从 t 到 1 的平均速度"（和 MeanFlow 的 u(x, t, 1) 是同一个东西）。 这是 Song et al. 的 c_skip / c_out 参数化在本课时间方向下的最简单版本。 """ return x + (1 - t) * model(x, t) def train_consistency(model, oracle, gmm, n_disc=32, steps=6000, batch=512, lr=1e-3, seed=0, ema=0.99): """ 一致性蒸馏（Song et al. 2023 的 CD），老师是精确的边缘向量场（"完美的老师"）。 把 [0, 1] 分成 n_disc 段：t_0 = 0 < t_1 < ... < t_N = 1。每一步： z ~ p_data，ε ~ N(0, I)，随机选 k，x_{t_k} = (1 − t_k) ε + t_k z （先落到路径上） 老师走一步欧拉：x̂_{t_{k+1}} = x_{t_k} + (t_{k+1} − t_k) u_{t_k}(x_{t_k}) （沿 ODE 往数据端走一小段） 损失 = ‖ f_θ(x_{t_k}, t_k) − f_{θ⁻}(x̂_{t_{k+1}}, t_{k+1}) ‖² θ⁻ 是 θ 的指数滑动平均（EMA），不求梯度——"目标网络"。 【直觉】同一条 ODE 轨迹上相邻两点，应该预测出同一个终点（"一致"）。t_{k+1} = 1 时 f(x, 1) = x 是已知的，于是正确的终点从数据端一段一段地"传"回噪声端。 【参数】n_disc：时间离散的段数；ema：目标网络的滑动平均系数 oracle：老师（精确的边缘向量场）。传 None 时是**一致性训练**（CT，不要老师）： x̂_{t_{k+1}} 直接取同一条**条件**路径上的点 (1 − t_{k+1}) ε + t_{k+1} z—— 用条件速度 z − ε 代替边缘速度走那一小步，和流匹配用条件向量场代替边缘向量场是同一个道理 （n_disc → ∞ 时期望一致，Song et al. 2023 定理 2 的思路）。 【返回】(logs, elapsed) 【边界】n_disc 太小：老师的欧拉步误差大；太大：每一步只传一点点信息，收敛慢（练习 5）。 """ target = copy.deepcopy(model) # 目标网络 θ⁻：一开始和 θ 相同 for p in target.parameters(): p.requires_grad_(False) grid = torch.linspace(0, 1, n_disc + 1) def loss_fn(g): z = gmm.sample(batch, g).float() eps = torch.randn(batch, 2, generator=g) k = torch.randint(0, n_disc, (batch, 1), generator=g) t, t_next = grid[k], grid[k + 1] xt = (1 - t) * eps + t * z with torch.no_grad(): if oracle is not None: u = oracle.vector_field(xt.double(), t.double()).float() x_next = xt + (t_next - t) * u # 老师沿边缘 ODE 走一小步 else: x_next = (1 - t_next) * eps + t_next * z # 沿条件路径走一小步（同一个 ε、z） tgt = consistency_fn(target, x_next, t_next) return ((consistency_fn(model, xt, t) - tgt) ** 2).sum(dim=1).mean() # 训练循环需要在每步之后更新 EMA，所以这里不用 _loop，而是写出完整的循环 gen = torch.Generator().manual_seed(seed) opt = torch.optim.Adam(model.parameters(), lr=lr) sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda k: min(1.0, (steps - k) / (0.2 * steps))) t0 = time.perf_counter() model.train() for _ in range(steps): loss = loss_fn(gen) opt.zero_grad() loss.backward() opt.step() sched.step() with torch.no_grad(): for p_t, p in zip(target.parameters(), model.parameters()): p_t.mul_(ema).add_(p, alpha=1 - ema) model.eval() return [], time.perf_counter() - t0 def consistency_sample(model, x0, ts=(0.0,), generator=None): """ 一致性模型的（多步）采样： 1 步：x_1 = f(x_0, 0)； 多步：先 x̂ = f(x_0, 0)，再"退回"到时刻 t_k：x_{t_k} = (1 − t_k) ε + t_k x̂（重新加噪），再 x̂ = f(x_{t_k}, t_k)…… 【参数】ts：依次使用的时刻，第一个必须是 0；例如 (0, 0.5) 是两步 【返回】(n, 2) float64 """ with torch.no_grad(): x = consistency_fn(model, x0.float(), ts[0]) for t in ts[1:]: eps = torch.randn(x0.shape, generator=generator) x = consistency_fn(model, (1 - t) * eps + t * x, t) return x.double() def meanflow_u(model, x, t, s): """MeanFlow 网络：u_θ(x, t, s) = 从时刻 t 的 x 出发、走到时刻 s 的平均速度。""" return model(x, t, s) def train_meanflow(model, gmm, steps=6000, batch=512, lr=1e-3, seed=0, frac_diff=0.25): """ MeanFlow（Geng et al. 2025），翻译成本课的时间方向（t = 0 噪声，t = 1 数据）。 平均速度：u(x_t, t, s) = 1/(s − t) ∫_t^s v(x_τ, τ) dτ（s ≥ t），所以 x_s = x_t + (s − t) u(x_t, t, s)。 对 t 求导（x_t 沿轨迹移动）得到 **MeanFlow 恒等式**： u(x_t, t, s) = v(x_t, t) + (s − t) · d/dt u(x_t, t, s)， d/dt u = ∂_x u · v + ∂_t u。 训练目标（stop-gradient）：u_tgt = v_cond + (s − t)(∂_x u_θ · v_cond + ∂_t u_θ)， v_cond = z − x_0 是 CondOT 的条件速度（和流匹配一样，用条件速度代替边缘速度，期望不变）。 导数项用一次"雅可比–向量积"（JVP）算：沿方向 (v_cond, 1, 0) 对 (x, t, s) 求方向导数。 【参数】frac_diff：s ≠ t 的样本比例；其余样本 s = t，此时目标就是 v_cond，退化成普通流匹配 【返回】(logs, elapsed) 【Python 小知识】torch.func.jvp(f, primals, tangents) 一次前向同时算出 f(primals) 和方向导数， 比"先求整个雅可比矩阵再乘向量"便宜得多。 【边界】s − t 大时目标里有网络自己的导数（"自举"），训练早期目标不准；frac_diff 太大会不稳定。 """ import warnings # 某些 PyTorch 版本在第一次调用 torch.func 时，内部会触发一条"torch.jit.script 已弃用"的 FutureWarning； # 它和本课的代码无关，这里只屏蔽这一条 warnings.filterwarnings("ignore", message=".*torch.jit.script.*", category=FutureWarning) from torch.func import jvp def loss_fn(g): z = gmm.sample(batch, g).float() x0 = torch.randn(batch, 2, generator=g) a = torch.rand(batch, 1, generator=g) b = torch.rand(batch, 1, generator=g) t, s = torch.minimum(a, b), torch.maximum(a, b) same = torch.rand(batch, 1, generator=g) >= frac_diff s = torch.where(same, t, s) xt = (1 - t) * x0 + t * z v = z - x0 u, du_dt = jvp(lambda xx, tt, ss: model(xx, tt, ss), (xt, t, s), (v, torch.ones_like(t), torch.zeros_like(s))) tgt = (v + (s - t) * du_dt).detach() return ((u - tgt) ** 2).sum(dim=1).mean() return _loop(model, loss_fn, steps, lr, seed) def meanflow_sample(model, x0, n_steps=1): """MeanFlow 采样：把 [0, 1] 分成 n_steps 段，每段 x ← x + (s − t) u(x, t, s)。【返回】(n, 2) float64""" with torch.no_grad(): x = x0.float() for k in range(n_steps): t, s = k / n_steps, (k + 1) / n_steps x = x + (s - t) * model(x, t, s) return x.double()