Technical Summary: Wasserstein Residuals: Learning Gradient Flows from Population Dynamics

Problem Statement

The paper addresses the inverse problem of reconstructing population dynamics from a sparse set of observational snapshots. In many scientific fields, such as computational biology and crowd dynamics, population evolution is modeled as a Wasserstein Gradient Flow (WGF): a curve of probability distributions $\rho_t$ driven by the steepest descent of an energy functional $F$. The goal is to recover the underlying energy functional $F$ such that its gradient flow fits the observed marginal distributions $\{q_t\}$ at discrete observation times $t \in \mathcal{T}_{obs}$.

Existing dominant approaches rely on the Jordan–Kinderlehrer–Otto (JKO) scheme, which defines WGFs as a sequence of proximal steps minimizing an energy plus an optimal transport (OT) penalty. However, JKO-based methods suffer from two primary limitations:

1. Inflexibility to Time Discretization: They require solving costly OT problems between consecutive snapshots.
2. Sensitivity to Gaps: They struggle with large temporal gaps between observations, often failing to capture curved trajectories (e.g., particles moving along a sinusoidal valley) because they rely on linear chord interpolation between snapshots.

Methodology: The Wasserstein Residual Framework

The authors propose a residual-based approach that bypasses the JKO scheme and optimal transport couplings. Instead of enforcing the JKO condition, they enforce the fundamental differential equations characterizing gradient flows directly via a non-negative loss function.

1. The Residual Formulation

The paper utilizes the Tangent and Energy Dissipation Equality (EDE) formulations of gradient flows. Specifically, it focuses on the velocity constraint:
$$v_t(x) = -\nabla_x \frac{\delta F}{\delta \rho_t}(x)$$
where $v_t$ is the velocity field and $\frac{\delta F}{\delta \rho_t}$ is the first variation of the energy functional.

The authors define a velocity residual $R_{vel}$ which is zero if and only if $\rho$ is a gradient flow of $F$:
$$R_{vel}[F, (\rho, v)] = \int_0^T \int_{\mathbb{R}^D} \left\| v_t(x) + \nabla_x \frac{\delta F}{\delta \rho_t}(x) \right\|^2 \rho_t(x) \, dx \, dt$$

The global objective combines this residual with a data-fitting divergence $D$ at observed times:
$$\mathcal{L}(F, \rho, v) = \lambda R_{vel}[F, \rho, v] + \sum_{t \in \mathcal{T}_{obs}} D(\rho_t, q_t)$$
This framework unifies existing methods like Path-Finding and Action Matching as instantiations of a single residual principle.

2. The "Stitching" Algorithm

To minimize this objective, the authors introduce Stitching, a simulation-free, particle-based method with the following characteristics:

• Learnable Trajectories: Unlike JKO methods where the curve $\rho$ is implicit, Stitching explicitly parametrizes the curve $\rho_\theta$ as a Kernel Density Estimate (KDE) over a set of $N$ learnable particle trajectories $x_{t,k}^\theta$.
• Simulation-Free: The method does not require solving Neural ODEs or integrating SDEs. The particle positions and velocities are direct learnable parameters.
• Empirical Approximation: To make computation efficient, the KDE is approximated by an empirical measure at the particle centers. This simplifies the velocity residual to a sum over particles:
  \mathcal{L}_{stitch} \approx \int_0^T \sum_{k=1}^N w_k \left\| \dot{x}_{t,k} + \nabla_x \frac{\delta F_\theta}{\delta \rho_\theta_t}(x_{t,k}) \right\|^2 dt + \sum_{t \in \mathcal{T}_{obs}} D(\rho_\theta_t, q_t)
• Functional Parametrization: The energy functional $F_\theta$ is parametrized as a sum of a potential term $V_\theta$, an entropy term (log-density), and an interaction term $W_\theta$, all represented by neural networks.

Key Contributions

1. Residual Framework: The authors recast WGF learning as minimizing a non-negative residual term, providing a unified perspective that subsumes Path-Finding and Action Matching.
2. Stitching Algorithm: They introduce a novel, simulation-free particle method that treats the trajectory curve $\rho$ as a first-class learnable variable. This allows the method to tolerate large gaps between snapshots and recover curved dynamics that JKO-based chord predictors miss.
3. State-of-the-Art Performance: The method achieves superior results on trajectory inference benchmarks, particularly in scenarios with sparse or unpaired data.

Experimental Results

The paper evaluates Stitching on three main benchmarks:

1. Illustrative Example (Continuous Flow vs. JKO)

In a synthetic "wavy valley" potential, particles follow a sinusoidal trajectory.

• Result: JKO-based methods (specifically JKOnet*) fail to capture the curvature between sparse snapshots ($t=10$ to $t=20$), producing straight-line interpolations. Stitching correctly recovers the curved trajectory, leading to a significantly better recovery of the underlying potential ($R^2 = 0.62$ vs. $0.51$).

2. Synthetic Potential Recovery

The method is tested on 15 two-dimensional potentials under paired (correlated trajectories) and unpaired (decorrelated snapshots) regimes.

• Result: Stitching maintains consistent performance across both regimes. In contrast, JKO-based methods degrade significantly or collapse in the unpaired setting because they rely on OT couplings between consecutive snapshots. Stitching's errors remain within a factor of $\sim 2\times$ between regimes, whereas JKO methods often fail to recover informative gradient fields in the unpaired case.

3. Single-Cell Trajectory Inference (Embryoid Body Dataset)

Applied to human embryonic stem cell differentiation data (5 snapshots over 27 days).

• Result: Stitching achieves the best performance (lowest $W_1$ and $W_2$ distances) compared to baselines including Neural SDEs, TrajectoryNet, and various JKO-based methods. It outperforms competitors in both full-data training and a "leave-two-out" temporal interpolation protocol.

4. Interaction Dynamics Recovery

The method is used to disentangle environmental potentials ($V$) and interaction kernels ($W$) in a McKean–Vlasov system.

• Result: Stitching successfully recovers both the confining potential and the attractive interaction kernel, achieving high pattern $R^2$ scores ($0.84$ for $V$, $0.81$ for $W$). It is also significantly faster per epoch than JKO-based methods, which require solving OT problems at every step.

5. Non-Gradient Flows

The authors demonstrate that the residual framework can be extended beyond gradient flows by replacing the gradient term with a general vector field, successfully tracking chiral (non-conservative) dynamics.

Significance and Claims

The paper claims that the Wasserstein Residual perspective offers a more flexible and robust alternative to the JKO scheme for learning population dynamics. By decoupling the learning of the energy functional from the constraints of optimal transport couplings, Stitching enables:

• Robustness to Sparsity: The ability to learn dynamics from widely spaced observations without the "curvature loss" inherent in linear interpolation.
• Unpaired Data Compatibility: The ability to learn from datasets where particle identities are not tracked across time (unpaired snapshots), a common scenario in biological imaging.
• Simulation-Free Efficiency: Avoiding the computational overhead of solving ODEs or OT problems during training.

The authors position this work as a unification of existing paradigms and a practical advancement for scientific discovery in fields ranging from single-cell biology to collective behavior, while acknowledging limitations such as $O(N^2)$ computational complexity for interaction terms and the lack of theoretical convergence guarantees as $N \to \infty$.