Research map / Probabilistic and scientific modeling / Dynamical systems

Wasserstein residuals: Learning Gradient Flows from Population Dynamics

NeurIPS 2026 · Probabilistic and scientific modeling › Dynamical systems · approach: Stochastic population dynamics · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes a residual framework for learning Wasserstein gradient flows from population snapshots, instantiated as a simulation-free KDE-particle method called stitching that jointly learns the functional and the trajectory curve.

Problem

Learning the energy functional of a Wasserstein gradient flow from sparse, possibly widely spaced population snapshots; existing JKO-based methods require costly optimal transport solves and degrade under long time gaps.

Approach

Replace the JKO scheme with a nonnegative velocity residual enforcing the continuity equation, combined with a data-fitting divergence; parametrize the curve itself as a KDE over learnable particle trajectories so it is learned jointly with the functional, without simulation or OT couplings.

Key results

• On the full-data EB benchmark, stitching with time-varying potential achieves mean W1 of 0.54 vs 0.58 for iJKOnet_{t,V} and 0.73 for JKOnet*_{t,V}.
• On interaction dynamics, stitching reaches R^2=0.84 on V and 0.81 on W, vs 0.04 and 0.72 for JKOnet*, with about 4 min total runtime vs about 69 min for JKOnet*.
• On the leave-two-out EB protocol, time-varying stitching obtains mean W2 0.88 vs 0.92 for the best published competitor.

Surprising findings

• JKOnet* in the unpaired synthetic setting fails to recover an informative gradient field (L2-UVP above 100%), while stitching remains stable across paired and unpaired regimes.
• Stitching is worse than JKO-based methods on forward-prediction EMD in the paired synthetic regime.

Limitations

• O(N^2) cost per time step from entropy and interaction terms of the particle method.
• No theoretical convergence guarantees as N grows to infinity.
• Residual is a soft penalty, so the exact gradient-flow condition R=0 is not guaranteed.

Experimental setup

Datasets: embryoid body scRNA-seq synthetic 2D potentials benchmark (15 potentials) simulated interacting particle system

Benchmarks: EB scRNA-seq W1 (5D) EB leave-two-out W2 15 synthetic 2D potentials (EMD, L2-UVP, Bd2-UVP) interaction dynamics (EMD, W2, BW2, MMD, R2)

Metrics: W1 W2 EMD L2-UVP Bd2-UVP MMD R^2

Baselines: JKOnet JKOnet* iJKOnet Neural SDE TrajectoryNet SB-FBSDE NLSB OT-CFM WLF-OT WLF-SB

Builds on: JKOnet JKOnet* iJKOnet Action Matching Path-Finding

Closest papers by meaning

• Wasserstein Gradient Flows and Forward-Only Diffusion Are Not Enough for Multimodal Sampling (NeurIPS 2026 · similarity 0.88)
• Generalized Wasserstein Flow Matching: Transport Plans, Everywhere, All at Once (NeurIPS 2026 · similarity 0.87)
• Parametric model reduction of mean-field and stochastic systems via higher-order action matching (NeurIPS 2024 · similarity 0.87)
• Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport (ICML 2024 · similarity 0.86)
• Partially Observed Trajectory Inference using Optimal Transport and a Dynamics Prior (ICLR 2025 · similarity 0.86)
• Curly Flow Matching for Learning Non-gradient Field Dynamics (NeurIPS 2025 · similarity 0.86)