A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots¶

会议: ICML 2026
arXiv: 2605.08550
代码: https://github.com/guanton/WLM
领域: 扩散模型 / 动力学学习
关键词: 人口动力学, Wasserstein 力学, 最小作用原理, 二阶动力学, 神经势能

一句话总结¶

本文从最小作用原理出发，提出 Wasserstein 拉格朗日力学（WLM）框架，学习二阶人口动力学而非传统梯度流的一阶动力学，从而能够捕捉周期性、旋转等更丰富的群体现象，并可在不需要参考过程的情况下完成插值与未来预报。

研究背景与动机¶

领域现状：传统人口动力学建模（从分子扩散到细胞分化、生物群体行为）普遍采用 Wasserstein 梯度流范式，从自由能泛函出发，能够建模纯耗散的演化过程。

现有痛点：梯度流本质上是一阶系统，最终走向某个平衡态。但真实的人口动力学常表现出周期性、旋转运动、振荡等非平衡现象——例如涡旋、Boids 群体行为——这些动力学超越了能量最小化的框架。

核心矛盾：虽然梯度流数学基础坚实、优化算法成熟，但表达能力受限。当我们只有边际分布的时间快照（无法追踪单个轨迹）时，如何在不预先指定拉格朗日量的情况下，推断更广泛的二阶动力学？

本文目标：从最小作用原理重新审视人口模型，用二阶系统取代一阶梯度流，使框架统一涵盖经典力学、量子力学甚至梯度流。

切入角度：在 Wasserstein-2 距离空间 \(\mathcal{P}_2(\mathbb{R}^d)\) 中定义人口"坐标"与拉格朗日量，通过变分原理推导 Hamiltonian 运动方程，形成"群体驱动自身演化"的力学观点。

核心 idea：用群体势能泛函 \(\mathcal{U}[\rho_t]\) 和阻尼系数 \(\gamma\) 参数化二阶动力学，使其涵盖梯度流（过阻尼极限 \(\gamma\to\infty\)）、保守经典力学（\(\gamma=0\)）以及量子力学等多种情形。

方法详解¶

整体框架¶

方法分为理论层与算法层两层。

理论层：在 \(\mathcal{P}_2(\mathbb{R}^d)\) 中将人口动力学表示为连续方程 \(\dot{\rho}_t = -\nabla\cdot(\rho_t \nabla s_t)\)，其中 \(s_t\) 为时变势函数。引入带阻尼的 Wasserstein 拉格朗日量 \(\mathcal{L}[\rho_t, s_t, t] = e^{\gamma t}(\frac{1}{2}\int\|\nabla s_t\|^2 \rho_t dx - \mathcal{U}[\rho_t])\)，应用变分原理，推导出粒子层面的运动方程：\(\frac{d}{dt}x_t = \nabla s_t(x_t)\)，\(\frac{d}{dt}v_t = -\nabla\frac{\delta\mathcal{U}[\rho_t]}{\delta\rho_t}(x_t) - \gamma v_t\)。这是"广义牛顿定律"的群体版本：粒子加速度由群体势能梯度和速度阻尼共同驱动。

算法层：神经力学模型直接学习势能泛函 \(\Psi_\theta\)，通过 Proposition 3.1 将函数导数转换为神经网络对粒子坐标的梯度。时间推进使用 Leapfrog 积分器，预测边际与观测边际之间用 Sinkhorn 散度度量差异，端到端反向传播优化。

关键设计¶

Hamiltonian 框架下的群体势能参数化:
- 功能：将不可计算的函数梯度 \(\frac{\delta\mathcal{U}[p]}{\delta p}(x)\) 转化为可微的神经网络参数梯度。
- 核心思路：对经验分布 \(\hat{p} = \frac{1}{N}\sum_i\delta_{x^{(i)}}\)，有恒等式 \(\nabla_{x^{(j)}}\Psi(x^{(1)},\ldots,x^{(N)}) = \frac{1}{N}\nabla_x\frac{\delta\mathcal{U}[p]}{\delta p}(x^{(j)})\big|_{p=\hat{p}}\)。这样势能梯度可由自动微分直接获得，无需显式计算泛函导数。
- 设计动机：绕过 Wasserstein 空间中泛函导数难以直接求解的困难，利用经验测度的离散结构，让神经网络学到的标量函数 \(\Psi_\theta\) 自然定义出势能场。
混合批次与无参考学习:
- 功能：在保留二阶动力学的同时灵活处理多源数据，无需预先指定 reference SDE 或 OT plan。
- 核心思路：从任意初始条件 \((p_0, v_0)\) 出发，用单一拉格朗日量 \((\Psi_\theta, \gamma)\) 覆盖所有观测时间。损失函数在每对预测/观测边际间累积 Sinkhorn 散度。不同于 AM/JKOnet/流匹配，本方法不依赖任何参考过程，直接从数据学到力学系统。
- 设计动机：在缺乏领域先验的场景下提供通用建模框架，使方法能用于未知物理或生物系统。
可学习阻尼与统一表达:
- 功能：通过学习阻尼系数 \(\gamma\)，自动确定动力学在梯度流与经典力学之间的位置。
- 核心思路：\(\gamma=0\) 对应保守系统（经典力学）；\(\gamma\to\infty\) 对应过阻尼极限（梯度流）；中间值 \(\gamma>0\) 对应有粘性但仍保留惯性的二阶系统。在梯度流数据上模型能自动学到 \(\gamma\geq 500\)（强阻尼极限），从而精确复现梯度流。
- 设计动机：单一可学习参数即可让模型在多种力学范式之间自适应切换，无需手工选择。

损失函数与训练策略¶

针对每段观测时间 \(t_i\to t_{i+1}\)，用 Leapfrog 推进得到预测边际 \(\hat{p}_{t_{i+1}}\)，与观测边际之间累积 Sinkhorn 散度 \(\sum_i \mathcal{S}_\epsilon(\hat{p}_{t_{i+1}}, p_{t_{i+1}})\)。优化变量为势能网络参数 \(\theta\)、初速度场 \(v_0\) 以及阻尼系数 \(\gamma\)。

实验关键数据¶

主实验¶

任务	方法	指标	结果	说明
梯度流 SDE（配对）	nn-APPEX	Forecast \(W_1\)	0.131±0.006	传统梯度流方法
梯度流 SDE（配对）	WLM (learnable \(\gamma\))	Forecast \(W_1\)	0.137±0.012	性能近似，但无先验
梯度流 SDE（不配对）	JKOnet*	Train \(W_1\)	0.236±0.040	配对假设破坏，性能崩
梯度流 SDE（不配对）	WLM (learnable \(\gamma\))	Train \(W_1\)	0.068±0.004	显著优于一阶方法
海湾涡旋插值（小涡旋）	WLM	多时刻 \(W_1\)	0.060–0.068	无先验下优于 AM/UAM/sAM
海湾涡旋预报（大涡旋）	WLM	预报 \(W_1\)	0.567±0.014	唯一无参考方法能预报
胚胎发育 scRNA	WLM	插值 \(W_1\)	优于梯度流+OT	高维真实数据有效
Boids 群体行为	WLM	预报	全面胜出	捕捉集体振荡

消融实验¶

配置	不配对 SDE Forecast \(W_1\)	涡旋预报 \(W_1\)	说明
Full WLM（\(\Psi_\theta\)+可学 \(\gamma\)）	0.246±0.026	0.567±0.014	完整模型
WLM（\(\gamma=0\) 保守）	0.346±0.045	0.689±0.120	无阻尼差，尤其插值
WLM（\(\gamma\) 固定过大）	0.298±0.031	0.612±0.025	过阻尼偏离梯度流

关键发现¶

二阶 vs 一阶：WLM 在不配对数据上 \(W_1\) 从 1.618 降至 0.246，二阶框架对混淆轨迹更鲁棒。
可学习阻尼的自适应性：模型在梯度流数据上自动收敛至过阻尼极限，在涡旋/Boids 数据上保持中等阻尼以捕捉旋转。
预报 vs 插值：WLM 是唯一不使用参考过程仍能预报涡旋未来状态的方法。
高维可扩展：scRNA 数千维数据上仍有效，证实神经势能参数化的通用性。

亮点与洞察¶

理论高度：完整推导了 Wasserstein 空间中的 Hamiltonian 框架，将梯度流、经典力学、量子力学纳入统一视角。
方法巧妙：Proposition 3.1 把不可计算的函数梯度转化为神经网络参数梯度，使整个学习过程在标准 autograd 框架内可行。
无参考学习：不需要预指定参考 SDE 或 OT 计划，使方法在未知系统上具备真正的通用性。
真预报而非插值：得到的是动力学方程，可外推到训练时间窗外，相比仅插值方法是定性进步。

局限与展望¶

仍假设观测是同一群体在多个时刻的快照，无法处理来自完全不同人口的快照匹配。
神经势能网络的表达能力边界缺乏理论刻画。
实验粒子规模较小（约 1000），大规模应用（百万粒子级）的可扩展性未验证。
高维场景下势能学习的样本复杂度可能很高，需要更高效的近似（诱导点、稀疏表示）。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 从最小作用原理重写人口动力学，二阶+无参考是真正的范式扩展。
实验充分度: ⭐⭐⭐⭐☆ 覆盖合成 SDE、物理涡旋、生物 scRNA、群体行为四类数据，消融清晰，但样本规模偏小。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨，算法伪代码完整，图表表达力强。
价值: ⭐⭐⭐⭐⭐ 为科学建模和动力学学习开辟新框架，跨学科影响潜力大。