Paper deep dive
Wasserstein Residuals: Learning Gradient Flows from Population Dynamics
Markus Heinonen, Yair Shenfeld, Ricardo Baptista, Daniel Waxman, Dmitry Batenkov, Tim Cooijmans, Eli Bingham
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%
Last extracted: 7/7/2026, 3:32:01 PM
Summary
The paper introduces Stitching, a simulation-free particle-based method for reconstructing population dynamics modeled as Wasserstein gradient flows. By enforcing continuity equations via a non-negative velocity residual loss combined with a data-fitting divergence, Stitching unifies existing paradigms like Path-Finding and Action Matching. It overcomes limitations of the dominant Jordan-Kinderlehrer-Otto scheme, particularly regarding time discretization and large temporal gaps, achieving state-of-the-art performance on trajectory inference benchmarks.
Entities (8)
Relation Signals (6)
Markus Heinonen → authored → Wasserstein Residuals: Learning Gradient Flows from Population Dynamics
confidence 99% · Wasserstein Residuals: Learning Gradient Flows from Population Dynamics Markus Heinonen 1,2
Markus Heinonen → affiliatedwith → Basis Research Institute
confidence 95% · 1 Basis Research Institute
Stitching → uses → Wasserstein residuals
confidence 95% · We take a residual approach, enforcing the continuity equations via a non-negative loss function whose minimum is the WGF.
Stitching → outperforms → Jordan-Kinderlehrer-Otto scheme
confidence 90% · JKO-based methods are inflexible to time discretisation and require solving costly optimal transport problems... stitching method achieves state-of-the-art performance
Stitching → unifies → Path-Finding
confidence 85% · This perspective unifies several existing methods and leads to a new particle-based method... unifies existing paradigms — Path-Finding
Stitching → unifies → Action-matching
confidence 85% · This perspective unifies several existing methods... Action Matching (Neklyudov et al., 2023)
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Reconstructing population dynamics is a central problem in the physical and data sciences. Often, the dynamics are modeled as a Wasserstein gradient flow (WGF): a curve of distributions driven by an energy functional. Though there are multiple mathematical characterizations of a WGF, the dominant algorithmic approach relies on the Jordan--Kinderlehrer--Otto (JKO) scheme. JKO-based methods are inflexible to time discretisation and require solving costly optimal transport problems. We take a residual approach, enforcing the continuity equations via a non-negative loss function whose minimum is the WGF. Combined with a data-fitting divergence, this gives a single global objective. This perspective unifies several existing methods and leads to a new particle-based method, stitching, that is simulation-free and robust to large gaps between observations. We demonstrate that the stitching method achieves state-of-the-art performance across trajectory inference benchmarks. For code see this http URL.
Tags
Links
- Source: https://arxiv.org/abs/2607.04738v1
- Canonical: https://arxiv.org/abs/2607.04738v1
Trouble viewing inline? Open PDF directly →
Full Text
137,560 characters extracted from source content.
Expand or collapse full text
Wasserstein Residuals: Learning Gradient Flows from Population Dynamics Markus Heinonen 1,2 Yair Shenfeld 1,3 Ricardo Baptista 1,4 Daniel Waxman 1,5 Dmitry Batenkov 1 Tim Cooijmans 1 Eli Bingham 1 1 Basis Research Institute 2 Aalto University 3 Brown University 4 University of Toronto 5 MIT Correspondence to markus@basis.ai Abstract Reconstructing population dynamics is a central problem in the physical and data sciences. Often, the dynamics are modeled as a Wasserstein gradient flow (WGF): a curve of distributions driven by an energy functional. Though there are multiple mathematical characterizations of a WGF, the dominant algorithmic approach relies on the Jordan–Kinderlehrer–Otto (JKO) scheme. JKO-based methods are inflexible to time discretisation and require solving costly optimal transport problems. We take a residual approach, enforcing the continuity equations via a non-negative loss function whose minimum is the WGF. Combined with a data-fitting divergence, this gives a single global objective. This perspective unifies several existing methods and leads to a new particle-based method, stitching, that is simulation-free and robust to large gaps between observations. We demonstrate that the stitching method achieves state-of-the-art performance across trajectory inference benchmarks. For code see github.com/BasisResearch/wasserstein-residuals. 1 Introduction Reconstructing how a population evolves from a handful of snapshots is a central challenge across many scientific fields, ranging from computational biology (Schiebinger et al., 2019) to crowd dynamics (Maury et al., 2010). A common language for such problems is that of Wasserstein gradient flows: curves of probability measures t↦ρt _t driven by the steepest descent of an energy functional ℱF (Ambrosio et al., 2005, Santambrogio, 2015). We study the inverse problem: recovering the energy ℱF from observed snapshots so that its gradient flow ρ:=(ρt)ρ:=( _t) fits the snapshots qt\q_t\ for a collection of observation times t∈obst _obs. The dominant line of work (Bunne et al., 2022, Terpin et al., 2024, Persiianov et al., 2026) on this problem learns the functional ℱF by relying on the Jordan–Kinderlehrer–Otto (JKO) definition of Wasserstein gradient flows (Jordan et al., 1998). This requires solving multiple optimal transport (OT) problems, which can be costly, and lead to inaccuracies under long temporal gaps; see Figure 1. We take a different perspective: a residual loss that vanishes exactly when ρ is a gradient flow of ℱF. Combined with a data-fitting divergence at observed times, the objective directly enforces both the gradient flow constraint and the data fit. The perspective unifies existing paradigms — Path-Finding (Liu and Zhou, 2026) and Action Matching (Neklyudov et al., 2023) — as instantiations of one residual framework. The residual framework leads us to a new particle-based method which achieves state-of-the-art results on real-world datasets. Contributions. We summarize our contributions along the following three axes: • We recast the problem as minimizing a nonnegative residual term, a framework subsuming Path-Finding (Liu and Zhou, 2026) and Action Matching (Neklyudov et al., 2023). • We introduce stitching, a simulation-free KDE-based method, which promotes the curve ρ to a first-class learnable variable alongside ℱF and tolerates large gaps between snapshots. • We achieve state-of-the-art performance on single-cell RNA trajectory inference, and interaction dynamics recovery. Outline. In Section 2 we present background for learning WGF from population dynamics. Section 3 presents the general Wasserstein residuals framework, and Section 4 introduces our stitching method. Section 5 describes our numerical experiments. Figure 1: Stitching recovers curvature that JKO chord predictors miss. Sparse SDE snapshots (N=20N=20) on a sinusoidal valley at t∈0,5,10,20,30t∈\0,5,10,20,30\. (a) True potential. (b) Stitching’s trajectory tracks the valley. (c) Lightspeed (Terpin et al., 2024) does not capture the curved trajectories of the particles; see Section 5.1. 2 Background We work in the Wasserstein space (2(ℝD),W2)(P_2(R^D),W_2) of probability measures 2(ℝD)P_2(R^D) over real space ℝDR^D with finite second moment, equipped with the 2-Wasserstein distance W2W_2 (Ambrosio et al., 2005, Villani, 2016, Santambrogio, 2015); see Appendix A for definitions. An absolutely continuous curve ρ:=(ρt)t∈[0,T]ρ:=( _t)_t∈[0,T] admits a unique minimal velocity field v satisfying the continuity equation ∂tρt()=−div(ρt()t()), _t _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=-div_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ( _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) ), (1) (Santambrogio, 2015, Thm 8.3.1). The equivalent “Lagrangian” description is via particle trajectories ˙t=t(t),0∼ρ0, [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t=v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t), [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_0 _0, (2) satisfying t∼ρt [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t _t for all t∈[0,T]t∈[0,T]. A Wasserstein gradient flow of a functional ℱ:2(ℝD)→ℝF:P_2(R^D) is a curve ρ with velocity t()=−∇δℱδρt(),v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x), (3) where δℱδρt δ _t is the first variation of ℱF at ρt _t. We focus on functionals of the form ℱ[ρt]=∫ℝDV()dρt()+∫ℝDlogρt()dρt()+∫ℝD×ℝDW(′−)dρt()dρt(′),F[ _t]= _R^DV( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _R^D _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _R^D×R^DW( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x - [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ), (4) where V:ℝD→ℝV:R^D is a potential term, and W:ℝD→ℝW:R^D is a symmetric W(−)=W()W(- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=W( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) interaction term (Villani, 2016, §5.2.2). Besides their ubiquity, functionals of the form (4) have the advantage of having an analytic Wasserstein gradient ∇δℱδρt()=∇V()+∇logρt()+∇(W∗ρt)(). _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)= _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xV( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x(W _t)( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x). (5) A few standard examples of Wasserstein functionals appear in Appendix A. Let us now state the focus of our work: Goal Given a collection of discrete snapshots at times obs⊆[0,T]T_obs [0,T] — where, for simplicity, we assume 0∈obs0 _obs — and samples from marginal distributions qtt∈obs\q_t\_t _obs, find ℱF whose gradient flow ρ satisfies ρt=qt _t=q_t for all t∈obst _obs. 3 Wasserstein residuals Following the classical reference Ambrosio et al. (2005, Ch. 11), gradient flows in metric spaces admit four equivalent formulations: a Tangent condition (the continuity equation), an Energy Dissipation Equality (EDE), the JKO scheme, and an Evolution Variational Inequality (EVI). In Appendix A we provide the Euclidean intuition for this four-way equivalence. For the purpose of numerical computations with Wasserstein gradient flows, the EVI is an inequality and is thus not amenable to residual minimization. Whereas the JKO scheme has been widely used (Bunne et al., 2022, Terpin et al., 2024, Persiianov et al., 2026), we focus on the Tangent and EDE formulations which have been largely overlooked. Both the Tangent and EDE formulations are based on (1) with velocity t=−∇δℱδρtv_t=-∇ δ _t. The tangent formulation requires ∂tρt()=div(ρt()∇δℱδρt()), _t _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=div_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ( _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) ), (density constraint) while the EDE is formulated in terms of the velocity field tv_t, t()=−∇δℱδρt().v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x). (velocity constraint) Each form is a pointwise constraint that holds if and only if ρ is a gradient flow of ℱF. 3.1 Density and velocity residuals The (density constraint) and (velocity constraint) produce corresponding residuals: ℛdens[ℱ,ρ] _dens[F,ρ] :=∫0T∫ℝD‖∂tρt()−div(ρt()∇δℱδρt())‖2ρt()ddt := _0^T _R^D \| _t _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)-div_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ( _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) ) \|^2 _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt (density residual) ℛvel[ℱ,(ρ,)] _vel[F,(ρ,v)] :=∫0T∫ℝD‖t()+∇δℱδρt()‖2ρt()ddt. := _0^T _R^D \|v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2 _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt. (velocity residual) Both residuals equal zero if and only if ρ is a Wasserstein gradient flow of ℱF. However, the density residual requires higher-order differentiation (due to the divdiv_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x operator), so instead we focus on the velocity residual which avoids this issue. 3.2 Coupling residuals to data The (velocity residual) enforces that ρ is a gradient flow of ℱF, but it does not couple ρ to the data qtt∈obs\q_t\_t _obs. To this end we add a statistical divergence (ρt,qt)D( _t,q_t) at observed times, where D satisfies (ρt,qt)=0⇒ρt=qtD( _t,q_t)=0 _t=q_t. This yields the global objective ℒ(ℱ,ρ,)=λℛvel[ℱ,ρ,]+∑t∈obs(ρt,qt),L(F,ρ,v)= _vel[F,ρ,v]+ _t _obsD( _t,q_t), (6) with λ>0λ>0. Standard choices for D include the Kullback–Leibler divergence (i.e., likelihood maximization), Fisher divergence (i.e., score matching (SM) (Hyvärinen and Dayan, 2005)), and denoising score matching (DSM) (Vincent, 2011); see Appendix B. 4 Stitching We assume that the true functional ℱF is of the form (4), and parametrize an approximating functional ℱθF^θ with neural networks for the potential VθV^θ and the interaction kernel WθW^θ, so that the Wasserstein gradient can be written as ∇δℱθδρtθ()=∇Vθ()+∇logρtθ()+∇(W∗ρtθ)(). _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)= _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xV^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x(W _t^θ)( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x). (7) Curve parametrization. The parametrization of the curve ρθρ^θ is done by a kernel density estimate (KDE) over moving differentiable trajectories [0,T]∋t↦t,kθ[0,T] t [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ, ρtθ():=∑k=1Nwkθϕ(−t,kθ),wkθ≥0,∑k=1Nwkθ=1, _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x):= _k=1^Nw^θ_k\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ), w^θ_k≥ 0, _k=1^Nw^θ_k=1, (8) with ϕφ a smooth probability kernel. Given the parametrization (8), the associated velocity field is (cf. Claim 1 in Appendix C) tθ()=∑k=1Nwkθϕ(−t,kθ)˙t,kθ∑l=1Nwlθϕ(−t,lθ),v_t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)= _k=1^Nw^θ_k\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ)\, [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ _l=1^Nw^θ_l\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,l^θ), (9) which is defined in terms of the particles t,kθ\ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ\ and their velocities ˙t,kθ\ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ\. Continuous objective. With this KDE parametrization, the objective (6) reads ℒstitch(θ)=∫0T∫ℝD‖tθ()+∇δℱθδρtθ()‖2ρtθ()ddt+∑t∈obs(ρtθ,qt),L_stitch(θ)= _0^T _R^D \|v_t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2 _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt+ _t _obsD( _t^θ,q_t), (10) with tθv_t^θ as in (9) and ∇δℱθδρtθ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ as in (7). Empirical-measure approximation. To make the computation of ℒstitchL_stitch more efficient we make the approximation ρtθ()=∑k=1Nwkθϕ(−t,kθ)≈∑k=1Nwkθδt,kθ(), _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)= _k=1^Nw^θ_k\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ)≈ _k=1^Nw^θ_k\, _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x), (11) i.e., we approximate the KDE by the empirical measure on its centers. We use the approximation (11) both when integrating in the velocity-residual term against ρtθ()d _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x, and when evaluating the velocity tθv_t^θ of (9). Specifically, the velocity at a particle simplifies to tθ(t,kθ)=˙t,kθv_t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ)= [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ (cf. Claim 2 in Appendix C), which leads to ℒstitch(θ)≈∫0T∑k=1Nwkθ‖˙t,kθ+∇δℱθδρtθ(t,kθ)‖2dt+∑t∈obs(ρtθ,qt).L_stitch(θ)≈ _0^T _k=1^Nw^θ_k \| [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ+ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ) \|^2dt+ _t _obsD( _t^θ,q_t). (12) Time discretization. We discretize time 0=t0<t1<⋯<tK−1=T0=t_0<t_1<·s<t_K-1=T and approximate ˙tj,kθ≈Δtj+1,kθ/Δtj+1 [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ≈ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ/ t_j+1, where Δtj+1:=tj+1−tj t_j+1:=t_j+1-t_j is the width of the jjth time interval and Δtj+1,kθ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ is a discretization rule, e.g., forward Euler Δtj+1,kθ:=tj+1,kθ−tj,kθ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ:= [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ. Stitching Objective ℒ^stitch(θ):=∑j=0K−2∑k=1Nwkθ1Δtj+1‖Δtj+1,kθ+Δtj+1∇δℱθδρtθ(tj,kθ)‖2+∑t∈obs(ρtθ,qt). L_stitch(θ):= _j=0^K-2 _k=1^Nw^θ_k 1 t_j+1 \| [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ+ t_j+1 _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ) \|^2+ _t _obsD( _t^θ,q_t). (13) Data divergence and functional. For the data divergence term (ρtθ,qt)D( _t^θ,q_t) we still use the KDE parametrization (8), which in turn allows us to use any of the divergences in Appendix B. For Equation (7), the term ∇logρtθ∇ _t^θ uses (8) (see Appendix C), while the interaction term is again approximated using the centers: ∇(Wθ∗ρtθ)(tj,kθ)≈∑l=1Nwlθ∇Wθ(tj,kθ−tj,lθ).∇(W^θ _t^θ)( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ)≈ _l=1^Nw^θ_l\,∇ W^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,l^θ). (14) Relation to prior velocity-residual instantiations. The (velocity residual) has previously been instantiated as a normalizing flow (Liu and Zhou, 2026), and can be seen as an instance of Action Matching (Neklyudov et al., 2023); see Appendix D. Stitching’s distinction is that the trajectory ρθρ^θ is itself a learnable particle cloud — rather than the output of a learned flow that requires ODE integration. 5 Experiments 5.1 Illustrative example: continuous flow vs. JKO chord Figure 1 on page 1 contrasts continuous flow against JKO chord interpolation on the wavy valley potential V(x1,x2)=0.6(x2−sin(πx1/2))2−0.3x1V(x_1,x_2)=0.6(x_2- (π x_1/2))^2-0.3\,x_1: particles flow along a sinusoidal floor under an SDE dx=−∇Vdt+2βdWdx=-∇ V\,dt+ 2β\,dW with β=0.00625β=0.00625. We observe N=20N=20 particles at five irregular times t∈0,5,10,20,30t∈\0,5,10,20,30\; the t=10→20t=10→ 20 gap is too wide for a first-order JKO predictor to interpolate the curve faithfully. Stitching and JKOnet⋆ (Terpin et al., 2024) jointly learn VθV^θ and the diffusion coefficient (full setup in Appendix E). The purpose of the experiment is to highlight the flexibility of our explicitly parametrized trajectory curves. Whereas JKOnet⋆’s first-order chord interpolation is restricted to straight-line segments between consecutive snapshots, the trajectories found by our stitching algorithm correctly curve according to the wavy valley. As a consequence, we more accurately recover the underlying potential, with R2=0.62R^2=0.62 vs. 0.510.51, where R2R^2 is the coefficient of determination between learned and ground-truth potentials evaluated on the data support. 5.2 Synthetic potential recovery We benchmark on the 1515 two-dimensional potentials of Terpin et al. (2024)Nsim=2,000N_sim=2,000 ODE particles are integrated for T=5T=5 steps of Δt=0.01 t=0.01, and split 50/5050/50 into train / held-out test snapshots. Under the paired setup we observe the particle trajectories, while in the unpaired setup, snapshots are decorrelated (See Figure 6). All methods use the same (64,64)(64,64) MLP architecture for VθV^θ; stitching uses N=1,000N=1,000 particle trajectories of length K=50K=50. We report the metrics of Persiianov et al. (2026): (i) the EMD measuring forward-prediction accuracy; (i) the L2L^2-UVP measuring recovery of potential gradient; and (i) the BdW22Bd^2_W_2-UVP measuring distributional moment match (See Appendix F). Table 1: Comparison on the 66 potentials of Persiianov et al. (2026, Tab. 3) most sensitive to the paired→ transition. Three methods (J = JKOnetV⋆ _V, I = iJKOnetV, S = stitching) on paired and unpaired regimes. EMD measures forward-prediction accuracy; L2L^2-UVP measures recovery of the underlying potential’s gradient; BdW22Bd^2_W_2-UVP measures distributional moment match. Stitching’s potential-recovery and moment-match errors are effectively unchanged when snapshots are decorrelated, while JKO-based methods degrade or collapse. JKO methods exploit paired snapshots best on forward prediction (EMD). paired unpaired EMD ↓ L2L^2-UVP ↓ BdW22Bd^2_W_2-UVP ↓ EMD ↓ L2L^2-UVP ↓ BdW22Bd^2_W_2-UVP ↓ # potential J I S J I S J I S J I S J I S J I S 1 flowers 0.010.01 0.010.01 0.300.30 149149 −- 0.000.00 0.0000.000 0.0000.000 0.180.18 0.390.39 0.300.30 0.310.31 151151 −- 0.010.01 1.71.7 0.350.35 0.240.24 4 zigzag_ridge 0.040.04 0.020.02 0.290.29 40414041 −- 0.030.03 0.050.05 0.0030.003 0.170.17 0.380.38 0.290.29 0.300.30 38703870 −- 0.040.04 1.71.7 0.370.37 0.290.29 6 watershed 0.0020.002 0.0060.006 0.300.30 8787 −- 0.000.00 0.0000.000 0.0000.000 0.190.19 0.380.38 0.300.30 0.300.30 8989 −- 0.010.01 1.51.5 0.350.35 0.240.24 7 ishigami 0.010.01 0.020.02 0.300.30 719719 −- 0.010.01 0.0000.000 0.0010.001 0.170.17 0.380.38 0.300.30 0.300.30 733733 −- 0.010.01 1.61.6 0.350.35 0.250.25 8 friedman 0.090.09 0.060.06 0.290.29 38993899 −- 0.060.06 0.020.02 0.0020.002 0.200.20 0.400.40 0.290.29 0.310.31 37473747 −- 0.060.06 1.81.8 0.330.33 0.270.27 11 wavy_plateau 0.530.53 0.140.14 0.270.27 34113411 −- 0.020.02 7.47.4 0.400.40 0.510.51 0.480.48 0.280.28 0.290.29 34583458 −- 0.040.04 4.94.9 0.630.63 0.540.54 Figure 2: Stitching’s potential recovery is qualitatively unchanged when consecutive snapshots are decorrelated. Top: ground-truth V. Middle / bottom: stitching’s VθV^θ trained on paired / unpaired snapshots. Per-panel labels: Rpattern2R^2_pattern. Full 1515-potential galleries in Appendix F. Stitching recovers V on 1111 of 1414 informative landscapes in both regimes; failures are the angular rotational and the degenerate flat (full galleries in Figure 7 and Figure 8). Table 1 reports absolute numbers on the 66 potentials Persiianov et al. (2026, Tab. 3) flagged as most sensitive to the paired→ transition. Stitching is consistent across both regimes: paired and unpaired errors stay within ∼ 2× on every metric and every potential, and stitching never collapses. JKO-step methods, by contrast, exploit paired snapshots and beat stitching on forward prediction (EMD) in the paired regime, but degrade or collapse on the unpaired UVPs — JKOnetV⋆ _V in particular fails to recover an informative gradient field at all (L2L^2-UVP >100%>100\%) at this default training budget. The underlying reason is that stitching’s velocity residual is evaluated against a learnable trajectory rather than an OT coupling between consecutive snapshots, so its quality does not depend on snapshot-to-snapshot pairing. 5.3 Single-cell trajectory inference We apply stitching to the embryoid body (EB) single-cell RNA sequencing dataset of Moon et al. (2019), which captures a population of human embryonic stem cells over 2727 days of differentiation as five snapshots at days 1–3, 6–9, 12–15, 18–21, 24–27\1--3,\,6--9,\,12--15,\,18--21,\,24--27\, indexed t∈0,1,2,3,4t∈\0,1,2,3,4\. We follow the preprocessing of Tong et al. (2020) and reduce each cell to its first 55 principal components, the same setup used by Terpin et al. (2024) and Persiianov et al. (2026). Recovering an underlying energy landscape from population snapshots in single-cell biology is a recurring problem (Huizing et al., 2026); unlike JKO methods that require well-estimated marginals at every proximal step, stitching optimizes the full trajectory globally and handles long temporal gaps naturally. We parametrize ℱθ[ρtθ]=cVρtθ[Vθ()]+cHρtθ[logρtθ]F^θ[ _t^θ]=c_V\,E_ _t^θ[V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)]+c_H\,E_ _t^θ[ _t^θ] with a (64,64)(64,64) MLP for Vθ()V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x), matching the architecture of Persiianov et al. (2026). The time-varying variant Vθ(,t) V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x,t) concatenates t to the input. The curve ρθρ^θ uses N=100N=100 particle trajectories of length K=50K=50. Full details in Appendix G. Figure 3: Both static and time-varying VθV^θ drive the population evolution accurately. Contours of VθV^θ in (PC1, PC2) with PC3–5 fixed at the global mean; half-integer columns are unseen at training. The standard EB benchmark trains on all five observed snapshots and reports W1W_1 between predicted and observed marginals at each transition ρk→ρk+1 _k→ _k+1. Table 2 compares stitching against the baselines collected in Persiianov et al. (2026, Table 5); stitching achieves the best performance. Figure 3 visualises the time-varying Vθ(,t)V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x,t): its basin migrates with the differentiating cell population, while the static Vθ()V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) is forced to compromise across the trajectory. Table 2: Full-data EB single-cell benchmark (5D), W1W_1 (↓ ). Baseline results from Persiianov et al. (2026, Table 5, no standard deviations reported). Our results averaged over 5 seeds. Method t=1t=1 t=2t=2 t=3t=3 t=4t=4 Mean Citation Neural SDE 0.690.69 0.910.91 0.850.85 0.810.81 0.820.82 Li et al. (2020) TrajectoryNet 0.730.73 1.061.06 0.900.90 1.011.01 0.930.93 Tong et al. (2020) SB-FBSDE 0.560.56 0.800.80 1.001.00 1.001.00 0.840.84 Chen et al. (2022) NLSB 0.680.68 0.840.84 0.810.81 0.790.79 0.780.78 Koshizuka and Sato (2023) OT-CFM 0.780.78 0.760.76 0.770.77 0.750.75 0.770.77 Tong et al. (2024) WLF-OT 0.650.65 0.780.78 0.760.76 0.750.75 0.740.74 Neklyudov et al. (2024) WLF-SB 0.630.63 0.790.79 0.770.77 0.740.74 0.730.73 Neklyudov et al. (2024) JKOnet 1.531.53 1.271.27 1.131.13 1.411.41 1.341.34 Bunne et al. (2022) Static potential JKOnetV∗^*_V 0.990.99 1.111.11 1.061.06 1.301.30 1.121.12 Terpin et al. (2024) iJKOnetV 0.920.92 1.111.11 0.950.95 1.211.21 1.051.05 Persiianov et al. (2026) Stitching, Vθ()V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) 0.460.46±0.01\,± 0.01 0.600.60±0.01\,± 0.01 0.600.60±0.01\,± 0.01 0.650.65±0.01\,± 0.01 0.580.58±0.01\,± 0.01 This paper Time-varying potential JKOnett,V∗^*_t,V 0.690.69 0.770.77 0.690.69 0.780.78 0.730.73 Terpin et al. (2024) iJKOnett,V 0.510.51 0.580.58 0.570.57 0.640.64 0.580.58 Persiianov et al. (2026) Stitching, Vθ(,t)V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x,t) 0.440.44±0.01\,± 0.01 0.560.56±0.01\,± 0.01 0.570.57±0.01\,± 0.01 0.600.60±0.02\,± 0.02 0.540.54±0.01\,± 0.01 This paper Under the harder leave-two-out protocol of Shen et al. (2025) (train on t∈0,2,4t∈\0,2,4\, evaluate at held-out t∈1,3t∈\1,3\, W2W_2 metric), stitching also outperforms every published baseline: mean W2W_2 of 0.880.88 for the time-varying variant against 0.920.92 for the best published competitor (iJKOnett,V) and 1.121.12 for iJKOnetV (full table in Appendix G, Table 5). 5.4 Recovering interaction dynamics The mean-field limit (N→∞N→∞) of the WGF associated with (4) is the McKean–Vlasov process (McKean, 1966, Jabin and Wang, 2017). For finite N the corresponding interacting particle system is dXti=−∇V(Xti)dt−1N∑i≠j∇W(Xti−Xtj)dt+σdBti,i=1,…,N.dX_t^i=- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xV(X_t^i)dt- 1N _i≠ j _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xW(X_t^i-X_t^j)dt+σ dB_t^i, i=1,…,N. (15) Such agent-based interacting particle models and their generalizations and extensions (for example, adding self-propulsion and alignment terms) have been shown to approximate a wide variety of collective behaviors across organisms and scales (Vicsek and Zafeiris, 2012, Ouellette, 2022, Couzin et al., 2002, D’Orsogna et al., 2006, Cucker and Smale, 2007). In (15), ∇V _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xV represents some environmental or intrinsic “force”, while ∇W _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xW encodes social interaction forces between the agents. The term dBtdB_t is the standard Wiener process. Simultaneous recovery of both V and W from the marginals Xt∼qtX_t q_t is a notoriously difficult problem (Wei and Lu, 2026, Guan et al., 2024, Carrillo et al., 2025), since these terms are mixed in the transient dynamics. Here we showcase the ability of our method to disentangle these terms, by simulating a 2D system with a confining “Mexican hat” potential V(x)=α(‖x‖2−β)2V(x)=α(\|x\|^2-β)^2 and an attractive Gaussian kernel W(r)=ηe−r2W(r)=η\, e^-r^2 (full setup in Appendix H). The attractive interaction results in formation of one or more clusters (Rainer and Krause, 2002, Motsch and Tadmor, 2014), while the presence of V confines the particles to a ring or radius β β. We jointly learn an MLP Vθ(x)V^θ(x) and radial Wθ(r)W^θ(r), reporting scale-invariant pattern R2R^2 (Persiianov et al., 2026) for the functionals and standard distributional metrics on held-out particles. Figure 4: Stitching captures the Gaussian → ring → cluster phase transition and recovers V,WV,W at pattern R2=0.84,0.80R^2=0.84,0.80. Top: observed data at eight times; rightmost cell shows learned Vθ(x)V^θ(x) filled, with V contours dashed. Bottom: Stitching KDE samples; rightmost cell shows Wθ(r)W^θ(r) (blue) vs. W (dashed) over the data pair-distance band (shaded). Results. Figure 4 shows that stitching reproduces the phase transition from snapshots alone, and the rightmost column compares the recovered Vθ,WθV^θ,W^θ against the truth: pattern R2=0.84R^2=0.84 on V and 0.810.81 on W, on the same order as the strongest landscapes of Table 4. Table 3 reports the distributional metrics; JKOnet⋆ has failed to recover W completely, and is ∼10× 10× slower per epoch since each JKO step requires an optimal transport coupling per snapshot pair (200200 pairs in our example). Table 3: Interaction dynamics results on learning potential V, entropy and radial interaction W. Method EMD ↓ W2 ↓ BW2 ↓ MMD ↓ R2(V)↑R^2(V)\, R2(W)↑R^2(W)\, per iter ↓ total ↓ JKOnetV,W⋆ _V,W (Terpin et al., 2024), radial W 50.0650.06 82.6282.62 1012010120 0.1730.173 0.040.04 0.720.72 ∼4.1 4.1 s ∼69 69 min Stitching, Vθ(),Wθ(r)V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x),W^θ(r) (ours) 0.490.49 0.850.85 0.190.19 0.0170.017 0.840.84 0.810.81 0.240.24 s 4 min 5.5 Recovering non-gradient flows As we discuss in further detail in Section 7 below, our method can in principle be applied to learning more general flows beyond WGFs. Here we showcase this ability by fitting a chiral dynamics (Liebchen and Levis, 2022) with the stitching loss. The dynamics is interaction-only, such that the kernel has a nonzero curl. In Figure 5 we show the original data and the reconstructed stitching trajectories over a few time snapshots. Full details are provided in Appendix I. Figure 5: Stitching tracks the chiral orbiting dynamics. Samples from the learned curve ρθρ^θ (model) and the data q (observed) at particular time snapshots over the training window. 6 Related work In this section, we discuss some related approaches in learning potential functions for population dynamics in the context of Wasserstein gradient flows. Appendix D provides a more detailed review. JKO-based methods. The dominant algorithmic framework to tackle our main goal is the Jordan–Kinderlehrer–Otto (JKO) scheme (Jordan et al., 1998), which solves a sequence of proximal problems with a Wasserstein penalty, cf. Appendix A. Currently, there are two strategies to learn the potential ℱF from data qtt∈obs\q_t\_t _obs using the JKO scheme. JKOnet (Bunne et al., 2022) learns ℱF by backpropagation through the inner loop that solves the proximal problem. The optimal transport is represented in terms of a transport map ∇ψθ∇ψ^θ where ψθ:ℝD→ℝψ^θ:R^D is an input-convex neural network (ICNN; Amos et al. (2017)). iJKOnet (Persiianov et al., 2026) avoids the need for ICNNs through a min-max formulation. JKOnet⋆ (Terpin et al., 2024) algorithms solve optimal transport of the empirical data once upfront, and then fit the gradient of the functional to the observed displacements. The above strategies directly represent ℱF, but leave the curve ρ defined implicitly through the JKO scheme on ℱF. In contrast, we parametrize both ℱF and ρ, coupled by a residual loss to enforce consistency. Directly representing ρ allows us to choose a temporal discretization independent from that of the observations. Residual Losses. Concurrent work of Liu and Zhou (2026) considers the problem of solving a WGF given a known functional ℱF using a form of (velocity residual), resulting in GenWGP. In particular, considering a K-step discretization of the interval [0,T][0,T], GenWGP parameterizes the flow-map Φθ(t,) ^θ(t, [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) as a normalizing flow. This flow is applied to a population of particles, and finite difference methods are then used to approximate the velocity residual, which becomes the objective in gradient descent. GenWGP requires solving a costly neural ODE via application of the learned flow-map, but sidesteps side-steps the KDE approximation that stitching must take. Action Matching. Neklyudov et al. (2023) introduces action matching (AM) for learning the velocity field of a population dynamic, and Neklyudov et al. (2024) extends the framework to a broader class of Wasserstein Lagrangian flows beyond pure gradient flows. AM fits a velocity field to data without access to ground-truth velocities. Adapted to our setting, the goal is to learn a functional whose gradient matches the velocity of the data. Neklyudov et al. (2024) effectively integrate (velocity residual) by parts to eliminate the dependence on the velocity. Unlike the divergence-based methods, AM fuses residual and data fitting into a single objective. The price is that AM requires data at the temporal endpoints and learns only the functional, leaving the trajectory ρ implicit. 7 Discussion Summary We introduce the Wasserstein residual framework as way to learn Wasserstein gradient flows from snapshots of population dynamics. We instantiate the Wasserstein residuals framework by the stitching algorithm, which is a simulation-free particle method robust against long temporal gaps in the observed time-series. The stitching achieves state-of-the-art results on the embryoid body (EB) single-cell RNA sequencing dataset. Limitations Being a particle method, stitching inherits the O(N2)O(N^2) quadratic computation per time step for the entropy and interaction terms, while the potential term is O(N)O(N) linear. In addition, theoretical convergence guarantees as N→∞N→∞ remain open. Finally, the residual framework enforces a constraint using a residual regularizer, which does not guarantee the WGF condition ℛ=0R=0. In practice the learned flows have residuals close to zero. Broader impacts Stitching enables learning population dynamics from sparse snapshots in domains ranging from single-cell biology and trajectory inference to crowd and collective behaviour. The most immediate benefits are in scientific discovery with dual-use risks generic to dynamical-system inference: applied uncritically to social or behavioural data, the recovered potentials may be taken as causal explanations. Practitioners should validate the gradient-flow assumption against domain knowledge before drawing conclusions from the learned functional. 7.1 Future work Density residuals. In this work we focus on (velocity residual) and instantiate it with our stitching method. Future research should explore the use of (density residual), which can be exploited using neural networks or particle-based methods. Beyond Wasserstein gradient flows. The residual approach applies in fact to a more general family of flows, beyond Wasserstein gradient flows. We demonstrate this point in Section 5.5, but the framework applies much more generally. Suppose our goal is to match the data qtt∈obs\q_t\_t _obs to a curve ρ=(ρt)ρ=( _t) of the form ∂tρt()=−div(ρt()()) _t _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=-div_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ( _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)u( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) ), where u is a vector field for which we assume some structure. For example, in this work we assume t()=−∇δℱδρt() u_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) for some functional ℱF. To find u we can use either (density residual) or (velocity residual) by replacing ∇δℱδρt() _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) by a parametrized family tu_t. Such non-gradient-flows arise in the study of transformers (Geshkovski et al., 2025), chiral active matter (Liebchen and Levis, 2022), non-reciprocal collection systems (Fruchart et al., 2021), and ocean currents (Petrović et al., 2025). Disclosure Yair Shenfeld’s and Ricardo Baptista’s contributions to this work resulted in part from their affiliation with Basis Research Institute (an outside organization with respect to Brown University and University of Toronto). References M. S. Albergo and E. Vanden-Eijnden (2023) Building normalizing flows with stochastic interpolants. In ICLR, Cited by: §D.5. L. Ambrosio, N. Gigli, and G. Savaré (2005) Gradient flows in metric spaces and in the space of probability measures. 2nd edition, Lectures in Mathematics ETH Zürich. Cited by: §A.3, §A.3, §A.3, Appendix A, §1, §2, §3. B. Amos, L. Xu, and J. Z. Kolter (2017) Input convex neural networks. In ICML, Cited by: §6. C. Bunne, L. Meng-Papaxanthos, A. Krause, and M. Cuturi (2022) Proximal optimal transport modeling of population dynamics. In AISTATS, Cited by: §D.1, §1, §3, Table 2, §6. J. A. Carrillo, G. Estrada-Rodriguez, L. Mikolás, and S. Tang (2025) Sparse identification of nonlocal interaction kernels in nonlinear gradient flow equations via partial inversion. Mathematical Models and Methods in Applied Sciences 35 (05), p. 1073–1131. Cited by: §D.6, §5.4. R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud (2018) Neural ordinary differential equations. In NeurIPS, Cited by: §D.5. T. Chen, G. Liu, M. Tao, and E. Theodorou (2023) Deep momentum multi-marginal Schrödinger bridge. In NeurIPS, Cited by: §D.5, Table 5. T. Chen, G. Liu, and E. A. Theodorou (2022) Likelihood training of Schrödinger bridge using forward-backward SDEs theory. In ICLR, Cited by: Table 2. I. D. Couzin, J. Krause, R. James, G. D. Ruxton, and N. R. Franks (2002) Collective memory and spatial sorting in animal groups. Journal of Theoretical Biology 218 (1), p. 1–11. Cited by: §5.4. F. Cucker and S. Smale (2007) Emergent behavior in flocks. IEEE Transactions on Automatic Control 52 (5), p. 852–862. Cited by: §5.4. M. R. D’Orsogna, Y. L. Chuang, A. L. Bertozzi, and L. S. Chayes (2006) Self-propelled particles with soft-core interactions: Patterns, stability, and collapse. Physical Review Letters 96 (10), p. 104302. Cited by: Appendix I, §5.4. V. De Bortoli, J. Thornton, J. Heng, and A. Doucet (2021) Diffusion Schrödinger bridge with applications to score-based generative modeling. In NeurIPS, Cited by: §D.5. M. Fruchart, R. Hanai, P. B. Littlewood, and V. Vitelli (2021) Non-reciprocal phase transitions. Nature 592 (7854), p. 363–369. Cited by: Appendix I, §7.1. B. Geshkovski, C. Letrouit, Y. Polyanskiy, and P. Rigollet (2025) A mathematical perspective on transformers. Bulletin of the American Mathematical Society 62 (3), p. 427–479. Cited by: §7.1. V. Guan, J. Janssen, H. Rahmani, A. Warren, S. Zhang, E. Robeva, and G. Schiebinger (2024) Identifying drift, diffusion, and causal structure from temporal snapshots. arXiv:2410.22729. Cited by: §D.6, §5.4. V. Guan, H. Rahmani, J. Janssen, A. Warren, E. Robeva, and G. Schiebinger (2026) Gradient-flow SDEs have unique transient population dynamics. In AISTATS, Cited by: §D.6. M. Hua, E. Vanden-Eijnden, and R. T. Q. Chen (2025) Simulation-free differential dynamics through neural conservation laws. In UAI, Cited by: §D.4. G. Huizing, J. Samaran, D. Capocefalo, A. Audit, L. Cantini, and G. Peyré (2026) STORIES: learning cell fate landscapes from spatial transcriptomics using optimal transport. Nature Methods 23, p. 522–531. Cited by: §5.3. A. Hyvärinen and P. Dayan (2005) Estimation of non-normalized statistical models by score matching.. JMLR 6 (24), p. 695–709. Cited by: Appendix B, §3.2. P. Jabin and Z. Wang (2017) Mean field limit for stochastic particle systems. In Active Particles, Volume 1, p. 379–402. Cited by: §5.4. R. Jordan, D. Kinderlehrer, and F. Otto (1998) The variational formulation of the Fokker–Planck equation. SIAM Journal on Mathematical Analysis 29 (1), p. 1–17. Cited by: §A.3, §D.1, §1, §6. P. Kidger, J. Foster, X. C. Li, and T. Lyons (2021a) Efficient and accurate gradients for neural SDEs. NeurIPS. Cited by: §D.5. P. Kidger, J. Foster, X. Li, H. Oberhauser, and T. J. Lyons (2021b) Neural SDEs as infinite-dimensional GANs. In ICML, Cited by: §D.5. T. Koshizuka and I. Sato (2023) Neural Lagrangian Schrödinger bridge: Diffusion modeling for population dynamics. In ICLR, Cited by: §D.5, §D.5, Table 2. C. Léonard (2014) A survey of the Schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems 34, p. 1533–1574. Cited by: §D.5. X. Li, T. L. Wong, R. T. Q. Chen, and D. Duvenaud (2020) Scalable gradients for stochastic differential equations. In AISTATS, Cited by: §D.5, Table 2. B. Liebchen and D. Levis (2022) Chiral active matter. Europhysics Letters 139 (6), p. 67001. Cited by: Appendix I, §5.5, §7.1. Y. Lipman, R. T.Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le (2023) Flow matching for generative modeling. In ICLR, Cited by: §D.5. C. Liu and X. Zhou (2026) Generative path-finding method for Wasserstein gradient flow. arXiv:2604.11519. Cited by: 4th item, §D.2, 1st item, §1, §4, §6. B. Maury, A. Roudneff-Chupin, and F. Santambrogio (2010) A macroscopic crowd motion model of gradient flow type. Mathematical Models and Methods in Applied Sciences 20 (10), p. 1787–1821. Cited by: §1. H. P. McKean (1966) A class of Markov processes associated with nonlinear parabolic equations. PNAS 56 (6), p. 1907–1911. Cited by: §5.4. P. Mokrov, A. Korotin, L. Li, A. Genevay, J. M. Solomon, and E. Burnaev (2021) Large-scale Wasserstein gradient flows. In NeurIPS, Cited by: §D.1. K. R. Moon, D. van Dijk, Z. Wang, S. Gigante, D. B. Burkhardt, W. S. Chen, K. Yim, A. van den Elzen, M. J. Hirn, R. R. Coifman, N. B. Ivanova, G. Wolf, and S. Krishnaswamy (2019) Visualizing structure and transitions in high-dimensional biological data. Nature Biotechnology 37 (12), p. 1482–1492. Cited by: Appendix G, §5.3. S. Motsch and E. Tadmor (2014) Heterophilious dynamics enhances consensus. SIAM Review 56 (4), p. 577–621. Cited by: §5.4. K. Neklyudov, R. Brekelmans, D. Severo, and A. Makhzani (2023) Action matching: Learning stochastic dynamics from samples. In ICML, Cited by: 4th item, §D.3, §D.3, 1st item, §1, §4, §6. K. Neklyudov, R. Brekelmans, A. Tong, L. Atanackovic, Q. Liu, and A. Makhzani (2024) A computational framework for solving Wasserstein Lagrangian flows. In ICML, Cited by: §D.3, Table 2, Table 2, §6. N. T. Ouellette (2022) A physics perspective on collective animal behavior. Physical Biology 19 (2), p. 021004. Cited by: §5.4. G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan (2021) Normalizing flows for probabilistic modeling and inference. JMLR 22 (57), p. 1–64. Cited by: §D.5. M. Persiianov, A. Korotin, and E. Burnaev (2026) Learning of population dynamics: Inverse optimization meets JKO scheme. In ICLR, Cited by: §D.1, Figure 6, Figure 6, Appendix F, Appendix F, Appendix F, Table 4, 2nd item, 2nd item, Appendix G, Table 5, Table 5, Table 5, Appendix H, §1, §3, §5.2, §5.2, §5.3, §5.3, §5.3, §5.4, Table 1, Table 2, Table 2, Table 2, §6. K. Petrović, L. Atanackovic, V. Moro, K. Kapuśniak, İ. İ. Ceylan, M. Bronstein, A. J. Bose, and A. Tong (2025) Curly flow matching for learning non-gradient field dynamics. In NeurIPS, Cited by: §7.1. H. Rainer and U. Krause (2002) Opinion dynamics and bounded confidence: Models, analysis and simulation. Journal of Artificial Societies and Social Simulation 5 (3). Cited by: §5.4. J. Richter-Powell, Y. Lipman, and R. T. Q. Chen (2022) Neural conservation laws: a divergence-free perspective. In NeurIPS, Cited by: §D.4. F. Santambrogio (2015) Optimal transport for applied mathematicians: Calculus of variations, PDEs, and modeling. Springer. Cited by: §A.1, §A.1, Appendix A, §D.1, §1, §2, §2. G. Schiebinger, J. Shu, M. Tabaka, B. Cleary, V. Subramanian, A. Solomon, J. Gould, S. Liu, S. Lin, P. Berube, et al. (2019) Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell 176 (4), p. 928–943. Cited by: §1. Y. Shen, R. Berlinghieri, and T. Broderick (2025) Multi-marginal Schrödinger bridges with iterative reference refinement. In AISTATS, Cited by: Table 5, §5.3. Y. Shi, V. De Bortoli, A. Campbell, and A. Doucet (2023) Diffusion Schrödinger bridge matching. In NeurIPS, Cited by: §D.5. Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole (2021) Score-based generative modeling through stochastic differential equations. In ICLR, Cited by: §D.5. A. Terpin, N. Lanzetti, M. Gadea, and F. Dörfler (2024) Learning diffusion at lightspeed. In NeurIPS, Cited by: §D.1, Appendix E, 2nd item, Appendix F, Appendix F, Table 4, 2nd item, Table 5, Table 5, Appendix H, Figure 1, §1, §3, §5.1, §5.2, §5.3, Table 2, Table 2, Table 3, §6. A. Tong, K. Fatras, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, G. Wolf, and Y. Bengio (2024) Improving and generalizing flow-based generative models with minibatch optimal transport. TMLR. Cited by: §D.5, Table 2. A. Tong, J. Huang, G. Wolf, D. Van Dijk, and S. Krishnaswamy (2020) TrajectoryNet: A dynamic optimal transport network for modeling cellular dynamics. In ICML, Cited by: Appendix G, Table 5, §5.3, Table 2. F. Vargas, P. Thodoroff, A. Lamacraft, and N. D. Lawrence (2021) Solving Schrödinger bridges via maximum likelihood. Entropy 23 (9), p. 1134. Cited by: Table 5. T. Vicsek and A. Zafeiris (2012) Collective motion. Physics Reports 517 (3-4), p. 71–140. Cited by: §5.4. C. Villani (2016) Topics in optimal transportation. 2nd edition, Graduate Studies in Mathematics, Vol. 58, American Mathematical Society. Cited by: Appendix A, §2, §2, Example 2. P. Vincent (2011) A connection between score matching and denoising autoencoders. Neural computation 23 (7), p. 1661–1674. Cited by: Appendix B, §3.2. M. P. Wand and M. C. Jones (1994) Kernel smoothing. Chapman & Hall/CRC. Cited by: Appendix C. D. Wang, Y. Jiang, Z. Zhang, X. Gu, P. Zhou, and J. Sun (2025) Joint velocity-growth flow matching for single-cell dynamics modeling. In NeurIPS, Cited by: §D.5. V. Wei and F. Lu (2026) Learning interacting particle systems from unlabeled data. arXiv:2604.02581. Cited by: §D.6, §5.4. Appendix A Mathematical background: Wasserstein gradient flows This appendix collects standard definitions and results for Wasserstein gradient flows used in the main text. We refer the reader to Ambrosio et al. (2005), Villani (2016), Santambrogio (2015) for comprehensive treatment. A.1 The Wasserstein space The space of probability measures with finite second moment is 2(ℝD):=ρ probability measure on ℝD:∫ℝD‖2dρ()<∞,P_2(R^D):= \ρ probability measure on R^D: _R^D \| [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x \|^2\,dρ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)<∞ \, (16) endowed with the 2-Wasserstein distance W22(ρ,ν):=infγ∈Π(ρ,ν)∫ℝD×ℝD‖−′‖2dγ(,′),ρ,ν∈2(ℝD),W_2^2(ρ,ν):= _γ∈ (ρ,ν) _R^D×R^D \| [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x \|^2\,dγ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x, [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ), ρ,ν _2(R^D), (17) where Π(ρ,ν) (ρ,ν) is the set of couplings (probability measures on ℝD×ℝDR^D×R^D with marginals ρ and ν). An absolutely continuous curve ρ:[0,T]→2(ℝD)ρ:[0,T] _2(R^D) admits a unique minimal velocity field v satisfying the continuity equation ∂tρt()+div(ρt()t())=0, _t _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+div_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ( _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) )=0, (18) (Santambrogio, 2015, p. 167); (18) is the Eulerian description. The equivalent Lagrangian description is via particle trajectories ˙t=t(t),0∼ρ0, [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t=v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t), [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_0 _0, (19) which satisfies (Santambrogio, 2015, Theorem 8.3.1) t∼ρt,t∈[0,T]. [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t _t, t∈[0,T]. (20) A.2 Functionals on the Wasserstein space A functional on the Wasserstein space is a map ℱ:2(ℝD)→ℝF:P_2(R^D) . A general family of functionals, common in applications, is ℱ[p]=∫ℝDV()dp()+∫ℝDU(p())d+∫ℝD×ℝDW(′−)dp()dp(′),F[p]= _R^DV( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,dp( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _R^DU(p( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x))\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x+ _R^D×R^DW( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x - [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,dp( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,dp( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ), (21) where p∈2(ℝD)p _2(R^D), V:ℝD→ℝV:R^D , U:ℝ≥0→ℝU:R_≥ 0 , and W:ℝD→ℝW:R^D are sufficiently regular. Let us consider some concrete examples. Example 1 (Kullback–Leibler divergence). Let V be such that ν=e−Vν=e^-V is a probability measure in 2(ℝD)P_2(R^D), let U(r)=rlogrU(r)=r r, and let W=0W=0. Then (21) is the Kullback–Leibler divergence functional ℱ[p]=KL[p∥ν]=∫ℝDlog(ρ()ν())dρ().F[p]=KL[p\|ν]= _R^D ( ρ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)ν( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) )\,dρ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x). (22) Example 2 (Aggregation). Let V=0V=0, U=0U=0, and let W be a symmetric interaction kernel. Then ℱ[p]=∫ℝD×ℝDW(′−)dp()dp(′)F[p]= _R^D×R^DW( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x - [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,dp( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,dp( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ) (23) Functionals of the form (23) model swarming, chemotaxis, and granular media (Villani, 2016, §5.4). A.3 Gradient flows The classical reference Ambrosio et al. (2005, Ch. 11, p. 279–280) identifies four equivalent formulations of gradient flows in metric spaces: the Tangent condition, the Jordan–Kinderlehrer–Otto (JKO) scheme, the Evolution Variational Inequality (EVI), and the Energy Dissipation Equality (EDE). To build intuition, we begin in Euclidean space, then map each formulation to its Wasserstein-space analogue. Let f:ℝD→ℝf:R^D be a smooth function and let (xt)t∈[0,T](x_t)_t∈[0,T] a smooth curve in ℝDR^D. Tangent condition. The most direct definition of (xt)t∈[0,T](x_t)_t∈[0,T] being a gradient flow of f is if the equation x˙t=−∇f(xt)for all t∈[0,T], x_t=-∇ f(x_t) all t∈[0,T], (TangentE) holds. The terminology comes from x˙t x_t lying in the tangent space at xtx_t. While intuitive, (TangentE) cannot accommodate non-differentiable f. JKO scheme. A more general definition of (xt)t∈[0,T](x_t)_t∈[0,T] being a gradient flow of f if it satisfies xt+h=argminx∈ℝD[f(x)+12h‖x−xt‖2]for all h>0.x_t+h= *arg\,min_x ^D [f(x)+ 12h\|x-x_t\|^2 ] all h>0. (JKOE) For differentiable f, the first-order optimality condition of (JKOE) reads xt+h−xth=−∇f(xt+h)for all h>0, x_t+h-x_th=-∇ f(x_t+h) all h>0, (JKO-FOE) which is the implicit-Euler discretization of (TangentE) and recovers it as h→0h→ 0. Equation (JKOE) is the basis of proximal algorithms, and the terminology comes from the seminal work of Jordan–Kinderlehrer–Otto (JKO) (Jordan et al., 1998) who used this formulation to define gradient flows in Wasserstein space. Evolution Variational Inequality (EVI). The EVI definition for (xt)t∈[0,T](x_t)_t∈[0,T] being a gradient flow of f requires the existence of α∈ℝα such that 12dt‖xt−y‖2≤f(y)−f(xt)−α2‖xt−y‖2for all y∈ℝD. 12 ddt\|x_t-y\|^2≤ f(y)-f(x_t)- α2\|x_t-y\|^2 all y ^D. (EVIE) For smooth α-convex f, the convexity inequality ∇f(xt)⋅(y−xt)≤f(y)−f(xt)−α2‖xt−y‖2∇ f(x_t)·(y-x_t)≤ f(y)-f(x_t)- α2\|x_t-y\|^2 combined with (EVIE) forces x˙t=−∇f(xt) x_t=-∇ f(x_t). The EVI form is central to the gradient flow theory in metric spaces (Ambrosio et al., 2005). Energy Dissipation Equality (EDE). The EDE definition for (xt)t∈[0,T](x_t)_t∈[0,T] being a gradient flow of f requires the validity of the identity f(xT)−f(x0)=−∫0T[12‖∇f(xr)‖2+12‖x˙r‖2]dr.f(x_T)-f(x_0)=- _0^T [ 12\|∇ f(x_r)\|^2+ 12\| x_r\|^2 ]dr. (EDEE) To see why (EDEE) characterizes gradient flows, note that for any smooth curve, the chain rule and successive applications of Cauchy–Schwarz and AM–GM inequalities give f(xT)−f(x0) f(x_T)-f(x_0) =∫0T∇f(xr)⋅x˙rdr≥−∫0T‖∇f(xr)‖⋅‖x˙r‖dr = _0^T∇ f(x_r)· x_r\,dr\;≥\;- _0^T\|∇ f(x_r)\|·\| x_r\|\,dr ≥−∫0T[12‖∇f(xr)‖2+12‖x˙r‖2]dr. \;≥\;- _0^T [ 12\|∇ f(x_r)\|^2+ 12\| x_r\|^2 ]dr. The first inequality is an equality if and only if x˙r=−cr∇f(xr) x_r=-c_r\,∇ f(x_r) for some cr≥0c_r≥ 0, while the second inequality is an equality if and only if ‖x˙r‖=‖∇f(xr)‖\| x_r\|=\|∇ f(x_r)\|. Together (when ∇f(xr)≠0∇ f(x_r)≠ 0) the two equalities force cr=1c_r=1, i.e. x˙t=−∇f(xt) x_t=-∇ f(x_t). It follows that (EDEE) holding with equality is equivalent to the tangent condition. From Euclidean to Wasserstein. The four formulations transfer to the Wasserstein space 2(ℝD)P_2(R^D) by replacing the Euclidean structure with the W2W_2 metric. Given a curve (ρt)t∈[0,T]( _t)_t∈[0,T] in the Wasserstein space a functional ℱ:2(ℝD)→ℝF:P_2(R^D) the following four formulations define what it means for (ρt)t∈[0,T]( _t)_t∈[0,T] to be a gradient flow of ℱF. • (TangentE) becomes the continuity equation ∂tρt()=−div(ρt()t())wheret()=−∇δℱδρt(), _t _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=-div_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x( _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)) where _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x), (TangentW) which is the constraint underlying the density residual (density residual). • (JKOE) becomes the Wasserstein proximal scheme ρt+h=argminν∈2(ℝD)[ℱ(ν)+12hW22(ρt,ν)], _t+h= *arg\,min_ν _2(R^D) [F(ν)+ 12hW_2^2( _t,ν) ], (JKOW) which is the basis of the JKO methods; cf. Section D.1. • (EVIE) becomes a W2W_2-EVI 12dtW22(ρt,ν)≤ℱ(ν)−ℱ(ρt)−α2W22(ρt,ν)for all ν∈2(ℝD), 12 ddtW_2^2( _t,ν) (ν)-F( _t)- α2W_2^2( _t,ν) all ν _2(R^D), (EVIW) which is an inequality rather than an equality, so is not amenable to residual minimization. • (EDEE) becomes the Wasserstein EDE, ℱ[ρT]−ℱ[ρ0]=−∫0T[12|ρ˙t|W22+12∫ℝD‖∇δℱδρt()‖2ρt()d]dt,F[ _T]-F[ _0]=- _0^T\! [ 12| ρ_t|^2_W_2+ 12 _R^D\! \| _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2\! _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ]dt, (EDEW) where |ρ˙t|W22:=limh→0W22(ρt,ρt+h)h2.| ρ_t|^2_W_2:= _h→ 0 W_2^2( _t, _t+h)h^2. Lemma 1 below shows how to reformulate (EDEW) in a way which is amenable to a residual formulation, which is the basis for the velocity residuals of Section 4 and the works Liu and Zhou (2026) and Neklyudov et al. (2023). Lemma 1 (EDE form of ℛvelR_vel). Let (ρ,)(ρ,v) be a curve satisfying the continuity equation ∂tρt()=−div(ρt()t()), _t _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)=-div_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ( _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) ), (24) such that, for all t∈[0,T]t∈[0,T], ρt∈C1(ℝD) _t∈ C^1(R^D) vanishes at infinity, and t∈C1(ℝD,ℝD)v_t∈ C^1(R^D,R^D). Let ℱ:2(ℝD)→ℝF:P_2(R^D) be a functional such that, for all t∈[0,T]t∈[0,T], ℱ(ρt)<∞F( _t)<∞, δℱδρt∈C1(ℝD) δ _t∈ C^1(R^D) vanishes at infinity, and ℱF satisfies the chain rule ℱ(ρT)−ℱ(ρ0)=∫0T∫ℝDδℱδρt()∂tρt()ddt.F( _T)-F( _0)= _0^T\!\! _R^D δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\, _t _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt. (25) Then, ℛvel[ℱ,(ρ,)]=∫0T∫ℝD‖t()+∇δℱδρt()‖2ρt()ddt=2(ℱ[ρT]−ℱ[ρ0])+∫0T[|ρ˙t|W22+∫ℝD‖∇δℱδρt()‖2ρt()d]dt. splitR_vel[F,(ρ,v)]&= _0^T\!\! _R^D \|v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2 _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt\\ &=2 (F[ _T]-F[ _0] )+ _0^T\! [| ρ_t|^2_W_2+ _R^D\! \| _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2\! _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ]dt. split (26) Proof. By Ambrosio et al. (2005, Theorem 8.3.1), |ρ˙t|W22=∫ℝD‖t()‖2ρt()d | ρ_t |^2_W_2= _R^D \|v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2 _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x. Expanding the squared norm in ℛvel[ℱ,(ρ,)]R_vel[F,(ρ,v)], and integrating by parts, ℛvel[ℱ,(ρ,)] _vel[F,(ρ,v)] =∫0T∫ℝD‖t‖2ρtddt+2∫0T∫ℝDt⋅∇δℱδρtρtddt+∫0T∫ℝD‖∇δℱδρt‖2ρtddt = _0^T\!\! _R^D \|v_t \|^2 _t\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt+2 _0^T\!\! _R^Dv_t· _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t\, _t\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt+ _0^T\!\! _R^D \| _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t \|^2 _t\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt =∫0T|ρ˙t|W22dt−2∫0T∫ℝDδℱδρtdiv(ρtt)ddt+∫0T∫ℝD‖∇δℱδρt‖2ρtddt. = _0^T | ρ_t |^2_W_2dt-2 _0^T\!\! _R^D δ _tdiv_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x( _tv_t)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt+ _0^T\!\! _R^D \| _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x δ _t \|^2 _t\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt. By (24) and the chain rule (25), −∫0T∫ℝDδℱδρtdiv(ρtt)ddt=∫0T∫ℝDδℱδρt∂tρtddt=ℱ(ρT)−ℱ(ρ0),- _0^T\!\! _R^D δ _tdiv_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x( _tv_t)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt= _0^T\!\! _R^D δ _t\, _t _t\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt=F( _T)-F( _0), which establishes (26). ∎ Appendix B Divergences In this section we expand on the divergence options D used in the data-fitting term. Kullback–Leibler (KL) Likelihood maximization leads to KL(ρt,qt):=−∫ℝD(logρt())qt()d=KL(qt∥ρt)+∫ℝD(logqt())qt()d,D_KL( _t,q_t):=- _R^D\!( _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x))q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x=KL(q_t\| _t)+ _R^D\!( q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x))q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x, (27) which can be approximated as KL(ρt,qt)≈−∑i=1Ntlogρt(t,i),wheret,ii=1Nt∼i.i.d.qt.D_KL( _t,q_t)≈- _i=1^N_t _t(y_t,i), where \y_t,i\_i=1^N_t i.i.d. q_t. (28) The divergence KL(ρt,qt)D_KL( _t,q_t) can be used whenever ρt _t can be evaluated pointwise. Score matching (Hyvärinen and Dayan, 2005). The score matching cost ∫ℝD‖∇logρt()−∇logqt()‖2qt()d, _R^D\|∇ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)-∇ q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\|^2q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x, (29) whose minimization over ρ is equivalent to minimizing over ρ, SM(ρt,qt):=∫ℝD[‖∇logρt()‖2+2Δlogρt()]qt()d,D_SM( _t,q_t):= _R^D\! [ \|∇ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2+2 _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) ]q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x, (30) leads to the divergence SM(ρt,qt)≈∑i=1Nt[‖∇logρt(t,i)‖2+2Δlogρt(t,i)],wheret,ii=1Nt∼i.i.d.qt.D_SM( _t,q_t)≈ _i=1^N_t[ \|∇ _t(y_t,i) \|^2+2 _t(y_t,i)], where \y_t,i\_i=1^N_t i.i.d. q_t. (31) The divergence SM(ρt,qt)D_SM( _t,q_t) can be used whenever ∇logρt∇ _t is available, potentially with the extra cost computing ∇logρt∇ _t is only logρt _t is parametrized. Denoising score matching (Vincent, 2011). Let qtσ:=∫κσ(⋅|)qt()dq_t^σ:= _σ(·|z)q_t(z)dz with κσ(⋅|)=(,σ2ID) _σ(·|z)=N(z,σ^2I_D). Minimizing over ρt _t the score matching cost between ρt _t and qtσq_t^σ ∫ℝD‖∇logρt()−∇logqtσ()‖2qt()d, _R^D\|∇ _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)-∇ q_t^σ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\|^2q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x, (32) is equivalent to minimizing over ρt _t, DSM(ρt,qt):=∫ℝD∫ℝD∥∇logρt()−∇logκσ(|)∥2κσ(|)qt()d.D_DSM( _t,q_t):= _R^D\!\! _R^D\!\|∇ _t(z)- _z _σ(z| [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\|^2 _σ(z| [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)dzd [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x. (33) The denosing score matching divergence DSM(ρt,qt)D_DSM( _t,q_t) can be approximated by DSM(ρt,qt)≈∑i=1Nt∥∇logρt(t,i)−∇logκσ(t,i|t,i)∥2,t,i∼κσ(⋅|t,i),t,ii=1Nt∼i.i.dqt.D_DSM( _t,q_t)≈ _i=1^N_t\|∇ _t(z_t,i)- _z _σ(z_t,i|y_t,i)\|^2,\,\,z_t,i _σ(·|y_t,i),\,\y_t,i\_i=1^N_t i.i.d q_t. (34) In contrast to (30), denoising score matching DSMD_DSM avoids the computation of the Laplacian in SMD_SM at the cost of fitting the smoothed qtσq_t^σ rather than qtq_t. Appendix C Stitching: details In this section we derive the velocity expression (9) and explain the reasoning behind (13). We start with Equation (9). Claim 1. Let ρ=(ρt)t∈[0,T]ρ=( _t)_t∈[0,T] be of the form ρt()=∑k=1Nwkϕ(−t,k) _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)= _k=1^Nw_k\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k), where for each k=1,…,Nk=1,…,N, the trajectory [0,T]∋t↦t,k[0,T] t [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k is differentiable. Then, (ρ,)(ρ,v) satisfies the continuity equation (1) with t()=∑k=1Nwkϕ(−t,k)˙t,k∑l=1Nwlϕ(−t,l).v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)= _k=1^Nw_k\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k)\, [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k _l=1^Nw_l\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,l). (35) Proof. Since ∇t,kϕ(−t,k)=−∇ϕ(−t,k) _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,kφ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k)=- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xφ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k) we have ∂tρt() _t _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) =∑k=1Nwk∂tϕ(−t,k)=∑k=1Nwk∇t,kϕ(−t,k)⋅˙t,k = _k=1^Nw_k\, _tφ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k)= _k=1^Nw_k\, _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,kφ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k)· [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k =−∇(∑k=1Nwkϕ(−t,k)˙t,k)=−∇(ρt()∑k=1Nwkϕ(−t,k)˙t,k∑l=1Nwlϕ(−t,l)). =- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ( _k=1^Nw_k\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k) [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k )=- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ( _t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) _k=1^Nw_k\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k) [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k _l=1^Nw_l\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,l) ). ∎ Next we explain the reasoning behind (13). Claim 2. Let ρt=∑k=1Nwkδt,k _t= _k=1^Nw_k _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k with differentiable trajectories t↦t,kt [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k, and let v be the unique minimal velocity satisfying (1) in the distributional sense. Then ˙t,k=t(t,k) [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k=v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k) for each k. Proof. For any test function η:[0,T]×ℝD→ℝη:[0,T]×R^D , weak satisfaction of the continuity equation gives ∫s+h∫[∂tη+∇η⋅t]dρtdt=∫η(s+h,⋅)dρs+h−∫η(s,⋅)dρs. _s^s+h\!\! [ _tη+∇η·v_t]d _tdt= η(s+h,·)d _s+h- η(s,·)d _s. Substituting ρt=∑k=1Nwkδt,k _t= _k=1^Nw_k _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k, and choosing η localized around a single t,k [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k, yields ˙t,k=t(t,k) [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k=v_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k). ∎ From the velocity residual to the boxed objective. Substituting the KDE parametrization (8) into the velocity residual (velocity residual) and approximating the spatial expectation at the centers (exact as ϕ→δφ→δ) gives ℛvel[ℱθ,ρθ,θ]≈∫0T∑k=1Nwkθ‖tθ(t,kθ)+∇δℱθδρtθ(t,kθ)‖2dt.R_vel[F^θ,ρ^θ,v^θ]\;≈\; _0^T _k=1^Nw^θ_k\, \|v_t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ)+ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ) \|^2\,dt. By Claim 2, tθ(t,kθ)=˙t,kθv_t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ)= [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ in this limit. Discretising time on 0=t0<⋯<tK−1=T0=t_0<·s<t_K-1=T with forward Euler ˙tj,kθ≈Δtj+1,kθ/Δtj+1 [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ≈ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ/ t_j+1 (where Δtj+1,kθ:=tj+1,kθ−tj,kθ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ:= [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ), replacing the time integral by the left-Riemann sum ∫0Tf(t)dt≈∑jΔtj+1f(tj) _0^Tf(t)\,dt≈ _j t_j+1\,f(t_j), and applying the algebraic identity h‖a/h+b‖2=(1/h)‖a+hb‖2h\,\|a/h+b\|^2=(1/h)\,\|a+h\,b\|^2 on each summand yields the boxed displacement form (13). Computation of the score term. The KDE (8) has analytic score ∇logρtθ()=∑k=1Nwkθ∇ϕ(−t,kθ)∑l=1Nwlθϕ(−t,lθ), _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\;=\; _k=1^Nw^θ_k\, _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xφ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ) _l=1^Nw^θ_l\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,l^θ), (36) which we need to evaluate at the centers t,jθ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,j^θ. The numerator self-term wjθ∇ϕ()=w^θ_j∇φ(0)=0 is harmless. The denominator self-term wjθϕ()w^θ_jφ(0) is not: it is the kernel’s peak value, and it stays large regardless of where the other particles sit. When the neighbours are far compared to the bandwidth, the numerator (a sum of tiny neighbour kernels and their gradients) is already small, and dividing by an over-strong self-normaliser collapses the score toward 0. The entropy contribution then drops out of the velocity residual and particles collapse onto minima of VθV^θ instead of spreading. We drop particle j from both numerator and denominator when evaluating the score at t,jθ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,j^θ: ∇logρtθ(t,jθ)≈∑k≠jwkθ∇ϕ(t,jθ−t,kθ)∑l≠jwlθϕ(t,jθ−t,lθ). _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,j^θ)\;≈\; _k≠ jw^θ_k\, _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xφ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,j^θ- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ) _l≠ jw^θ_l\,φ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,j^θ- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,l^θ). (37) The two small quantities now divide to give a finite vector pointing toward the nearest neighbours, restoring the diffusion pressure of the entropy term (Wand and Jones, 1994). Centers vs. KDE-Monte-Carlo quadrature. The boxed objective evaluates the residual integrand at the particle centers t,kθ\ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ\. The implementation also supports a stochastic alternative: for each particle, draw a perturbation εt,k∼ϕ(,h2I) _t,k φ(0,h^2I) from the KDE kernel and evaluate the residual at t,k=t,kθ+εt,ky_t,k= [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t,k^θ+ _t,k, with the velocity at t,ky_t,k given by the Nadaraya–Watson estimator (35). The two estimators are combined convexly through a parameter α∈[0,1]α∈[0,1]: ℛ^velα:=(1−α)ℛ^velcenters+αℛ^velKDE-MC. R_vel^α\;:=\;(1-α)\, R_vel^centers\;+\;α\, R_vel^KDE -MC. (38) At α=0α=0 the residual is the deterministic centers approximation: exact in the small-bandwidth limit, but blind to off-trajectory information. At α=1α=1 it is an unbiased Monte-Carlo estimator under the KDE measure, at the cost of MC variance and a Nadaraya–Watson regression bias of order O(h2)O(h^2). Intermediate α trades the two biases against each other. We use α=0α=0 by default in all experiments. We found the α=0.5α=0.5 to give slightly better results in the low-data regime wavy-valley illustration. Time-discretisation schemes. The forward-Euler discretisation ˙tj,kθ≈Δtj+1,kθ/Δtj+1 [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ≈ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ/ t_j+1 used above is one of several supported choices. The implementation also exposes: • Backward (implicit) Euler: Evaluates ∇δℱθδρtθ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ at the right endpoint tj+1,kθ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ rather than tj,kθ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ. • Midpoint: Evaluates both the displacement and ∇δℱθδρtθ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ at tj+1/2t_j+1/2, ˙tj+1/2,kθ≈Δtj+1,kθ/Δtj+1 [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1/2,k^θ≈ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j+1,k^θ/ t_j+1. O(h2)O(h^2) accurate per step. • Trapezoidal: Averages the forward and backward Euler residuals, evaluating ∇δℱθδρtθ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ at both endpoints. O(h2)O(h^2) accurate; this is the default in our experiments and the natural choice on non-uniform grids. On uniform grids the schemes differ only in higher-order error terms; the trade-off is between accuracy (midpoint, trapezoidal) and per-step cost (forward Euler avoids re-evaluating ∇δℱθδρtθ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ at the right endpoint). Appendix D Extended Related Work In this section we outline the details for the related work to our approach, including JKO-based methods in Section D.1, methods based on residual losses in Section D.2, and action matching in Section D.3. Further methods are discussed in Section D.4– Section D.6. D.1 JKO-based methods The dominant algorithmic framework to tackle our main goal is the Jordan–Kinderlehrer–Otto (JKO) scheme (Jordan et al., 1998). In this scheme we fix a discretization step h>0h>0 and sequentially solve ρt+h:=argminν∈2(ℝD)[ℱ(ν)+12hW22(ρt,ν)]. _t+h\;:=\; *arg\,min_ν _2(R^D) [F(ν)+ 12hW_2^2( _t,ν) ]. (JKO) As h→0h→ 0, this iteration scheme converges to the Wasserstein gradient flow of ℱF (Santambrogio, 2015, Ch. 11). There are currently two strategies for the purpose of learning the potential ℱF from the data qtt∈obs\q_t\_t _obs using the JKO scheme, which we outline next. Forward simulation. The JKOnet algorithm (Bunne et al., 2022) iterates (JKO), starting from ρ0=q0 _0=q_0, by solving a bi-level optimization. JKONet uses Brenier’s theorem to characterize the solution to (JKO) as an optimization problem over a convex function ψ. In particular, given a parametrized functional ℱξF^ξ, the inner optimization problem minν∈2(ℝD)[ℱξ(ν)+12hW22(ρt,ν)] _ν _2(R^D) [F^ξ(ν)+ 12hW_2^2( _t,ν) ] is rewritten as, minν∈2(ℝD)[ℱξ(ν)+12hW22(ρtξ,ν)]=min∇ψθ:ψθ:ℝD→ℝ convx[ℱξ(∇ψ♯θρtξ)+12h∫ℝD‖−∇ψθ()‖2ρtξ()d], split& _ν _2(R^D) [F^ξ(ν)+ 12hW_2^2( _t^ξ,ν) ]\\ &= _∇ψ^θ:ψ^θ:R^D convx [F^ξ(∇ _ ^θ _t^ξ)+ 12h _R^D\| [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x-∇ψ^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\|^2 _t^ξ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ], split (39) where ψθ\ψ^θ\_θ is taken to be a class of convex neural networks (ICNN); see Mokrov et al. (2021) for using ICNNs for the proximal step. Then, given a minimizer θ set ρt+hξ:=(∇ψθ)♯ρtξ _t+h^ξ:=(∇ψ^θ)_ _t^ξ, and then minimize over ξ, W22(ρt+hξ,qt+h)W_2^2( _t+h^ξ,q_t+h), (in fact a Sinkhorn divergence). The iJKOnet (Persiianov et al., 2026) algorithm avoids the usage of ICNN by using a min-max formulation. First-order linearization. The JKOnet⋆ (Terpin et al., 2024) algorithms use the optimality condition of (JKO), ′−h=−∇δℱδρt+h(′)∀(,′)∈suppγt, [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x - [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xh\;=\;-∇ δ _t+h( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ) ∀( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x, [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ) _t, (JKO-FO) where γt _t is the optimal transport plan between consecutive marginals in qtt∈obs\q_t\_t _obs. In this style of algorithms the optimal transport problem is first solved between any two consecutive marginals in qtt∈obs\q_t\_t _obs. Given these optimal transport plans γt\ _t\ the functional ℱF is parametrized by a neural network ℱθF^θ, and the optimization problems becomes minimizing over θ the residual of (JKO-FO). In both strategies the only learnable object is the functional ℱF. The curve ρ can be recovered only by iterating the scheme (JKO). Therefore, trajectory expressivity is coupled to ℱF’s, and evaluating ρt _t at unobserved times requires post-hoc simulation. Instead, our approach promotes the parametrization ρθρ^θ of the curve ρ to a first-class learnable object on equal footing with ℱθF^θ, coupled only through a residual loss. Both are optimized at the same level, and the curve’s capacity is determined by its own parametrization, not by what the JKO operator can produce. D.2 Residual Losses Concurrent work of Liu and Zhou (2026) considers the problem of solving a WGF given a known functional ℱF using a form of (velocity residual), resulting in GenWGP. In particular, considering a K-step discretization of the interval [0,T][0,T], GenWGP parameterizes the flow-map Φθ(t,) ^θ(t, [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) as a normalizing flow, i.e., Φθ(tk,⋅)=Ψk∘⋯Ψ1∘ρ0 ^θ(t_k,·)= _k ·s _1 _0, for invertible neural networks Ψk _k (k=1,…,Kk=1,…,K). This flow is applied to a population of particles, xk(i)x_k^(i); densities can then be empirically evaluated via the normalizing flow, which enables an approximate evaluation of velocity fields using finite difference methods. The resulting approximation of (velocity residual) is then optimized via gradient descent. The GenWGP approach is stated only for a known functional ℱF, but can be adapted to our setting where ℱF is unknown. The main drawback, compared to stitching, in this inverse problem setting is that GenWGP requires simulation; in particular, the normalizing flow comprises solving a neural ODE. In contrast, stitching does not require an ODE solve, as the particles θ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x^θ are directly parameterized as part of the resulting optimization problem. Notably, however, GenWGP does not require a further approximation of the density ρt _t, e.g., in the form of a KDE. D.3 Action Matching Neklyudov et al. (2023) introduced action matching (AM) for learning the velocity field of a population dynamic, and Neklyudov et al. (2024) extends the framework to a broader class of Wasserstein Lagrangian flows beyond pure gradient flows. The main idea of Neklyudov et al. (2023) is to fit a velocity field to data without access to ground-truth velocities. Adapted to our setting, the goal is to learn ℱθF^θ such that −∇δℱθδρtθ- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ matches the (unknown) true velocity t⋆v_t of the data. The naive objective ∫0T∫ℝD‖∇δℱθδρtθ()+t⋆()‖2qt()ddt _0^T\!\! _R^D \| _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)+v_t ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt (40) cannot be optimized as written due to the unknown velocity. By proceeding as in Neklyudov et al. (2023, Theorem 2.2), using integration-by-parts, the θ-dependent part of the objective is given by ℒAM(θ) _ AM(θ) =∫ℝDδℱθδρ0θ()q0()d−∫ℝDδℱθδρTθ()qT()d = _R^D ^θδ _0^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)q_0( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- _R^D ^θδ _T^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)q_T( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x (41) +∫0T∫ℝD[12‖∇δℱθδρtθ()‖2+∂tδℱθδρtθ()]qt()ddt, + _0^T\!\! _R^D\! [ 12 \| _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2+ _t ^θδ _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) ]q_t( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt, (42) which can be evaluated using only samples from qtq_t. In our taxonomy, AM is an example of the velocity residual evaluated at data samples, with the integration-by-parts trick eliminating the need for t⋆v_t . Unlike the divergence-based methods, AM fuses residual and data fitting into a single objective. The price is that AM requires data at the temporal endpoints (for the boundary terms in (41)) and does not directly constrain ρtθ _t^θ at intermediate times—it learns ℱθF^θ such that its gradient flow has the right velocity, leaving the trajectory ρθρ^θ to be recovered post-hoc. D.4 Consistency by construction. An alternative to the residual-loss approach we take is to bake the continuity equation into the parametrization itself, so that any candidate (ρθ,θ)(ρ^θ,v^θ) satisfies it by construction. Neural Conservation Laws Richter-Powell et al. (2022) parametrize divergence-free vector fields via differential forms, and Hua et al. (2025) managed simplify this construction. Their setting is closely related to ours: like stitching, it parametrizes the trajectory and the dynamics jointly, but it enforces the PDE via the architecture rather than via a residual loss, which trades off architectural flexibility against exact constraint satisfaction. D.5 Neural SDEs, Schrödinger bridges, and flow matching Neural SDEs (Li et al., 2020, Kidger et al., 2021a, b) are a popular way to learn dynamical systems: they parameterize the drift and diffusion terms of an SDE, trained by optimization of a variational objective. Contrary to our setting, however, neural SDEs are typically formulated and trained with respect to trajectory/pathwise data. Applying them to population data is non-trivial, accomplished via a Sinkhorn divergence in Koshizuka and Sato (2023). Schrödinger bridges (Léonard, 2014, De Bortoli et al., 2021, Shi et al., 2023, Koshizuka and Sato, 2023, Chen et al., 2023), diffusion models (Song et al., 2021) and flow matching (Lipman et al., 2023, Tong et al., 2024, Albergo and Vanden-Eijnden, 2023, Wang et al., 2025) learn velocity fields between distributions without imposing a gradient-flow structure. They are thus solving a strictly weaker problem: they recover dynamics, but not an underlying energy. Similarly to flow matching, normalizing flows parametrize a velocity network tθv_t^θ and let ρtθ _t^θ be the continuous normalizing flow (Chen et al., 2018, Papamakarios et al., 2021) obtained by transporting ρ0 _0 along tθv_t^θ. This provides exact normalized density evaluation, and permits learning the model parameters by maximizing the likelihood (i.e., minimizing the KL divergence) of the data at observed times under the model ρtθ _t^θ. D.6 Other methods We also mention several additional related methods from the literature which do not fall into the above categories. We leave detailed comparison of our method to those in a future work. Guan et al. (2024, 2026) study identifiability of the drift and diffusion of an SDE from temporal marginals, complementary to our setting where the drift is constrained to be the gradient of a learned functional. Carrillo et al. (2025) considers recovery of interaction kernel from gridded density data using regularized basis pursuit. In Wei and Lu (2026), the authors derive a self-test loss based on the weak form of the stochastic evolution equation for the empirical measure. Appendix E Wavy valley: details This appendix collects all hyperparameters and protocol details for the wavy-valley experiment summarised in Section 5.1 and Figure 1. Potential and SDE. The wavy valley is the 2D landscape V(x1,x2)=K(x2−sin(πx1/2))2−τx1,K=0.6,τ=0.3,V(x_1,x_2)\;=\;K (x_2- (π x_1/2) )^2\;-\;τ\,x_1, K=0.6,\ τ=0.3, whose minimum-energy curve is the sinusoid x2=sin(πx1/2)x_2= (π x_1/2). The first term confines particles to the valley floor; the second tilts the floor downward in +x1+x_1. We simulate the SDE dx=−∇V(x)dt+2βdWdx=-∇ V(x)\,dt+ 2β\,dW with β=0.00625β=0.00625 from a tight Gaussian initial distribution centred at (−3,1)(-3,1) with std 0.10.1, using Euler–Maruyama with at least 1010 substeps per integration interval. Snapshot protocol. We observe the population at the irregular times t∈0,5,10,20,30t∈\0,5,10,20,30\. Each snapshot draws N=20N=20 particles from the same coupled simulation pool (matched particle identities across times), yielding T=5T=5 paired snapshots and 2020 training points per snapshot. The t=10→t=20t=10→ t=20 gap of width 1010 is intentional: a first-order JKO predictor cannot interpolate a curved trajectory across a step that long. Stitching configuration. The stitching model uses N=20N=20 particles with K=50K=50 trajectory nodes spaced linearly between t0=0t_0=0 and tT=30t_T=30. The KDE bandwidth is initialised at 0.50.5 and trained; the entropy coefficient is initialised at 0.050.05 and trained (matching lightspeed’s setting); particle mixture weights are trainable softmax variables. The velocity residual uses an α=0.5α=0.5 convex combination of centers quadrature and Nadaraya–Watson Monte Carlo (Appendix C), with one MC sample per particle per snapshot. The VθV^θ network is a (64,64)(64,64) MLP. Training: Adam at lr=5⋅10−3lr=5· 10^-3 for 10K10K iterations. Lightspeed configuration. We use the published JKOnetV⋆ _V recipe (Terpin et al., 2024) with one modification: the entropy coefficient is trainable rather than frozen at zero, matching stitching for fairness. Other hyperparameters: (64,64)(64,64) hidden, lr=10−3lr=10^-3, 10K10K iterations, OT-Hungarian coupling between consecutive snapshot pairs. Metrics. We report (i) per-snapshot W2W_2 between trained-model samples and held-out particles; (i) coefficient of determination R2(V)R^2(V) between learned and ground-truth potentials, evaluated on the data support (the values in Figure 1); (i) pattern R2R^2 between gradient fields, the scale-invariant variant. Appendix F Synthetic potential recovery: details Metric definitions. We use the three metrics of Persiianov et al. (2026, Eqs. 17–20), all lower-is-better. Let ρtk+1θρ^θ_t_k+1 denote the model’s predicted marginal at observation time tk+1t_k+1 and qtkq_t_k the observed marginal at tkt_k. • One-step EMD: the Earth Mover’s (1-Wasserstein) distance between predicted and observed marginal, averaged over consecutive transitions, EMD:=1T−1∑k=0T−2W1(ρtk+1θ,qtk+1).EMD\;:=\; 1T-1 _k=0^T-2W_1 (ρ^θ_t_k+1,\,q_t_k+1 ). • L2L^2-UVP (gradient-level error): the residual squared L2(qtk)L^2(q_t_k) error in the learned gradient ∇Vθ∇ V^θ as a fraction of the variance of the true gradient, L2-UVP:=∼qtk‖∇Vθ()−∇V()‖2Var∼qtk[∇V()]× 100%.L^2-UVP\;:=\; E_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x q_t_k \|∇ V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)-∇ V( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2Var_ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x q_t_k [∇ V( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) ]\;×\;100\%. L2L^2-UVP measures recovery of the gradient field independently of additive constants in V, which is the right quantity for a Wasserstein gradient flow (the dynamics is invariant to such constants). • BdW22Bd^2_W_2-UVP (Bures–Wasserstein UVP): the squared Bures–Wasserstein distance between Gaussian approximations of the predicted and observed marginals, normalized by the trace of the observed covariance. With μθ,Σθμ^θ, ^θ the mean and covariance of ρtk+1θρ^θ_t_k+1 and μ,Σμ, those of qtk+1q_t_k+1, BdW22-UVP:=‖μθ−μ‖2+tr(Σθ+Σ−2(Σ1/2ΣθΣ1/2)1/2)tr(Σ)× 100%.Bd^2_W_2-UVP\;:=\; \|μ^θ-μ\|^2\;+\;tr\! ( ^θ+ -2( ^1/2 ^θ\, ^1/2)^1/2 )tr( )\;×\;100\%. This captures how well the predicted first two moments match the observed ones. Pattern and raw R2R^2 for V. The galleries (Figure 2, Figure 7, Figure 8) report two coefficients of determination between the learned and ground-truth potentials, both evaluated on a uniform grid covering the data support. Rraw2R^2_raw is the standard coefficient of determination, Rraw2(Vθ,V):= 1−∑i(Vθ(i)−V(i))2∑i(V(i)−V¯)2,R^2_raw(V^θ,V)\;:=\;1- _i (V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_i)-V( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_i) )^2 _i (V( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_i)- V )^2, where V¯ V is the mean of V over the grid. Rraw2R^2_raw is sensitive to additive and multiplicative shifts of VθV^θ. Rpattern2R^2_pattern replaces the residual-sum-of-squares ratio with the squared Pearson correlation, Rpattern2(Vθ,V):=corr(Vθ,V)2,R^2_pattern(V^θ,V)\;:=\;corr (V^θ,V )^2, again evaluated on the same grid; this is scale- and shift-invariant. Because Wasserstein gradient flow is invariant to additive constants in V and identifiable only up to a scale by β (Persiianov et al., 2026, App. A), Rpattern2R^2_pattern is the natural similarity score for the recovered potential, while Rraw2R^2_raw additionally penalises any residual scale or offset. Stitching configuration. • Trajectory: N=1,000N=1,000 particles per run, length K=50K=50, identity-coupled at initialization. • Functional: ℱθ[ρtθ]=cVρtθ[Vθ()]F^θ[ _t^θ]=c_V\,E_ _t^θ[V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)] with VθV^θ a (64,64)(64,64) MLP and cV>0c_V>0 a softplus-parametrised scalar (entropy and interaction terms disabled — the Terpin et al. (2024) benchmark is deterministic gradient flow of a single potential). • Density model: per-dimension Gaussian KDE bandwidth set by Silverman’s rule, frozen uniform mixture weights wk=1/Nw_k=1/N. • Loss: velocity residual evaluated at particle centers (α=0α=0, no KDE–MC perturbation), trapezoidal time scheme, plus a KDE-KL data divergence at the five observed snapshots. • Optimizer: full-batch Adam at learning rate 5⋅10−35· 10^-3 with cosine decay, 2,0002,000 steps. • Configuration is identical across all 3030 runs (1515 potentials × \paired, unpaired\); each run takes a few minutes on a single CPU. Baseline configuration. JKOnetV⋆ _V and iJKOnetV are run directly from the upstream codebases of Terpin et al. (2024) and Persiianov et al. (2026) respectively, both at default hyperparameters: a (64,64)(64,64) MLP VθV^θ matching stitching, no entropy or interaction terms, single seed, on the same train/test splits as stitching. JKOnetV⋆ _V uses Hungarian-OT couplings between consecutive snapshots; iJKOnetV uses its inverse-JKO solver with K=5K=5. We train JKOnetV⋆ _V for 100100 epochs and iJKOnetV for 2,0002,000 epochs (the upstream defaults). Metrics are pulled from each method’s own evaluation pipeline (parsed from upstream stdout for iJKOnetV, saved-params ++ a forward SDE rollout in our pipeline for JKOnetV⋆ _V); iJKOnetV’s CLI does not expose L2L^2-UVP, so those cells are marked ‘–’ in Table 1. Because of compute, the baseline runs are restricted to the 66 potentials Persiianov et al. (2026, Tab. 3) flag as most paired→ sensitive; stitching’s appendix numbers below cover all 1515. Per-potential numbers. Table 4 reports stitching’s three metrics on each of the 1515 landscapes of Terpin et al. (2024), in both regimes. The main-text Table 1 compares stitching against the JKO baselines on the 66 sensitive potentials. The consistency claim of the main text is visible row-by-row in Table 4: paired and unpaired columns differ by less than ∼ 2× on every metric and every potential except the four most delicately structured (sphere, bohachevsky, rotational, and to a lesser extent relu), where stitching’s L2L^2-UVP and BdW22Bd^2_W_2-UVP increase but never collapse. Table 4: Stitching’s per-potential metrics on the 1515 two-dimensional landscapes of Terpin et al. (2024), in the paired (correlated trajectories, original protocol) and unpaired (independent snapshots, Persiianov et al. (2026) protocol) regimes. Same configuration across all 3030 runs (1,0001,000 particles, 2,0002,000 steps, single seed). Metrics follow Persiianov et al. (2026): one-step EMD, L2L^2-UVP (gradient-level error), and BdW22Bd^2_W_2-UVP (Bures–Wasserstein UVP); L2L^2-UVP and BdW22Bd^2_W_2-UVP in percent. Definitions in the metric paragraph above. Lower is better. paired (correlated) unpaired (independent) # potential EMD ↓ L2L^2-UVP ↓ BdW22Bd^2_W_2-UVP ↓ EMD ↓ L2L^2-UVP ↓ BdW22Bd^2_W_2-UVP ↓ 1 flowers 0.30 0.00 0.18 0.31 0.01 0.24 2 styblinski_tang 0.26 0.02 0.15 0.29 0.04 0.29 3 holder_table 0.29 0.06 0.18 0.30 0.06 0.24 4 zigzag_ridge 0.29 0.03 0.17 0.30 0.04 0.29 5 oakley_ohagan 0.24 0.01 0.20 0.25 0.02 0.28 6 watershed 0.30 0.00 0.19 0.30 0.01 0.24 7 ishigami 0.30 0.01 0.17 0.30 0.01 0.25 8 friedman 0.29 0.06 0.20 0.31 0.06 0.27 9 sphere 0.62 0.86 1.45 0.64 0.91 1.24 10 bohachevsky 0.39 7.39 13.99 0.38 7.46 12.93 11 wavy_plateau 0.27 0.02 0.51 0.29 0.04 0.54 12 double_exp 0.34 0.02 0.22 0.35 0.06 0.29 13 relu 0.35 0.04 0.18 0.38 0.10 0.33 14 rotational 0.41 0.82 1.60 0.45 0.82 1.80 15 flat 0.30 0.00 0.19 0.31 0.01 0.24 Figure 6: Paired vs. unpaired snapshot evolution on the synthetic potentials following Persiianov et al. (2026). Top: In ‘paired’ setting we observe same particles driven by a potential gradient over time. Bottom: In ‘unpaired‘ trajectory structure is lost, as if the observations remove the corresponding particles from the system. Both cases have the same number of observations. Figure 7: Stitching recovers the level-set geometry of V on 1111 of 1414 informative landscapes. Synthetic potential recovery, paired regime: true V (left) vs learned VθV^θ (right) for each potential, mean-centered with shared per-potential color scale; titles report scale-invariant Rpattern2R^2_pattern and on-support Rraw2R^2_raw. Figure 8: Stitching remains robust to independently-sampled snapshots; detectable degradation only on the four most delicately structured potentials. Synthetic potential recovery, unpaired regime; same conventions as Figure 7. Appendix G Single-cell trajectory inference: details Dataset and preprocessing. The embryoid body (EB) dataset of Moon et al. (2019) comprises ∼17,000 17,000 cells observed across five 3-day windows of human embryonic stem cell differentiation. We use the PCA-reduced version distributed by Tong et al. (2020): each cell is represented by its first 55 principal components, and time labels are scaled to [0,1][0,1]. We use a 70/30 particle-level train/test split, deterministic per seed. Stitching architecture. • Static Vθ()V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x): MLP, input dim 55, hidden (64,64)(64,64), SiLU activations with residual skips, scalar output. • Time-varying Vθ(,t)V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x,t): same MLP but input dim 66, with t concatenated to [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x. Matches the design of iJKOnett,V in Persiianov et al. (2026, Sec. 5) and JKOnett,V∗^*_t,V in Terpin et al. (2024). • Trajectory: K=50K=50 learnable snapshots between tmint_ and tmaxt_ , N=100N=100 particles per snapshot, OT-coupled at initialization (Hungarian assignment + linear interpolation between consecutive observed marginals). • Density model: per-dimension Gaussian KDE bandwidth (learnable, softplus-parametrized) and learnable mixture weights (softmax-normalized). • Functional: ℱθ[ρtθ]=cVρtθ[Vθ()]+cHρtθ[logρtθ]F^θ[ _t^θ]=c_V\,E_ _t^θ[V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)]+c_H\,E_ _t^θ[ _t^θ] with cV,cH>0c_V,c_H>0 softplus-parametrized scalars. Training. • Loss: ℒ=wklKL(ρtobsθ,qtobs)+wvelℛvel[ℱθ,(θ,ρθ)]L=w_kl\,D_KL(ρ^θ_t_obs,q_t_obs)+w_vel\,R_vel[F^θ,(v^θ,ρ^θ)] with the trapezoidal-scheme velocity residual (Equation 13), α=0α=0 (no KDE-MC perturbation), kinetic factor 11. • Optimizer: Adam, batch size 256256, learning rate 5⋅10−35· 10^-3 with cosine decay to 5%5\% of the initial value. • Epochs: 10,00010,000. • Seeds: 0,…,40,…,4 for the leave-two-out comparison; a single seed for the full-data benchmark and Figure 3. • Wall time: ∼2 2 minutes per run on a CPU laptop (M-series). Evaluation protocols. • Full data (Table 2): train on all five marginals; report W1W_1 between the stitching KDE marginal at each observed t and the held-out test cells at that t. Forward-rollout from the learned trajectory; no JKO chain. • Leave-two-out (Table 5): restrict training to t∈0,2,4t∈\0,2,4\; evaluate at held-out t∈1,3t∈\1,3\ via KDE marginals of the learned particle cloud at the target time. No retraining or architectural change; the model is continuous-time by construction. Metric W2W_2 to match Persiianov et al. (2026, Table 1). • All distances computed with the POT library: EMD via emd2(Euclidean) and W2W_2 via emd2(sqEuclidean) emd2(sqEuclidean). Leave-two-out results. Table 5 compares stitching against the full set of baselines reported in Persiianov et al. (2026, Table 1), all using the same W2W_2 metric. Stitching with a static V already outperforms every published baseline (mean W2W_2 0.920.92); the time-varying Vθ(,t)V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x,t) variant further reduces the mean to 0.880.88, with the best published competitor (iJKOnett,V) at 0.920.92. Table 5: Leave-two-out temporal interpolation on the EB single-cell dataset (5 D), W2W_2 (↓ ). Models trained on t∈0,2,4t∈\0,2,4\; evaluated at held-out t=1t=1 and t=3t=3. Baseline results from Persiianov et al. (2026, Table 1); our results averaged over 5 seeds. Method t=1t=1 t=3t=3 Mean Citation TrajectoryNet 2.032.03±0.04\,± 0.04 1.931.93±0.08\,± 0.08 1.981.98 Tong et al. (2020) Vanilla-SB 1.491.49±0.06\,± 0.06 1.551.55±0.03\,± 0.03 1.521.52 Vargas et al. (2021) MMSB 1.271.27±0.03\,± 0.03 1.571.57±0.05\,± 0.05 1.421.42 Shen et al. (2025) DMSB 1.131.13±0.08\,± 0.08 1.451.45±0.16\,± 0.16 1.291.29 Chen et al. (2023) Static potential JKOnetV∗^*_V 1.151.15±0.03\,± 0.03 2.532.53±0.01\,± 0.01 1.841.84 Terpin et al. (2024) iJKOnetV 1.081.08±0.01\,± 0.01 1.151.15±0.00\,± 0.00 1.121.12 Persiianov et al. (2026) Stitching, Vθ()V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) 0.850.85±0.06\,± 0.06 0.990.99±0.03\,± 0.03 0.920.92±0.03\,± 0.03 This paper Time-varying potential JKOnett,V∗^*_t,V 4.414.41±1.50\,± 1.50 2.772.77±0.20\,± 0.20 3.593.59 Terpin et al. (2024) iJKOnett,V 0.980.98±0.04\,± 0.04 0.850.85±0.02\,± 0.02 0.920.92 Persiianov et al. (2026) Stitching, Vθ(,t)V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x,t) 0.840.84±0.05\,± 0.05 0.920.92±0.03\,± 0.03 0.880.88±0.04\,± 0.04 This paper Appendix H Interaction dynamics: details This appendix collects the full setup, model, training, and metric definitions deferred from Section 5.4. Dataset. The interaction dynamics data is a single-trajectory simulation of a 2D SDE (15) dXti=−∇V(Xti)dt−1N∑j≠i∇xW(Xti−Xtj)dt+σdBti,,i=1,…,N,dX_t^i=- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7xV(X_t^i)\,dt- 1N\! _j≠ i _xW(X_t^i-X_t^j)\,dt+σ\,dB_t^i, ,i=1,…,N, with V(x)=α(‖x‖2−β)2V(x)=α(\|x\|^2-β)^2 (α=0.1,β=4α=0.1,β=4, minimum on the ring ‖x‖=2\|x\|=2), W(r)=ηe−r2W(r)=η\,e^-r^2 (η=−2.2η=-2.2), σ=0.045σ=0.045, N=310N=310 particles, integrated with Euler–Maruyama at dt=0.05dt=0.05 over [0,200][0,200] (T=201T=201 snapshots) from a Gaussian initial condition. We use a 50/5050/50 particle-level train/test split applied uniformly across all T snapshots. Energy parametrization. We learn ℱθ[ρtθ]=cVρtθ[Vθ()]+cWρtθ⊗ρtθ[Wθ(‖−′‖)]+cHρtθ[logρtθ]F^θ[ _t^θ]=c_V\,E_ _t^θ[V^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)]+c_W\,E_ _t^θ _t^θ[W^θ(\| [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x \|)]+c_H\,E_ _t^θ[ _t^θ] with Vθ:ℝ2→ℝV^θ:R^2→R and the radial kernel Wθ:ℝ≥0→ℝW^θ:R_≥ 0→R, both parametrized as two-layer (64,64)(64,64) MLPs (zero-init last layer), and softplus-parametrized positive coefficients cV,cW,cH>0c_V,c_W,c_H>0. The radial parametrization Wθ(‖−′‖)W^θ(\| [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x \|) enforces translational and rotational symmetry of the kernel without restricting its functional form. Stitching model. The trajectory ρθρ^θ is represented by 100100 learnable particles at K=50K=50 snapshots between t0=0t_0=0 and t200t_200, OT-coupled to the data at initialization, with a Gaussian KDE marginal. The KDE bandwidth is initialized by Silverman’s rule and learned. Mixture weights are held uniform (wk=1/Nw_k=1/N). Training. Full-batch Adam at lr=5⋅10−3lr=5·10^-3 for 1,0001,000 epochs, optimizing the centers-quadrature velocity residual (velocity residual) (trapezoidal scheme, α=0α=0 in the convex combination) plus a KDE KL data-fit term, identical to the synthetic recovery setup of Section 5.2. CPU wall time on a 2023 MacBook is ≈4≈4 minutes per run. Baseline. We compare against JKOnet⋆ (Terpin et al., 2024) on the same train/test split. Pattern R2R^2 for V and W. WGF inference from snapshots is identifiable only up to a joint (V,W,σ)(V,W,σ)-scale ambiguity (Persiianov et al., 2026, App. A): only ∇V/σ∇ V/σ and ∇W/σ∇ W/σ are determined by the data, not the absolute scale of V,WV,W. We therefore report the scale- and shift-invariant pattern R2R^2, equal to the squared Pearson correlation between learned and true fields on the relevant support: a uniform grid covering the data support for V, and the band r∈[r5%,r95%]r∈[r_5\%,r_95\%] of observed pair distances at t=100t=100 for W. This is the same metric used in Table 4. Appendix I Beyond gradient flows: non-conservative interactions The residual framework of Section 3 only requires a target velocity field; it does not require that target to be the Wasserstein gradient of any functional. This is relevant in applications, as many interacting particle systems of practical interest have non-conservative pairwise interactions, including chiral active matter (Liebchen and Levis, 2022) and the broader class of non-reciprocal collective systems (Fruchart et al., 2021), none of which arise as Wasserstein gradient flows. Concretely, for any model curve ρθρ^θ with Eulerian velocity θv^θ and any model velocity field θu^θ, ℛvel[θ,(ρθ,θ)]:=∫0T∫ℝD‖tθ()−tθ()‖2ρtθ()ddtR_ vel[u^θ,(ρ^θ,v^θ)]\;:=\; _0^T\!\! _R^D \|v_t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)-u_t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) \|^2\, _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x\,dt (43) is a well-defined nonnegative loss whose minimum is the curve whose velocity matches θu^θ. Choosing tθ=−∇δℱθδρtθu_t^θ=- _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ recovers (velocity residual) and the WGF stitching of Section 4; any other choice of θu^θ instantiates stitching for a different class of dynamics. In particular, the KDE parametrization (8), the centers approximation (11), the time discretization leading to (13), and the data divergence (ρtθ,qt)D( _t^θ,q_t) are all independent of any gradient structure of θu^θ. Here we illustrate the extension on a non-conservative interacting particle system. We replace the symmetric scalar interaction W in (4) by a vector-valued pairwise kernel :ℝD→ℝDK:R^D ^D. The corresponding SDE reads dXti=−∇V(Xti)dt+1N∑j≠i(Xti−Xtj)dt+σdBti,dX_t^i\;=\;-∇ V(X_t^i)\,dt\;+\; 1N _j≠ iK(X_t^i-X_t^j)\,dt\;+\;σ\,dB_t^i, (44) in direct analogy with the conservative SDE of Appendix H. When =−∇WK=-∇ W for a scalar W, (44) is a Wasserstein gradient flow of (4); when K has nonzero curl, no scalar W satisfies =−∇WK=-∇ W, and (44) cannot be written as a WGF. A canonical non-conservative case in two dimensions is the chiral kernel (′−)=α(−∇W(′−))+ωRπ/2(−∇W(′−)),Rπ/2=(0−110),K( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x - [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\;=\;α\, (-∇ W( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x - [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) )\;+\;ω\,R_π/2 (-∇ W( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x - [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x) ), R_π/2\!=\! pmatrix0&-1\\ 1& -0 pmatrix, (45) where α,ω∈ℝα,ω are scalars and W:ℝ2→ℝW:R^2 is a radial scalar potential. The first summand is a standard conservative attraction/repulsion; the second is a 90∘ rotation of the same gradient, inducing a circulating component. Whenever ω≠0ω≠ 0 the kernel has nonzero curl. Parametrization. We parametrize the learned kernel as an unrestricted vector-valued MLP θ:ℝ2→ℝ2,K^θ:R^2 ^2, (46) that ingests the displacement =′−r= [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x - [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x directly and outputs a 2-vector, with no structural prior — in particular, no built-in decomposition into radial and chiral parts. We omit both the confinement VθV^θ and the entropy term in this experiment, so the target velocity field reduces to the convolution of θK^θ against the curve, tθ()=(θ∗ρtθ)(),(θ∗ρtθ)()=∫ℝDθ(′−)ρtθ(′)d′.u_t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\;=\;(K^θ _t^θ)( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x), (K^θ _t^θ)( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)= _R^DK^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x - [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x)\, _t^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x )\,d [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x . (47) The stitching loss (13) is structurally unchanged: substitute −tθ-u_t^θ for ∇δℱθδρtθ _ [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x ^θδ _t^θ and approximate θ∗ρtθK^θ _t^θ at the particle centers analogously to (14), (θ∗ρtθ)(tj,kθ)≈∑l=1Nwlθθ(tj,lθ−tj,kθ).(K^θ _t^θ)( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ)\;≈\; _l=1^Nw^θ_l\,K^θ( [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,l^θ- [rgb]0,0.3,0.7 [named]pgfstrokecolorrgb0,0.3,0.7x_t_j,k^θ). (48) Setup. We simulate (44) on [0,T][0,T] with T=800T=800, V≡0V≡ 0, σ=0σ=0, N=500N=500 particles, and the chiral kernel (45), where W is a generalised Morse potential (D’Orsogna et al., 2006) W(r)=Cre−r/ℓr−Cae−r/ℓa,(Cr,ℓr,Ca,ℓa)=(1.0, 0.5, 0.375, 1.5),W(r)\;=\;C_r\,e^-r/ _r-C_a\,e^-r/ _a, (C_r, _r,C_a, _a)=(1.0,\,0.5,\,0.375,\,1.5), with chirality ω=1.5ω=1.5 and radial scale α=0.2α=0.2. The initial condition is a mixture of two horizontal Gaussian blobs centred at (0,±2)(0,± 2) with covariance diag(1.52,0.22)diag(1.5^2,0.2^2); with ω≠0ω≠ 0 the two lumps orbit each other while the kernel’s radial component would otherwise relax them to the rotationally-symmetric Morse equilibrium. Particles are saved at Δt=0.5 t=0.5 via Euler–Maruyama with internal step 0.050.05, yielding 16011601 marginal snapshots. Stitching configuration. θK^θ is a (64,64)(64,64) MLP with SiLU activations and a 2-D output; the final-layer weights are rescaled by 10−210^-2 at initialization so that θ≈0K^θ\!≈\!0 at start while gradients still propagate through every layer. We use 200200 particles. The confinement and entropy terms are not part of the model. The learning setup otherwise follows Appendix H. Result. Figure 5 in the main text shows the stitching marginals tracking the rotation of the two-lump pattern across snapshots, with the learned curve ρθρ^θ following the data q at the displayed times.