One-Step Generative Surrogate Models via Block-Triangular Joint Drifting
Nicholas Geissler, Shreya Jha, Ricardo Baptista, Benjamin Peherstorfer
cs.LG, math.NA
2026-09-22
对相邻状态联合分布做块三角投影漂移，一次网络评估采样下一步。二维湍流能量误差2.14e-2，约为单步ReFlow蒸馏的四分之一。
随机系统的廉价代理要能从当前状态采样下一步的条件分布，再自回归滚出整条轨迹。扩散和流模型能拟合这种转移律，但每一步物理时间往往要 K 次去噪或积分，T 步 rollout 就是 KT 次网络评估，集合预报会很贵。
蒸馏路线（MeanFlow、ReFlow、consistency）能把 K 压到 1，通常要先训一个多步老师再压缩，精度换速度。只拟合时间边缘的方法（DICE、边际扩散）更快，但边缘定不了轨迹级统计。确定性步进器从同一初值永远走出同一条路径，表示不了固有随机性。
Drifting 本来就是为单步生成设计的：训练时用分布相关的漂移场把模型样本推向目标，收敛后噪声到样本只需一次前向。麻烦在于漂移场要靠目标分布的经验样本来构造。轨迹数据对每个观测到的当前状态只有一个实现的下一步，条件分布的经验近似就是点质量，没有宽度、形状、多峰。扩散和流匹配的回归目标不需要同一条件状态的多次实现，漂移需要。
BTJD 换了漂移目标。不去漂条件转移律，去漂相邻状态的联合分布。轨迹对本身就是联合分布的样本，只要轨迹条数大于 1，每个时刻都有一堆对可以用。
生成器钉成块三角形式：输出第一块恒等于当前状态、与参考噪声无关，所以当前状态边缘永远保持为真实边缘；第二块才是可学习的条件映射 g。联合分布一旦匹配上，g 就把噪声推到条件转移律上。这个结构来自块三角传输的既有结果。推理时给定当前状态抽一个噪声，一次网络评估得到下一步。
漂移场对联合分布的 Sinkhorn 散度求一阶变分，再正交投影到第二块：只移动下一步分量，当前状态钉死。损失因此退化成只对 g 的平方回归。时间上参数共享，每个 batch 抽若干时刻和若干转移对，用经验 Sinkhorn 重心投影估场。
理论给了一个无正则（熵系数 ε=0）的理想化定理：紧支撑、密度上下有界、两边第一边缘相同，投影后的漂移场为零当且仅当模型联合分布等于目标。实际训练全是 ε>0 的 Sinkhorn。Duffing 上 ε=0.01、每步 50 次迭代；PDE 上 ε 约 0.07 到 0.1、20 次迭代。高维场先压进自编码器潜空间再漂。
四组实验。低维用 sliced Wasserstein-2 看边缘，再用路径相关的兴趣量看是不是真学到了转移律。
Duffing 振子（噪声强度 0.5，1200 步，5000 条训练轨迹），随机初值：
| 方法 | sliced-W2 | 轨迹 QoI 误差 |
| 确定性步进器 | 3.48e-1 | 3.68e-2 |
| DICE | 5.80e-1 | 8.30e-2 |
| 边际扩散 | 8.10e-2 | 1.07e-1 |
| SDE learning | 8.30e-2 | 5.70e-3 |
| SDE matching | 5.49e-1 | 4.47e-2 |
| BTJD | 4.80e-2 | 2.75e-3 |
固定同一初值时，BTJD 的 W2 是 5.80e-2、QoI 2.49e-3，仍低于 SDE learning 的 1.08e-1 和 6.80e-3。同一点出发能散开一簇轨迹，随机性来自转移律，不是初值抖动。
9 维 Rayleigh–Bénard，训练控制参数取 13.5 到 14.2 的离散点，测试未见参数 13.65：BTJD 的 W2 为 4.60e-2，旋转电流 QoI 为 3.00e-4。次优 SDE learning 是 5.40e-2 和 2.20e-2。边缘对了不算完，旋转方向这种轨迹量也对上了。
高维才跟多步生成模型比。随机 Burgers，64 格点压到 16 维潜空间：
| 方法 | 能量误差 | 涡量平方误差 | 每步 NFE |
| Operator learning | 2.21e-2 | 2.55e-1 | 1 |
| ARDM 100 步 | 1.24e-2 | 2.11e-1 | 100 |
| CFM 20 步 | 2.71e-3 | 1.53e-1 | 20 |
| MeanFlow 1 步 | 8.22e-1 | 3.59e+2 | 1 |
| ReFlow+蒸馏 1 步 | 2.78e-3 | 1.40e-1 | 1 |
| BTJD | 2.22e-3 | 5.87e-2 | 1 |
涡量平方误差不到次优方法的一半。MeanFlow 单步在这组上崩掉了。
二维随机强迫湍流，训练分辨率 64×64：
| 方法 | 能量误差 | 涡量平方误差 | 每步 NFE |
| ARDM 100 步 | 1.16e-1 | 1.34e-1 | 100 |
| CFM 20 步 | 1.39e-1 | 1.04e-1 | 20 |
| MeanFlow 4 步 | 1.22e-1 | 6.74e-2 | 4 |
| ReFlow+蒸馏 1 步 | 8.60e-2 | 7.16e-2 | 1 |
| BTJD | 2.14e-2 | 2.52e-2 | 1 |
能量误差大约是次优 ReFlow 的四分之一，涡量平方大约三分之一。同一初值多次 rollout 会分叉成不同湍流实现；确定性算子学习从同一初值永远走出同一条。
集合预报、不确定性量化和外循环设计需要又准又便宜的随机轨迹。BTJD 把单步采样做成训练目标本身，不经过老师蒸馏。实验室和仿真日志里的轨迹，本来就是「每个状态只有一个后继」，这个设定对得上。
适合已经有一批随机轨迹、要快速再采样路径级统计（能量、涡量、过势垒次数）的场景，而不是只画某一时刻的直方图。训练阶段每步要跑 Sinkhorn，推理才是真便宜。
Burgers 上 MeanFlow 单步崩掉，说明「先训 CFM 再蒸馏到 1-NFE」不是免费午餐。这仍是科学计算向的生成式降阶模型，不是通用图像生成器。
固定点结论只覆盖 ε=0、紧支撑、密度有界的理想情况。实验全是 ε>0 的 Sinkhorn。投影会不会引入额外零点，论文只在这个理想设定里排除。
训练开销几乎没报。每步 20 到 50 次 Sinkhorn，batch 大到 8192 粒子，低维 10 万步、湍流 6.5 万步。1-NFE 是推理账，不是训练账。
Burgers 和湍流都先压进自编码器，表格误差是解码后的场量，潜空间重构误差没有单独拆开。低维实验没放 CFM、ARDM 和蒸馏，高维才放，跨设定的全面最优要打折。MeanFlow 在 Burgers 上的灾难性失败，对比强度取决于蒸馏实现有没有调到位。
条件输入还人为加了 1% 到 2% 的高斯噪声。长程误差累积、超出训练时间窗的外推，没有专门压力测试。