Double-Stitch：免模拟方法学习 Wasserstein 空间拉格朗日力学
Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals
Basis Research 提出的 Double-Stitch 免去训练时反复跑数值求解器，比 WLM 快 4-14 倍，还覆盖梯度流搞不定的周期性动力学，JAX 代码已开源。
这篇论文提出 Double-Stitch，通过 Clebsch 变分原理导出运动方程残差，沿学习到的种群路径惩罚残差来学习 Wasserstein 拉格朗日力学，无需在每个训练步运行数值求解器。相比基于模拟的 WLM 方法，训练速度快 4-14 倍。在合成数据、单细胞和海洋涡流数据集上，Double-Stitch 在多数任务上达到或超过梯度流方法和 WLM。该方法能覆盖保守和周期性动力学，而 Wasserstein 梯度流无法描述这两类动力学。JAX 实现已开源在 GitHub。
Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals
The dynamics of cells, organisms, and fluids are often modeled as probability distributions evolving over time. Reconstructing and extrapolating this evolution from unpaired snapshots requires assumptions about the underlying process. Wasserstein gradient flows are a common choice, but they cannot describe conservative or periodic dynamics. Lagrangian mechanics in Wasserstein space covers both, but existing methods for learning it are simulation-based: they run a numerical solver at every training step, which makes training expensive. We propose Double-Stitch, a simulation-free method that learns these mechanics by penalizing the residual of the equation of motion along a learned population path. We derive this equation from a Clebsch variational principle that does not require gradient velocities, and show that the residual vanishes exactly when the equation holds. We test Double-Stitch on synthetic, single-cell and ocean vortex datasets and find that it matches or outperforms gradient-flow methods and simulation-based WLM on most tasks, while training $4$-$14$ times faster than WLM. We provide a JAX implementation of Double-Stitch at https://github.com/BasisResearch/stitching.
