Simulation-Free Learning of Population Dynamics with Wasserstein Lagrangian Residuals

• Flow methods are generally unsuited to extrapolation, gradient flows cannot effectively model conservative or periodic behavior, and WLM simulates trajectories at every training step.Source: §1 · pp. 1–2
• Lagrangian methods model acceleration rather than velocity, and Double-Stitch extends this approach with simulation-free learning and time extrapolation.Source: §1 · pp. 1–2
• Double-Stitch fits a latent population path and potential by penalizing a discretized Lagrangian residual; its Clebsch derivation admits non-gradient velocities.The residual vanishes if and only if the dynamics are satisfied.Source: §1 · p. 2
• Double-Stitch matches or improves on competing methods across most synthetic, single-cell, and ocean-vortex tasks while training more cheaply than WLM.Source: §1 · p. 2

2.1 PROBLEM SETTING

• Unpaired snapshots do not identify a continuous-time evolution, so the task requires assumptions about the generating process.The model predicts distributions at selected times, with prediction quality evaluated using the p-Wasserstein distance.Source: §2.1 · p. 3

2.2 DYNAMICAL ASSUMPTIONS

• Wasserstein gradient flows model steepest descent of a time-independent free energy, but cannot represent periodic dynamics.The section motivates a second-order assumption in which velocity can evolve independently of particle positions.Source: §2.2 · p. 3

2.3 EQUATIONS OF MOTION

• The Clebsch formulation derives second-order motion for the full velocity, which need not be a gradient, while continuity constrains density evolution.In the canonical class, zero damping gives Wasserstein Hamiltonian flow and conserves K+U for time-independent energy; the overdamped limit recovers Wasserstein gradient flow.Source: §2.2 · p. 3, §2.3 · p. 4
• Using observed velocities directly avoids canonical projection, which can alter subsequent density evolution when velocity is part of the second-order state.The derivation’s particle labels and Lagrange multipliers are not learned.Source: §2.3 · p. 4

3 DOUBLE-STITCH

• The nonnegative residual vanishes exactly when the equation of motion holds, and can be evaluated directly on a fitted path without simulation.It squares the equation discrepancy, averages over the population, and integrates over time along particle trajectories.Source: §3 · p. 5
• The objective combines snapshot-fitting data loss with residual loss for the assumed dynamics, balanced by λ.The data loss sums distances between learned and observed snapshots at observation times; D is a divergence, typically forward KL.Source: §3 · p. 5
• The learned distribution path uses a KDE, while neural networks parameterize external potentials and particle interactions and KDE approximates the entropy term.Here the entropy term acts as isothermal pressure with force −cH∇log ρ.Source: §3 · p. 5
• The numerical residual uses KDE-based density forces, center-path acceleration, quadrature, and finite differences to approximate its population and time terms.The particle approximation differs from the exact material acceleration of the KDE mean velocity.Source: §3 · p. 6
• With M time steps and B parallel force evaluations, one Double-Stitch iteration costs O(⌈M/B⌉C(P)), versus O(M C(P)) for WLM, apart from data loss.The comparison assumes the same task, time steps, particles, and data-loss evaluation time.Source: §3 · p. 6

4 EXPERIMENTS

• The evaluation uses direct learned-particle scoring, push-forward, and rollout, alongside EB holdout protocols and ocean interpolation and forecasting splits.The ocean interpolation task uses nine snapshots of 111 particle positions; the forecasting task uses 11 snapshots of 400 positions and tests on the final snapshot.Source: §4 · p. 6, §4.2 · p. 7, §4.3 · p. 7

4.1 SYNTHETIC DEMONSTRATION

• Test W1 was 1.324 vs 0.389 on conservative flow, 0.711 vs 0.366 on damped flow, and 0.273 vs 0.292 on gradient flow (Stitching vs Double-Stitch).The comparison used one seed; the study considered conservative harmonic motion, damped harmonic motion with γ = 1, and Wasserstein gradient flow.Source: §4.1 · pp. 6–7

4.2 SINGLE-CELL DYNAMICS

• On EB, Double-Stitch matches or outperforms WLM with both static and time-dependent potentials, and matches or beats Stitching except for time-dependent L2O.Both comparisons use the authors’ reruns.Source: Table 2 · p. 7, §4.2 · p. 7

4.3 OCEAN VORTEX DYNAMICS

• For ocean interpolation, static Double-Stitch matches WLM and beats all but Curly-FM; forecasting favors WLM under push-forward, while static Double-Stitch performs poorly under rollout.Under rollout, time-dependent potential brings its error close to WLM, though static Stitching performs best; at t = 10, the potential is evaluated beyond the last observed snapshot, assuming its learned time dependence extrapolates.Source: Table 3 · p. 8, §4.3 · p. 8

4.4 COMPUTATIONAL COST

• On matched synthetic WHF training, WLM took 7.6 times longer (225 vs 30 seconds) to reach test W1 of about 0.09; EB and ocean training times were 4–12 and 7–10 minutes for Double-Stitch versus 16–54 and 97–145 minutes for WLM.The EB and ocean measurements were taken on identical hardware.Source: §4.4 · p. 8

• Flow-based methods learn time-dependent velocity fields or stochastic bridges to transport snapshots by an optimality criterion, but generally do not extrapolate beyond the observed interval without additional assumptions.Reference processes can extend them to non-gradient flows.Source: §5 · pp. 8–9
• Wasserstein gradient-flow methods learn a free-energy potential, but common JKO approaches unroll the scheme, while Stitching optimizes trajectories and minimizes a Wasserstein residual.GenWGP uses the same residual for known energy but integrates a Neural ODE.Source: §5 · p. 9
• Prior Lagrangian approaches use action matching or a CK Lagrangian with Lagrange multipliers; Double-Stitch derives the dynamical law through a Clebsch variational formulation.The cited Clebsch formulation is used for Euler–Arnold and Euler–Poincaré equations.Source: §5 · p. 9
• Double-Stitch penalizes dynamics through a residual while enforcing continuity by construction; it and Stitching are collocation methods, unlike shooting methods that simulate trajectories forward.The passage notes that multiple shooting is not represented in the field.Source: §5 · p. 9

6 DISCUSSION

• 4–14 times faster than WLM, Double-Stitch matches or outperforms Stitching and WLM on single-cell data and ocean interpolation, although WLM remains better on ocean forecasting.Source: §6 · p. 9
• Because it penalizes a residual rather than simulating the law, Double-Stitch does not guarantee exact satisfaction and may be hard to optimize with poor initialization or competing residuals.Source: §6 · p. 9
• Future work could examine alternative Lagrangians, active matter and viscous flows, or multiple-shooting hybrids combining Double-Stitch with WLM.Source: §6 · p. 10

B DERIVATION OF THE EQUATIONS OF MOTION

• The Caldirola–Kanai Lagrangian weights kinetic minus potential energy by e^γt, with γ≥0 representing friction, and yields equations usable in simulators or residual losses.Source: B DERIVATION OF THE EQUATIONS OF MOTION · p. 16

B.1 LAGRANGE MULTIPLIERS

• The Lagrange-multiplier derivation recovers prior damped Hamilton–Jacobi equations but restricts velocity to an irrotational, minimum-kinetic-energy canonical field.Source: §B.1 · pp. 16–17

B.2 CLEBSCH ACTION PRINCIPLE

• The Clebsch action transports particle labels and permits non-gradient initial velocities while yielding the same particle equation.Source: §B.2 · pp. 17–18

B.3 RECOVERED FLOWS

• With canonical initial velocity, the law recovers WLM; γ=0 conserves energy for time-independent U, while γ>0 dissipates it.Setting γ=0 in canonical dynamics gives Wasserstein Hamiltonian flow; the formal overdamped limit recovers Wasserstein gradient flow.Source: §B.3 · pp. 18–19

B.4 GENERAL AND CANONICAL VELOCITIES

• Rotational velocity can change later density evolution even when its instantaneous density change is unchanged, a trajectory effect canonical velocity may miss.Source: §B.4 · p. 19

C.1 VELOCITY AND ACCELERATION

• Differentiating the KDE path gives a continuity-equation velocity from responsibility-weighted center velocities, whose material derivative includes their local dispersion.Source: §C.1 · p. 19

C.2 ENERGY VARIATION AND WASSERSTEIN GRADIENT

• For symmetric interaction kernels, the energy variation combines external-potential, entropy, and interaction gradients; an asymmetric kernel contributes only its even part.Source: §C.2 · p. 20

C.3 LAGRANGIAN RESIDUAL AND EMPIRICAL APPROXIMATION

• The residual vanishes exactly when the equation of motion holds under the stated assumptions, while its empirical particle form is an approximation rather than an exact KDE-mean-velocity evaluation.Source: §C.3 · p. 20

C.4 TIME DISCRETIZATION

• The time-discrete loss uses central differences at interior nodes and trapezoidal integration, with one-sided endpoint stencils of the same order.Source: §C.4 · pp. 20–21

C.5 PSEUDOCODE

• Algorithm 2 initializes particle paths, KDE, potential and friction, then optimizes data fit plus weighted residual and optional initial-velocity losses.Source: §C.5 · p. 22

C.6 COMPUTATIONAL COMPLEXITY

• Up to min(B, M) times faster per update than WLM, Double-Stitch parallelizes residual force evaluations instead of sequentially unrolling a simulated trajectory.Apart from data loss, the stated costs are O(ceil(M/B) C(P)) for Double-Stitch and O(M C(P)) for WLM.Source: §C.6 · pp. 22–23

D.1 DATA

• Synthetic data comprise conservative harmonic, damped Lagrangian, and Ornstein–Uhlenbeck gradient flows; training uses snapshots 0, 2, and 4, with 1 and 3 held out.Source: §D.1 · p. 23
• The embryoid-body dataset contains 16,819 cells across five time bins, reduced to five standardized principal components and evaluated with LPO, L1O, and L2O protocols.Figure 3 shows the snapshots over time in a shared UMAP.Source: §D.1 · p. 23, Figure 3 · p. 24
• The ocean data include a 111-particle small-vortex interpolation task and a 400-particle large-vortex forecasting task, using unpaired snapshots.Figure 4 marks training points as circles and test points as diamonds, with color indicating time.Source: §D.1 · p. 24, Figure 4 · p. 24

Evaluation and timing

• Accuracy is reported as mean ± standard deviation over five EB seeds and three ocean seeds, while ablations and timing comparisons use three seeds.The training-time comparisons use different seed ranges for Double-Stitch and WLM, and exclude final evaluation.Source: §D.2 · pp. 24–25

Double-Stitch setup

• Double-Stitch uses OT-path trajectory initialization, a coarsest-scale bandwidth, and learned friction; EB and ocean experiments differ in particle counts, time grids, and training updates.The EB setup uses 1,000 centers and 10,000 updates; ocean setups use 100 or 400 centers and 40,000 updates.Source: §D.2 · pp. 24–25

Baseline setups

• Stitching reruns use endpoint OT initialization, while WLM uses released architectures and solver settings, with ocean blur estimated from training snapshots in the main setting.The original WLM ocean setting estimates blur from all snapshots.Source: §D.2 · p. 25

Initialization sensitivity

• Initialization sensitivity tests vary bandwidth initialization with OT-path fixed, then vary trajectory initialization with Silverman bandwidth fixed.The compared trajectory initializations are endpoint OT, OT-path, and McCann-path.Source: §D.2 · p. 25

Prediction

• Direct prediction evaluates the learned KDE, pushforward integrates from the latest observation, and rollout integrates from the initial snapshot without later restarts.For second-order dynamics, pushforward uses the observed velocity at its starting snapshot, while rollout uses only the global initial velocity.Source: §D.2 · p. 25

Synthetic demonstration

• The synthetic demonstration compares velocity-residual Stitching with Lagrangian-residual Double-Stitch on Gaussian-mixture flows using held-out W1 between fitted and analytic position densities.Densities are smoothed for visualization, but scores use unsmoothed densities.Source: §D.2 · p. 26

E.1 IDENTIFIABILITY

• Finite position snapshots need not identify Lagrangian dynamics: equal-density velocity fields, distinct initial velocities, and distinct potentials can produce indistinguishable observations.The examples include trajectories passing through the same observed points and nearly intersecting trajectories from different potentials.Source: §E.1 · pp. 26–28

E.2 SINGLE-CELL DYNAMICS

• Per-snapshot results are tabulated for EB leave-particles-out, leave-one-out, and leave-two-out evaluations.Double-Stitch predictions for the leave-one-out and leave-two-out scenarios are also shown.Source: §E.2 · p. 29, Table 6 · p. 29, Table 7 · p. 29
• Double-Stitch prediction plots show the EB leave-one-out and leave-two-out scenarios.The figures use a shared UMAP, with ground truth and predictions distinguished by pale and dark points.Source: §E.2 · p. 30, Figure 8 · p. 30, Figure 9 · p. 30

E.3 OCEAN VORTEX DYNAMICS

• Per-snapshot results cover small-vortex interpolation and large-vortex forecasting across three protocols.Double-Stitch predictions are also shown for small-vortex interpolation, with learned trajectories and pushforward predictions displayed separately.Source: §E.3 · p. 30, Table 9 · p. 31, Table 10 · p. 31

E.4 SYNTHETIC DEMONSTRATION

• Double-Stitch succeeds on conservative, damped, and gradient Gaussian-mixture flows, whereas Stitching fails on the conservative flow and underperforms on the damped flow.Both methods use the same training snapshots and fitting budget.Source: Figure 11 · p. 32, §E.4 · p. 32

Prediction mode and settings

• Double-Stitch performs similarly under direct and pushforward ocean evaluation, while rollout is worse but improves with a time-dependent potential; Stitching scores well only under direct evaluation.The authors suggest Stitching's first-order condition is likely misspecified for this data.Source: §E.5 · p. 32, Table 11 · p. 33

Training time and number of particles

• A tenfold increase in particle count raises training time 4–5 times across tasks; on oceans, P = 1000 gives the lowest forecasting error but the highest training time.For ocean interpolation, error increases with P; the authors suggest P = 100 already matches the dataset's 111 particles per snapshot.Source: §E.5 · pp. 33–34, Table 12 · p. 33, Table 17 · p. 34

Initialization

• Endpoint OT has the largest error on all five tasks, while McCann-path is generally close to OT-path except for ocean interpolation.Coarsest bandwidth initialization is better on ocean interpolation, with similar performance between bandwidth choices elsewhere.Source: §E.5 · p. 33, Table 13 · p. 33

Time-dependent potential

• A time-dependent potential outperforms a static potential on EB LPO and ocean forecasting, and matches it on the other reported tasks..Source: §E.5 · p. 33, Table 14 · p. 34

Residual weight and initial friction

• Smaller residual weights are slightly better on LPO, L1O, and ocean interpolation, while larger weights perform better on L2O and forecasting; γ0 = 0.1 gives slightly lower errors than γ0 = 1 across tasks.The authors hypothesize that larger gaps between observed snapshots make dynamics more important on L2O and forecasting; γ0 = 10 increases errors on EB and ocean interpolation.Source: §E.5 · p. 34, Table 15 · p. 34, Table 16 · p. 34