Wasserstein Residuals: Learning Gradient Flows from Population Dynamics¶

会议: NeurIPS 2026
arXiv: 2607.04738
代码: https://github.com/BasisResearch/wasserstein-residuals
领域: 物理与科学计算
关键词: Wasserstein梯度流, 群体动力学, 速度残差, 逆问题, 核密度估计

一句话总结¶

针对传统 JKO 方案依赖高耗时最优传输求解且难以应对大时间跨度稀疏快照的瓶颈，本文提出 Wasserstein 速度残差最小化框架与免模拟的拼接（stitching）算法，将分布曲线作为一阶可学习粒子变量，实现了对未知能量泛函与群体动力学轨迹的全局联合恢复。

研究背景与动机¶

从离散时间切片中重构群体动力学演化过程，是计算生物学、群体运动学以及统计物理等领域的核心逆问题。在连续介质理论中，这类问题常被建模为 2-Wasserstein 空间上的梯度流（Wasserstein Gradient Flow, WGF），即概率测度在某个未知能量泛函（通常包含外部势能、内部熵以及粒子间对称相互作用能）驱动下的最速下降曲线。长期以来，从观测快照反演能量泛函的主流算法几乎完全建立在 Jordan–Kinderlehrer–Otto（JKO）离散近端步进格式之上，如 JKOnet 与 JKOnet⋆ 等方法。

然而，基于 JKO 的逆问题方案存在固有的结构性缺陷。JKO 框架将连续曲线切分为相邻时间步之间的离散极小化问题，不仅需要频繁求解代价昂贵的最优传输（OT）耦合，而且其一阶离散本质上是在相邻观测点间做“弦截（chord）”线性近似。当快照时间间隔较大、轨迹在复杂势能面上呈现显著弯曲时，JKO 往往因忽略时间段内部的流形曲率而发生严重失真或梯度场崩溃；此外，JKO 仅将泛函参数化，而将连续分布曲线交由事后模拟隐式生成，使得观测时间点外的边缘分布推断高度受限。

针对这一矛盾，本文跳脱出 JKO 离散近端框架，回归到度量空间梯度流的切向量（连续性方程）与能量耗散等价性（EDE）本源。核心 idea：将满足连续性方程的速度场与能量泛函变分梯度之间的失配构造成非负速度残差损失，结合观测时间步的数据匹配散度，并引入以粒子轨迹为一阶优化变量的免模拟核密度拼接（stitching）算法，实现能量势能与弯曲轨迹的全局联合学习。

方法详解¶

整体框架¶

本文方法的目标是从离散时间集合 $\mathcal{T}_{\mathrm{obs}}$ 处的观测样本群 $\{q_t\}$ 中，联合反演参数化能量泛函 $\mathcal{F}^\theta$ 与连续测度演化曲线 $\rho^\theta$。整个框架无需对神经微分方程进行前向数值仿真，而是构建了一个统一的全局连续优化目标：一方面在全时间区间上最小化速度残差 $\mathcal{R}_{\mathrm{vel}}$，强约束连续性方程与梯度流物理定律；另一方面在观测快照时刻通过统计散度 $\mathcal{D}$ 将模型边缘分布对齐到观测数据点云。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["离散时序群体快照<br/>{q_t | t in T_obs}"] --> B["一阶显式曲线参数化<br/>N 条粒子连续轨迹 x_k(t) + KDE 构建 rho^theta_t"]
    B --> C["经验测度与速度场逼近<br/>中心点速度 v_t(x_k)=dx_k/dt 与解析 Score 评估"]
    C --> D["全局残差联合优化<br/>时间梯形离散速度残差 R_vel + 观测时刻散度 D_data"]
    D --> E["输出物理量<br/>能量势能 V_theta、相互作用核 W_theta 及完整平滑轨迹"]

关键设计¶

1. 速度残差形式化：规避高阶微分的梯度流约束

在 Wasserstein 空间 $(\mathcal{P}_2(\mathbb{R}^D), W_2)$ 中，绝对连续测度曲线 $\rho_t$ 的连续性方程与其能量耗散等价性（EDE）均要求速度场 $\mathbf{v}_t$ 严格等于能量泛函 $\mathcal{F}$ 的负第一变分梯度 $-\nabla_{\mathbf{x}} \frac{\delta \mathcal{F}}{\delta \rho_t}$。直接使用密度层面的连续性方程残差 $\mathcal{R}_{\mathrm{dens}}$ 需要计算散度算子 $\operatorname{div}_{\mathbf{x}}$，涉及对神经网络的高阶导数，计算代价极高且数值极不稳定。

为此，方法采用加权 $L^2(\rho_t)$ 速度残差项，避开对密度直接求散度： $$ \mathcal{R}{\mathrm{vel}}[\mathcal{F}, (\rho, \mathbf{v})] := \int}^{T} \int_{\mathbb{R}^D} \left| \mathbf{vt(\mathbf{x}) + \nablat $$ 该残差非负且当且仅当 }} \frac{\delta \mathcal{F}}{\delta \rho_t}(\mathbf{x}) \right|^2 \rho_t(\mathbf{x}) \, \mathrm{d}\mathbf{x} \, \mathrm{d$\rho$ 为 $\mathcal{F}$ 的 Wasserstein 梯度流时精确为零。将其与观测时刻的数据拟合散度 $\mathcal{D}(\rho_t, q_t)$（如负对数似然 KL 散度或去噪分数匹配 DSM）结合，构建出单一全局可微损失函数 $\mathcal{L} = \lambda \mathcal{R}_{\mathrm{vel}} + \sum_{t \in \mathcal{T}_{\mathrm{obs}}} \mathcal{D}(\rho_t, q_t)$。

2. 拼接（Stitching）算法：粒子轨迹曲线的一阶参数化与免仿真求解

以往的残差法（如 GenWGP）依赖基于连续归一化流或神经常微分方程（Neural ODE）的求解，单步反传计算开销巨大。拼接方法（stitching）将演化曲线 $\rho^\theta$ 本身提升为与泛函 $\mathcal{F}^\theta$ 平等的一阶直接学习对象，采用动态粒子轨迹支撑的核密度估计（KDE）进行参数化： $$ \rho_t^\theta(\mathbf{x}) := \sum_{k=1}^N w_k^\theta \, \phi\big(\mathbf{x} - \mathbf{x}{t,k}^\theta\big), \quad \mathbf{v}t^\theta(\mathbf{x}) = \frac{\sum}^N w_k^\theta \, \phi(\mathbf{x} - \mathbf{x}^\theta) \, \dot{\mathbf{x}{t,k}^\theta}{\sum $$ 通过引入经验测度近似 }^N w_l^\theta \, \phi(\mathbf{x} - \mathbf{x}_{t,l}^\theta)$\rho_t^\theta \approx \sum_{k=1}^N w_k^\theta \delta_{\mathbf{x}_{t,k}^\theta}$，在粒子中心点处速度场严格退化为可微轨迹的导数 $\mathbf{v}_t^\theta(\mathbf{x}_{t,k}^\theta) = \dot{\mathbf{x}}_{t,k}^\theta$。这一设计将复杂的连续介质流形积分简化为沿粒子节点的时间求和，彻底摆脱了数值微分方程求解器，即使在快照间隙极大时亦能自由学习弯曲轨迹。

3. 中心点 Score 剔除自作用项与二阶梯形时间积分

在计算泛函变分中的扩散熵贡献项 $\nabla_{\mathbf{x}} \log \rho_t^\theta$ 时，若直接在自身中心点 $\mathbf{x}_{t,j}^\theta$ 计算 KDE 梯度，分母包含的核峰值常数项 $w_j^\theta \phi(\mathbf{0})$ 将在近邻稀疏时主导分母，导致数值迅速衰减为零，使粒子丧失向外扩散的压力并塌缩在势阱底。

为恢复正确的扩散流体物理行为，设计在评估点 $j$ 处剔除粒子自身对密度及其梯度的贡献： $$ \nabla_{\mathbf{x}} \log \rho_t^\theta(\mathbf{x}{t,j}^\theta) \approx \frac{\sum} w_k^\theta \, \nabla_{\mathbf{x}} \phi\big(\mathbf{x{t,j}^\theta - \mathbf{x}}^\theta\big)}{\sum_{l \neq j} w_l^\theta \, \phi\big(\mathbf{x{t,j}^\theta - \mathbf{x} $$ 在时间轴离散上，算法在节点网格 }^\theta\big)$0 = t_0 < \dots < t_{K-1} = T$ 上应用二阶精度的梯形积分法则（Trapezoidal rule），对前后两端点的物理变分梯度取平均，以 $O(\Delta t^2)$ 精度严密保持动力学守恒性与跨步平滑性。

损失函数 / 训练策略¶

离散化后的总优化目标形式如下： $$ \mathcal{L}{\mathrm{stitch}}(\theta) = \sum}^{K-2} \frac{\Delta t_{j+1}}{2} \sum_{k=1}^N w_k^\theta \left( \left| \frac{\Delta \mathbf{x{t},k}^\theta}{\Delta t_{j+1}} + \nabla_{\mathbf{x}} \frac{\delta \mathcal{F}^\theta}{\delta \rho_{t_j}^\theta}(\mathbf{x{t_j,k}^\theta) \right|^2 + \left| \frac{\Delta \mathbf{x}},k}^\theta}{\Delta t_{j+1}} + \nabla_{\mathbf{x}} \frac{\delta \mathcal{F}^\theta}{\delta \rho_{t_{j+1}}^\theta}(\mathbf{x{t},k}^\theta) \right|^2 \right) - \lambda_{\mathrm{data}} \sum_{t \in \mathcal{T{\mathrm{obs}}} \sum) $$ 其中泛函包含潜在势能神经网络 }^{N_t} \log \rho_t^\theta(\mathbf{y}_{t,i$V^\theta$（两层 64 维 MLP）以及粒子间对称相互作用网络 $W^\theta(r)$。参数通过 Adam 优化器配合余弦退火学习率进行全量端到端联合训练。

实验关键数据¶

主实验¶

在单细胞轨迹推断（Embryoid Body, EB 数据集，5个观测时间步）基准测试中，评估模型在各时间步的前向推断分布与真实观测测度之间的 Wasserstein-1 距离（$W_1$，越低越好）：

方法	类型	t=1 ($W_1$↓)	t=2 ($W_1$↓)	t=3 ($W_1$↓)	t=4 ($W_1$↓)	平均 $W_1$↓
Neural SDE	随机微分方程	0.69	0.91	0.85	0.81	0.82
TrajectoryNet	连续最优传输流	0.73	1.06	0.90	1.01	0.93
OT-CFM	条件流匹配	0.78	0.76	0.77	0.75	0.77
WLF-SB	薛定谔桥动力学	0.63	0.79	0.77	0.74	0.73
JKOnet	离散近端最优传输	1.53	1.27	1.13	1.41	1.34
JKOnet⋆$_V$	JKO 一阶弦截法	0.99	1.11	1.06	1.30	1.12
iJKOnet$_V$	逆 JKO 优化	0.92	1.11	0.95	1.21	1.05
Stitching (静态 $V^\theta(\mathbf{x})$)	本文残差拼接方法	0.46	0.60	0.60	0.65	0.58
JKOnet⋆$_{t,V}$	JKO 时变势能	0.69	0.77	0.69	0.78	0.73
iJKOnet$_{t,V}$	逆 JKO 时变势能	0.51	0.58	0.57	0.64	0.58
Stitching (时变 $V^\theta(\mathbf{x}, t)$)	本文残差拼接方法	0.44	0.56	0.57	0.60	0.54

消融实验与稳健性分析¶

在二维合成势能面反演（含配对 paired 与独立未配对 unpaired 快照机制）以及相互作用粒子系统（Confining Mexican Hat 势能 + 径向吸引核 $W$）基准中，对比算法在数据配对退化及多项相互作用解耦下的性能表现：

评估任务与配置	评估指标	JKOnet⋆	iJKOnet	Stitching (本文)	说明与增益分析
合成势能 (Paired)	梯度误差 $L^2$-UVP (%)↓	149.0	-	0.00	Stitching 几乎完全拟合势能真实梯度场
合成势能 (Unpaired)	梯度误差 $L^2$-UVP (%)↓	151.0	-	0.01	快照去相关时 JKO 梯度彻底崩溃，Stitching 保持稳定
合成势能 (Unpaired)	矩匹配 $\mathrm{Bd}^2_{W_2}$-UVP (%)↓	1.70	0.35	0.24	边缘分布前两阶统计矩拟合精度最优
相互作用核恢复	径向势能拟合 $R^2(W)$↑	0.72	-	0.81	精确分离环境势能与粒子间吸引作用力
相互作用动力学	单次迭代耗时 / 总收敛时间↓	4.1s / 69min	-	0.24s / 4min	Stitching 免前向 OT 求解，训练加速超 17 倍

关键发现¶

对无配对独立快照具有极强稳健性：JKO 依赖前后时间步样本对的最优传输位移耦合，在快照未配对（unpaired）时梯度恢复误差（$L^2$-UVP）普遍超过 100% 出现发散；而 Stitching 基于内生参数化轨迹与残差优化，在配对与未配对场景下的各项误差均保持在 $2\times$ 以内，几乎完全不受快照相关性丢失的影响。
大时间步长曲率重构能力完胜 JKO 弦截：在波状山谷（Wavy Valley）非线性势能测试中，长观察间隔使得 JKOnet 仅能生成直线割线而错失势能盆底；Stitching 的显式连续轨迹能精准沿着正弦山谷底部自适应弯曲，反演势能的拟合优度达到 $R^2 = 0.62$（相比 JKOnet 的 0.51）。
计算效率带来数量级飞跃：在解耦 McKean–Vlasov 多智能体相互作用动力学时，免除了单步 200 对微批次 OT 耦合的 Stitching 将单步耗时由 4.1 秒压至 0.24 秒，训练总耗时从 69 分钟剧缩至 4 分钟。

亮点与洞察¶

统一理论视角：揭示了 Path-Finding（生成流路径搜索）与 Action Matching（动作匹配）本质上均是连续介质速度残差在特定参数化下的特例，确立了 Wasserstein 残差最小化作为 WGF 逆问题的一般性范式。
轨迹作为第一类对象（First-Class Learnable Variable）：突破了“参数化势能 $\rightarrow$ 事后积分推断轨迹”的单向绑定惯性思维，通过将带权移动粒子曲线直接参数化并用 KDE 解析链接，使得曲线表达力与泛函表达力双向解耦且能独立选择离散步长。
可复用的 Score 计算去自环技巧：在基于粒子的粒子系统数值优化中，KDE 分数导数剔除自身节点贡献的设计巧妙解决了分子梯度为零而分母常数占优导致的“扩散压崩溃”数学瑕疵，可直接推广至各类基于粒子滤波与扩散生成模型的高维流形计算中。

局限与展望¶

粒子规模的二次方复杂度：由于采用全连接交互计算粒子间相互作用与核密度熵项，Stitching 单时间步计算复杂度关于粒子数 $N$ 呈 $O(N^2)$ 增长，在大规模高维粒子云应用中需引入树码或快速多极子（FMM）加速。
弱松弛约束下的精确性边界：残差最小化在数值优化层面是通过软惩罚项实现的，并不在物理层面强保证残差恒等于零，且无限粒子极限下的理论收敛性（$N \to \infty$）在数学上仍待严格证明。
非梯度流推广前景：论文初步验证了在旋度不为零的手性流体（Chiral Active Matter）中泛化拟合的能力，未来可进一步拓展至非保守场、主动物质系统以及大语言模型表示轨迹等非互易动力学建模。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 颠覆了 JKO 主导的传统范式，从速度残差本源切入提出一阶粒子拼接法，理论视野宽广且自洽。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖合成解析势能、单细胞高维分化基因轨迹及多粒子非线性相互作用力学，评测指标严谨全面。
写作质量: ⭐⭐⭐⭐⭐ 数学推导干净利落，从欧氏梯度流到 Wasserstein 空间的四重等价映射阐述极具启发性。
价值: ⭐⭐⭐⭐⭐ 为动力学逆问题、单细胞推断及生成扩散模型提供了崭新的免模拟、抗稀疏快照技术工具箱。

评估任务与配置	评估指标	JKOnet⋆	iJKOnet	Stitching (本文)	说明与增益分析
合成势能 (Paired)	梯度误差 \(L^2\)-UVP (%)↓	149.0	-	0.00	Stitching 几乎完全拟合势能真实梯度场
合成势能 (Unpaired)	梯度误差 \(L^2\)-UVP (%)↓	151.0	-	0.01	快照去相关时 JKO 梯度彻底崩溃，Stitching 保持稳定
合成势能 (Unpaired)	矩匹配 \(\mathrm{Bd}^2_{W_2}\)-UVP (%)↓	1.70	0.35	0.24	边缘分布前两阶统计矩拟合精度最优
相互作用核恢复	径向势能拟合 \(R^2(W)\)↑	0.72	-	0.81	精确分离环境势能与粒子间吸引作用力
相互作用动力学	单次迭代耗时 / 总收敛时间↓	4.1s / 69min	-	0.24s / 4min	Stitching 免前向 OT 求解，训练加速超 17 倍