Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme
Authors: Mikhail Persiianov
Jiawei Chen
Petr Mokrov
Alexander Tyurin
Evgeny Burnaev
Alexander Korotin
Venue: ICLR 2026 Poster (2026)
Keywords: population dynamics
JKO scheme
inverse problem
Primary area: generative models
TL;DR: We learn population dynamics by combining the JKO scheme with an inverse optimization framework.

Abstract: Learning population dynamics involves recovering the underlying process that governs particle evolution, given evolutionary snapshots of samples at discrete time points. Recent methods frame this as an energy minimization problem in probability space and leverage the celebrated JKO scheme for efficient time discretization. In this work, we introduce ``iJKOnet``, an approach that combines the JKO framework with inverse optimization techniques to learn population dynamics. Our method relies on a conventional *end-to-end* adversarial training procedure and does not require restrictive architectural choices, e.g., input-convex neural networks. We establish theoretical guarantees for our methodology and demonstrate improved performance over prior JKO-based methods.

Published as a conference paper at ICLR 2026
LEARNING OF POPULATION DYNAMICS: INVERSE OPTIMIZATION MEETS JKO SCHEME
Mikhail Persiianov
Jiawei Chen
Applied Al Institute, Moscow, Russia persiianov mi 16 gma com
Applied AI Institute, Moscow, Russia MIRAI,, Moscow, Russia
Petr Mokrov
Alexander Tyurin
Applied AI Institute, Moscow, Russia
AXXX, Russia
Applied AI Institute, Moscow, Russia
Evgeny Burnaev
Alexander Korotin
Applied Al Institute, Moscow, Russia AXXX, Russia
Applied AI Institute, Moscow, Russia
AXXX, Russia
iamal in gmai com
ABSTRACT
Learning population dynamics involves recovering the underlying process that governs particle evolution, given evolutionary snapshots of samples at discrete time points. Recent methods frame this probability space and leverage the celebrated JKO scheme for efficient time dis- cretization. In this work, we introduce iJKOnet, an approach that combines the JKO framework with inverse optimization techniques to learn population dynam- ics. Our method relies on conventional end-to-end adversarial training procedure and does not require restrictive architectural choices, e.g., input-convex neural net- works. We establish thcoretical guarantees for our methodology and demonstrate improved performance over prior JKO-based methods. Source code:
energy minimization problem
INTRODUCTION
Modeling population dynamics is a funda- mental challenge in many scientific domains, including biology (Schiebinger et al.] 20199]
peedisted evolation
tne evolution
lnverse JKO loes
Moon et al. 2019). ecology (Ayala et (SigriSt 2015 Verma et al. [024]] Price al.
(1)
1973).
meteorology (Fisher et al. [2009] [Sigrist
P0
Wa space
2025 transportation flows in urban networks (MedinaSSalaaoo et al., 2022), and epidemiol- ogy (Wang al. [222 Kosma et al. 2023), among others. The task to infer the under- lying stochastic dynamics of system typi- cally modeled by stochastic differential equa- tions (SDEs) from observed marginal distri- butions at discrete time points. While this prob-
Jol Sr dG, (p
(eve) curves
P216
Figure 1: iJKOnet working scheme: our method minimizes the gap between the optimal values of parametric (0) JKO functional and suboptimal values obtained at ground truth population measures.
lem has been studied extensively in settings where individual trajectories available (Krishnan et al. 2017 et al. 2020 Brogat-Motte et al.) 2024), such data are often unavailable in practice. In many real-world scenarios, we only observe population-leve data, where is infeasible continu- ously track individual entities. Instead, we are forced to rely on temporally separated and mutually independent snapshots of the population.
Moscoo Independent Research Institute of Artificial Intelligence

---

Published as a conference paper at ICLR 2026
In single-cell genomics (Macosko et al. 2013). for example, measuring the state of individual cells typically involves destructive sampling a process where cells are destroyed during measurement, preventing any further observation of their future behavior. As a result, the data consist only of isolated profiles of cellular populations taken at discrete time points. Reconstructing the continuous developmental trajectories of cells from such fragmented data poses a challenge. A similar situation arises in financial markets (Gontis et al. 2010 Yang et al., (023), where analysts often have access only to marginal distributions of asset prices specific times, making necessary to infer the un- derlying dynamics that govern these distributions. In crowd dynamics (Maury et al., 2010; 2011] Wan et al. 2023), modeling the temporal evolution of population densities crucial for understand- ing and managing pedestrian flows. In this case, data on individual trajectories are also typically unavailable, and only aggregate distributions at various time points are observed.
promising approach to address the absence of particle trajectories is the Bondan-Kinderlehee---O (JKO) scheme (Jordan al. 1998), which models the evolution of particle system as a sequence of distributions that gradually approach the minimum of a total energy functional while remaining close to the previous distributions. However, implementing the JKO scheme involves solving an optimization problem over the space of probability measures, which is computationally demanding. The first attempt to leverage JKO for learning population dynamies was introduced in (Bunne al., 2022b JKOnet). While innovative, this approach relies on a complex learning objective and is limited potential energy functionals- meaning it cannot capture stochasticity in the dynamics. A recent work. (Terpin et al. 2024] JKOnet"), ), proposes replacing the JKO optimization step with its first-order optimality conditions. This relaxation allows modeling more general energy functionals and reduces computational complexity, but it does not support end-to-end training. Instead, it re- quires precomputing optimal transport couplings (Cuturi 2013) between subsequent time snapshots, which limits its scalability and generalization.
Contributions. In this work, we present novel approach for recovering the system dynamics based on observed population-level data. Our key contributions are as follows:
1. Methodology: In $3.2 cast the problem of reconstructing energy functionals within the JKO framework as an inverse optimization task. This perspective leads to a novel min-max optimization objective for population dynamics recovery.
2. Algorithm: We equip our population dynamics methodology with a conventional end-to-end ad- versarial learning procedure (33.3). Importantly, our practical scheme does not pose restrictions on the architectures of the utilized neural networks, which contributes to the scalability.
Theoretical Guarantees: Under suitable assumptions, we show that our method can accurately recover the underlying energy functional governing the observed population dynamics ($) (.44.).
In $5)and Appendix B] we evaluate our approach on range of synthetic and real- weal-world datasets, in- cluding single-cell genomics. The results show that our method demonstrates improved performance over previous JKO-based approaches for learning population dynamics.
Notation. Let X' be a compact convex subset of TR equipped with the Euclidean norm || 112. Let P(Z) denote the set of probability measures on X, and let Px)(X) denote its subset of probability measures absolutely continuous with respect to the Lebesgue measure. For م E (2), we use p denote both the measure and its density function with respect to the Lebesgue measure. For measures μ1, VE E P(x'), we denote the set of couplings (transportation plans) between them by [1(, v). For a measure o E P(X) and measurable map T X + X, we denote by Ttp the associated push-forward measure, and idx X + X' the identity mapping. F. ===x =Z Zza denotes the divergence operator for a continuously differentiable vector field F +x
an
BACKGROUND
To describe the evolution of population measures, we use the theoretical framework of Wasserstein gradient flows (WGFs). In what follows, we first introduce the preliminary concept of Optimal Transport between probability measures (Villani et al. 2008 /Santambrogio 2015). Then we provide sufficient theoretical background on WGFs; see (Santambrogio 2017, Figalli & Glaudo, 2023) for an overview and (Ambrosio et al. (2008) for a comprehensive study of WGFs theory. Finally, we get the reader acquainted with the JKO scheme, which is the cornersoone of our approach.
Optimal Transport. The (squared) Wasserstein-2 distance dw, between two probability measures v P(x) is defined as the solution to the Kantorovich problem (Kantorovich 1942;
2

---

Published as a conference paper at ICLR 2026
Liv, (v, v) def
min e) xxxx
| 21 d 2),
where the distribution x² delivering the minimum to is called the optimal coupling (or optimal plan). p (a) then the Brenier's theorem (Brenier 1991) establishes equivalence between
the Kantorovich formulation ① and Monge's 's problem (Monge 1781:
/ || (xx) du(x),
dig (1,
min T:Thumu
(2)
where the optimal transport map T known the Monge map. Moreover, there exists a unique (up to an additive constant) convex potential v/* X' - IR such that T = Vib" and (Vi)*)
(McCann 1995). In this setting, the optimal coupling ① is given by x [idxx 1/.
Wasserstein Gradient Flows. For clarity, we first consider the concept of gradient flows in Eu- clidean space RD and then move to the space of probability measures P(t).. For Euclidean space RP, a gradient flow is an absolutely continuous curve xxt) starting Zo F2 that minimizes an energy functional J I2 + R "as fast as possible". To find such a curve, one needs to solve ordinary differential equation (ODE) (Teschl 2012) of the form (t) (xx(tt) with initial condition x0) = Zo- The same idea is applicable to the space probability measures P(X). If we equip P(X) with the Wasserstein-2 distance dz we obtain the Wasserstein space (2(X)= = (P()), d complete and separable metric space (Bogachev & Kolesnikov, 2012). In this case, for an energy functional
P(X) + R. the gradient flow in (x')'), called the Wasserstein gradient flow (WGF), is an absolutely continuous curve R+ - P(X) starting at Po that follows the steepest descent direction of J, i.e., solves
0,pt = - V P(t=o) Po:
(3)
where (pe) denotes the Wasserstein gradient in W2 given by -V 87
8p with
63 denoting the first variation (Chewi et al. 2025) energy J. Thus, equation (3) can be rewritten Ap
in the form of the continuity equation, expressing mass conservation under the velocity field 0y:
8J dept + V (p(er) = 0, = 8p (p.).
There exists intriguing connection between WGFs and certain partial differential equations (PDEs) (Evans 2022). In particular, different energy functionals in the Wasserstein space corre-
spond to distinct PDEs (Santambrogio 2015), some of which are associated with diffusion processes
appearing in practice (astro--Catrro 2024] Bailo et al. 2024b).
Examples of PDEs as WGFs. Consider the free energy functional (Carrillo et al. 2003 Eq. (1.3)):
[PE( = / ((x) dp(x))
/ Xxx
W(x - u) dp(z))p(u)+ + / (((()))d
(5)
V(p)
W(p)
u((p)
where V. W, and correspond to the system's potential, interaction, and internal energies, respec- tively. The PDE corresponding to the WGF driven by this functional is known the aggregation- diffusion equation (ADE). It describes the evolution of density Pe under a corresponding velocity field v, This evolution reflects balance of three effectt: drift driven by the potential V (xx), modeling an external field (e.g., gravity, electric potential); interaction governed by the symmetric interaction kernel W (x - y), accounting for non-local effects (e.g.. particle interactions, long-range forces); and diffiusinn modeled by the internal energy U(()) representing the energy associated with the local state of the system (e.g., thermal energy, chemical energy). A weak solution to the equation exists if the no-flux conditions: VV - γ O on 0.Y' holds, where 7ī is the outward normal vector on the boundary, and W E C(X x X¹), i.e. W is a continuous function on x X. These conditions ensure that no mass crosses the boundary of X;; mass only be redistributed within X. For a detailed discussion of the existence and classification solutions, see (Gómez-Castte 2024).
Such energy functionals are ubiquitous in real-world applications. including physics (Carrillo & Gvalani 2021) biology (Keller & Segel 1971 Potts & Painter 2022 Potts 2024), economics (Fi- aschi & Ricci 2025), machine learning (Suzuki al.] 2023 Chizat et al. 2024: Nitandal 2024). and nonlinear optimization (Bailo et al. 2024a), to name a few. For further discussion and references,
|024
the recent surveys (Carrillo et al. 2019a Gomez-Castro
2024 Bailo et al. 2024b).
3

---

Published as a conference paper at ICLR 2026
To demonstrate the flexibility of the energy formulation in ( we examine few representative cases. When the system's 's energy is purely internal and given by the negative differential entropy,
u(p) - -o)(p) def
[ p(x) log ((x), continuity equation 4 reduces to the classical heat equation: p.p- V2p (Vázquez 2017 Alternatively, when the energy is Jip(p) V(p) (H(()), the
resulting PDE Is the linear Fokker-Planee equation with diffusion coefficient 3:
Tep:= = V (VV(z) pe) + VV2 Pt.
(6)
A fruitful branch of theoretical and practical research stems from the connection between the Fokker- Planck PDE and stochastic differential equations (SDEs) (Risken & Frank) 1996; Bogachev et al.)
2022), the latter describes the stochastic evolution of particles. particular, equation (6) is equiva-
lent (Weinan et al. (202I) to the following Itô SDE:
dx, dx, 22
where X = (so is RD-valued stochastic process and W 120 is standard Wiener process (Särkkä & Solin 2019). In other words, if X, ~ Pt evolves according to (7). then the density Pe evolves according to the Fokker-Plance equation (6) in the space of probability measures.
WGFs provide a compelling framework for modeling PDEs and their associated SDEs across various domains, giving rise to range of methods for approximating solutions to WGFs. These include deep
caaming-based approaches (Mokrov et al. [022 Alvarez-Melis et al.) |022 Alteküüger et al.. 2023).
particle-based methods such as Stein Variational Gradient Descent (SVGD) (Liu & Wang 2016)
and its extensions (Das & Nagaraj 2023; Tankala et al. 2025), ) as well as classical discretization
techniques Wasserstein space (Carrillo et al. 2022] see 91.2). Many of these advances build on
the celebrated Jordan Miaaaan-Knnnreeeee-OOOtet (JKO) scheme (Jordan et al. 1998), which we review next.
The JKO Scheme. A classical method for solving ODEs in Euclidean space the implicit Eu- ler scheme. For an (Euclidean) gradient flow (t) = -v)(((t)) and T >> 0. this scheme ap- proximately solves this ODE by iteratively applying the proximal operator (Parikh
al,, [014]] 51.1): Fx+11 prox, (#*) = arg m [(x) + 27 || I& 12] Compared to the explicit Eu- ler scheme, this approach offers improved numerical stability (Butcher 2016). Jordan, Kinderlehrer, and Otto (Jordan et al. 1998) extended this idea to the space of probability measures, introducing variational time discretization of the Fokker-Planck equation (6). now known as the JKO scheme:
Pk+1 = -arg min { J(p) + dr (p. (T) (KO,)(PH), PEP(X)
1
P = Po-
(8)
27
where > 0 is the time step. As 0. the sequence PE. k E N converges to the continuous solution Pr of the Likker-Planck equation (6). Later, the JKO scheme's convergence was gener- alized (Ambrosio et al. 2008 Thm. 4.0.4) for the free energy functional JrE (). Note that since
di (P.P%) 2 0, the energy is non-increasing along the sequence: J(p +4 S J (pk). This mono-
tonicity property is often utilized in JKO-based algorithms (Salim al. 2020 Wasserstein space.
3 IJKONET METHOD
We begin by formally stating the problem addressed in our work ($3.I). then we develop our method- Finally, we discuss the practical ($3.3) and theoretical aspects ($3.4) of our approach.
ology ( 3.2
3.1
PROBLEM STATEMENT
In our work, we address the problem of recovering the underlying energy functional J that gov- emns the evolution of density Pr (P) in (4) based marginal population measures (Bunne et al. 2022b Lavenant et al. 2024). Specifically, we are given independent samples from marginals {px k= of the evolving distribution Pt at corresponding time points to tt < < tx- Impor- tantly, each distribution Pk: may be represented by a different number of samples. As discussed in 22 the JKO scheme approaches the continuous-time dynamics the step sizes Ale la+1 i small enough. This motivates our modeling assumption that the ground truth sequence of measures K follows the scheme Pk+1 = JKO Ana. (pk).
{p*)
Although the time intervals Al between observations may vary, corresponding to non-uniform step sizes in the JKO scheme, we assume equal spacing At& 7 in the remaining text for simplicity.
4

---

Published as a conference paper at ICLR 2026
3.2 METHODOLOGY AND LOSS DERIVATION
In this section, we assume Pk P (X) and minimization is always taken over p (a) ) in order to utilize Brenier's 's theorem. The key idea of our method builds on inverse optimization techniques (Chan et al. 2025) and is demonstrated in Figure [ Thanks to assumption Pk+1 = JKO+5+ (pk), we can derive an inequality that becomes an equality a candidate functional J matches the ground truth functional J. since PK++1 is obtained via the JKO step associated with J":
{ min p²
T(0))
1 des (px: (0x:)00))) $ ((k+11) + der, (pkr Pk+1)) 25
1
2T
Moving the right-hand side to the left yields an expression that is always upper-bounded by zero, regardless of the choice of J. Maximizing the resulting gap with respect to I encourages the candidate functional to approximate the true functional J This yields the following objective:
K-1
1 _ 2T
- _dlns dlws (pu, pk++t) 2T
+ min [(p*)
- I [px+1))
max
k=0
indepenndent of J.
K-1
max
Σ min [[(p)) 1 J(p+t1) + die, (pk p) + Const. 27
1
(10)
1991) ensures that each distribution p² appearing in (TO) be written a pushforward pM 7 for trans-
Thanks to our assumption Ps E (a)(()), Brenier's theorem (Brenier
port map 1 X' + X. In addition, the Wasserstein distance (2) admits the upper-bound (x))
dip (Pk, S - T²) are independent across k. the sum and the minimization can be interchanged. Applying these observations, we end up with our final loss objective equivalent to (TO):
(z). Since the minimizations over p² (or equivalently over
K-1
max max L(J, T def min
Σ
[(( - [ 27 / | T (z) dz). (11) R=
max
3
T
J
T
The key idea is that the inner minimization over maps T approximates JKO step for given energy functional J; that is, each optimal map T = arg min [(T) + 27 Sx ||
T (x)11ppx(x) dz) pushes P to the JKO-updated distribution Px+1 JKO,J The outer maximization then drives J toward the true functional J causing the pushforward distributions Pk+1 Tytpk to approach Pk-1, see illustration in Figure
3.3 PRACTICAL ASPECTS: METHOD PARAMETRIZATION AND LEARNING PROCEDURE
We denote by θ and y E & the parameters of sufficiently expressive function classes (c.g., neural networks) used to approximate the candidate functional I and the transport maps T re- spectively. Specifically, To and Th are their parameterized counterparts.
Mapping Parameterization. Building upon [Benamou et al., 2016), prior works (Mokrov et al., Alvarez-Melis et al., 2022, Bunne et al., 2022b) typically parameterize the transport maps where V.C are modeled using input-convex neural networks (ICNNs) (Amos et al. However, ICNNs suffer from poor scalability in high- high-dimensional settings (Korotin et al, In contrast, since objective (TT) imposes convexity constraints, parameterize T directly using standard architectures like MIPs or ResNets (He et al. (0166) This relaxation simplifies optimization and yields improved empirical stability. Alternatively, one employ more task-specific parametrizations, such as triangular maps (Baptista et al. (024)..
2021
as 7
VV
2017
2021bb.
Energy functional parametrization. Following (Terpin et al. 2024). we parameterize the energy
functional J as specific instance of the free energy formulation in (5) where the internal energy
term is chosen to be scaled negative differential entropy:
Wo, dp(x))p(() 0g / log p(x) dp(xx)
To(p)
Vo (x)
(12)
xxx
(0)
(oy
(p))=-0-H10
5

---

Published as a conference paper at ICLR 2026
Here 8 (01. 82. f2} represents the set of learnable parameters: 01 and 0g are parameters of neu- ral networks that define the potential Ve, : X' + IR and the interaction kernel Wo, : 2' + R, respectively; 0g R is learnable scalar diffusion coefficient, see (7).
Entropy evaluation. All the terms in the objective function (II) can be estimated using Monte Carlo integration (Metropolis & Ulam 1949), except for the internal energy term Ug To handle this remaining component, we follow (Mokrov et al. [202]] Theorem 1), (Alvarez- Melis et al. 2022
Eq. (II)) and apply the change-of -vrriaalles formula, yielding:
Ue (TT)pk) Uog (px) - 0g J log del (x))
(13)
where T denotes the Jacobian of the invertible transport map T We note that, during training, the computation of log det even for non-invertible mappings Th be effectively handled by modern programming tools; see Appendix [C.3] for details. To efficiently compute the gradient of the log det. term in (13) with respect sp, one can use the Hutchinson trace estimator (Hutchinson) 1989). as proposed by (Finlay et al. [2200] Huang et al. (021). However, in our experimental setting ( we find that computing the full Jacobian is tractable and sufficient. For estimating the negative differential entropy M((s)) in (13). we employ the sranchhnnko-LLnennne nearest seirest- estima- tor (Kozachenko) 1987, Berrett et al. 2019). Notably, entropy values for the population measures can be precomputed prior to training.
{pk k==)
Learning Objective. To facilitate gradient- based optimization, we employ Monte Carlo estimation
to approximate the loss in III. At each time step &x (for & = O.....K) K), we draw a batch of Bj samples X = (x) .
13 Pa and optimize the following empirical loss objective:
K
Be ToDi1B
[00 - Σ
-o(T)(Xa)) -x((x)1)) +
| 2T j=1
1
(14)
where Th (Xx) denotes the batched pushforward used compute the empirical estimates To of the functional To- We optimize the objective in (14) with respect to 0 and şp using standard
gradient descent-ascent scheme; implementation details are provided in $5 and Appendix C with the complete training procedure summarized in Algorithm therein.
3.4 THEORETICAL ASPECTS
The central idea behind our loss formulation (:3.2) is to minimize the inverse gap in which measures the discrepancy between the optimal value (left-hand side) and the expected value (right- hand side) of the JKO step for a candidate functional J. While this objective is intuitively justified, it remains to be shown whether the minimizer Jmin of (I]) truly approximates the ground-truth energy functional J Our Theorem 3.1 addresses this question, demonstrating that, under suitable assumptions, optimizing (1) indeed recovers J up an additive constant that does not affect the
dynamics governed by (4) Full details and proofs are provided in Appendix D
We state theorem for K = 1. When considering K > 1, the statement of the theorem holds true independently for each timestep k = 0, 1,..., K - 1, and the results extend straightforwardly to this more general setting. We further assume that both the ground-truth functional J' and the candidate functional J are purely of potential energy form, i.c., J* = (p) = Sx Vx) dp(x); (p) = (p)(p) fxx V(x) dp(z)) For notational simplicity, we denote the loss as L(V, T) in place of [(JoE, T). For technical purposes, we introduce the modified version of a potential as V, == -V + Il 11 X + R; subscript stands for "quadratic".
Theorem 3.1 (Quality bounds for recovered potential energy). Let e(V) def L(V*, L(V*,Tv.) C(V, Tv) be the gap between the optimal and optimized value of inverse JKO loss ①1] with internal minT problem solved exactly, i.e., Tv def miny L(V,T), Let X be convex set; (modified) potentials V be strictly convex and 1 smooth (see the definition in Appendix D). Then there exists C = C(r, 3) such that following inequality holds:
f [|V - (V())||dpx00 (0) Ce(V).
(15)

---

Published as a conference paper at ICLR 2026
Notably, equation (15) compares gradients of the recovered and ground truth potentials, hiding the appearance of a redundant additive constant. Importantly, the assumptions of the theorem are not particularly restrictive in practice. Specifically, the smoothness of potentials can be ensured by employing smooth activation functions such as CELU (Barron (017), SiLU (Hendrycks & Gimpel 2016), SoftPlus, and others. Strict convexity, in turn, be enforced through architectural design choices (Amos et al., 2017). Furthermore, the strict convexity of the (modified) potentials Va can often be assumed when the step size T is sufficiently small. In our experiments ($5) we parameterize V using standard MLPs and observe that this approach performs adequately.
To the best of our knowledge, this work is the first to provide a quality analysis (Theorem 3.1 of JKO-based solvers for population dynamics. At present, our analysis focuses exclusively on the potential energy component. Extending this framework incorporate interaction and internal ener- gies presents an interesting direction for future research. In contrast, (Wu & Wang 2025) analyzes the JKO scheme in terms of convergence to the true density of the Tokker-Plance equation (ō. ac- counting for parameter uncertainty in iterative updates. Our work differs by focusing on the ability to restore energy components from data, rather than reconstructing the full density.
4 RELATED WORKS
This section reviews research directions most relevant to our work. For an extended discussion of related works, see Appendix A.I First, we discuss existing methods for modeling the dynamics of P given a known energy functional J ($4.1). Next, we focus on approaches closely related to our setting specifically, learning population dynamics from observed data ($4.2) via the JKO scheme.
4.1 SOLVING AGGREGATION DIFFUSIIO EQUATION WITH GIVEN ENERGY
This line of research focuses on methods for solving aggregation- -ifffuuoonn equations derived from WGFs with given energy functionals of the form
(5) A variety of numerical approaches exist,
including mesh-based schemes (Cancès et al., 2023 Jüngel et al. 2024), J. particle methods (Craig &
Bertozzi 2016 Campos Pinto et al.) 2018 [ Carrillo et al. [201b)), and JKO-based schemes (Carrillo ct al. 2022). For a comprehensive overview, see (Bailo et al./ [024b] 85). Recently, the research hield was empowered by deep learning techniques. In particular, (Alvarez-Melis et al., 2022 Mokrov et al. 2021 Fan et al., [022] Park et al.] 2023 Lee et al.) (2024) employ gradient- bradee optimization
and neural-network parametrization solve WGFs with given energy as special cases of (5)
4.2 LEARNING ENERGY FOR SAAELIGAAATIO-DIUUSSIIE EQUATION VIA JKO SCHEME
This line of research focuses on methods for learning population dynamics (3.11 using the JKO scheme, as in JKOnet (Bunne et al. 2022b) and JKOnet" (Terpin et al. 2024). We discuss these methods detail below, they are the most closely related our approach. Additional methods for learning population dynamics are reviewed in Appendix A.I.I
JKOnet (Bunne et al., 2022b) formulates the task of population dynamics recovery as a bi-level optimization problem aimed at minimizing the discrepancy between observed distributions Pk and model predictions PN:
K-1 Loxonet (0,4) = Σ k=0
de Pk). s.t. po = po.
Pu+1 = VeL
(16)
del 1/
arg min To (ppptpt) + / || poott ECVx
(1) dpp(
2r x
where CVX denotes the set of continuously differentiable convex functions from X' to R. The trans- port maps the are parametrized either using ICNNs (Bunne et al., 2022a) or through the 'Monge gap' approach (Uscidda & Cuturi 2023). However, the energy functional is limited to the potential energy term (p(p), excluding interaction and internal components. This framework has two key limitations: (i) solving the bi-level optimization problem in (6)) is computationally challenging and requires specific techniques like unrolling optimizer's 's steps, and (fi) extending the method to more realistic settings with richer energy structures is questionable due computational complexity.
JKOnet" Terpin et al. 2024) addresses the limitations of the original JKOnet by replacing the full JKO optimization problem with its first-order optimality conditions. This reformulation allows
7

---

Published as a conference paper at ICLR 2026
The method minimizes the following objective, utilizing a precomputed optimal transport plans Ty between Pk and Pk+1:
the use of more expressive energy functionals (22
K-
/ VUo, (#x+1 yk++) dp (3++1)
Lucnnet (0)= -
Σ k=0
V Vo. (xx+1))
xxx
(17)
V (P + - -1x++1 - -200| do Ix+1).
+0g
P++1(I0+1)
The authors emphasize several advantages over JKOnet, including reduced computational costs and support for more general energy functionals. However, their method requires an additional op- timization round to precompute optimal transport plans wk, introducing extra sources of inaccuracy and rendering the approach non-end-to-end. Moreover, the authors of (Terpin et al., 2024) employ discrete OT solvers to compute Tę. which may fail to accurately represent OT maps between the underlying distributions, particularly in high-dimensional settings (Deb et al. 2021).
iJKOnet (Ours) integrates the strengths of both JKOnet and JKOnet". From JKOnet, it inherits parameterized models for the energy functional To and the transport map T.. avoiding reliance on precomputed discrete optimal transport plans Tg. which can hinder flexible mappings between Pk and Pk+1- From JKOnet", ", it adopts a more expressive formulation of the energy functional ([2) capable of capturing complex structures. Unlike JKOnet", lytical expressions of the JKO step (8) optimality conditions; instead, it solves the step directly via the inner minimization in III This approach makes the method more readily extendable to general energy forms, such as porous medium energies (Alvarez-Melis et al. 2022), although exploring such extensions is beyond the scope of the current work.
formulation does not rely ana-
5 EXPERIMENTS
Our inplemenntaionn is primarily based on the publicly available code from JKOnet" (erpin et al. 2024]. written JAX (Bradbury et al., 2018), and available at the following repository Since our approach builds on the JKO framework, in I.. we first reproduce the experimental setup from JKOnet" * for learning a known potential energy. Although we found the original codebase well-structured and accessible, we identified several inconsistencies in the data generation process, which we discuss in detail there. Finally, in $5.22 we compare our method against popular non-JKO baselines single-cell RNA-seq datasets. Extended comparisons be found in Appendix B
JKO-based Models. As demonstrated in Appendix 1.1.2 and B.1.3 learning the interaction and internal energy components directly from samples is challenging likely because their estimation requires computing integrals over functions that must themselves be estimated. Consequently, for large-scale experiments, we restrict our model to the potential energy parametrization, imposing the inductive bias 02 03 = O. We refer to this variant as jKKnetvv. For moderate-dimmnnsiona tasks, we explore all combinations of parametrizations recovery. For fair comparison with (Teppin et al. 2024). we also implemented a time-varying potential energy parametrization, as described in Appendix B of the original paper, and denote this variant by i JJOOnet,r However, as discussed in Appendix (A.2] this parametrization does not cor- respond to single, consistent energy functional governing the dynamics. Instead, it actually solves sequence of K independent optimal transport problems between consecutive snapshots (Bunne) cet al., 2023). Moreover, because JKOnet (Bunne et al., 2022b) is computationally expensive, we limit comparisons to experiments with moderate dimensionality, with the corresponding results reported in Appendix B.2 Technical details of model training are provided in Appendix(C)
illustrate the challenges of full energy
Metrics. We evaluate the next-step distribution Pk+1 from Pk: and compare it to the ground-truth distribution Pi++1 using the following metrics: (a) Earth Mover's Distance (EMD) (or Wi) (Rub- ner et al. 1998) and W2, (b) Wasserseess Unexplained Variance Percentage (Bal -UVP) (Korotin et al. 2021c), (c) L2-based Unexplained Variance Percentage (C,-UVP) (Korotin et al.) 202 which measures the discrepancy between the ground-truth functional J and its reconstruc- tion F, and (d) Maximum Mean Discrepancy (MMD) (Gretton et
2012). Formal definitions and implementation details for all metrics are provided in Appendix C.T
staar
http::///gtttub..cmm///.a//
8

---

Published as a conference paper at ICLR 2026
1X0nntt-2KlOursi
IJKOnet, Ours
JKonet 222
1K0e1;-10K
BC
01
WP
W5
K5
W5
F30
ST
ST
PN
(a) EMD
(b) Bdlly -UPP
(c) L2-UVP
Numerical results from 5..1 for the unpaired setup. The reported absolute values show that, while increasing the number of samples generally improves performance across metrics, certain potentials remain challenging, highlighting the difficulty of this setup.
Figure 2:
5.1 LEARNING POTENTIAL ENERGY
As discussed in T the population dynamics setting assumes that particle trajectories are not tracked across time- i.c., each time snapshot should consist of independently sampled particles without temporal correlation. This implies that data at each time step should be regenerated from seratch. However, identified inconsistency in the original codebase particle trajectories across time steps, resulting in temporally correlated samples. We refer to this as the paired setup, where each particle Vi is directly linked to 2x++1 along trajectory. In contrast, the intended, tem- porally uncorrelated setting is referred to as the unpaired setup. Table 3 in Appendix B.1.1 shows that switching between setups has a substantial impact on performance. detailed discussion of this phenomenon is provided in Appendix A.3
preserved
84.1), we begin by examining how our method learns potentials the 2D paired setup. which allows for direct visual comparison with JKOnet". These experiments are conducted on the synthetic dataset from (Terpin et al. (2024] Appendix B), with results presented in Figure 4 and further details provided in Appendix B. We then repeat the same procedure in unpaired setup, analyzing how performance scales with the number of samples (2K and 10K). The corresponding results are shown in Figure 2 They indicate that our method outperforms JKOnet" on nearly all potentials, and that for certain cases, increasing the number of samples does not improve performance, highlighting the difficulty of the unpaired setup.
Following Terpin
2024
5.2 LEARNING POPULATION DYNAMICS FOR SINGLE-CEE CELL DATA
Following (Chen al. 2023) and (Shen al. 2025). in this section, we apply our method mod- cling the dynamics of single-cell RNA data, with extended comparisons provided in Appendix B.2
Dataset. We use the Embryoid Body (EB) dataset (Moon al. 2019) and follow the preprocessing pipeline described in (Tong et al.) 2020). The EB dataset comprises a collection of 5 timepoints and describes the differentiation human embryonic stem cells over period of 30 days.
Non-JKO Models. We evaluate our method against classic algorithms for trajectory inference such 22020), MIOFLOW (Huguet al. 2022). DMSB (Chen et al.
as Tra TrajectoryNet (Tong
2023 NLSB (Koshizuka & Sato 2023) and the recently proposed MMSB (Shen et al. 2025). For these methods, the reported metrics are taken from the corresponding referenced papers.
Leave-two-out in 5D. Following (Shen et al. (025), we perform experiments using leave-two-out setup. Since the EB dataset contains five timesteps ((0-4). we remove the first (t₂) and third (t3) timesteps from training data (i.e.. it consists of 2,,2,4 timesteps), and then evaluate how well our method can reconstruct the data for t1 and t3 from the remaining to and t2 timesteps, respectively. One see that our method with the "t, V parameter-
The results are demonstrated in Table ization yields the best metrics, whereas the use of alternative parameterizations (with interaction and internal energies) results in poorer metrics. Vani 1 1 a- corresponds the IPF (GP) method

---

Published as a conference paper at ICLR 2026
Vargas et al. 2021;; further details are provided in Appendix B.6 of (Shen al. 2025). Our method outperforms previous approaches in terms of the provided metric.
Leave-one-out in 100D. Following (Chen et al. 2023),
conduct experiments using leave-one- out setup. One of the timesteps 61. 122 or t3 is omitted, and we evaluate the method's ability to reconstruct the distribution of the left-out timestep. The results shown in Table 2 Our method iJKOnety achieves results on par with DMSB in the all-time (w/o LO) setting while using a simpler, simulation-free optimization routine that requires no trajectory caching. Consequently, outperforms DMSB in execution time, as shown by comparing Table 12 in Appendix E of (Shen
al. 2025) (for 5D) with Table ( (for 100D).
DISCUSSION
Contributions. We introduce the novel iJKOnet method, which pro- vides general framework for re- covering any type of energy func- tional governing the evolution of a system; in this work, we focus the free energy parametrization (12). Our work bridges the gap between in- verse optimization theory and com- putational gradient flows. The pro- posed method outperforms previous JKO-based approaches in all com- parisons and achieves comparable re- sults to non-JKO methods on real- world tasks. We also establish theo- retical guarantees for recovering dy- namics driven by potential energy.
Table 5D experiment. dw, distance (L) comparison across t and 1z. Results for baselines (Shen et al. 2025).
Method Vanilla-SB
1.49 0.063
1.55 + 0.034
DMSB
1.13 + 0.082
1.45 +: 0.16
TrajectoryNet
2.03 0.04
1.93 0.08
MMSB
1.27 + 0.028
1.57 + 0.048
Static
JKOnet
1.145 0.033
2.52
JKOnet /+U
1.099 = 0.119
2.537 +.055
JKOnety +WW
1.419 0.173
2.510
JKOnet W+U
1.887 = 0.017
1.739 #739=#00 0.037
JKOnet
1.361 + 0.257
2.557 0.042
Static (Ours)
iJKOnety
1.082 0.011
1.147 + 0.001
Limitations. Our method can not re- cover other types of internal energy functionals except for the entropy, it handle time-dependent interaction energies, although theo- retical results for such cases exist (Ferreira & Valencia-Guevara 2018). Furthermore, it does not account for birth-death dynamics as in recent tra- jectory inference methods (Zhang al. 20256). The reliance on entropy estimation may also lead to performance degradation in higher-diment linenioai settings. and we find that learning interaction energies remains particularly challenging. As result, jointly optimizing all energy parameters (0,, 0z, 03) can cause instability and convergence inaccurate potential estimates. Finally, from theoretical standpoint, our analysis currently provides guarantees only for potential energy (Theo- rem 3.1 extending it other energy types is avenue for future work.
iJKOnety+U
1.065 + 0.018
1.150 0.004
iJKOnety
2.865 0.166
1.376 + 0.015
iJKOnetw+U
1.649 = 0.005
+.600
iJKOnet
3.577 = 0.166
1.395 = 0.032
Time-varying JKOnet iJKOnet,,v (Ours)
4.414 = 1.499
2.771 + 0.197 0.849 + 0.021
0.983 + 0.037
Table 100D experiment. MMD distance (4). Comparison of methods across different leave-one- out splits. The results are averaged across 3 runs. Results for baseline methods: (Chen et al. 2023
Method
LO-ta
L0
LO-ts
w/o LO
NISS Koshizuka & Sato,
2023
0.38
0.37
01.7
0.66
MIOFLOW [uuguet
2027
0.23
0.90
0.23
0.23
DMSB Chen al. 2023
0.042
0.020
0.033 + 0.003
0.040
0.020
0.032 + 0.003
JKOnet; (Terpin et al. 2024
0.220 + 0.025
0.293 + 0.018
0.235
0.006
0.229 + 0.052
iJKOnety (Ours)
0.137 + 0.001
0.123 1 0.001
0.097 + 0.002
0.085 + 0.024
JKOnet, Terpin 2024
0.575
0.119
0.619 + 0.157
0.456
0.056
0.477 + 0.098
iJKOnet. (Ours)
0.848 1: 0.043
0.370 + 0.233
0.055 + 0.007
0.124 + 0.243
10

---

Published as a conference paper at ICLR 2026
Reproducibility. We provide the experimental details in Appendix C and the code required produce the conducted experiments is available in this repository (see README. md for details).
LLM Usage. Large Language Models (LLMs) were used exclusively to assist with sentence rephrasing and improving text clarity. All scientific content, results, and interpretations in this paper were produced solely by the authors.
ACKNOWLEDGMENTS
The work was supported by the grant for research centers in the field of Al provided by the Ministry of Economic Development of the Russian Federation in accordance with the agreement W00000C313925P4F0000 and the agreement 1313--0--005--033
REFERENCES
Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery data mining, pp. 2623-2631, 2019.
Vitalii Aksenov and Martin Eigel. An eulerian approach to regularized jko scheme with low-rank tensor decompositions for bayesian inversion. Journal Scientific Computing, 105(1):3, 2025.
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference Learning Representations, 2023. URL
https://openrewt
net
Fabian Altekrüger, Johannes Hertrich, and Gabriele Steidl. Neural Wasserstein gradient flows for discrepancies with riesz kernels. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Bar- bara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 664-690. PMLR, 23-29 Jul 2023. URL https: altekruger23a 1tmmmmmm......11.111..mmm
mlr press/v202
David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing functionals the space of probabilities with input convex neural networks. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2008.
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In International conference on machine learning, pp. 146-155. PMLR, 2017.
Francisco Andrade, Gabriel Peyré, and Clarice Poon. Learning from samples: Inverse problems over measures via sharpened fenchel-young losses. arXiv preprint arXiv:2505.0711 007124, 2025.
Arip Asadulaev, Alexander Korotin, Vage Egiazarian, Petr Mokrov, and Evgeny Burnaev. Neu- ral optimal transport with general cost functionals. In The Twelfth International Conference Learning Representations, 2024. URL
gIiz7tBtYZ
Lazar Atanackovic, Xi Zhang, Brandon Amos, Mathieu Blanchette, Leo J Lee, Yoshua Bengio, Alexander Tong, and Kirill Neklyudov. Meta flow matching: Integrating vector fields on the wasserstein manifold. In The Thirteenth International Conference on Learning Representations, 2025. URL /ittpss//ooeneeee....eeee/eoeo wwtp:eeeetttonane/e111 95Yc22230000
Francisco J Ayala, Michael E Gilpin, and Joan G Ehrenfeld. Competition between species: theoret- ical models and experimental tests. Theoretical population biology, 4(3)331-356, 1973.
Jean Baccou and Jacques Liandrat. Subdivision scheme for discrete probability measure-valued data. Applied Mathematics Letters, 158: 109233, 2024.

---

Published as a conference paper at ICLR 2026
Rafael Bailo, Alethea Barbaro, Susana N Gomes, Konstantin Riedl, Tim Roith, Claudia Totzeck, and Urbain Vaes. Cbx: Python and julia packages for maassenuus-basee interacting particle methods. Journal of Open Source Software, 9(98);661,, 2024a.
Rafael Bailo, José A Carrillo, and David Gómez Canez-Castro, Aggregation- diffuusson equations for col- lective behaviour in the sciences. arXiv preprint arXiv:2405.1667,, 2024b.
Amartya Banerjee, Harlin Lee, Nir Sharon, and Caroline Moosmiller. Efficient trajectory infer-
ence in wasserstein space using consecutive averaging. In International Conference Artificial Intelligence and Statistics, pp. 2260-2268. PMIR, 2025.
Arindam Banerjee, Srujana Merugu, Inderjit S Dhillon, and Joydeep Ghosh. Clustering with breg- man divergences. Journal of machine learning research, 6(0ct):17055 170--1749, 2005.
Ricardo Baptista, Youssef Marzouk, and Olivier Zahm. On the representation and learning of mono- tone triangular transport maps. Foundations of Computational Mathematics, 24(6):203322108 2024.
Jonathan T Barron. Continuously differentiable exponential linear units. arXiv:1704.0748,, 2017.
arXiv preprint
Nico Battich, Joep Beumer, Buys De Barbanson, Lenno Krenning, Chloé S Baron, Marvin E Tanen- baum, Hans Clevers, and Alexander Van Oudenaarden. Sequencing metabolically labeled tran- scripts in single cells reveals mrna turnover strategies. Science, 36764882 2020.
Amir Beck. First-Order Methods in Optimization, volume 25 of MOS-SIAM Series on Optimization. SIAM, Philadelphia, 2017. ISBN 9781611974980.
doi: 10.1137/1.
9781611974997.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge- kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393 2000.
Jean-David Benamou, Guillaume Carlier, Quentin Mérigot, and Edouard Oudet. Discretization of functionals involving the monge-ampal operator. Numerische mathematik, 134(3):611-636, 2016.
Jean-David Benamou, Thomas Gallouët, and François-Xavier Vialard. Second-order models for optimal transport and cubic splines on the wasserstein space. Foundations of Computational Mathematics, 1935::1113-114 2019.
Jules Berman and Benjamin Peherstorfer. CoLoRA: Continuous low-rank adaptation for reduced implicit neural modeling of parameterized partial differential equations. In Ruslan Salakhutdi-- nov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, vol- ume 235 of Proceedings of Machine Learning Research, pp. 3565-3583. PMLR, 21-27 Jul 2024. URL http:::sspppocceeciiicc mlr www.prressaeeee sateesss/v23/bbeee rman2 html
Jules Berman, Tobias Blickhan, and Benjamin Peherstorfer. Parametric model reduction of mean- field and stochastic systems via higher-order action matching. Advances in Neural Information Processing Systems, 17:56588-56611 2024.
Thomas B Berrett, Richard J Samworth, and Ming Yuan. Efficient multivariate entropy estimation via k-nearest neighbour distances. The Annals of Statistics, 47(1):288-318, 2019.
Vladimir Bogachev and Aleksandr V Kolesnikov. The boonge-katorroiichhrbbee achievements, connections, and perspectives. Russian Mathematical Surveys, 67(5):785, 2012.
Vladimir Bogachev, Nicolai V Krylov, Michael Röckner, and Stanislav Shaposhnikov. Porker-PPanck-Kolnngggn Equations, volume 207. American Mathematical Society, 2022.
Clément Bonet, Christophe Vauthier, and Anna Korba. Flowing datasets with wasserstein over wasserstein gradient flows. arXiv preprint arXiv:2506.07534, 2025.
12

---

Published as a conference paper at ICLR 2026
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milmg and Qiao
Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//githube com foom/jax-ml/jal
Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions.
Communications on pure and applied mathematics, 44(4):375-417, 1991.
Luc Brogat-Motte, Riceardo Bonalli, and Alessandro Rudi. Learning controlled stochastic differen- tial equations. arXiv preprint Ttt::2111019980 2024.
Charlotte Bunne. Andreas Krause, and Marco Cuturi. Supervised training of conditional monge
maps. Advances in Neural Information Processing Systems, 35:6859-6872. 2022a.
Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal optimal trans- port modeling of population dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 3511-6528 PMLR, 2022b.
Charlotte Bunne, Stefan Stark, Gabriele Gut, Jacobo Sarabia Del Castillo, Mitch Levesque, Kjong-Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single- cell perturbation responses using neural optimal transport. Nature methods, 20(11):1759-1768, 2023.
Daniel Burkhardt, Jonathan Bloom, Robrecht Cannoodt, Malte D Luecken, Smita Krishnaswamy, Christopher Lance, Angela O Pisco, and Fabian Theis. Multimodal single-cell integration across time, individuals, and batches. NeurIPS Competitions, pp. 12, 2022.
Ivan Butakov, Aleksandr Tolmachev, Sofia Malanchuk, Anna Neopryatnaya, and Alexey Frolov. Mutual information estimation via normalizing flows. Advances in Neural Information Processing Systems, 37:3027-3057 2024.
John Charles Butcher. Numerical methods for ordinary differential equations. John Wiley & Sons, 2016.
Martin Campos Pinto, José A Carrillo, Frédérique Charles, and Young-Pil Choi. Convergence of linearly transformed particle method for aggregation equations. Numerische Mathematik, 139: 743-793, 2018.
Clément Cancès, Maxime Herda, and Annamaria Massimini. Finite volumes for generalized yoisson-nernst-Po planck system with cross -iffuudif and size exclusion. In International Conference on Finite Volumes for Complex Applications, pp. 57-73. Springer, 2023.
Xavier Aramayo Carrasco, Grigoriy Ksenofontov, Aleksei Leonov, Iaroslav Sergeevich Koshelev, and Alexander Korotin. Entering the era of discrete diffusion models: A benchmark for schrödinger bridges and entropic optimal transport. In The Fourteenth International Conference on Learning Representations, 2026. URL. /ttes:://ppereviii ew. net. forum?
XcPDT615Gd
José Carrillo and Rishabh Gvalani. Phase transitions for nonlinear nonlocal aggregation- diffusion equations. Communications in mathematical physics, 382(1):485-544 2021.
José A Carrillo, Robert J McCann, and Cédric Villani. Kinetic equilibration rates for granular media and related equations: entropy dissipation and mass transportation estimates. Revista Matematica Iberoamericana, 19(3)-911-0018 2003.
José A Carrillo, Katy Craig, and Yao Yao. Aggregation-diiffsiin equations: dynamics, asymptotics, and singular limits. Active Particles, Volume 2: Advances in Theory, Models, and Applications, pp. 65-108, 2019a.
José A Carrillo, Katy Craig, Li Wang, and Chaozhen Wei. Primal dual methods for wasserstein gradient flows. Foundations of Computational Mathematics, pp. 1-55, 2022.
José Antonio Carrillo, Katy Craig, and Francesco S Patacchini. A blob method for diffusion. Calculus of Variations and Partial Differential Equations, 58:1-53, 2019b.
13

---

Published as a conference paper at ICLR 2026
Timothy CY Chan, Rafid Mahmood, and Ian Yihang Zhu. Inverse optimization: Theory and appli- cations. Operations Research, 1302);14461107, 2025.
Ricky T. Q. Chen and Yaron Lipman. Flow matching on general geometries. The Twelfth
International Conference Learning Representations, 2024. htttttt:://ppereeeeeee
net//orumm?Id=q7ohDiIr
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018a.
Tianrong Chen, Guan-Horg Liu, and Evangelos Theodorou. Likelihood training of schrödinger bridge using forward-backward SDEs theory. In International Conference on Learning Representations, 2022a. URL /ttps://openrevieww nwtp:/eee//eoeeeeee http://wwww.......ooomoomioooooo
Tianrong Chen, Guan-Horng Liu, Molei Tao, and Evangelos Theodorou. Deep momentum multi- marginal schrödinger bridge. Advances in Neural Information Processing Systems, 36:57058- 57086, 2023.
Yongxin Chen, Giovanni Conforti, and Tryphon T Georgiou. Measure- spline curves: An optimal transport viewpoint. SIAM Journal on Mathematical Analysis. 50(6))5947- -566688, 2018b.
Yongxin Chen, Tryphon T Georgiou, and Michele Pavon. The most likely evolution of diffusing and vanishing particles: Schrodinger bridges with unbalanced marginals. SIAM Journal on Control and Optimization, 60(4):2016-2033) 2022b.
Xiuyuan Cheng, Jianfeng Lu, Yixin Tan, and Yao Xie. Convergence of flow-based generative models via proximal gradient descent in wasserstein space. IEEE Transactions on Information Theory, 2024.
Ziheng Cheng, Shiyue Zhang, Longlin Yu, and Cheng Zhang. Particle-based variational inference with generalized wasserstein gradient flow. Advances Neural Information Processing Systems, 36:66583-6666 666099, 2023.
Andrei Chertkov and Ivan Oseledets. Solution of the fokker- equation by cross approximation method in the tensor train format. Frontiers in Artificial Intelligence, 4:668215, 2021.
Sinho Chewi, Julien Clancy, Thibaut Le Gouic, Philippe Rigollet, George Stepaniants, and Austin Stromme. Fast and smooth interpolation wasserstein space. In International Conference Artificial Intelligence and Statistics, pp. 3061-3069. PMLR, 2021.
Sinho
Chewi, Jonathan Niles-Weed, and Philippe Rigollet.
Statistical Optimal Transport:
Eicole d' Été de Probabilités de Saint-Flour XLIX - 2019, volume 2364 of Lecture Notes in doi: 10.1007/
Mathematics. Springer Cham, 2025. ISBN 978-3-031-85159-99 978-3-031-85160-
Lenaic Chizat, Gabriel Peyré. Bernhard Schmitzer, and François-Xavier Vialard, An interpolat- ing distance between optimal transport and fisher-rao metrics. Foundations of Computational Mathematics, 18(1):1 2018a.
Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and Trançois-Xavier Vialard. Unbalanced opti- mal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 17)(1)): 3090-3123,
20180. 2018b.
Lénaïc Chizat, Maria Colombo, Xavier Fernández-Real, and Alessio Figalli. Infinite-width limit of deep linear neural networks. Communications on Pure and Applied Mathematics, 77(10): 3958-4007, 2024.
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Generative modeling through the semi-dua for- mulation of unbalanced optimal transport. Advances in Neural Information Processing Systems, 36:42433-42443 2023.
Jaemoo Choi. Jaewoong Choi, and Myungjoo Kang. Scalable wasserstein gradient flow for gener- ative modeling through unbalanced optimal transport. In International Conference on Machine Learning, pp. 8629 8650. PMLR, 2024.
14

---

Mathematics of Published as a conference paper at ICLR 2026
Julien Clancy and Felipe Suarez. Classesstinn-fssher. splines. arXiv e-prints, pp. 2022.
Katy Craig and Andrea Bertozzi. A blob method for the aggregation equation. computation, 85(300):1611-771) 2016.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Marco Cuturi, Laetitia Meng- Pana-Papaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint urXiv:2001.23222 2022.
Aniket Das and Dheeraj Nagaraj. Provably fast finite particle variants of svgd via virtual particle stochastic approximation. Advances in Neural Information Processing Systems, 36:49748 49760, 2023.
Constantinos Daskalakis, Andrew Ilyas, Vasilis Syrgkanis, and Haoyang Zeng. Training gans with optimism. In International Conference on Learning Representations, 2018.
Valentin De Bortoli, James Thomton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in neural information processing systems, 34: 34:17695-17709, 2021.
Nabarun Deh, Promit Ghosal, and Bodhisattva Sen. Rates of estimation of optimal transport maps using plug-in estimators via barycentric projections. Advances in Neural Information Processing Systems, 34:29736-29777 297533, 2021.
DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Clau- dio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King. Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojeiech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http:
//cithub.com http://///.ppooeeeeeee.em.m..eee
Michael Ziyang Diao, Krishna Balasubramanian, Sinho Chewi, and Adil Salim. Forward-backward gaussian variational inference via jko the wassersesserttein space. In International Conference on Machine Learning, pp. 7960-7991. PMLR, 2023.
Nira Dyn and Nir Sharon. Subdivision schemes in metric spaces. arXiv preprint arXiv:2509.08070 2025.
Rossin Erbe, Genevieve Stein-O' Oriien, and Flana Fertig. Transcriptomic forecasting with neural ordinary differential equations. Patterns, 4(8). 2023.
Lawrence Evans. Partial differential equations. volume 19. American Mathematical Society, 2022.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Scalable computations of wasserstein barycenter via input convex neural networks. International Conference on Machine Learning, pp. 1581- PMLR, 2021.
Jiaojiao Fan, Qinsheng Zhang, Amirhossein Taghvaei, and Yongxin Chen. Variational wasserstein gradient flow. In International Conference on Machine Learning, pp. 6185-6215, PMLR, 2022.
Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estima- tion and its applications. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. net URL http:: http:://ppeeevviiiee / wwwp:rewwwwww.eeeeeeeeeeeeeeeee f d==mmmml11 [sc Featured Certification.
Lucas CF Ferreira and Julio C Valencia-Guevarz Gradient flows dinendepend functionals in
metric spaces and applications to pdes. Monatshefte für Mathematik, 185(2):231-266 2018.
15

---

Published as a conference paper at ICLR 2026
Davide Fiaschi and Cristiano Ricci. The spatial evolution of economic activities and the emergence of cities. Journal Economic Theory, pp. 105971, 2025.
Alessio Figalli and Federico Glaudo. An Invitation to Optimal Transport, Wasserstein Distances, and Gradient Flows. EMS Press, Berlin, 2nd edition, 2023. ISBN 978-3-98547-050-1 doi:
10.4171/ETB/25. URL https: / /mmm.. prwss/eesss.pp.. http://booo///ee///e///
Chris Finlay, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neu- ral ode: the world of jacobian and kinetic regularization. In International conference on machine learning, pp. 3154-3164, PMLR, 2020.
Mike Fisher, Jorge Nocedal, Yannick Trémolet, and Stephen J Wright. Data assimilation in weather forecasting: case study in pde-constrained optimization. Optimization and Engineering, 10(3): 409-426 2009.
Rémi Flamary, Cédric Vincent-Cuaz, Nicolas Courty, Alexandre Gramfort, Oleksii Kachaiev, Huy Quang Tran, Laurène David, Clément Bonet, Nathan Cassereau, Théo Gnassounou, Eloi Tanguy, Julie Delon, Antoine Collas, Sonia Mazelet, Lactitia Chapel, Tanguy Kerdoncuff, Xizheng Yu, Matthew Feickert, Paul Krzakala, Tianlin Liu, and Eduardo Fernandes Montesuma. Pot python optimal transport (version 0.9.5), 2024.
/tte::/giitttt com/ wwt.eeteettyyttee PythooT/PooTT
Nicolas Fournier and Arnaud Guillin. On the rate of convergence in wasserstein distance of the empirical measure. Probability theory and related fields, 162(3):707-73, 2015.
Milena Gazdieva, Arip Asadulaev, Alexander Korotin, and Evgeny Burnaev. Light unbalanced optimal transport. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024a. URL http::s:///pnnrriieeeeeetet.
wttprrmmmmmcccoooo//oo YK
Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Extremal domain translation with neural optimal transport. Advances Neural Information Processing Systems, 36, 2024b.
Milena Gazdieva, Jaemoo Choi, Alexander Kolesov, Jaewoong Choi, Petr Mokrov, and Alexan- der Korotin. Robust barycenter estimation using cami-unbalanced neural optimal transport. In The Thirteenth International Conference on Learning Representations, 2025. URL https:
//openeview. wwtp:/rwwwwwww...ee.eeeeeeeeeeee wttp::/toooommmmddccccoooom
David Gómez-Castro. Beginner's guide to aggregation-diifusion equations. SeMA Journal, 81(4);
531-587, 2024.
V Gontis, J Ruseckas, and A Kononovičius. A rangerange memory stochastic model of the return in financial markets. Physica A: Statistical Mechanics and its Applications, 389(1):100-100 2010.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(1 (3(11))139-144, 2020.
Erhard Gramer. Probability measures with given marginals and conditionals: I-projections and conditional iterative proportional fitting. Statistics & Risk Modeling. 18(3):311-33) 2000.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf. and Alexander Smola. A kernel two-sample test. The journal of machine learning research, 13(1):723-773, 2012.
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Im- proved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning, 2024a.
Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry P Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. Advances in Neural Information Processing Systems, 36, 2024b.
16

---

Published as a conference paper at ICLR 2026
Nikita Gushchin, Alexander Kolesov, Petr Mokrov, Polina Karpikova, Andrei Spiridonov, Evgeny Burnaev, and Alexander Korotin. Building the bridge of schr" optimal tramsport benchmark. Advances in Neural Information Processing Systems, 36, 2024c.
schr" odinger: A continuous entropic
Nikita Gushchin, Daniil Selikhanovych, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Adversarial schr" schr"odinger bridge matching. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024d. URL https: /ttp///ppeeeeveeeeeee /
net fww.mmrrre
1.3Knn i cuu
Fuqun Han, Stanley Osher, and Wuchen Li. Tensor train based sampling algorithms for ap- proximating regularized wasserstein proximal operators. SIAM/ASA Journal on Uncertainty Quantification, 13(2):7558044, 2025.
Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-level diffusions with generative rnss. In International Conference on Machine Learning, pp. 2417-2426. PMLR, 2016.
Doron Haviv, Aram-Alexandre Pooladian, Dana Pe'er, and Brandon Amos. Wasserstein flow matching: Generative modeling over families of distributions. In Forty-second International Conference on Machine Learning, 2025. URL
ew. net forummrrmm
MRRI6BX3gg 6833gd
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.084155 2016.
Wanli Hong, Yuliang Shi, and Jonathan Niles-Weed, Trajectory inference with smooth schrödinger bridges. In Forty-second International Conference on Machine Learning, 2025. URL https:
//openreview. wtt.:eeeeeooomeeentocootoutuuuue
Chin-Wei Huang, Ricky T. Q. Chen, Christos Tsirigotis, and Aaron Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. In International Conference on Learning Representations, 2021.
Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy. Manifold interpolating optimal-transpon flows for trajectory inference. Advances in neural information processing systems, 15:29705-29718. 2022.
Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in 1989.
and Computation, 18(3):1059-1076
Jan-Christian Hütter and Philippe Rigollet. Minimax estimation of smooth optimal transport maps.
The Annals of Statistics, 49(2):1166 1194, 2021. doi: 10. 10.1214/0-AASS1997 URL https:
Haoming Jiang, Zhehui Chen, Minshuo Chen, Feng Liu, Dingding Wang, and Tuo Zhao. On com- putation and generalization of gans with spectrum control. arXiv preprint arXiv: 1012192. 2018.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker- planck equation. SIAM journal mathematical analysis, 29(1):1-17, 1998.
Ansgar Jüngel. Stefan Portisch, and Antoine Zurek. A convergent finite-volume scheme for nonlo- cal cross-diffusion systems for multi-species populations. ESAIM: Mathematical Modelling and Numerical Analysis, 58(2):759-75 2024.
Jorge Justiniano, Martin Rumpf, and Matthias Erbar. Approximation of splines in wasserstein spaces. ESAIM: Control, Optimisation and Calculus of Variations, 30:64, 2024.
17

---

Published as a conference paper at ICLR 2026
Leonid V Kantorovich. On the translocation of masses. In Dokl. Akad. Nauk. USSR (NS), vol- ume 37, pp. 199-201, 1942.
Kacper Kapusniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael Bronstein, Joey Bose, and Francesco Di Giovanni. Metric flow matching for smooth interpolations on the data manifold. Advances in Neural Information Processing Systems, 37:13501 1335044 2024.
Keisuke Kawano, Takuro Kutsuna, Naoki Hayashi, Yasushi Esaki, and Hidenori Tanaka. Ct-ot flow: Estimating continuous-time dynamics from discrete temporal snapshots. arXiv preprint arXiv:2505.17354 2025.
Evelyn Keller and Lee Segel. Model for chemotaxis. Journal of theoretical biology, 30(2):
225- 234, 1971.
Sergei Kholkin, Grigoriy Ksenofontov, David Li, Nikita Kornilov, Nikita Gushchin, Alexandra Suvorikova, Alexey Kroshnin, Evgeny Burnaev, and Alexander Korotin. Diffusion &amp: ad- versarial schrödinger bridges via iterative proportional markovian fitting. In The Fourteenth International Conference on Learning Representations. 2026. URL htttt::://pppeeeee///
net het//oorum2io==38fCChh3
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint ar1iv:1412.6980, 2014.
Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Genot: stein flow matching with applications to single-cell genomics. Advances in Neural Information Processing Systems, 37:103897-103944, 2024.
Entropic (gromov) wasser-
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Evgeny Burnaev, and Alexander Korotin. Estimating barycenters of distributions with neural optimal transport. In International Conference on Machine Learning, pp. 25016-25041. PMLR, 2024a.
Alexander Kolesov, Petr Mokrov, Igor Udovichenko, Milena Gazdieva, Gudmund Pammer, Anas- tasis Kratsios, Evgeny Burnaev, and Aleksandr Korotin. Energy-guided continuous entropic barycenter estimation for general costs. Advances in Neural Information Processing Systems, 37:107513-107546, 2024b.
Nikita Kornilov and Alexander Andreevich Korotin. On the equivalence of optimal transport prob- lem and action matching with optimal vector fields. Uspekhi Matematicheskikh Nauk, 80(5): 185-186, 2025.
Nikita Kornilov, Petr Mokrov, Alexander Gasnikov, and Aleksandr Korotin. Optimal flow match- ing: Learning straight trajectories in just one step. Advances in Neural Information Processing Systems, 37: 37:104180-104204, 2024.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. International Conference on Learning Representations, 2021a. URL /tLp:tttss:////eeeee/e
net. htt.:/ooommmmmddbbooodeoe EXsa
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in neural information processing systems,
2021b.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021c.
Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative net- works for barycenter estimation. In Thirty-Sixth -ixth Conference on Neural Information Processing Systems, 2022. htt.ss:::: www.oeeevvviieeeeeoooo//e//eeeee www./nee//eee ftt./formmmmm/ d= Enzx EnzxTnaMN
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernll neural optimal transport. In International Conference on Learning Representations, 2023a. URL https: //ttopts///ot///eeeeeeee //tpp//eve/....///// net/forru d=zu MHt.Uma 4
18

---

Published as a conference paper at ICLR 2026
Alexander Korotin, Daniil Selikhanovych, and Evgeny Burneev. Neural optimal transport. In
The Eleventh International Conference on Learning Representations, 2023b. URL https:
//openzeview. wwwt//eeteeenneeeeeeeeeeem http://oouum211-EEEBBB.L......
Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light achrodinger In The Twelfth International Conference on Learning Representations, 2024. URL ltttps://openrevicee
Takeshi Koshizuka and Issei Sato. Neural lagrangian schr schr\"(()) "(ogern bridge: Diffusion modeling for population dynamics. In The Eleventh International Conference on Learning Representations, 2023.
Chrysoula Kosma, Giannis Nikolentzos, George Panagopoulos, Jean-Mare Steyaert, and Michalis Vazirgiannis. Neural ordinary differential equations for modeling epidemic spreading. Transactions on Machine Learning Research, 2023.
Leonenko Kozachenko. Sample estimate of the entropy of a random vector. Probl. Pered. Inform., 23:9, 1987.
Rahul Krishnan, Uri Shalit, and David Sontag. Structured inference networks for nonlinear state space models. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017.
Grigoriy Ksenofontov and Alexander Korotin. Categorical schrödinger bridge matching. In Forty-second International Conference Machine Learning, 2025. URL https://
Marc Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via wasserstein gradient flows. Advances in Neural Information Processing Systems,
35:14434-14447, 2022.
Jeffrey M Lane and Richard F Riesenfeld. A theoretical development for the computer genera- tion and display of piecewise polynomial surfaces. IEEE Transactions on Pattern Analysis and Machine Intelligence, (1):35 46, 1980.
Hugo Lavenant, Stephen Zhang, Young-Heon Kim, Geoffrey Schiebinger, et al. Toward a math- ematical theory of trajectory inference. The Annals of Applied Probability, 34(1A))428-500 2024.
Wonjun Lee, Li Wang, and Wuchen Li. Deep jko: time-implicit particle methods for general non- linear gradient flows. Journal of Computational Physics, 514:113187, 2024.
Ruikun Li, Jiazhen Liu, Huandong Wang, Qingmin Liao, and Yong Li. Weightflow: Learning
stochastic dynamics via evolving weight of neural network. arXiv preprint S.Ii:22508000555 2025a.
Xuechen Li, Ting-Kan Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 3870-3882, PMLR, 2020.
Zihao Li, Zhichen Zeng, Xiao Lin, Feihao Fang, Yanru Qu, Zhe Xu, Zhining Liu, Xuying Ning, Tianxin Wei, Ge Liu, al. Flow matching meets biology and life science: survey. arXiv e-prints, pp. arXiv-2507, 2025b.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In 11th International Conference on Learning Representations, ICLR 2023. 2023.
Guan-Horng Liu, Yaron Lipman, Maximilian Nickel, Brian Karrer, Evangelos Theodorou, and Ricky T. Q. Chen. Generalized schrödinger bridge matching. In The Twelfth International Conference on Learning Representations, 2024. URL/https: httphhtpps//ppeeeviieeeee /wwp/reveeneeeeeee/oe/.e wwt. formm?
id=Soi smgox7zgeee
Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems, 29, 2016.
19

---

Published as a conference paper at ICLR 2026
Xingchao Liu and Lemeng Wu. Learning diffusion bridges constrained domains. In international conference on learning representations (ICLR), 2023.
Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023. URL httthttp:////
Zewen Liu, Xiaoda Wang, Bohan Wang, Zijie Huang, Carl Yang, and Wei Jin. Graph odes and be- yond: A comprehensive survey integrating differential equations with graph neural networks. In Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. pp. 61 68-6688, 2025.
Evan Z Macosko, Anindita Basu, Rahul Satija, James Nemesh, Karthik Shekhar, Melissa Goldman, Itay Tirosh, Allison R Bialas, Nolan Kamitaki, Emily M Martersteck, et al. Highly parallel genome-wide expression profiling of individual cells using nanoliter droplets. Cell, 161(5))1202- 1214, 2015.
Bertrand Maury, Aude Roudneff-Chuppe and Filippo Santambrogio. macroscopic crowd motion model of gradient flow type. Mathematical Models and Methods in Applied Sciences, 20(10): 1787-1821, 2010.
Bertrand Maury, Aude Goudneff-Chuping Filippo Santambrogio, and Juliette Venel. Handling con- gestion in crowd motion modeling. Networks and Heterogeneous Media, 6(3):485-519, 2011.
Robert McCann. Existence and uniqueness of monotone measure--peesrrving maps. Duke
Mathematical Journal, 80(2):309, 1995.
Boris Medina-Salgado, Eddy Sánchez- Beachez-Deaarrz, Pilar Pozos-Parra, and Javier E Sierra. Urban traffic flow prediction techniques: A review. Sustainable Computing: Informatics and Systems, 35:100739, 2022.
Nicholas Metropolis and Stanislaw Ulam. The monte carlo method. Journal of the American statistical association, 44(247):3353441 1949.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, 2018.
URL /BRLLttpps//opeeeeeeee.
net forum? 20-50
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, and Evgeny Burnaev. Large-scale wasserstein gradient flows. Advances in Neural Information Processing Systems, 34:15243-15256, 2021.
Petr Mokrov, Alexander Korotin, Alexander Kolesov, Nikita Gushchin, and Evgeny Burnaev. guidgy-guided entropic neural optimal transport. In The Twelfth International Conference on Learning Representations, 2024. URL bttps:://penreviiew net/ net/forum?id=
@6tUs2eVV>
Gaspard Monge. Mémoire sur la théorie des déblais et des remblais. Mem. Math. Phys. Acad. Royale Sci., pp. 666-704, 1781.
Kevin R Moon, David Van Dijk, Zheng Wang, Scott Gigante, Daniel Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, et al. Visualizing structure and transitions in nigh-dimensional biological data. Nature biotechnology, (7(11):1482 2019.
Alex Morehead, Lazar Atanackovic, Akshata Hegde, Yanli Wang. Frimpong Boadu, Joel Sel- varaj. Alexander Tong. Aditi Krishnapriyan, and Jianlin Cheng. How to go with the flow: flow matching in bioinformaties and computational biology. Authorea Preprints, August 2025.
doi: 11533////..11782400.. 175382408 8 0944666666770
URL latp::// dx.
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Lear- ing stochastic dynamics from samples. In International conference on machine learning, pp. PMLR, 2023.
20

---

Published as a conference paper at ICLR 2026
Kirill Neklyudov, Rob Brekelmans, Alexander Tong, Lazar Atanackovic, Qiang Liu, and Alireza Makhzani. computational framework for solving wasserstein lagrangian flows. In Proceedings of the s1st International Conference Machine Learning, pp. 37461-37485, 2024.
Atsushi Nitanda. Improved particle approximation error for mean field neural networks. In
The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL
net/ forum? wedddd7666 76Niiddcct
Ivan V Oseledets. Tensor-train decomposition. SIAM Journal on Scientific Computing, 33(5):
2295-233 2317, 2011.
F Otto. The geometry of dissipative evolution equations: the porous medium equation. Comm Partial Differential Equations, 26:101-174, 2001.
Victor M Panaretos and Yoav Zemel. Statistical aspects of wasserstein distances. Annual review of statistics and its application, 6(1):405-431, 2019.
Neal Parikh, Stephen Boyd, et al. Proximal algorithms. Foundations and trends@ in Optimization, 1()):127-239) 2014.
Min Sue Park, Cheolhyeong Kim, Hwijae Son, and Hyung Ju Hwang. The deep minimizing move- ment scheme. Journal of Computational Physics, 494: 12518, 2023.
A Paszke. Pytorch: An imperative style, high-performance deep learning library. arXiv preprint arXiv:1912.01703, 2019.
Stefano Peluchetti. Diffusion bridge mixture transports, schrödinger bridge problems and generative modeling. Journal of Machine Learning Research, 24(374):1-51, 2023a.
Stefano Peluchetti. Non-denoising forward-time diffusions. 2023b.
arXiv preprint arXiv:2312.14589,
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ Machine Learning, 11(5-6):35-6077, 2019.
Simon Plazotta and Jonathan Zinsl. High-frequeng, limit of non-auuonomous gradient flows of functionals with time-periodic forcing. Journal of Differential Equations. 66166)66806-6 2016.
Jonathan R Potts Aggregation-ddiftssion in heterogeneous environments.
arxiv preprint
arXiv:2409.1014 10147, 2024.
Jonathan R Potts and Kevin J Painter. Distinguishing between long- transiransient and asymptotic states in a biological aggregation model. Bulletin of Mathematical Biology, 86(3):28, 2024.
Ilan Price, Alvaro Sanchez-Gonzalez, Ferran Alet, Tom R Andersson, Andrew El-Kadi, Dominic Masters, Timo Ewalds, Jacklynn Stott, Shakir Mohamed, Peter Battaglia, et al. Probabilistic weather forecasting with machine learning. Nature, 637(8044):84-90) 2025.
Hannes Risken and Till Frank. The Fokker-Plande Equation: Methods of Solution and Applications, volume 18. Springer Science & Business Media, 1996.
R. Tyrrell Rockafellar. Convex analysis. Princeton Mathematical Series. Princeton University Press, Princeton, N. J., 1970.
Martin Rohbeck, Charlotte Bunne, Edward De Brouwer, Jan-Christian Huetter, Anne Biton, Kelvin Y. Chen, Aviy Regev, and Romain Lopez. Modeling complex system dynamics with flow matching across time and conditions. In The Thirteenth International Conference on Learning Representations, 2025. URLhttps: http::////ppeeeiiie/e net wtwpareewweweeee/oouuuueeme. htw.uummiiidwwwwwwmmooooo
Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. In Proceedings of the 10th International Conference on Learning Representations (ICLR), 2022. UR..https://pperrve ew. net htt.:/oormmm--0mm wett://oormmm..055tti7......
21

---

Published as a conference paper at ICLR 2026
Yossi Rubner, Carlo Tomasi, and Leonidas J Guibas. A metric for distributions with applica- tions to image databases. Sixth international conference on computer vision (IEEE Cat. No. 98CH36271), pp. 59 66. IEEE, 1998.
Adil Salim, Anna Korba, and Giulia Luise. The wasserstein proximal gradient algorithm. Advances in Neural Information Processing Systems, 33:12356-12366 2020.
Filippo Santambrogio. Optimal transport for applied mathematicians, volume 87. Springer, 2015,
Filippo Santambrogio. {Euclidean, metric, and Wasserstein) gradient flows: an overview. Bulletin of Mathematical Sciences, 7:87-154, 2017.
Daniel Saragih, Deyu Cao, and Tejas Balaji. Flows and diffusions on the neural manifold. arXiv e-prints, pp. arXiv-2507, 2025.
Simo Särkkä and Arno Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Christopher Scarvelis and Justin Solomon. Riemannian metric learning via optimal transport. In The Eleventh International Conference on Learning Representations, 2023. URL https:/7
openreview.woo net htte://ooommmmmd1d3mma333a WEz
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transport analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928 943, 2019.
Yutong Sha, Yuchi Qiu, Peijic Zhou, and Qing Nie. Reconstructing growth and dynamic trajectories from single-cell transcriptomics data. Nature Machine Intelligence, 6(1):25 39, 2024.
Yunyi Shen, Renato Berlinghieri, and Tamara Broderick. Multi-marginn schrödinger bridges with itcrative reference refinement. The 28th International Conference on Artificial Intelligence and Statistics, 2025. URL httpt::t:///pseeeeeee
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnad Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36:62183-66223,, 2023.
Fabio Sigrist, Hans Künsch, and Werner A Stahel. Stochastic partial differential equation based modelling of large space-time data sets. Journal of the Royal Statistical Society Series B: Statistical Methodology, 77(1):3-33 2015.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
Yuhao Sun, Zhenyi Zhang, Zihan Wang, Tiejun Li, and Peijic Zhou. Variational regularized unbal- anced optimal transport: Single network, least action. arXiv preprint 11X1::2505511823 2025.
Taiji Suzuki, Atsushi Nitanda, and Denny Wu. Uniform-in iintamm-ii-timee propagation of chaos for the mean-field gradient langevin dynamics. In The Eleventh International Conference on Learning Representations, 2023. URL https:: /
form JSCU 9TBUn
Chandan Tankala, Dheeraj M Nagaraj, and Anant Raj. Beyond propagation of chaos: A stochastic algorithm for mean field optimization. arXiv preprint arXiv:2503.13115, 2025.
Roman Tarasov, Petr Mokrov, Milena Gazdieva, Evgeny Burnaev, and Alexander Korotin. A statistical learning perspective semi-dual adversarial neural optimal transport solvers. In
The Fourteenth International Conference on Learning Representations, 2026. URL https:
Antonio Terpin, Nicolas Lanzetti, Martín Gadea, and Florian Dorfler. Learning diffusion light-
speed. Advances Neural Information Processing Systems, 37:6797-6832, 2024.
Gerald Teschl. Ordinary differential equations and dynamical systems, volume 140. American Mathematical Soc,, 2012.
22

---

Published as a conference paper at ICLR 2026
Panagiotis Theodoropoulos, Augustinos D Saravanos, Evangelos A Theodorou, and Guan- Horng Liu. Momentum multi-marginal schr\" odinger bridge matching. arXiv preprint arXiv:2506.10166. 2025.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference on machine learning, pp. 9536-9536. PMLR, 2020.
Alexander Tong, Manik Kuchroo, Shabarni Gupta, Aarthi Venkat, Beatriz P San Juan, Laura Rangel, Brandon Zhu, John G Lock, Christine L Chaffer, and Smita Krishnaswamy, Learning transcrip- tional and regulatory dynamics driving cancer cell plasticity using neural ode-based optimal trans- port. bioRxiv, pp. 2023 03, 2023.
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, pp. 1-34, 2024a.
Alexander Tong, Nikolay Malkin, Kilian Fatras, Lazar Atanackovic, Yanlei Zhang, Guillaume Huguet, Guy Wolf, and Yoshua Bengio. Simulation-fin schrödinger bridges via score and flow matching. In The 27th International Conference Artificial Intelligence and Statistics, pp. 1279- 1287. Journal of Machine Learning Research- Track, 2024b.
Théo Uscidda and Marco Cuturi. The monge gap: A regularizer to learn all transport maps. In International Conference on Machine Learning, pp. 34709-3440 34733. PMLR, 2023.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9): 1134, 2021.
Juan Luis Vázquez. Asymptotic behaviour methods for the heat equation. convergence to the gaus- sian. arXiv preprint arXiv:17066 10034, 2017
Yogesh Verma, Markus Heinonen, and Vikas Garg. ClimODE: Climate and weather forecast- ing with thysics-informen neural ODEs. In The Twelfth International Conference on Learning Representations, 2024. URL
net. htt.:pawww.a.a.e....e...oooooaan
Alexander Vidal. Samy Wu Fung, Luis Tenorio, Stanley Osher, and Levon Nurbekyan. Taming hyperparameter tuning in continuous normalizing flows using the jko scheme. Scientific reports, 13(1):4501, 2023.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2008
Wei Wan, Yuejin Zhang, Chenglong Bao, Bin Dong, and Zuoqiang Shi. A scalable deep learning approach for solving high-dimensional dynamic optimal transport. SIAM Journal Scientific Computing, 2023.
Dongyi Wang, Yuanwei Jiang, Zhenyi Zhang, Xiang Gu, Peijie Zhou, and Jian Sun. Joint velocity- growth flow matching for single-cell dynamics modeling. arXiv preprint arXiv:2505.13413 2025.
Xiaomei Wang, Qi An, Zilong He, and Wei Fang. A literature review of social network analysis in epidemic prevention and control. Complexity, 2021()):381221, 2021.
Weinan, Tiejun Li, and Eric Vanden-Eijnden. Applied stochastic analysis, volume 199. American
Mathematical Soc., 2021.
Shang Wu and Yazhen Wang. Computational and statistical asymptotic analysis of the jko scheme for iterative algorithms to update distributions. arXiv preprint arXiv:2501.0644 2025.
Chen Xu, Xiuyuan Cheng, and Yao Xie. Normalizing flow neural networks by jko scheme. Advances in Neural Information Processing Systems, 36:47379-47400 2023.
Luxuan Yang, Ting Gao, Yubin Lu. Jinqiao Duan, and Tao Liu. Neural network stochastic dif- ferential equation models with applications to financial data forecasting. Applied Mathematical Modelling, 115:279-299, 2023.
23

---

Published as a conference paper at ICLR 2026
Stephen Zhang, Suryanarayana Maddu, Xiaojie Qiu, and Victor Chardès. Inferring stochastic dy- namics with growth from cross-sectional data. arXiv preprint arXiv:2505.13197, 2025a.
Yanbo Zhang and Michael Levin. Equilibrium flow: From snapshots to dynamics. arXiv preprint AxKi::5509.19990 2025.
Zhenyi Zhang, Tiejun Li, and Peijie Zhou. Learning stochastic dynamics from snapshots through regularized unbalanced optimal transport. In The Thirteenth International Conference on Learning Representations, 2025b. URL U ner/forum we=iid==gdddl0 ru
Zhenyi Zhang, Yuhao Sun, Qiangwei Peng, Tiejun Li, and Peijic Zhou. Integrating dynamical systems modeling with spatiotemporal scrna-seq data analysis. Entropy, 27(5):453, 2025c.
Zhenyi Zhang, Zihan Wang, Yuhao Sun, Tiejun Li, and Peijie Zhou. Modeling cell dynamics and interactions with unbalanced mean field schr odinger bridge. arXiv preprint arXiv:2505, 2025d.
Zhenyi Zhang, Zihan Wang, Yuhao Sun, Jiantao Shen, Qiangwei Peng, Tiejun Li, and Peijie Zhou. Deciphering cell-fate trajectories using spatiotemporal single-cell transcriptomic data. Authorea Preprints, September 2025e. doi: 10.22541/au 77tp://2200177775334444449999999
45794422244499 URL http://dx.
doi Goi.org/10 2054222241aa //all/aauu 11t...175733333 457954334447792227977
Qinqing Zheng, Matt Le. Neta Shaul, Yaron Lipman, Aditya Grover, and Ricky TQ Chen. Guided flows for generative modeling and decision making. arXiv preprint arXiv:2311.138 133443, 2023.
Huminhao Zhu, Fangyikang Wang, Chao Zhang, Hanbin Zhao, and Hui Qian. Neural sinkhorn gradient flow, arXiv preprint arXiv:2401.14069 2024.
24

---

Published as a conference paper at ICLR 2026
CONTENTS
A Extended Discussions
25
Related works
25
A.2 Time-varying potentials discussion
28
Paired vs. unpaired setups
28
Extended Comparisons
29
B. Synthetic comparisons
30
B.2
Learning single-cell dynamics
Experimental Details
35
Metric computation details
35
C.2
Training details
36
Entropy estimation details
37
Optimizer
38
Network architecture
38
C.6
Hardware
38
7 Functionals
38
D Proofs
38
A EXTENDED DISCUSSIONS
We first review related work in Appendix and then discuss the validity of vtre-varying potentials proposed by Terpin al. (2024) in Appendix A.2 in contrast the stationary setting considered in this paper. Finally. Appendix A.3] we discuss potential reasons for the performance gap between the paired and unpaired setups introduced in Appendix 5.1
A.1 RELATED WORKS
In this section, we first provide an extended discussion of related work on learning population dy- namies (also known trajectory inference) in Appendix followed by an overview of appli- cations of the JKO scheme (Jordan et al. 1998) in deep learning in Appendix A.1
A.1.1 LEARNING POPULATION DYNAMICS
In contrast our approach, which models population dynamics as Wasserstein gradient flow (WGF) (Ambrosio et al. 2008) of the free-energy functional (5) (Gómez- Castro 2024), several other fruitful approaches have been proposed the literature. We review these altematives below.
RNNs and Neural ODEs. One of the first works on learning population dynamics is (Hashimoto at. 2016), which used recurrent neural networks (RNNs) to learn the SDE (7) as WGF of poten- tial energy functional. Later, (Chen et al.) 2018a CNF) introduced neural ODEs as the continuous- time limit RNNs for learning ODE-based dynamics, and (Li et al. 2020 subsequently extended this approach to handle SDEs. Building these ideas, (Erbe et al. 2023 RNAForecaster ap- plied neural ODEs to predict future transcriptomic states single cells (Battich et al. 2020). More recently, (Li al. 2025a) proposed Wei Weight Foww which models stochastic population dynamics by learning the continuous evolution of neural network weights that parameterize distributions at each time point, rather than modeling the dynamics directly in the state space. Their key idea is
25

---

Published as a conference paper at ICLR 2026
to project the evolving distribution into the weight space of backbone neural network trained to represent the probability distribution at each snapshot.
Static OT. Another line of work models population dynamics by directly learning K independent transport maps from time-snapshot data, using the static optimal transport formulation T or 2,. without attempting recover the unifying energy functionals guiding the system evolution. For example, (Schiebinger et al. 2019 Waddingt
employs discrete transport maps, whereas
(Bunne et al. 2023 CellOT) learns neural transport maps, as in (Fan et al. 2023 (6), (Rout
et al. [022] (14), and (Korotin et al. 2023b (15)) see Appendix A.2 for more details. Such
approaches are typically called Neural OT solvers (Korotin et al. adl,2222bb, 2023a
2
Mokrov et al. 2024]
| Gushchin et al. al.]| [2244]] Tarasov et al. 2026] Asadulaev et al. (0224) and have recently obtained many
different extensions to certain extended OT settings, in particular to the barycenter problem (Kolesov) et al. 2024a Gazdieva et al., 2025 Korotin et al. 2022, 2021c; Fan et al. 2021) which can also be relevant for modeling data evolution time (Kolesov et al.) 2024b) Appendix C.5). However, they are not that widespread in applications to population dynamics.
Dynamic OT. (Tong et al., 2020, TrajectoryNett) models population dynamics using the dy-
namic optimal transport formulation of (Benamou & Brenier 2000), parameterized by regularized CNF. The work (Huguet et al. 2022, MIOFIow later extended this approach by incorporating a geodesic autoencoder to better capture manifold structure. More recently, (Wan et al. 2023) pro-
posed computational techniques that make dynamic OT scalable to higher-dimensional settings.
Flow Matching. The works in this paragraph extend Flow Matching (FM) techniques (Lipman
et al. 2023: Albergo & Vanden- Eiijnden 2023) Liu et al.) 2023) in several directions. Conditional
Flow Matching (Tong et al. 20244a CFM) and Optimal Flow Matching (Kornilov al. 2024 OFM) introduce simulation-free objectives for learning deterministic and optimal flows. Building on the celebrated work of (Otto 2001), the Wasserstein space can be nominally (i.c., heuristicallly) viewed as Riemannian manifold. Motivated by this observation, Wasserstein Flow Matching (Haviv et al., 2025 WFM) applies the idea of Riemannian Flow Matching (Chen & Lipman 2024 RFM) to the Wasserstein space, i.e., performing FM over a distribution of distributions. Further developments, such as Meta Flow Matching (Atanackovic et al. [2025] Meta FM), embed the population of sam- ples with Graph Neural Network (GNN) (Liu et al. (025) and use these embeddings as con- ditioning inputs for the learning vector field. Finally, Multi-Marginal Flow Matching (Rohbeck et al. 2025 MMFM) constructs flow using smooth spline- based interpolation across time points and conditions, and regresses it with a neural network under the classifier free guided Flow Matching framework (Zheng et al., 2023). For broader overview of FM for learning population dynamics, see the recent surveys (Morehead et al.
2025 Li et. 2025b).
AM) introduced Action Matching (AM), method that optimizes an action-gap objective. In contrast to Flow Matching, AM learns a vector field expressed as the gradient of an action, Vst, which uniquely defines the velocity field that transports particles optimally in the sense of optimal transport. Building on this idea, (Berman el al. 2024) conditioned the action $1 on physical parameters μu. used (Berman & Peherstorfer 2024 CoLoRa) for efficient action parametrization, and demonstrated its utility for surrogate modeling of classical numerical solvers, enabling fast prediction of system behavior across different physics parameter settings. Most recently, (Neklyudov et al. 2024], WLF) proposed Wasserstein Lagrangian Flows, a unifying framework that minimizes Lagrangian action functionals over the space of probability densities rather than the ground space, thereby encapsulating AM as special case. Interestingly, action matching can be modified just like the flow matching (Kornilov et al. 2024 OFM) to obtain the solutions to optimal transport after training, (Kornilov & Korotin 2025 OAM).
Action Matching. (Neklyudov et al., 2023
Schrödinger bridges. Two seminal works, Vargas et al. 2202 IPF (G2)) and De Bortoli et al. (2021 IPF (NN) ), proposed solving the Schrödinger bridge problem using the Iterative Propor- tional Fitting (IPF) algorithm (Gramer) 2000) in application to generative modeling, which al- ternates between forward and backward processes. The first approach employed Gaussian pro- cesses as parametrization, whereas the second relied on score-based neural networks, see (Gushchin ct al.] 2024c) for a quick survey and benchmark. Later, bridge matching methods were developed [Peluchetti] 2023b 202b3b Ksenofontov & Korotin 2025) Kholkin et al. [2026] Tong et al. [0024]] Carrasco et al. (026). Ad-
Liu & 2023 Shi et al. [2023] 10023/oone Korotin et al., 22224] 2024//ushhhhnn Gushchin et al. 2024d a
ditional contributions include generalized setups for Schrödinger bridges (Koshizuka & Sato 2023
Liu et al. 2024). multi- multi-marginal generalizations of the problem (Chen et al.) 2023
DMSB),
(Shen
26

---

Published as a conference paper at ICLR 2026
et al.
2025
MMSB), (Hong et al. 2025), and commentum-accelenann formulations ( (Thoodoropoulos 3MSBM) which are particularly suitable for trajectory inference problems.
et al. 2025
Unbalanced OT. Several works jointly model marginal transitions and growth dynamics by mini- mizing the action in the Wasserstein-Fistee Fisher-Raoo (WFR) metric (Chizat et al, 2018a), i... by solving the dynamical unbalanced optimal transport problem (Chizat et al. 2018b). (Tlng et al. 2023) pro- posed BEMIOFI fIMw, flow, an extension of (Huguet et al. [022] MIOf low), that incorporates a neural network to predict cell growth and death rates continuously in time, thereby augmenting optimal transport with population-size dynamics. Later, (Chen et al., 2022b) formulated an unbalanced Schrödinger Bridge problem, which was applied practice by (Zhang et al., 2025b, DeepRUOT), following the developments in (Sha et al.) 2024] TIGON). DeepRUOT was further simplified into single-network formulation, Var-RUOT (Sun et al., 2025), by exploiting optimality conditions; extended to include interaction modeling in (Zhang et al. 2025d] CytoBric idg);; and adapted to use two independent networks to separately model distributional drift and mass growth (Wang et al.)
2025 VGFM), (Zhag et al. 2025a). In different direction, (Klein al.) 2024 GENOT) generalized OT to simultaneously handle stochasticity, entropy regularization, and unbalanced mass transport, enabling applications cross-modal and heterogeneous data. For broader survey of trajectory inference methods and their applications to biology data, see (Zhang et al.] 2025e c). For com- pleteness, we note that there exists line of works developing efficient solvers for unbalanced OT Gazdieva et al. 2024b a Choi et al. 2023), but they do not consider applications to biological data.
The concurrent work (Andrade et al. 2025 iJKO) also connects inverse optimization with the JKO scheme (8) but in the unbalanced setting. However, their focus is on sample complexity, the method addresses only potential energy, and experiments are restricted to low-dimensional cases.
Riemannian perspective. (Scarvelis & Solomon, 2023) proposed learning a metric tensor A(x) that minimizes the average I-Wasserstein distance on the learned manifold between pairs of con- secutive population snapshots. Furthermore, (Kapusniak et al. 2024 MFM) introduced Metric Flow Matching (MFM), which learns interpolants Ft.n = (1-ttl1a t)xg + tæ + t(1 (-1-t)ee t), .1). where n are the parameters of a neural network ftty providing a nonlinear "correction" to straight-line inter- polants (Albergo& & Vanden Vanden-Eijnden, (023). The resulting velocity field minimizes date-dependent Riemannian metric, assigning lower transport cost to regions with higher data density.
Splines Wasserstein Space. In the seminal work of (Schiebinger et al. 2019), a piecewise lin- ear OT interpolation method was proposed to infer cell trajectories. Subsequent works introduced higher-order piecewise polynomials (c.g., cubic splines) in Wasserstein space: (Chen et al. 2018b) later
Benamou et al. (2019) formulated global cubic spline minimization problem, which tended by (Chewi et al. 2021) to use Euclidean interpolation algorithm in TR: D after a finding optimal
Monge map 2 between samples from consecutive measures, with further refinements by (Clancy & Suarez 2022 Justiniano et al.] 2024). More recently, (Banerjee et al.) 2025 WLR) proposed the
Wasserstein Lane-Riesenfeld method, leveraging the classical subdivision algorithm of (Lane & Riesenfeld 1980). In parallel, (Dyn & Sharon 2025) developed subdivision schemes for general
metric spaces, including the Wasserstein space. (Baccou & Liandrat 2024 Kawano al.) 2025)
combines subdivision schemes with OT.
Dynamics from stationary distributions. A recent line of research aims to recover meaningful dynamical structure from static data alone. The Equilibrium Flow framework of (Zhang & Levin 2025) addresses this problem by learning vector field whose flow preserves the observed stationary distribution p. Their method enforces the invariance condition V [p(xr, f)u(a, ()] 0 using score- based estimates of V logp, p. enabling the recovery plausible dynamics even without temporal information or trajectories. Despite this minimal supervision, the learned flows can be used for inverse design, where one constructs dynamics that yield a desired target distribution.
1.1.2 JKO SCHEME MEETS DEEP LEARNING
Generative Modeling. (Vidal et al. 2023 JKO- F low) and (Xu et al. 2023 JKO iFlcw) rein-
terpret the JKO scheme g through the framework of CNTs (Chen et al. 2018b), applying it to 2D synthetic data and image generation tasks, respectively. Subsequently, (Cheng et al. 2024) provided theoretical convergence guarantees for this approach. (Choi et al. 2024 S-JKO) further acceler- ated WGF-based generative modeling by leveraging a semi-dual anbalanced OT formulation. They constructed the WGF of an f-divergence D, and used an (Song et al. 2021 NCSN++) backhone
27

---

Published as a conference paper at ICLR 2026
similar to (Zhu et al. 2024 NSGF), who designed generative model as the WGF with respect to the Sinkhorn divergence (Peyré et al. 2019).
Tensor Train (TT) (Oseledets 2011). (Aksenov & Eigel 2025 introduces method for approxi- mating probability distributions in Bayesian inversion by minimzing KL-based functional via an entropically regularized JKO scheme. The approach discretizes the resulting coupled heat equations on a high-dimensional grid and solves them using accelerated fixed-point methods combined with low-rank TT representations. Related work includes (Chertkov & Oseledets) (021)) which applies TT approximations to the Fokker-Planck equation with drift and diffusion, and (Han et al. 2025). which employs -egroyy-regularized proximal steps with TT-cross for particle- based evolution.
Variational Inference. (Lambert et al., 2022 and (Diao et al., 2023) study variational inference through the lens of Soures-Wassersteer gradient flows. (Cheng al., [023 GWG) introduce gener- alized minimizing movement scheme on the space Per. (I.".. where the transport cost is defined as en (I, y) = 9( = h with g continuously differentiable Young function, and apply this framework to particle- particle-based variational inference.
Datasets/ /eeistss Learning, (Bonet et al. 2025 propose explicit scheme that is computationally more efficient in practice than the implicit JKO scheme on the space (22 (R))), and apply this approach to modeling flow datasets viewed as random measures. (Saragih et al. (2025) train a meta-model, based on JKO as well as other flow- and diffusion- based approaches, that generates dataset-conditioned classifiers by producing the neural network weights in single forward pass.
A.2 TIME-VARYING POTENTIALS DISCUSSION
In this section, we examine the validity of using vime-varying potentials, as proposed in (Terpin et al, 2024 $4.4). We demonstrate that under this formulation, the original problem of reconstructing energy functionals (Bunne et al., 2022b) via JKO Scheme reduces to learning K independent neural optimal transport maps (Fan al. 2023 (6)) (Rout al. 2022 (1.(14) between the data snapshots Pk and Pk++1 To illustrate this, consider the loss in (II) and assume that, instead of a single
functional J, we now assign a separate only potential energy functional JrK Sx V V*(x) dp for each time step k:
max min C(J T def Jí
K-1
(p* - (PE()) + 1 / ||x ||x- I (z) (2)) dz) 27
max min T
T
JE
then becomes clear that the loss optimization decomposes into independent terms L (Tpr. T*:;
1 (TT,, del The (T* ( + 25 / 1| (x))p((()) dx + max min.
(18)
Jr.
By replacing V with the dual potential fú and T with To. we recover Equation (6) from Fan
et al. 2023), which aligns with the objective proposed in (Bunne et al. 2023 Eq. (9)) for modeling single-cell dynamics. Thus, (Terpin al.) 2024) introduces new strategy for minimizing this loss
rather than directly addressing the recovery of ground-truth energy functionals (Bunne et al. 2022b)
Fortunately, the theoretical framework developed by Ferreira & Valencia-Guevara (2018) and later extended by Pllazotta & Zinsl (2016) establishes the validity of this approach even for time-varying free energy functionals (5), not only for the potential energy case, in the limit the time discretiza- tion tends to zero. However, this connection is not acknowledged in (Terpin et al. 2024).
A.3 PAIRED VS. UNPAIRED SETUPS
In this section, we discuss the substantial performance gap observed between the paired and unpaired setups, as illustrated in Table We suggest that this discrepancy stems from the intrinsic connection between the considered solvers, iJKOnet and JKOnet., ", and the underlying OT problem (
Both solvers are designed to recover the optimal transport mappings k. between consecutive dis- tributions Pk and Pk-1, satisfying the pushforward relation The 4Pp Pk+11 In JKOnet", these mappings are obtained explicitly by precomputing the OT maps between Pk and Pk+1, whereas in our method, they are learned implicitly via neural networks T
28

---

Published as a conference paper at ICLR 2026
In practice, the continuous distributions Pk replaced by their empirical counterparts, constructed from training samples { ......, , x = X ~ Pk- This is precisely where the distinction be- tween the paired and unpaired setups becomes critical:
Paired. When samples are paired, they follow the ground-truth transport mappings, i.e.,
(x))) = =i, (*+((+)= = zz
In this case, estimating T becomes substantially simpler and can be viewed as supervised regression problem over the given sample pairs. Unpaired. In contrast, when the samples Xs. k {o..... K}. are mutually independent, the quality of the estimated OT maps degenerate significantly. This degradation stems from the high sample complexity of Wasserstein-2 distances and maps (known to be relatively poor, see Hütter & Rigollet (001)). ) as well as from the increased variance of the estimated transport mappings.
deeper theoretical analysis of this gap between paired and unpaired setups remains an interesting direction for future work. To quantify how performance changes when transitioning from unpaired to paired data, we conducted the following ablation study.
Experimental Setup. Following Appendix B.I.I. we used the same set of potentials and focused solely on the potential parametrization. In the 2D setting with K = 5, we trained for 5K iterations using N 2K training samples per time step and a 40% test split. The fraction of full trajectories in the training set was gradually increased from 0,0 (unpaired) to 1.0 (paired) in steps of 0.25.
Results. Figure 3 shows the results, As the fraction of paired samples decreases, the performance of JKOnet" degrades more rapidly than that of our method. For the Bohachevsky potential, both methods show reduced accuracy in both paired and unpaired scenarios.
JK0nnt
OJet (Ours)
Buhuchorahi
Double exponeriall
Fix
Flowers
Fredrrn
3:56
0.00
6:15
3:4
0:0
0025
0.05
0000
3:0
0.00
0.00
8:75
0.3
3:58
0.00
C,00P
0.0-
0.15
3.8
60
823
0.75
022
03
0.9
1.0
025
035 0.19
0:23
0.79
09 023 05 075
Schure
Vhery pMMou
20-44g rdge
030
3.44
804
006
Figure 3: Ablation study comparing paired (1.0) and umpaired (0.0) setups. As the proportion of paired samples (i.e. full trajectories in the dataset) decreases, the performance of JKOnet" degenerate more rapidly than that of our method.
B EXTENDED COMPARISONS
In this section, we first present extended synthetic experiments in Appendix B.I followed by results on real-world single-cell data in Appendix B.2
29

---

Published as a conference paper at ICLR 2026
B.1 SYNTHETIC COMPARISONS
In this section, further investigate our method's ability to learn potential energy (Ap- pendix B.I.I). interaction energy (Appendix B.1.2). and internal energy (Appendix B.1.3).
Since our approach builds on the JKO framework for modeling population dynamics, we adopt the experimental setup from JKOnet" and compare our method against both JKOnet [unnee et al. 2022b) and JKOnet" "Terpin et (024)) on corresponding synthetic evaluation tasks.
Styblinski- Tang
Holder table
Flowers
Bakley-Ohagan
Watershed
Ishigami
Friedman
Sphere
Bohachevski
Wavy plateau
Zig-zag agg ridge
Double exponential
ReLU
Rotational
Flat
Figure 4: Level curves of the true (green) and estimated (blue) potentials reconstructed by our proposed iJKOnety method for the paired setup, following (Terpin et al.) 2024, Appendix F). These results can be directly compared with those in (Terpin et al.) [024] Figure 6). Note that for the flat potential, the value range is near zero, expected for the ground-truth potential.
B. LEARNING POTENTIAL ENERGY
Paired Setup. Following Terpin et al. (2024 $4.1), we evaluate our method in the paired setup (see details in $5.1). which enables direct visual comparison with JKOnet (Terpin et al. 2024). The experiments are conducted on the synthetic dataset from Terpin et al. (2024 Appendix B), using step size of T = 0.01, = 5 time steps, and N = 2000 samples per step with 40% of test samples. Figure 4 presents the ground-truth potentials V(x) (green), defined in Terpin et al. (2024, Appendix F. Eqs. (31)-(45)) alongside the reconstructed potentials (z(z) (blue) learned by our method. As shown, approach achieves performance comparable to JKOnet" in the paired setup.
30

---

Published as a conference paper at ICLR 2026
Unpaired Setup. We repeat the experiments corrected version of the synthetic dataset, where samples are not temporally correlated. We select the six most challenging potentials in terms of convergence. Table B reports the ratio of the final metrics between the unpaired and paired setups. As shown, switching to the unpaired setup significantly affects the method's performance.
Table 3: Comparison of paired and unpaired setups. Each value indicates the ratio of the final metric in the unpaired setup that the paired setup. Most ratios significantly exceed 1, highlighting the increased difficulty of the unpaired setting.
i iJKOnety (Ours)
JKÖnet
Alias
Potential Name
Bdh
-UP
EMD
Cv-UVP
B -UVP
EMD
Cv-UVP
FL
Flowers
8066
90
297
135325
305
5809
FN
Friedman
556
5
263
6
3
Isiggmii
2259
36
4276
210
WS
Watershed
18135
220
1509
512720
1014
102592
WP
Wavy plateau
2
ZR
Zigzag ridge
158
23
13
155
25
31
B.1.2 LEARNING INTERACTION ENERGY
In this section, we evaluate how accurately our method and JKOnet" (Terpin al., 2024) can approximate the interaction energy in the unpaired setup. We exclude JKOnet from comparison, as design does not support learning interaction energies.
We identified an inconsistency in the theoretical formulation presented in the JKOnet" paper (Ter- pin et al. 2024). Specifically, as discussed in b (see also (Santambrogio 201lbrooooooi15]]]] 97.2)), the ground- truth interaction functional must be symmetric; that is, in ( the interaction kernel should satisfy (2)) = W(- (-)). Incorporating this correction, we conducted an experiment to assess the ability of both methods to recover the interaction energy in 2D unpaired setting.
Experimental Setup. The ground-truth interaction kernels are defined as ((2) = 2 (Wo(2) +
Wi(-z)). where Wi denotes the base functionals listed in Table 3 Both methods were restricted use only the interaction energy component We, in the energy parameterization (12). The corre- sponding variants are denoted as iJKOnetw and JKOnet Similar to B..I.1 we use a step size of r 0.01, K
time steps, and two sample size settings: small-scale setup with 12K total samples (N 2K per step, with 40% reserved for testing) and large-scale setup with 60K total samples (N = 10K per step, with 40% reserved for testing).
Results. Qualitative results are shown in Figure for the paired setup and in Figure 6 for the unpaired setup, with quantitative C2-UVP metries reported in Table 4 (see SC.I for details). In the paired setup, most potentials are accurately restored, although Waby Plateau, Friedmann, and Flowers are not. In the unpaired setup, neither method successfully recovers the interaction energy, likely due to biases introduced by the batched estimation of the interaction term We, in (12) which involves squaring the input. This approximation seems to cause both models to converge to batch- specific minima.
Table 4: Quantitative L2-UVP results for the unpaired setup in 2D interaction energy learning. Increasing the number of samples improves performance, though the effect is relatively limited.
Sample size: 2K
Sample size: 10K
Potential
iJKOnetw (Ours)
JKOnet
iJKOOeew (Ours)
JKCnet
Flowers
0,0032
0,0101
0.0009
0,0069
Friedman
0.0087
0.0010
0.0024
0.0093
Ishigami
0.0046
0.0087
0.0008
0.0066
Watershed
0.0043
0.0086
0.0009
0.0081
Zigzag Ridge
0.0060
0.0104
0.0017
0.0034
31

---

Published as a conference paper at ICLR 2026
Flower
Friedman
Ishigami
Watershed
Wavy plateau
mis
2010
2010
14.00
1150
1120
Zig-zag ridge
(a) Ground-truth W
(b) iJKOnetw (Ours)
(c) JKOnet Tr
Figure 5: Qualitative results from B.2 in the paired setup, using 10K samples per time step (with 40% reserved for testing). Compared to Figure [ which shows the unpaired setup. the model per- forms noticeably better: nearly all potentials are accurately reconstructed, highlighting the relative simplicity of the paired scenario.
32

---

Published as a conference paper at ICLR 2026
Flower
400
400
Friedman
Ishigami
Watershed
Wavy plateau
D500
Zig-zag ridge
8
80
(a) Ground-truth 7
(b) iJKOnetw (Ours)
(c) JKOnet Wr
Figure 6: Qualitative results from B.2 in the unpaired setup with 10K samples in total (40% reserved for testing). Neither model is able accurately reconstruct the ground-truth interaction energy W likely due to biases introduced in the estimation of the integral in [12
33

---

Published as a conference paper at ICLR 2026
B.1.3 LEARNING INTERNAL ENERGY
In this section, we assess how accurately our method and JKOnet" (Terpin et al. 2024) can estimate the diffusion coefficient 83 in (12) corresponding the internal energy term under the unpaired setup. We exclude JKÖnet from comparison, as its design does not support learning internal energy.
Experimental Setup. Following the protocol of Terpin et al. (2024 f4.3), we estimate the diffusion coefficient 0g in 2D and 20D settings using both t and JKCnet", i.e., the full parameter- ization (2). The same functional from Appendix C.7 is used for both potential and interaction energies. The ground-truth diffusion levels are set to 3 E {0.0, 0.1,
Results. F shows that iJKOnet fails to recover the ground-truth 3 values, while JKOnet" provides a closer approximation. However, all predicted 03 values tend to converge toward 0, devi- ating from the true levels 3 E {0.0, 0.1, 0.2]. This indicates that accurately learning the internal energy remains challenging for both models. We assume that such behavior stems from difficulties of entropy estimation developed methods.
JKOnetiy
OJottt((Suss))
2.5
0.4
2.0
0%
0.3
1.5
10
1.0
0.1
0.5
0.0
0.0
dimensions: 2
dimensions: 20
Figure 7: Estimation eπror of the diffusion coefficient 0g relative to the ground-truth values 8 Blue bars correspond to iJKOnet and orange bars to JKOnet'. The Y-axis indicates the absolute deviation between the estimated fg and true f values.
B.2 LEARNING SINGLE-CEI DYNAMICS
In this section, we present extended comparisons on real-world single-cell datasets, starting with SD results and then considering 50D and 100D.
Datasets. Following 5..2
consider the Embryoid Body (EB) dataset (Moon et al. 2019) and
additionally the Multiome (Multi) dataset (Burkhardt et al. 2022). For both datasets, we apply the preprocessing pipeline of (Tong et al.) 2020). The Multi dataset contains single-cell measurements across four time points (days 2, 3, 4, and 7), which we use for high-dimensional experiments.
5D. Following (Neklyudov al., 2024 Table 2) and (Terpin et al., 2024 54.4), we conducted
experiments on the EB dataset in 5D. At each time step, we train for 1000 epochs on 60% of the data and compute the dw, distance between the observed distribution yr (remaining 40% the data) and the one-step-aheal prediction Au. The results are shown in Table 5 Our method outperforms several previous approaches when using a vime-narying potential parametrization.
50D and 100D data. Following Neklyudov et al.) (2024), we conducted experiments the Multi dataset in 50D and 100D using a leave-one-out setup, averaged over three runs. Models were trained marginals from timepoint partitions [2.4.7] and [2.3,7]. and evaluated the corresponding left- out marginals at timepoints [3] and [4]. Training was performed for 5000 epochs.
The results presented in Table 6 Our method achieves comparable performance for the left-out [3] marginal. However, performance on the [4] marginal is poor, leading to lower overall average. This is likely due the longer time intervals between the learned marginals, which violate
34

---

Published as a conference paper at ICLR 2026
Table 5: The results for the EB dataset in 5D are obtained by training on all time steps, evaluated for next-time-step prediction using the dæ, metric. Results for non-JKO methods are taken from (Neklyudov al. 2024 Table 2):
Model
to
t1
f2
f3
Mean
Neural SDE (Li al. 2020)
0.69
0.91
0.85
0.81
0.82
TrajectoryNet [Tog al.
2020)
0.73
1.06
0.90
1.01
0.93
SB-FBSDE CChn et al. 2022a
0.56
0.80
1.00
1.00
0.84
NLSB Koshizuka & Sato 2023
0.68
0.84
0.81
0.79
0.78
OT-CFM TTng et al. 20224
0.78
0.76
0.77
0.75
0.77
WLF-OT Neklyudov al.
2024
0.65
0.78
0.76
0.75
0.74
WLF-SB
Neklyudov et al.
2024
0.63
0.79
0.77
0.74
0.73
JKOnet Bunne et al.
2022b
1.53
1,27
1.13
1.41
1.34
1KOnett Terpin al.
(2024)
0.99
1.11
1.06
1.30
1.12
iJKOnety (Ours)
0.92
1.11
0.95
1.21
1.05
JKOnet (Terpin al. 2024
0.69
0.77
0.69
0.78
0.73
iJKOnet,,v (Ours)
0.51
0.58
0.57
0.64
0.58
assumption that consecutive marginals are obtained from JKO steps (see 33. accuracy of predicting the [4] marginal.
thereby reducing the
Table 6: The results for Multi dataset for direct comparison with from Nekklyddv et al.
2024
Table 1), averaged for runs for leave-one-out setup, dw, metric.
Dimension
50
100
Metric
BW-UVP
EMD
MMD
BW-UVP
EMD
MMD
IKOnet,
128.729 78.307
68.406
6.055
0.0003
92.512 + 38.075
72.639 + 3.161
0.0002
JKOOnct (Ours)
38.318 # 18.325
50.560 = 7.187
0.0003
38.991 = 16,918
59.216 = 6.869
0.0002
JKOnet,,v
89.264 1: 38.372
68. 6.6000 9.846
0.0003
93.429 + 34.332
78.674 + 8.236
0.0002
JKOnet,, (Ours)
36.156 16.515
50.026 7.727
0.0003
38.387 + 18.525
59.318
7.365
0.0002
C EXPERIMENTAL DETAILS
C.I METRIC COMPUTATION DETAILS
In this section, we provide details on the evaluation metrics used our experiments.
Estimation protocol. All metrics computed using mini- batch Monte Carlo estimation. For syn- thetic experiments, we use 250 samples per time step. For real-world datasets, we use all available samples (approximately 500 per time step).
L2-UVP. When the ground-truth functional F" (e.g., V or W) is available, we assess the discrepancy between it and its reconstruction / using the backward Ly-based Unexplained Variance Percentage (C2-UVP) metric introduced by Korotin et (2021a), defined as follows:
VF
C2-UVP(F* F) 100
PE11 %,
(19)
Var(ps)
where the norm 11 is computed with respect to the ground-truth distribution Pk++1 and denotes the time step size. Values close to 0% indicate that V closely approximates F*.
Baly -VV. The Bures- Wasserstin UVP introduced by Korotin et al. 2021c) is defined as
Bdy, (pk. Pk)
Bdy, -UVP(Pt. Px) def 100
%
(20)
2 (ar(pk)
35

---

Published as a conference paper at ICLR 2026
where the Bures- Wasserstein distance is given by
Bd (I, Q del dip, (W(p,, N Eg)),
(21)
with pip and Ep denoting the mean and covariance of distribution P, respectively.
EMD or dw., The Earth Mover's Distance (EMD) (Rubner et al.
1998) is defined as
EpD(p%, Pb)) = drw, (pks Pk))
min xell(pe.pe) xxx
> yl dz(x,,)
(22)
where Pk: and Pk denote the ground truth and predicted distributions at time step k.
In practice, we approximate the continuous formulation (22 using its discrete empirical counterpart
(Peyré et al. 2019) and compute it with the POT library (Flamary et al.) 2024). is well known that high dimensions (Foumier
empirical Wasserstein distances suffer from poor sample complexity
& Guillin 2015 Panaretos & Zemel, 2019), meaning that a large number of samples is required for accurate estimation. In our experiments (up to 100 dimensions), the available sample sizes were sufficient to obtain stable estimates.
MMD. We compute the Maximum Mean Discrepancy (MMD) (Gretton et al. 2012) using its kernel formulation. Let K X' X' + IR be positive definite kernel with associated reproducing kernel Hilbert space (RKHS) H. and feature map Φ. The squared MMD between distributions pk and Ps. is
MMD (pk, Pk:) E K K(x, x') 2Ezzpp. ((x, y) + Ey.'' K ((y, y')
(23)
Importantly, MMD can be computed directly via kernel evaluations without explicitly constructing feature maps. Given samples { ~ Pls and { M Pk. we use the unbiased estimator
N
MMD K = N(N - 1)
Σ itp""
K(T,)r)
2
K(xi,0;) + M(M - 1) 3#8
Σ
K(yj U3).
NM
(24)
This estimator has quadratic complexity in the number of samples. our experiments, we use Gaussian kernel K (xx, 1y) exp
||x- ||--v11
with the bandwidth parameter o 10.
2n
For experiments involving the DMSB method (Chen et al. 2023). we instead use an adaptive band- width provided in the official implementation" In this case, is computed as the average squared pairwise distance between all distinct samples in the batch.
C.2 TRAINING DETAILS
Energy. To ensure stability, we accumulate gradients across all time steps = O,..., using mini-batch estimates Jo of the energy function To (see Eq. (2)). For each time step (x. we sample mini-batch X { , 2x Pk. with fixed batch size B. The energy is then estimated as:
1 B
1 Σ Wo, (z, 2zy) f((x)
Jo(pk) [X{:
Vo, (ox) +
(25)
B
t=
where (()xx) denotes the estimated entropy of Pk discussed the following section.
Map. In practice, we use two strategies to parameterize each transport map Th: for large-scale tasks, we assign a separate network to each time step; for moderate-dimensional tasks, we encode the time index & as additional input, i.e., T (x) = (T, k), which helps prevent overfitting. This approach performs better than using a network without any time embedding. Thanks to Optax (DeepMind al. 2020), the inference of all T be performed in parallel, treating them single generator step. We also experimented with initializing the maps during the early training epochs by aligning them with discrete OT maps computed using the OTT-JAX (Cuturi al. 2022)
Training. We aggregate gradients for Jo and T across all time steps. We also experimented with alternative aggregation strategies, but found that equal aggregation across all time steps was the
36

---

Published as a conference paper at ICLR 2026
Algorithm 1: iJKOnet Training
Input:
Sequence of measures { k=0 (accessible via samples); 1K
Mapping networks T X + X;
Potential network Ve, X' + IR, interaction network We, X + R:
Diffusion coefficient 03 E R=, time step max inner iterations Ir: Output:
Learned parameters 0 = (01, 0g, 03) and yo; while not converged do
for i <- 1 to Ly do
// Update transport maps TØ
for k to K - do
Sample batch X ~ Ph. Xb+1 ~ Pat1i
X41 + T (Xa);
c + Jo( kt1 + 27 k+1
X X
t p p- 4
Σx c.:
// Update energy parameters 0 for k 0 K do
Sample batch Xx= ~ Pk:
Xpred & Th (Xx);
C [o pred + To (X++1):
0
1 Vo Σx C%;
most effective, During optimization, we perform multiple update steps for the transport maps T%k parameterized by s., while updating the energy function Jg. parameterized by 0, only once. This follows standard practice in min-man max optimization (Goodfellow et al. 2020). The overall procedure summarized in the pseudocode in Algorithm
Stability. The main challenge was avoiding suboptimal energy functionals that did not converge zero. We found that the simplest setup (combined with careful hyperparameter tuning using the Optuna framework (Akiba et al.) (019h) was the most effective. We additionally experimented with common stabilization techniques from GAN training, including gradient penalties (Gulrajani et al. 2017), spectral normalization (Miyato et al. 2018 Jiang et al.) 2018). and extragradient up- dates Daskalakis et al. 2018). Overall, we perform optimization for single energy functional, this naturally acts regularizer, improving the stability of the training.
Scalability. Our method scales well with the number of time points. Leveraging JAX (Bradbury al. 2018) and Optax (DeepMind al.) 2020), avoid memory issues even for large batches and many time steps, since we do not backpropagate through time as in JKOnet (Bunne et al. 2022b). Each step is processed independently, and the loss is computed by averaging outputs from separate networks, avoiding large computation graphs. Training time on GPU is comparable to JKOnet but our approach eliminates the costly precomputation of OT couplings. In high-dimensional settings (50D-100D), JKOnet" exceeds GPU memory limits and requires CPU-based precomputations, making it substantially slower.
C.3 ENTROPY ESTIMATION DETAILS
Prior to training, we estimated H(pp) for each k: = 0, K using the Kozachenko-Leoome Leonenko nearest-neighhor estimator (Kozachenko/ 1987 Berrett et al. 2019) with 5 nearest neighbors. We used the publicly available implementation from (Butakov et al. 2024
As discussed in 8 3.3
37

---

Published as a conference paper at ICLR 2026
the following equation for estimating H((TTtp)):
F((T)to)) = F(((w) / log det V (x)) dox
(26)
Since our map required to invertible during training, the case for gradient-based ICNN AAmos et al. 2017) parameterizations (Bunne et al. (202bb), we compute the sign and the natural
logarithm of the absolute value of the Jacobian determinant, both of which are supported by modern computational frameworks (Bradbury et al. 2018 Paszke 2019).
C.4 OPTIMIZER
We optimize two loss functions with respect to To and T using the Adam optimizer Kingmaa 2014). For Tg. we use parameters 31 = 0.9, 32 0.999, é = x 10-*, and constant learning rate of 10 4, with gradient clipping applied using maximum global norm of 10. For T use parameters 31 = 0.5, B2 = 0.9, and E = x 10-8, with a learning rate of 1 x 10 Training performed with mini-batches of size 500.
C.5 NETWORK ARCHITECTURE
The neural networks for the potential Vo, and interaction W6, energies are multi-layer percep- trons (MLPs) with two hidden layers of size 64, using softplus activation functions, and one-dimensional output layer. The neural network for the optimal transport maps Th is also an MLP with two hidden layers of size 64. The time step k is concatenated with the input xt. and the network uses selu activations, with an output layer that matches the input dimension of x
C.6 HARDWARE
Experiments were conducted on a CentOS Linux 7 (Core) system with NVIDIA A100 GPUs. Most of the computation time spent metric evaluation. Leveraging JAX's just-in-time compilation, our method completed 100 training epochs in approximately one minute.
just'jjus-
Table Training time (in hours) for EB Moon et al. 2019) dataset for 5 and 100 dimensions.
Solver
Dimension
mean
iJKOnety
100
0.333
0.020
iJKOOne,,
100
0.338
0.023
iJKOnety
0.321
0.005
iJKOnet,
0.332
0.001
C.7 FUNCTIONALS
For easier representation of our experimental results, we use abbreviations for each potential, each abbreviation being shown parentheses and the definitions can be found in Teerpin et al. 2024] Appendix F): Wavy Plateau (WP). Double Exponential (DE), Rotational (RO), ReLU (RE), Flat (FT), Friedman (FN), Watershed (WS), Ishigami (IG), Flowers (FL), Bohachevsky (BC), Sphere (SP), Styblinski-Tang (ST), Oakley-Ohagan (00), Zigzag Ridge (ZR), and Holder Table (HT).
D PROOFS
To begin with, we recall the notion of strong smoothness, see (Beck 2017 $5.1). A functional F: X called Y/4 strongly smooth it continuously differentiable on and it holds:
B VF)() - 2 s ||x - -112 Va, E X.
Now proceed auxiliary results.
38

---

Published as a conference paper at ICLR 2026
Lemma D.1 (Solution to the JKO problem with potential energy is unique). Let (p) Ix. V (x)dp(r)) Pu E (2.ac()) and p1 = JKO+J- (Po). Then Tv. = arg minrxx CVV*, is the (unique) optimal transport map between Pa and pt. In particular, Ty-1pp- = Pr.
Proof. Consider the functional p + dey, (Po. p), where E P2(2). In what follows, we prove that this functional is strictly convex.
Consider Pa. Pb E P2(x)); a > 0. Let ape + (1 a)p.. Note that Po is absolutely continuous. By the Brenier's s theorem, holds:
/ || (==-T(())(dppo(()) dwr (po, pb) = / ||x [b==7((+)12ddot
dR,, (po, (po,Poo) Pa)
for the corresponding (unique) OT maps To Tatpo = pai Tb Toipo = pb. Now consider the weighted sum:
advy (po, Pa) + (1 - ai ai) Pb) / (a - Ta(x))12 +1 a)]l= z d))=))()ddp(()))
f x ||x | yl2d
where plan Tab E II(po has the following conditionals:
" (y)a=) ad (x))) (1 a)o Tb(x)).
Since the OT map between and is unique (and deterministic), we have
min / ((=))())))ppz)) < / ||. vl vl Tipo=p
Pr,, (po. p)
which yields the strict convexity of p dby, (po.
The main statement of the Lemma follows from the strict convexity + fx (x)dp(() + dey, (Pos p) and Brenier's theorem.
of
Now we are ready to prove our main quality bound theorem.
Proof. (Theorem 3.1
Prelimiary Note: The facts from convex analysis (properties of convex functions and their conju- gates) which use below could be found in (Kornilov et al. 2024 Lemma 1).
At first, we simplify the expression for the JKO loss:
C(V.Tv)= / ((Iv(=)) / (()) 1 2r / 10 (v[](())(3dpo(a) (=(Tv(=)))po(e) - / (o(o)dpp()) 11112 dpr(y)+ /
- 1 -
(27)
| dpo(x),
2r
2r
#c(po.P)
where Vg =V + 111 11.
From 27 we note that Tv = arg miny (x) (T. From the convex
analysis, it follows that (z( = vV,(x) delivers the minimum; V, is the (Fenchel) conjugate of Vg. Note that Tv measurable since it is continuous almost surely (Rockafellar 1970 Theorem 25.5). Substituting VV, into (77 on par with Fenchel- Young equality yields:
((V,Tv) - (a(o)ppo(a)) - + / V,(())dp(()) + C(po.Pt).
(28)
39

---

Published as a conference paper at ICLR 2026
Note that Tv- P1 (Lemma D.1 Therefore, eq- (27 for L(V", Tv-) could be simplified:
- / Vi (1))dp - - /
(z) (1) + C(po. =
Vi (TV-
=5 (v)dpx(
1 (v-))))))) + C(po, pi). (29)
Now we analyze the gap E(V) between optimal and optimized JKO losses. Leveraging (28 and yields:
29
- = vy(x)dpo(a)+ + / (V(W)ddi(V) - 1 (2, Tv- (x)))po()))
TE(V) = TL(V", - -I.
(30)
Now we note that Tv. = VV; X' - RDD From the properties of convex functions follows that VV2 :' + Rd defines the inverse Optimal Transport mapping. In particular (we assume Vi to be convex) VV 5p1 = po- Then, changing the variables in (30) results in:
TE(V)= / V VV (v) (v) + / V.(u)dp1()) (VVVT())
Vo (v)) +v())) (VV dp)()). D DVg (v), (w)
(31)
(32)
where ( :) is the Bregman divergence, see (Banerjee et al. 2005 Def. 1) for the definition
and (Banerjee et al. 2005, Appendix A) for the "Dual Divergences" property used in the transition
from (31) to 32)
Since V, is 3 -moooth then V,, is B-strongly convex. Therefore, by the property of strongly convex functions (we pick y E supp(e1)):
2 VV -V
(v)11 & -WV -V)
D (VV) (v), VV,())).
(33)
Combining
(32 and 33 yields:
TE(1) (V) 2 8 / (v) (v) (VG(v)//dd110)))
x
To complete the proof, are left to note that VV - - -V()))
40


Reviews:
[{"review_id": "xdDpqWweo0", "reviewer": "Reviewer_XRFW", "summary": "This paper introduces iJKOnet, an approach to learning population dynamics by bridging inverse optimization techniques with the Jordan-Kinderlehrer-Otto (JKO) scheme. The method frames the task as a min-max optimization problem for identifying the underlying energy functional that governs observed distributions, enabling end-to-end adversarial training without restrictive neural architectures. Theoretical guarantees are provided for recovery of potential-driven dynamics, and empirical results on both synthetic and real datasets show improved performance over prior JKO-based baselines.", "questions": "- In which regimes or datasets is JKO-based modeling justified or superior to flow-matching, neural ODEs, or Schrödinger bridge-based methods?\n- How sensitive is iJKOnet to missing, irregularly spaced, or noisy population snapshots? Do learned energy functionals overfit to sampling artifacts in sparse settings?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- iJKOnet is different from previous methods that either require potential-only energies or precomputed OT couplings. The formulation of energy functional recovery using inverse optimization within the JKO scheme provides a conceptually clear route to modeling population level dynamics from discrete snapshots, directly leveraging optimal transport geometry.\n- Theorem 3.1 gives a non-trivial quality guarantee, explicitly bounding the distance between the gradients of learned and ground-truth potential functions in terms of the inverse optimization gap. This is a concrete advance compared to many contemporary works that lack precise recovery guarantees.\n- iJKOnet dispenses with input-convex neural networks for parameterizing transport maps, allowing the use of more scalable and widely applicable neural network architectures such as MLPs and ResNets", "weaknesses": "- The main theoretical guarantee (Theorem 3.1) exclusively addresses potential energy functionals with strongly convex smooth potentials, while the practical method is claimed to work for broader energy functionals (e.g., including interaction and internal terms). This leaves a theoretical gap between what the analysis covers and what is demonstrated empirically. Specifically, the inability to provide quality bounds for non-potential and more general functionals limits the rigor of the claims. (This could explain why the performance of iJKOnet is much worse than baseline in Table 2)\n- While iJKOnet is, in principle, applicable to full free-energy functionals, all main experiments revert to only learning potential energy, explicitly setting interaction and entropy terms to zero. As a result, the empirical support for learning richer dynamics is missing, and the scalability to very high-dimensional or truly multimodal interaction-dependent data is unproven.\n- Some recent and directly related works [1,2,3,4] on similar task with unknown parameters, alternative population-level diffusion modeling, and advanced regularization techniques are not discussed.\n\nref:\n\n[1] Computational and Statistical Asymptotic Analysis of the JKO Scheme for Iterative Algorithms to update distributions\n\n[2] An Eulerian approach to regularized JKO scheme with low-rank tensor decompositions for Bayesian inversion\n\n[3] WeightFlow: Learning Stochastic Dynamics via Evolving Weight of Neural Network\n\n[4] Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schrodinger Bridge", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OrHm6aAEN5", "reviewer": "Reviewer_tCBu", "summary": "In this paper, the author propose  a method to learn the population dynamics, which combine JKO framework with inverse optimization techniques. In addition, this method relays on adversarial training procedure. Theoretical guarantees are also provided in this paper.", "questions": "- Could you explain about the potential stability of the proposed method resulted by the adversarial training procedure, and what further improvements could be made?\n\n- Could you more explicitly summarize the key algorithmic or theoretical innovations that distinguish your approach from its predecessors?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper is well written.\n\n2. The figures are well-prepared and greatly facilitate the understanding of the content.\n\n3. The theoretical proofs is provided, which give solid theoretical guarantees of the proposed method.\n\n4. The background section is clearly presented and offers helpful context.", "weaknesses": "1. The interchange of the min and Σ operators between Equations (10) and (11) lacks justification or analysis. It is unclear under what conditions this exchange is mathematically valid.\n\n2. The method relies on an adversarial training procedure, which is known to introduce instability. It would be important to discuss whether any measures were taken to mitigate this issue.\n\n3. The work builds heavily on existing studies (e.g., Terpin et al., 2024; JKOnet∗). The specific contributions and novelty of the present approach should be more clearly emphasized.\n\n4. Several relevant references appear to be missing from the related work section, such as:\n\n[1] Zhang Z, et al. Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schr\\\" odinger Bridge[J]. arXiv preprint arXiv:2505.11197, 2025.\n\n[2] Li R, et al. WeightFlow: Learning Stochastic Dynamics via Evolving Weight of Neural Network[J]. arXiv preprint arXiv:2508.00451, 2025.\n\n[3] Zhang Y, Levin M. Equilibrium flow: From Snapshots to Dynamics[J]. arXiv preprint arXiv:2509.17990, 2025.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Byjv3poN7g", "reviewer": "Reviewer_vGio", "summary": "The paper is an attempt to use the (famous) JKO optimization scheme to learn a dynamical system from a time- series $\\rho_1, \\cdots, \\rho_M$  where $\\rho_i$ represent the data of the system at time $T_1< T_2 < \\cdots < T_M$.  The JKO scheme are expressed as proximal operator of the $W_2^2$ optimal cost with a functional (typically of the form energy plus entropy plus possibly interaction terms described by a kernel).   The approach taken here is to parametrize a transport map (for the $W^2$ term) and to parametrize  the energy terms using neural architectures and to estimate the entropy terms via an entropy estimator.", "questions": "1) Can you clarify what are the obstacles to implement an interaction kernel and the entropy term in the algorithm?\n\n2) Can you comment on the difficulty of computing the transport maps?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The JKO formulation of gradient flow as a proximal optimzation problem is a fundamental tool in probbaility theory.  It has been found difficult to use it effectively as a generating tool and and the reformulation of the JKO functional in this paper  is interesting and could be promising.", "weaknesses": "1) The main weakness is that the paper does not fully implement the proposed loss functional. Indeed the the interaction terms and the entropy terms are ignored and the same architecture as in JKONet is used.  This reduces the class of PDEs considered to a very small class of linear PDE. The mismatch between the theory which considers general PDEs and the implementation is too great to make the paper convincing.   Ignoring the entropy term,  which I suspect is very difficult to estimate,  is a serious defect of the approach.  \n\n2) The proposed architecture requires the construction of a transport map which is both expensive and only works theoretically when the measures have density.  The authors propose to replace input convex neural networks (since the transport maps are gradient of a convex functions) by a more general architecture but I do not see why this will make the problem more scalable or stable.  Maybe a approach using triangular map as in the Yousef Mazrouk group at MIT may be more scalable?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)