Learning diffusion at lightspeed
Authors: Antonio Terpin
Nicolas Lanzetti
Martín Gadea
Florian Dorfler
Venue: NeurIPS 2024 oral (2024)
Keywords: optimal transport
diffusion processes
gradient flows
Primary area: optimization

Abstract: Diffusion regulates numerous natural processes and the dynamics of many successful generative models. Existing models to learn the diffusion terms from observational data rely on complex bilevel optimization problems and model only the drift of the system.
We propose a new simple model, JKOnet*, which bypasses the complexity of existing architectures while presenting significantly enhanced representational capabilities: JKOnet* recovers the potential, interaction, and internal energy components of the underlying diffusion process. JKOnet* minimizes a simple quadratic loss and outperforms other baselines in terms of sample efficiency, computational complexity, and accuracy. Additionally, JKOnet* provides a closed-form optimal solution for linearly parametrized functionals, and, when applied to predict the evolution of cellular processes from real-world data, it achieves state-of-the-art accuracy at a fraction of the computational cost of all existing methods.
Our methodology is based on the interpretation of diffusion processes as energy-minimizing trajectories in the probability space via the so-called JKO scheme, which we study via its first-order optimality conditions.

Learning diffusion at lightspeed
Antonio Terpin ETH Zürich aterpinGethz.ch
Nicolas Lanzetti ETH Zürich Imicolas@ethz.cr ch
Martín Gadea ETH Zürich mgadeaßethz.cht ch
Florian Dörfler ETH Zürich dorfler@ethz.co
Abstract
Diffusion regulates numerous natural processes and the dynamics of many suc- cessful generative models. Existing models to learn the diffusion terms from observational data rely complex bilevel optimization problems and model only the drift of the system. We propose new simple model, JKOnet", which bypasses the complexity of existing architectures while presenting significantly enhanced representational capabilities: JKOnet" recovers the potential, interaction, and internal energy components of the underlying diffusion process. JKÜnet" mini- mizes a simple quadratic loss and outperforms other baselines in terms sample efficiency, computational complexity, and aceuracy. Additionally, JKOnet" pro- vides a closed-form optimal solution for linearly parametrized functionals, and, when applied predict the evolution of cellular processes from real-world data, it achieves state-of-the-art accuracy at fraction of the computational cost of all existing methods. Our methodology is based the interpretation of diffusion pro- cesses as energy-miningg meergy-miiimizing trajectories in the probability space via the so-called JKO scheme, which we study via its first-order optimality conditions.
Source code: https://gitee comm//anooniiorrrpp///ooooett--e
Introduction
Diffusion processes govern the homeostasis of biological systems [40], stem cells reprogramming [20, 36), and the learning dynamics of diffusion models [16, 22, 54] and transformers [19, 52]. The diffusion process of interest often originates from three quantities: a drift term due to potential field, the interaction with other particles, and a stochastic term. these three components are known, predictions follow from simple forward sampling [27] or the recent work in optimization in the probability space [1, 3, 11, 25, 34, 38, 41]. In this paper, we consider the case when the diffusion process is unknown, and we seek to learn its representation from observational data. The problem has been addressed when the trajectories of the individual particles are known [5, 32), but it is often the case that we only have "population data". For instance, single-cell RNA sequencing techniques enabled the collection of large quantities of data on biological systems [35], but the observer cannot access the trajectories of individual cells since measurements are destructive [20, 36). The most promising avenue to circumvent the lack particle trajectories the Jordan-Kinderaa Kinderlehrer-Out (JKO) scheme [24] which predicates that the particles while not deviating too much from the current configuration. However, the JKO scheme entails an optimization problem in the probability space. Thus, the problem of finding the energy functional that minimizes a prediction error (w.r.t observational data) takes the form of a computationally- challenging infiite-diimeniioall bilevel optimization problem, whereby the upper-level problem is the minimization of the prediction error and the lower-level problem is the JKO scheme. Recent work [1, 9] exploits the theory of optimal transport and in particular Brenier's theorem [7] to attack this bilevel optimization problem, a model henceforth referred to as JKOnet. Despite promising initial results in [9], this complexity undermines scalability, stability, and generality of the model.
whole move decrease aggregate energy,
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

1
M4-1 argmin Jo(n) + W2 pEP(
2r
causes mismalch: Σ; [VV[(( (x =F!)
P4
j22
μ12
A1
μT
A41
effects mismatch W2(11, j1)
Given sequence of snapshots (Ho.-. ................................ diffusion, we want to find the parameters 0 of the parametrized energy function Jg that best explains the particles evolution. Given 0, the effects mismatch the Wasserstein distance between the observed trajectory and the predicted trajectory obtained iteratively solving the JKO step with Jg. The first-order optimality condition in [30] applied to the JKO step suggests that the "gradient" of Ja with respect to each júr vanishes at optimality, i.e., for jår = μız. For Jo(u) Sn ((())d this condition is depicted on the right. The gradient (dashed blue arrows) of the true V (level curves in dashed blue) each observed particle x! (blue circles) in the next snapshot (4++1 opposes the displacement (dotted red arrows) from a particle x (red triangles) in the previous snapshot μ4, Instead, the gradient (solid green arrows) of the estimated Vo (level curves in solid green) each observed particle xr! 1++1 (square) does not oppose the displacement from particle xf in the previous snapshot μt.. This mismatch the causes of the diffusion process what JKOnet" minimizes.
Figure 1:
.M......TT) PT) of a population of particles undergoing
Furthermore, to be practical, it is limited learning only potential energies, modelling the underlying physics only partially. Alternatively, [10, 43] learn directly the transport map describing the evolution of the population (i.e., the effects), bypassing the representation of the underlying energy functional (i.e.. the causes). Motivated by robustness, interpretability, and generalization, here we seek method learn the causes. [23, 42], the authors try to learn geometry that explains the observed transport maps. Unfortunately, the cost between two configurations along a cost-minimizing trajectory is often not metric [47]. Other attempts include recurrent neural networks [21], neural ODEs [15], and Schrödinger bridges [12, 28].
Contributions. We study the first-order necessary optimality conditions for the JKO scheme, an optimization problem in the probability space, and show that these conditions can be exploited to learn the energy functional governing the underlying diffusion process from population data, effectively bypassing the complexity the infinite-dimnnse We provide a closed-form solution in the case of linearly parametrized energy functionals and simple, interpretable, and efficient algorithm for non-linear parametrizations. Via exhaustive nu- merical experiments, we show that, in the case of potential energies only, JKOnet" outperforms the state-of-the-an in terms of solution quality, scalability, and computational efficiency and, in the until now unsolved case of general energy functionals, allows us to also learn interaction and internal energies that explain the observed population trajectories. When applied to predict the evolution of cellular processes, it achieves state-of-the-art accuracy at fraction of the computational cost. Figure shows an overview of our method, detailed in Section
bilevel optimization problem.
2 Diffusion processes via optimal transport
2.1 Preliminaries
The gradient of p Rd + DR i V E Rd and the Jacobian of ó Rd R is V Raxd We say that f R - F. has bounded Hessian if | < C for some C (and some matrix

---

norm II-ID. The divergence of F R + is V F and its laplacian is V2F. The identity function Id : Ie + Rd Id(x) = T. We denote by P((R") the space of (Borel) probability measures over R with finite second moment. For E P(RR), supp(H) its support. The Dirac's delta measure at & E RR, is 6z. All the functions are assumed to be Borel, and for f I + R, Jae (x))d((()) is the (Lebesgue) integral of w.r.t. μ.. If is absolutely continuous w.r.t the Lebesgue measure, / < dr, then it admits density
Red Rzo. and the integral becomes Jaa f (z)))d The pushforward of μ via (Borel) map R + R" is the probability measure f# defined by (B# (B)); when μ is empirical with N, H = N E oz,, then f# = 1 Ei=1 ((xz)). Given p. P(RR), we say that probability measure 7 E P(R Rd) is a transport plan (or coupling) between μ and y if its marginals are μe and v. We denote the set of transport plans between I and v by [[u, v). The Wasserstein distance between μ and v is
W2(,, = min (eE((p))) gdxgd || ||-- dy,, $
/
(1)
When lr and are discrete, (1) is a linear program. If, additionally, they have the same number of particles, the optimal transport plan is 7 = (Id, T) for some (transport) map T: IRd + Rd [39].
When is absolutely continuous, y (Id, V x 2 for some convex function và [7].
2.2 The JKO scheme
Many continuous-time diffusson processes can be modeled by partial differential equations (PDEs) or stochastic differential equations (SDEs): Example 2.1 Bookker-Plancke, The Tokker-Plance equation.
Op(t. x) = V (VV((r))((t,)) x)) + p(t,(x ŌL
(2)
describes the time evolution of the distribution p of a set of particles undergoing drift and diffusion.
dx = - (() + (3)
where (t) is the state of the particle, (ιx) the driving potential, and ((t) the Wiener process.
The pioneering work of Jordan, Kinderlehrer, and Otto [24], related diffusion processes to energy- minimizing trajectories in the Wasserstein space (i.e., probability space endowed with the Wasserstein distance), providing discrete-time counterpart of the diffusion process, the JKO scheme.
argmin J(u) W2(,.p)) (EP(
M+++1
where J : P(Rd) R {+o} is an energy functional and r > is the time discretization. Wasserstein gradient flow). The Fokker- Bakker--laank equation (2) results
Example 2.2 (Fokker-Planl from the continuous-timg limit (i.e., T - 0) of the JKO scheme (3) for the energy functional
J(N) = / V(z)du(x) + 3 [ p(x) log(p(z))dz with du((x) = p(x)ddx
2.3 Challenges
Section 2.2 suggests that we can interpret the problem of learning diffusion processes as the problem of learning the energy functional J in (3). Specifically, the setting is follows: We have access to sample populations μu, #1..... Mr, and we want to learn the energy functional governing their dynamics. A direct approach to tackle the inverse problem others. for the model JKOnet in [9]. This approach bases on the following two facts. First, by Brenier's theorem, the solution of (3). Pt+1. can be expressed! the pushforward of μ4 via the gradient of a convex function 1/ RRd - R and, thus,
bilevel optimization, used, among
W2(M, 22+1)
(V-VV((x)) (d.x)))
Under an absolute continuity assumption.
3

---

Second, the optimization problem (3) is equivalently written as
1 2T
/
argmin J (V (4 +
|| (z))))) Rd
WIEC
where C' is the class of continuously differentiable convex functions from R to R. Therefore, the learning task can be cast into the following bilevel optimization problem, which minimizes the discrepancy between the observations (1+) and the predictions of the model (ft):
min
W2 14))
-
s.t. A0 = Po.
P++1 V V pec
(4)
V argmin (V (Veee())) 2r Jad
| ((((()))
$EC
practical implementation of the above requires a parametrization both the transport map and the energy functional. The former problem has been tackled via input convex neural network (ICNN) parametrizing 2% [2, 8] or via the "Monge gap" [51]. The second problem is only addressed for energy functional of the form [μ) JR Ve (z).. without interaction and internal energies, where Vo is a non-linear function approximator [9].
Challenges. This approach suffers from two major limitations. First, bilevel optimization problems are notoriously hard and we should therefore expect (4) to be computationally chal- lenging. Second, most energy functionals are not potential energies but include interactions and internal energy terms as well. Although it is tempting include other terms in the energy functional (e-g., parametrizing interaction and internal energies), the complexity of the bilevel optimization problem renders such an avenue viable only in principle.
3 Learning diffusion at lightspeed
Our methodology consists of replacing the optimization problem (3) with its first-order necessary conditions for optimality. This way, bypass its computational complexity which, ultimately, leads to the bilevel optimization problem (4). Perhaps interestingly, our methodology for learning diffusion based on first principles: whereas e.g. [9] minimizes an error on the effects (the predictions), minimize an еrror on the causes (the energy functionals driving the diffusion process); see Figure 1. As detail in Section the resulting learning algorithms are significantly faster and more effective.
3.1 Intuition
To start, we illustrate our idea in the Euclidean case (i.e., Rd) and later generalize it to the probability space (i.e., P)R Consider the problem of learning the energy functional J Rdd - R U {+oo} of the analog of the JKO scheme in the Euclidean space, the proximal operator
1 2T
T++1 argmin J(x +
(5)
DERd
Under sufficient regularity, we can replace (5) by its first-order optimality condition
1 V J(r+1) + (x x.)
6
Given dataset (ro, F1, ,T). seek the energy functional that best fits the optimality condition:
VJ(Ie1)) ( (0 21)
min
In the probability space, we can proceed analogously and replace (3) with its first-order optimality conditions. This analysis, which is based on recent advancements in optimization the probability space [30], allows us to formulate the learning task single-level optimization problem.

---

3.2 Potential energy
Consider initially the case where the energy functional a potential energy, for V R2 R.
J(m = .fe V(x)du(())
(8)
The following proposition is the counterpart of (6) in P(RR)::
Proposition 3.1 (Potential energy). Assume V is continuousty differentiable, lower bounded, and has bounded Hessian. Then, the JKO scheme (3) has an optimal solution px41 and, if p+41 is optimal for (3), then there is an optimal transport plan 7tt between 14 and 1441 such that
f J (2(4)) + 1 + - -aa) dreeIe, =.
Proposition 3.1 is by all means the analog of (6), since for the integral to be zero, (Z (xr+1 - Øt) = must hold for all (zz, Tt-11) 1 E supp(7=). Since the collected population data optimization variables in the learning task, the optimal transport plan % be computed beforehand. Thus, we can learn the energy functional representation by minimizing over class of continuously differentiable potential energy functions the loss function
No. M1. : H&r
T
E 2dx2d
1 (a++1 xx)) (1r,
V
(9)
3.3 Arbitrary energy functionals
Consider now the general case where the energy functional consists of potential energy (with the potential function V Rd + IR), interaction energy (with interaction kernel U IRd R), and internal energy (expressed the entropy weighted by 3 Rgo):
J(p) V(z)dp(x) + ff U(x y)d(H p) y)
/ p(x) log(p(z))dr.
JR
Jgd
(10)
The first-order necessary optimality condition for the JKO scheme then reads as follows. Proposition 3.2 (General case). Assume V and U are continuously disferentiable, lower bounded,
and have a bounded Hessian. Then, the JKO scheme (3) has an optimal solution pr41 which is absolutely continuous with density Pt+1 and, if dp)+1(() pt+ (x)((z) is optimal for (3), then there an optimal transport plan "tt between μt and pc+1 such that
0
V ((
+
VU - Deeee
Bd
Rd
(1+1(
3
+ ( xx) (2((1,)441))
P(41(1441)
Thus, Proposition 3.2 suggests that the energy functional can be learned by minimizing over a class of continuously differentiable potential and internal energy functions and / Rzo the loss function
VV
V((++) +
VU(++11 - -1/41 )duee (2(41)
1=0 R
Rd
(11)
Pr+ D (I++1)
+ (x++1 - x2) dr (rt
Ptt1 (I++1)
Remark 3.3. We generalize the setting to time-varying energies in Section 4.4 and in Appendix B.
3.4 Parametrizations
For our model JKOnet", we parametrize the energy functional at a measure du ((x)de follows:
JØ(N)= [ Ve, (x))d(x)) + / U2 y)d(p x n)(2,)) v) + 03 I log(((())d))
where 01, 02 1R 0g R. and we set 0 10T,03 03 E 2=
5

---

Generality
Model
FLOPS per epoch
Seq. op. per particle
V(x) B U(r)
JKOnet
o (T (DNd+ N² [og(N)
DNd+ + No log((N)
w/o TF
X
x
JKOnet
(T (DNd+ N2 (o
(NNd+ + N2 (og((N)
X
X
w/ TF
JKOnet
(TD Nd N2 (og(()
TD Nd No 1og((N)
w/ MG w/o TF
X
X
JKOnet
TD Nd (o
( Nd + N²
x
x
w/ MG, TF
JKOnet"
w/o (xx)
(TNdd)
(d)
x
JKOnet*
0 (TN2d)
O (Nd)
w/ (x))
JKOnet; w/o ((r)
0 (TNdn n²)
O NT + n²)
JKOnet
o (TN2dn + +3)
O (TN2dn + n²)
w/ (z)
Table 1: Per-epoch complexity (in FLOPs) and per per-particle minimum number of sequential operations (maximum parallelization) for the JKDnet and JKOnet' model families (we refer to the linear parametrization of our model with JKOnetī).. Here, T is the length of the population trajectory, N the number of particles in the snapshots of the population (assumed constant), d is the dimensionality of the system, n the numher features for the linear parametrization, D, TeacherForcing (TF) are JKOnet parameters: the number of inner operations (which may or not be constant), the accuracy required for the Sinkhorn algorithm, and a training modality (see [9] for details), respectively. MG stands for the Monge gap regularization [51].
Linear parametrizations. When the parametrizations are linear, i.e. Ve, (1) - 0J ((x), ) Ua, (x y) 82 p(z 1e--y) y) for feature maps 01, 62 R + R", the optimal 6* can be computed in closed-form: Proposition 3.4. Assume that the features Φ1,; and Φ2,5 are continuously differentiable, bounded, and have bounded Hessian. Define the matrix IRd - R(2++1)xd by
ye(x))= ( Voz(I, 21) (af), patx.)
Vellan)
and suppose that the data is sufficiently exciting that D Jex (x(x))) dup(ax)) is invertible. Then, the optimal solution of (11) is
[ Rd
0" Σ
(x( due(ze)) (dpe)))
Σ 1=0 RexRd
(((41))(2111 - (2))
(12)
Remark 3.5. The excitation assumption can be enforced via regularization terms Ae110,11 with X; > 0, in the loss (11). Another practical alternative is to use pseudoinverse solve the least- squares problem corresponding to (12) by gradient descent.
Non-linear parametrizations. When the parametrizations are non-linear, we minimize (11) by
gradient descent.
Inductive biases. By fixing any the parameters 01, 02, 03 to zero, the corresponding energy term dropped from the model. It is thus possible to inject into JKOnet" the proper inductive bias when additional information on the underlying diffusion process are known. For instance, if the process deterministic and driven by an external potential, one can assumed that the interaction between the particles is negligible, we can set 02 = 0.
03 0. Similarly, if it can be
3.5 Why first-order conditions
Here, we motivate the theoretical benefits of JKOnet" over JKOnet using the desiderata:

---

Styblinski-Taun (31)
Flowers (33)
Ishigami (36)
Friedman (37)
Figure 2: Level curves of the true (green-colored) and estimated (blue-colored) potentials (31), (33),
(36) and (37), see Appendix F. See also Figure 6 in Appendix A.
1. Total computational complexity per epoch (i.e., the cost to process the observed populations).
Per-particle computational complexity (i.e., the cost to process single particle when maximally parallelizing the algorithm, prior merging the results).
3. Representational capacity of the method (i.e., which energy terms the model can learn).
We collect this analysis in Table 1. Fundamentally, the first-order optimality conditions allow reformulation of the learning problem that decouples prediction of the population evolution and learning of the dynamics (such coupling is the of (4)). As result, JKOnet" enjoys higher parallelizability We also observe that the interaction energy comes with an increase in complexity, and in a way resembles the attention mechanisms in transformers [19, 52]. The linear dependence of JKOnet" on the size of the batch implies that our method can process larger batch sizes for free (to process the entire dataset we need fewer steps in an inverse relationship with the batch size). In practice, this actually increases the speed (less memory swaps). These considerations do not hold for the JKOnet family. Finally, JKOnet" enjoys enhanced representational power and interpretability. JKOnet; generally needs more computation per epoch (primarily related to the number of features) but requires a single epoch. Table we also report the computational complexity of the variants of JKOnet using a vanilla multi-layer perceptron (MIP) with Monge gap regularization [51] instead of a ICNN as a parametrization of the transport map. Despite the success in simplifying the training of transport maps over the of ICNN [51], the Monge gap requires the solution of an optimal transport problem at every inner iteration, an unbearable slowdown [33].
Remark 3.6. Unlike JKOnet, JKOnet" requires the construction of the optimal transport couplings beforehand. However, JKOnet constructs the current estimate of the potential, whereas JKOnet" needs to do so only once, at the beginning. Moreover, discussed in Section 4.1, in Section 4.2, in the application single-cell diffusion dynamics in Section 4.4, and in the ablations in Appendix C.2, this additional cost is minimal.
optimal transport plan each iteration depending
Experiments
The code for the experiments is available at https://githut com/antonioterpin ikonet-star. We include the training and architectural details for the JKOnet' models family in Appendix C. The settings of the baselines considered are the provided by the corresponding pa- pers, reported for completeness in Appendix E, and the hardware setup is described in Appendix C.7. In all the experiments, we allow the models a budget of 1000 epochs.
Our models. We use the following terminology for our method. JKOnet" is the most general non-linear parametrization in Section 3.4 and JKOnet;, introduces the inductive bias 02 = 0g 0. Similarly, we refer to the linear parametrizations by JKOnetř, and JKOnet
Metrics. To evaluate the prediction capabilities (EMD), minyeT(us.pl, f || y| dy(x, y). where μr, and üú are the observed and predicted populations. In particular, we consider the average and standard deviation over a trajectory.
the one-step-ahead earth-mover distance
4.1 Training at lightspeed
Experimental setting. We validate the observations in Section 3.5 comparing (i) the EMD error, (ii) the convergence ratio, and (iii) the time per epoch required by the different methods on synthetic dataset (see Appendix B) consisting of particles subject to non-linear drift, 3++1 = VV((r), with 7 = 0.01, 7 5, and the potential functions V (x) (31)-(45) Appendix F, shown Figure and in Figure 6 in Appendix A.
7

---

JKOnet"
JKOnet
JKOnet
JKOnet
JKOnet
JKOnet-vanilla
NaN
EMD (norm.)
10-1
O
10-3
10-6
(31)
(35)
(40)
(45)
Potential (z)), from left to right: (31)-(45) in Appendix F
Convergence ratio
JKOnet
JKOneti.v
JKOnet'
0.5
JKOnety
JKOnet
0
JKOnet-vanila
500
1,000
10-2 10
100 10
Epochs
Time per epoch [s]
Figure 3: Numerical results of Section 4.1. The scatter plot displays points (IA ys) where x; indexes the potentials in Appendix and y; are the errors (EMD, normalized so that the maximum error among all models and all potentials is 1) obtained with the different models. We mark with NaN each method that has diverged during training. The plot on the bottom-left shows the EMD error trajectory during training (normalized such that and are the minimum and maximum EMD), averaged over all the experiments. The shaded area represents the standard deviation. The box plot analyses the time per epoch required by each method. The statisties are across all epochs and all potential energies.
Results. Figure 3 summarizes our results. All our methods perform uniformly better than the baselines, regardless of the generality. The speed improvement of the JKOnet" models family suggests that theoretically guided loss may provide strong computational benefits on par with sophisticated model architectures. Our linearly parametrized models, JKOnetž and JKOnetī, v require a computational time per epoch comparable to the JKUnet family, but they only need one epoch solve the problem optimally. Our non-linear models, JKOnet" and JKOneti, instead both require significantly lower time per epoch and converge faster than the JKūnet family. In these experiments, the computational cost associated with the optimal transport plans beforehand amounts to as little as 0.03 = 0.01s, and thus has negligible impact on training time. The true and estimated level curves of the potentials are depicted Figure 2 and Figure 6 in Appendix A. Compared to JKOnet, our model also requires simpler architecture: we drop the additional ICNN used in the inner iteration and the related training details (e.g., the strong convexity regularizer and the teacher forcing). Notice that simply replacing the ICNN in JKOnet with vanilla MIP deprives the method of the theoretical connections with optimal transport, which, our experiments, appears to be associated with stability (NaN in Figure 3).
The results suggest orders of magnitude improvement also in terms of accuracy of the predictions. These performance gains can be observed also between the linear and non-lineat parametrization of JKDnet". In view of Proposition 3.4, this is not unexpected: the linear parametrization solves the problem optimally, when the features are representative enough. However, the feature selection presents problem in itself; see e.g. [4. $3 and $4]. Thus, whenever applicable, we invite researchers and practitioners to adopt the linear parametrization, and the non-linear parametrization as demanded by the dimensionality of the problem. We further discuss the known failure modes Appendix G.
4.2 Scaling laws
Experimental setting. We assess the performance of JKOnet;, to recover the correct potential energy given N (1000, 2500, 5000, 7500, 10000) particles in dimension E (10, 20, 30, 40, 50). generated as in Section 4.1.

---

Styblinski-Tane (31)
Holder table (32)
Flowers (33)
-10
N. particles N
10,000
7,500
0.8
10
5,000
0.6
0.5
5
2,500
0.4
1,000
0.2
10 20 30 40 50
10 20 30 40 50
10 20 30 40 50
Dimension d
Dimension
Dimension
Figure 4: Numerical results of Section 4.2, reported in full in Figure 7 Appendix A. The colors represent the EMD error, which appears to scale sublinearly with the dimension d.
Days
Days 12-15
Days
Days 23-23
PC
Figure 5: Visualizations of Section 4.4. The top row shows the two principal components of the scRNA-seq data, ground truth (green, days 1-3, 6-9, 12-15, 18-21, 24-27) and interpolated (blue, days 4-5, 10-1 11, 16-17, 22-23). The bottom row displays the estimated potential level curves over time. The bottom left plot superimposes the same three level curves for days 1-3 (solid), 12-15 (dashed), and 24-27 (dashed with larger spaces) highlight the time-dependency.
Results. We summarize findings in Figure 4 for the potentials (31)-(33) and in Figure 7 in Appendix A for all other the potentials. Since the EMD error related to the Euclidean norm, is expected grow linearly with the dimension d (i.e., along the rows); here, the growth is sublinear up to the point where the number of particles is not informative enough: along the columns, the error decreases again. The time complexity of the computation of the optimal transport plans is influenced linearly by the dimensionality d, and is negligible compared to the solution of the linear program, which depends only on the number of particles; we further discuss these effects Appendix C.2. We thus conclude that JKOnet" is well suited for high-dimensional tasks.
4.3 General energy functionals
Experimental setting. We showcase the capabilities of the JKOnet" models to recover the potential, interaction, and internal energies selected as combinations of the functions in Appendix F2 and noise levels E (0.0, 0.1, 0.2}. To our knowledge, this is the first model to recover all three energy terms. Results. We summarize
findings the right. Compared to the setting in Section 4.1, there are two JKOnet; additional sources of inaccuracies: (i) the noise, which introduces an inevitable sampling error, and the (ii) the estimation of the densities (see Appendix C for training details). Nonetheless, the low EMD errors demonstrate the capability of JKOnet" recover the energy compo- nents that best explain the observed populations.
JKUnet"
10-4
10-2
10
102
EMD error (not norm.)
4.4 Learning single-cell diffusion dynamics
Experimental setting. Understanding the time evolution of cellular processes subject to external stimuli is a fundamental open question in biology. Motivated by the intuition that cells differentiate minimizing some energy functional, deploy JKOnet" analyze the embryoid body single-cell
We exclude the functional (32) from the interaction energies due to numerical issues in the data generation.

---

RNA sequencing (scRNA-seq) data [35] describing the differentiation of human embryonic stem cells over period of 27 days. We follow the data pre- pro-esssssiing in [50. 49]: in particular, we use the same processed artifacts of the embryoid data, which contains the first 100 components of the principal components analysis (PCA) of the data and, following [49], we focus on the first five. The cells are sequenced in five snapshots (days 1-3, 6-9, 12-15, 18-21, 24-27); we visualize the first two principal components in Figure 5. The visualization suggests that the energy governing the evolution is time-varying, possibly due to unobserved factors. For this, we condition the non-linear parametrization in Section 3 on time & and minimize the loss
Σ 1= xd
1 (041 mnd dr d=l(01,
VV (V ++
To predict the evolution of the particles, then, we use the implicit scheme (see Appendix B)
T++1 = 2, V (V 1).
We train the time-varying extension of J JKOnet and JKOnet-vanilla for 100 epochs on 60% of the data at each time and we compute the EMD between the observed pt (40% remaining data) and one-step ahead prediction مa at each timestep. We then average over the trajectory and report the statistics for 5 seeds.
Results. We display the time evolution of the first two principal components of the level curves of the inferred potential energy in Figure 5, along with the cells trajectory (in green the data, in blue the interpo- lated predictions). As indicated by the table on the right, JKOnet" outperforms JKOnet. We also compare JKOnet" with recent work in the literature which fo- cuses on the slightly different setting, namely the infer- ence μt from the evolution at all other time steps pkx k / without train/test split of the data (the numerical values are taken directly from [12, 49] and our statistics computed over the timesteps). Since the experimen- tal setting slightly differs, we limit ourselves to observe that JKOnet" achieves state-of-the-art performance, but with significantly lower training time: JKOnet" trains in few minutes, while the methods listed take hours run. We further discuss these results in Appendix E.
Algorithm
EMD
JKOnet [9]
1.363 + 0.214
JKOnet-vanilly [9]
3,237 = 1.135
TrajectoryNet [50]
0.848 +
Reg. CNF [17]
0.825 + --
DSB [14]
0.862 = 0.023
I-CFM [49]
0.872 = 0.087
SB-CFM [49]
1.221 + 0.380
OT-CFM [49]
0.790 = 0.068
NLSB [28]
0.74 +
MIOFLOW [23]
0.79 +
DMSB [12]
0.67
JKOnet
0.624 + 0.007
(arminaay
Conclusion and limitations
Contributions. We introduced JKOnet", model which recovers the energy functionals governing various classes of diffusion processes. The model is based on the novel study of the first-order optimality conditions of the JKO scheme, which drastically simplifies the learning task. In particular, we replace the complex bilevel optimization problem with a simple mean-square error, outperforming existing methods terms of computational cost, solution accuracy, and expressiveness. In the prediction of cellular processes, JKOnet" achieves state-of-the-art performance and trains in less than minute, compared to the hours of all existing methods.
Limitations. Our work did not address a few important challenges, which we believe be exciting open questions. On the practical side, JKOnet" owns its performances to loss function motivated by deep theoretical results. However, its architecture still "vanilla" and we did not investigate data domains like images. Moreover, this work does not investigate in detail the choice of features for the linear parametrization, which in our experiments displays extremely promising results nonetheless. We further discuss the known failure modes in Appendix G.
We expect the approach followed in this work to apply to other exciting avenues of applied machine learning research, such as population steering [47], reinforcement learning [37, 44, 48), diffusion models [16, 22, 54] and transformers [19, 52].
Outlook.
10

---

Acknowledgments and Disclosure of Funding
We thank Elisabetta Fedele for the helpful and detailed comments. We also thank the reviewers for their constructive suggestions. This project has received funding from the Swiss National Science Foundation under the NCCR Automation (grant agreement 5INF40 180545).
References
[1] David Alvarez-Mens Melis, Yair Schiff, and Youssef Mrouch. Optimizing functionals the space of probabilities with input convex neural networks. Transactions on Machine Learning Research, 2022.
[2] Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pages 155- PMLR, 2017.
[3] Jean-David Benamou, Guillaume Carlier, and Maxime Laborde. An augmented Lagrangian approach to Wa Wasserstein gradient flows and applications. ESAIM: Proceedings and surveys, 54:1-17, 2016.
[4] Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, volume 4. Springer, 2006.
[5] Riccardo Bonalli and Alessandro Rudi. non-parametric learning stochastic differential equations with fast rates of convergence. arXiv preprint arX55:2305.15557, 2023.
[6] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[7] Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications Pure and Applied Mathematics, 44(4):375-417, 1991.
[8] Charlotte Bunne, Andreas Krause, and Marco Cuturi. Supervised training of conditional Monge maps. Advances in Neural Information Processing Systems, 35:6859-6872, 2022.
[9] Charlotte Bunne, Laetitia Meng- Paang-Ppaaaanths, Andreas Krause, and Marco Cuturi. Proximal optimal transport modeling of population dynamics. In International Conference on Artificial Intelligence and Statisties (AISTATS), 2022.
[10] Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia Del Castillo, Mitch Levesque, Kjong-Vas Van Lehmann, Lueas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single- cell perturbation responses using neural optimal transport. Nature Methods, pages 1-10, 2023.
[11] José A Carrillo, Katy Craig, Li Wang, and Chaozhen Wei. Primal dual methods for Wasserstein gradient flows. Foundations of Computational Mathematics, pages 1-55, 2022.
[12] Tianrong Chen, Guan-Horng Liu, Molei Tao, and Evangelos Theodorou. Deep momentum multi-marginal schrödinger bridge. Advances in Neural Information Processing Systems, 36, 2024.
[13] Marco Cuturi, Laetitia Meng-Papaxanthos. Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (OTT): JAX toolbox for all things Wasserstein. arXiv preprint 22Xiv:2201.12322 2022.
[14] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications score-based generative modeling. Advances in Neural Information Processing Systems, 14917695-17709 2021.
[15] Rossin Erbe, Genevieve Stein-O' Brien, and Elana Fertig. Transcriptomic forecasting with neural ordinary differential equations. Patterns, 4(8), 2023.

---

[16] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transform- ers for high- nigh-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
[17] Chris Finlay, Jöm--Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ode: the world of jacobian and kinetic regularization. In International Conference on Machine Learning pages 1154-3164, PMLR, 2020.
[18] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78))1-8, 2021.
[19] Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. A mathematical perspective transformers. arXiv preprint arXiv:2312.10799 2023.
[20] Jacob Hanna, Krishanu Saha, Bernardo Pando, Jeroen Van Zon, Christopher J Lengner, Menno P Creyghton, Alexander van Oudenaarden, and Rudolf Jaenisch. Direct cell reprogramming is a stochastic process amenable to acceleration. Nature, 462(7273):595--66 601, 2009.
[21] Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-leve diffusions with generative RNNs. In International Conference on Machine Learning, pages 2417-2426. PMLR, 2016.
[22] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[23] Guillaume Huguet, Daniel Sumner Magruder, Alexander Tong, Oluwadamilola Fasina, Manik Kuchroo, Guy Wolf, and Smita Krishnaswamy, Manifold interpolating optimal- flows for trajectory inference. Advances in Neural Information Processing Systems, 35:29705-2971 18, 2022.
[24] Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker- Planck equation. SIAM journal on mathematical analysis, 1911):1117 1998.
[25] Carson Kent, Jiajin Li, Jose Blanchet, and Peter W Glynn. Modified Frank Wolfe in probability space. In Advances in Neural Information Processing Systems, volume 34, pages Curran Associates, Inc., 2021.
[26] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[27] Peter E Kloeden, Eckhard Platen, Peter Kloeden, and Eckhard Platen. Stochastic differential equations. Springer, 1992.
[28] Takeshi Koshizuka and Issei Sato. Neural Lagrangian Schrödinger bridge: Diffusion modeling for population dynamics. arXiv preprint S2ti::2204.085555 2022.
[29] Nicolas Lanzetti, Saverio Bolognani, and Florian Dörfler. First-order conditions for optimization in the Wasserstein space. SIAM Journal on Mathematics of Data Science, 2024. In press.
[30] Nicolas Lanzetti, Antonio Terpin, and Florian Dörfler. Variational analysis in the Wasserstein space. arXiv preprint arXiv:2406. 10676, 2024.
[31] Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Scungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invarin neural networks. In International conference on machine learning, pages 3744 3753. PMLR, 2019.
[32] Xuechen Li, Ting-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradi- ents for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pages 3870-3882. PMLR, 2020.
12

---

[33] Jianzhou Luo, Dingchuan Yang, and Ke Wei. Improved complexity analysis of the Sinkhorn and Greenkhorn algorithms for optimal transport. arXiv preprint arXiv:2305. 2023.
[34] 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,
2021.
[35] Kevin R Moon, David van Dijk, Zheng Wang, Scott Gigante, Daniel Burkhardt, William Chen, Kristina Yim, Antonia van den Elzen, Matthew Hirn, Ronald R Coifman, et al. Vi- sualizing structure and transitions in high-dimensional biological data. Nature biotechnology, 37(12):1422-1422 2019.
[36] Rob Morris, Ignacio Sancho-Mand Martinez, Tatyana Sharpee, and Juan Carlos Izpisua Belmonte. Mathematical approaches to modeling development and reprogramming. Proceedings of the National Academy of Sciences, 111(14):5076555 000776655022200 2014.
[37] Mirco Mutti, Riccardo De Santi, Piersilvio De Bartolomeis, and Marcello Restelli. Challenging common assumptions in convex reinforcement learning. Advances in Neural Information Processing Systems, 35:4489-4502, 2022.
[38] Gabriel Peyré. Entropic approximation of Wasserstein gradient flows. SIAM Journal on Imaging Sciences, 8(4):2323 21452,333-3351,,,, 2015.
[39] Gabriel Peyré, Marco Cuturi, al. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 116--6):355-0777 2019.
[40] Jane Reece, Lisa Urry, and Michael Cain. Campbell biology. Pearson, 2017.
[41] Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58- 63):94, 2015.
[42] Christopher Scarvelis and Justin Solomon. Riemannian metric learning via optimal transport. In The Eleventh International Conference Learning Representations, 2023.
[43] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin. Peter Berube, et al. Optimal-transporr analysis of single-cell gene expression identifies developmental trajectories reprogramming. Cell, 176(4):928-944 943, 2019.
[44] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint drXiv:170..06347, 2017.
[45] Simon Fraser University. Test functions and datasets for optimization, 2023. Accessed: 2024- 05-15.
[46] Ole Tange. Gnu parallel, September 2024.
[47] Antonio Terpin, Nicolas Lanzetti, and Florian Dörfler. Dynamic programming in probability spaces via optimal transport. SIAM Journal on Control and Optimization, 62(2):1188) 18:-11068. 2024.
[48] Antonio Terpin, Nicolas Lanzetti, Batuhan Yardim, Florian Dorfler, and Giorgia Ramponi. Trust region policy optimization with optimal transport discrepancies: Duality and algorithm for continuous actions. Advances in Neural Information Processing Systems, 15:19786-19997, 2022.
[49] Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector Brookss, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based genera- tive models with minibatch optimal transport. Transactions on Machine Learning Research. 2024.
[50] Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajecto- rynet: A dynamic optimal transport network for modeling cellular dynamics. In International Conference Machine Learning, pages 9526-9536. PMLR, 2020.
13

---

[51] Théo Uscidda and Marco Cuturi. The Monge gap: A regularizer to learn all transport maps, In International Conference on Machine Learning. pages 34709-34733. PMLR, 2023.
[52] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[53] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Coumaapaa,, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C Carey, ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, António H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python, 2020.
[54] Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys, 2022.
[55] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barabaas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in Neural Information Processing Systems, 30, 2017.
14

---

A Eye candies
We collect the level curves of the ground-trune potentials and the ones recovered with each of our methods, together with the predictions of the particles evolution super-imposed to the ground-truth population data, in Figure 6. The plots discussed in Section 4.2 are in Figure 7.
Styblinski-Tand (31)
Holder table (32)
Flowers (33)
Oakley- Ohagan (34)
Watershed (35)
Ishigami (36)
Friedman (37)
Sphere (38)
Bohachevski (39)
Wavy plateau (40)
Zig-zu zag ridge (41)
Double exponential (42)
ReLU (43)
Rotational (44)
Flat (45)
Figure 6: Level curves of the true (green-colored) and estimated (blue-colored) potentials Ap- pendix F, from top-left to bottom-right, row-hy-row.
B Details on the data generation and prediction of the particles evolution
The prediction scheme for each particle,
- rt+1 = Ft 7 VV - / VU(Z--y)dpitl + / 2T
(13)
where n, is sampled from the d-dimensional Gaussian distribution at each 1, follows from the Fuler-Maruyama discretization (see [27]) of the diffusion process
dx(() VV f VU(X()) v)d (2)dW(I).
(14)
In particular, we sample 2N particles (in Section 4. 1, 1000: Section 4.2, N is indicated in the experiment) uniformly in [-4,4] 4] and update their state according to (13) for 5 timesteps, for a total of snapshots including the initialization. Then, the data of the first N particles is used for training. and the data of the remaining N particles is left out for testing.
15

---

Styblinski-Tand (31)
Holder table (32)
Flowers (33)
-10
N. particles N
10,000
7,500
0.8
1
10
5,000
0.6
2,500
0.4
0.5
5
1,000
0.2
Oakley-Ohan Ohagan (34)
Watershed (35)
Ishigami (36)
-10
-10
N. of particles N
10,000
0.3
7,500
5,000
0.2
2
2,500
1,000
0.1
Friedman (37)
Sphere (38)
Bohachevsky (39)
-10 -
-10
N. of particles N
10,000
1.8
0.8
7,500
1.6
10
1.4
5,000
0.6
1.2
2,500
5
1,000
0.4
Wavy Plateau (40)
Zig-zag ridge (41)
Double exp (42)
-10-
N. of particles N
10,000
1.4
0.3
7,500
1.2
2.5
0.2
5,000
2,500
1,000
0.8
0.1
1.5
ReLU (43)
Rotational (44)
Flat (45)
.10
.10 -10-3
N. particles N
10,000
0.8
7,500
0.6
5,000
0.4
2,500
0.2
1,000
10 20 30 40 50
10 20 30 40 50
10 20 30 40 50
Dimension d
Dimension d
Dimension
Figure 7: Additional numerical results for Section 4.2. Each heat-map corresponds to functional in Appendix F, from top-left to bottom-right, row-by-row. The r-axis corresponds to the dimension and the y-axis corresponds to the number of particles. The colors represent the EMD error. Thus, method that scales well to high-dimensional settings should display relatively stable color along rows: the error related to norm and, thus, is linear the dimension d; here, growth is sublinear.
16

---

ground-truth (3 trajectories)
predicted
2
1.5
0.25 0.5 0.75
0
0.25
0.5 0.75
0.25 0.5 0.75 0 0.25 0.5 0.75 1
Figure Comparisons between implicit and explicit losses and predictions for time-varying potentials for 3 different particles trajectories. From left to right: explicit loss (16) explicit prediction (13), explicit loss (16) explicit prediction (13) with more observations, implicit loss (17) explicit prediction (13), implicit loss (17) implicit prediction (15).
When only the potential energy is considered, (VV(1)+1 which is also suggested by the loss in Proposition 3.1. In particular, we adopt this scheme for the vary-ng potential in Section 4.4:
adopt the implicit scheme Fe+1 3t
I++1 = 2t
TV VVV
(15)
We refer this scheme as implicit. and to (13) explicit, since in (15) x++1 appears in both sides of the equation and its value is found solving an optimization problem, whereas in (13) it can be computed directly.
Implicit vs explicit prediction scheme with time-varying potential. We now discuss one potential pitfall when training JKOnet" with time-varying potentials. For this, work on R (instead of the space of probability measures) and consider the analysis in Section 3.1 for the time-varying function
0.2 < S 0.3 or 0.7 S 0.8,
V(I, = -0.75 x² otherwise.
Suppose one adapts the loss in (7) so that the value of the potential the previous timestep is used, i.c.,
T-1 E/
VV + !=U 34x3d
+ (0141 - 21) dx
(16)
This explicit adaptation, in our experiments, requires more observations to correctly recover the energy potential; the first two figures on the left in Figure 8. Since in Section 4.4 we given only timesteps, we explored the use of an implicit adaptation, i.c.
Σ
(
1 (2441 21) d (2t,
(17)
RdXRd
However, if the predictions are made with the explicit scheme in (13), the third figure in Figure 8 is obtained, which suggests that the predictions are shifted time. With the implicit prediction scheme in (15), instead, we able to recover the correct potential with only 10 timesteps; see the right-most figure in Figure We rely on these considerations to learn the time-varying energy potential of the scRNA-seq data in Section 4.4.
C Training and model details
This section describes the training details for our model JKDnet".
17

---

[(
( (3) +41 n (a,
(5)
M=--1
#111 Pe+1 i.g.t=
Compute Couplings
Fit Densities
Snapshot
T
Measurement System
Figure 9: Data pipeline for JKDnet"..
C.1 Data pre-processing
In principle, measurement system is used to take a trajectory of populations. Then, we compute the optimal couplings between consecutive snapshots and, if needed, infer the density of each snapshot. This information is then combined to get the training set (see Figure 9).
C.2 Computation of the optimal couplings
The raw population data No, Mi..... Hr typically empirical. Since they not optimization variables in the learning task, the optimal couplings "oo Ti.... "T can be computed beforehand. Differently from [9], do not need to track the gradients within the Sinkhorn algorithm compute the bac--propagation [13], and we do not need to solve any optimal transport problem during training. If the number of particles per snapshot is very large, one may consider splitting the single snapshot into multiple ones and computing the couplings of the sub-sampled population. For datasets larger than 1000 particles, compute the couplings in batches of 1000. In our experiments, the optimal transport problem via linear programming using the POT library [18]. In Figure 10, we compare the performances resulting from training JKOnet" to recover the potentials in Appendix F when the couplings are computed with plain linear programming or with Sinkhorn-type algorithms [13] with various degrees of regularization e. particular, we conclude that, long as the couplings are close to the correct one, the algorithm used to compute them does not impact the performance of JKDnet". However, small regularizers slow down the Sinkhorn algorithm and, thus, we prefer to directly solve the linear program without regularization. In general, the solver choice can be considered an additional knob that researchers and practitioners can tune when deploying JKOnet".
solve
To instead investigate how the dimensionality affects the computation of the optimal transport plan, recall that the optimal transport problem between two empirical measures μ1 = -En ((:;)52, and M #2 Σ/= p(y)) can be cast as the linear program
min Ni 20
Ciylis
s.t.
Σ Mi (2(U3)
Vj (1,..,,
i=1
M
Σ Mis = (m;) #=1
(1,
where Yi is the weight of the coupling between and Øý. (xx, Uj) Yi and Gj || x2,1? is the itth entry of the cost matrix C. Therefore, the dimensionality d of the data impacts the solving time only relation to the construction of the cost matrix C (in particular, linearly). However, this cost is negligible and dwarfed by the actual solution of the linear program. In particular, the dimensionality does not affect the size of the linear program. Moreover, the optimal plans different timesteps can be computed in parallel. As result, only the number of particles (i.c., the size of the
18

---

No reg.
0.01 DE 0.1 oE 1.0
Computation time [s]
600
Wasserstein eIror
0.6
Time [s]
400
0.4
200
0.2
F
An
10
10-1
10"
No reg.
1000 2500
5000
7500
10000
N
Figure 10: Comparison of the performances resulting from training JKOnet" to recover the potentials when the couplings computed with plain linear programming (No reg.) with Sinkhorn-type algorithms with regularization E. In the plot on the left, the left y-axis (in red, crosses) represents the EMD error (normalized that 0 and 1 are the minimum and maximum obtained with the different methods), and the right y-axis (in blue, circles) represents the computation time in seconds. In the plot on the left, we report the computational time for the different methods for different numbers of particles. The statistics are computed on all the data generated with the potentials listed in Appendix F with the same settings as Sections 4.1 and 4.2, with no batching for the optimal transport couplings computation and using the default configuration in the POT library [18] for linear programming and in the OTT JAX library [6] for the Sinkhorn algorithm.
dataset) is scaling bottleneck. In practice, as showcased by the real-world application in Section 4.4, this is often not an issue. For completeness, we report an ablation on the number of particles for the computation times of plain linear programming and Sinkhorn algorithm with different values of the regularizer in Figure 10, using the default configuration in the POT library [18] for linear programming and in the OTT-JAX library [6] for the Sinkhorn algorithm. This analysis suggests that, while the number of particles affects the computational time in an exponential fashion, the computation of optimal transport plans does not pose a problem for reasonably sized datasets. Also, highlight that this operation is done only once and independently from, c.g., the network architecture. Finally, as discussed above, batching strategies can help reduce the computational cost. Specifically, rather than computing the couplings between N particles, one can compute [N/b] couplings between b particles, sampled uniformly from the entire population at each timestep, and then aggregate them.
C.3 Estimation of and Vp
The case with Og non-zero imposes the estimation of the density Pe and its gradient Vp: from the empirical probability measures μr. To the extent of estimating Pt, there are many viable options (c.g.. see the SciPy package [53]). We use a mixture 10 gaussians for all the experiments in this paper. To compute V Pt. we rely on the autograd feature of JAX [6]. The complexity related to Pt and Pe. instead, is bypassed if there is no internal energy in (10) (i.e., 0g = 0).
C.4 Optimizer
We use the Adam optimizer [26] with the parameters B1 = 0.9, 82 = 0.999, 1e-8, and constant learning rate Ir = 1e-3. The model is trained with gradient clipping with maximum global norm for an update of 10. We process the data in batches of 250.
C.5 Network architecture
The neural networks of potential and interaction energies are multi-layer perceptrons with 2 hidden layers of size 64 with softplus activation functions and one-dimensional output layer (ef. [9]). Future work shall investigate different architectures for the interaction energy, for instance using Sret-TTranfformers [31, 52] or Deep Sets [55].
19

---

C.6
Features selection for the linear approximation
The features used in this work are polynomials up to degree and radial basis exponential functions ||--|l
exp
with 0.5 and in the discretized (10 points per dimension) grid [-4, 4
Finally, regularizer À 0.01 on the square norm of the parameters 0 [0T 03
C.7 Hardware
The empirical data was collected entirely an Ubuntu 22.04 machine equipped with an AMD Ryzen Threadripper PRO 5995WX processor and a Nvidia RTX 4090 GPU. Interestingly, both JKOnet" and J demonstrate comparable computational times when executed on a GPU, highlighting the method's strong parallelizability. Among the methods compared in Section 4.1, JKOnet" the only one that significantly benefits from GPU parallelization in our experiments, while all others exhibit similar computational times when run on CPU. When parallelizing over CPU cores, we used on parallel [46].
D Proofs
D.1 Preliminaries
We briefly collect and summarizes the notation, definitions, and results of [30] used in the proofs of this work. For more intuition and details, refer the reader to [30]. A sequence (Kn)enN S P(R") converges narrowly to if Jg. (zz)d)n(z) + Jg (x))du(x) for all bounded continuous functions á IRd + R. We say that the convergence is "in Wasserstein" if W2 μ)) + 0 or, equivalently, Jaa (())d(()))))) + fge ((z)dμ((z) for all continuous functions á Rd + R with at most quadratic
A(1 | for some A 0). For i, N, S í S m, denote by T, (Re)m Rdd the projection map the 7th component, i.e., (T1, Tm) I, We compose (z)).
growth (i.e., ((x)
projections as n. (z) =
Wasserstein calculus. We use two notions of subdifferentabiiity in the Wasserstein space.
Regular Wasserstein subdifferentiaaiiityy. A transport plan & P( Rd) belongs the regular subdifferential of functional P(RT) + i at p E P(Rd), denoted å.( if for all μ E P(Rd) and E (11, 72 ( we have
J(m) - J())
max QEP(R'XR' XR* xRd) Fing a
J
8 Vuda 0,0) (P,
Here, a can be interpreted as a coupling of the tangent vectors & (the subgradient) and & (the tangent vector μ (T).
(General) Wasserstein subdiffirrnntiabiliyy transport plan & P(Rd Rd) belongs to the (general) subdifferential of functional J P((R)) R jt, denoted aJ if there sequences (NN)REN C P(R.)) and (En)nEN C P(R Rd) so that (i) En E aJ(()). (ii) &n
converges in Wasserstein to E.
From these definitions, the analogous ones for the supergradient follows: The regular supergradients of J at j the elements of (á(-))(z) and the supergradients the elements of -O(-))(p)) We then call the grudient of J the unique element, if it exists, of -å(-J)(p) åJ and we say that is differentiable.
Necessary Conditions for Optimality. For J proper and lower semi-continuous w.r.t. convergence in Wasserstein (i.e., if μ4, converges in Wasserstein to μ. then lim inf J (ux) 2 (u)),, consider the optimization problem
inf J(n). WEP(R4)
(18)
20

---

Then, if E P(RR") is optimal for (18), i.e. J(2)) = inf (EP(RR) J(p,, it must satisfy the "Fermat's
rule Wasserstein space" (see [30, Theorem 3.3]):
14" x d0 E (J(p)). In particular, this implies that for some E aJ((() must hold
Oμ- =
(19)
=
de("",
(20)
RaxRd
Outline of the appendix and proof' strategy Although the statement of Proposition 3.1 follows directly from Proposition 3.2, we prove it separately in Appendix D.3 as it is the simplest setting and, thus, the easiest for the reader familiarize with the techniques used this work. Since the internal energy (perhaps surprisingly) simplifies the setting by restricting to absolutely continuous measures (i.e., {ur P(RR)
< ). we prove also the statement for potential and interaction energy only in Appendix D.4. We then provide the proof for the most general statement in Appendix D.5. All the solution. Second, we characterize
proofs follow the same recipe. First, we prove existence Wasserstein subgradients for the functional considered. Finally, we conclude deploying [30, Theorem 3.3]. Because the Wasserstein subgradients will be constructed using the subdifferential calculus rules (cf. [30, Proposition 2.17 and Corollary 2. 28)8), collect them all in Appendix D.2. This appendix ends with Appendix D.6, in which we prove Proposition 3,4 using standard gd optimization arguments.
D.2 Wasserstein subgradients
(Scaled) Wasserstein distance. By [30, Corollary 2.12]. the (Wasserstein) subgradients of W2(u, p2)² at P(R4) are all of the form (τı, ग - 22) for an optimal transport plan Ya between and y4- Then, by [30, Corollary 2.18), the subgradients of 2r (1, are
n2
7t-
(21)
#
Since the Wasserstein distance is in general not regularly subdifferentiable (cf. [30, Proposition 2.6)). it is not differentiable.
Potential energy. Under the assumptions in Propositions 3.1, 3.2 and 3.4, we deploy [30, Proposi- tion 2. 13] to conclude that the potential energy is differentiable at E P(IR") with gradient given by
(Id, VV) 4/.
(22)
Interaction energy. Under the assumptions in Propositions 3.2 and 3.4, we deploy [30, Proposition 2.15] to conclude that the interaction energy is differentiable at E P((Rd) with gradient given by
[dd, {
VU(Id=
μ.
(23)
Internal energy. Under the assumptions in Propositions 3.2 and 3.4, we deploy [29, Example 2.3] and the consistency of the tangent space (ef. [30, $2.2)) to conclude that the internal energy is differentiable at < dæ with gradient given by
Vp
ld.
(24)
In particular, here we consider μ < dr since the internal energy is otherwise +oo by definition and is certainly not a minimum.
D.3 Proof of Proposition 3.1
The JKO step in Proposition 3.1 is the optimization problem, resembling (18), 27
J(H) = [ v(z)dp(z) + We(m, We(
inf (E)(
Since lower bounded, up to replacing V by V mi-miin V(x), we can without loss of generality assume that V is non-negative. We now proceed in three steps.
21

---

Existence of solution.
J has closed level sets W.E.L. narrow convergence, As an intermediate step to prove compactness of the level sets prove their closedness, which is equivalent to lower semi-continuity of J. Since V is continuous and lower bounded, the functional μ + Je V (z))(((z) is lower semi- continuous [29]. The Wasserstein distance is well-known be lower cemi-continuous [30]. Thus, their sum J is also lower semi-continuous.
J has compact (w.rf. narrow convergence) level sets. Without loss generality, assume À is large enough that the level set is not empty. Otherwise, the statement is trivial. Closedness of the level sets follows directly from lower semi amiiccontinutt, thus suffices to prove that level sets are contained in a compact set, because closed subsets of compact spaces are compact. By non-negativity of V.
((R") J(p) S 1) c n E P(RR))
27 W2(H, p2) S =d}
= {u E P(RR) We),,m) & 2r
Since the Wasserstein distance has compact level sets [29] w.r.t. narrow convergence, we conclude.
Compact level sets imply existence of a solution. Let a = "nf (()) and let (JN)NEN be minimizing sequence so that ( - a. Then, by definition of convergence, Pn E { P(R"): J(p) S 2 for all sufficiently large. By compactness of the level sets, Mn, converges narrowly (up subsequences)
some μ E {u P(TRR)) [(() A}- Since compactness of the level sets implies closedness and a functional with closed level sets is lower semi-continuous (w.r.t. narrow convergence) we conclude that J (μ) lim inf J ( a. Thus, μ is minimizer of J.
Wasserstein subgradients. By [30, Proposition 2.17]. we combine (21) and (22) to conclude that any subgradient of at M4+1 must of the form (I1, T2 + T3) 0, for some "sum" coupling E P(RRd x I& IR"), 122 (Id, VV) MM++ and T13 = (m), -=="² # 4 (ef. [30, Definition 2.1)). with 7 being an optimal transport plan between Mt+1 and f42.
Necessary condition for optimality. Any minimizer Me+1 of satisfies Opees E (1(41)) Thus, for pr+1 to be optimal, there must exist 7/41 and a, so that (ef. (20))
0
||02 + vzll da( 01,
f JRIXRARR XB VV + 22 da
-fe JRExRd
[V + u22 v2l (1
[V +2l (
-|e Seeeee
क. 122
(1)+1,
T
#
2
VV(x++1)+ It+1 Ft
d7/4 (2441)
Finally, consider 7/- = (T2, 111 2 to conclude.
D.4 The case of potential and interaction energies
As preliminary step to Proposition 3.2, consider the case B = 0. The JKO step is the optimization problem, resembling (18),
J(p) [ v(z)du(z) + S U(x y)d(u x ()) We(H,
1
(25)
2T
Since and U are lower bounded, up to replacing J with inf ( J(p), we can without loss of generality assume that / is non-negative.
22

---

Existence of a solution. w.r.t. narrow convergence [29], the proof of existence is analogous to the one in Proposition 3.1.
Since also μ μ fRe U(z v)d(a x (z) y) is lower semi-continuont
Wasserstein subgradients. By [30, Proposition 2.17], combine (22) and (23) to conclude that each subgradient of the functional
# + J v(z)du(2)+ + J JR. U( - v)dp(y)du(x))
(26)
1441 is of the form (I1, 42 + +s) for some coupling a' E P(R Rd RR) such that
T12 = (Id, 4/0 and Ti = (Id, Sg (Uldd=z)dpra (z)) Mo-1. But then a'
(Id, Sg. VU -x -x)ppt1(())) # 144 and (26) has the unique subgradient
Id. VV
VU (I)dd=))d(e+1(s))
Hz+1-
(27)
24
#
Similarly, we conclude that (Id, (VV+ V =z))d(()))( # H is the unique supergra-
dient of (26) y and, thus, it differentiable there. We thus deploy again [30, Proposition 2.17] to combine (27) with (21) and conclude that the subgradients of (25) are of the form (T1, T2 + T for a E P(R.
(Id, VV + S&a VU -2)dde++1111 #/++1 and (I,, x2) a (mi, -===² # M41 (cf. [30, Definition 2.1]), with 7441 being an optimal transport plan between 14+11 and paa.
R Rdd). #12
Necessary condition for optimality. The steps are analogous to Appendix D.3. Any minimizer Me+1 of J satisfies Opeia E @d((1))). Thus, for pc+1 to be optimal, there must exist %441 and so that (cf. (20))
[ dexx<xRR x
102 + Uz V2 (a
/
VV + VU(x++1 (U(=+11=)dduua x)d (x) + 02 (1,
=/r
V V((I++)
VU(x++1 VU(=HI-2)d/ae (0ite+1-))ddu++11) (x) + d((TT13) + (x) (11+1, u2)
JR
B2d
VV(I++1) +
VU(I+11 z2)dp++1 R++1 - It d d%/+1 (
Bdd
Finally, consider 7/- (T2, TT ) conclude. Importantly, for 3 we do need restrict
to o & dr in Proposition 3.2, so that the statement holds regardless of 33 (see also Appendix D.5).
D.5 Proof of Proposition 3.2
For B = 0, see Appendix D.4. Let here. The JKO step Proposition 3.1 the optimization problem, resembling (18),
SEe V(x)dp(x) + Serxe U(x - y)d(H x (x) +B Sax p(x:) log(z)dæ + 27 (1),
J(p)= =
if E C, else.
(28)
+00
Since V.U, and Snu p(x) log(x)dæ are lower bounded, up to replacing J with inf (EP(E))))))) (()). we can without loss of generality assume that J non-negative.
U(r - v)d(H x 1) y) and Existence of a solution. Since also μ + Jee x
S2a p(x) +00
log())))gg()))))))) log(x)dr f & C,
else.
are lower semi-continuous w.r.t. narrow convergence [29], the proof of existence analogous the one in Proposition 3.1.
23

---

Wasserstein subgradients. When / > O, only p4+1 < dL can be optimal and we are thus interested in characterizing the subgradients of (28) only at these probability measures. Let P++1 be the density of 14+1- Consider first the functional
14+1 + { V(x)dumn(0)) Jg. U(x (ud--)dde JRd
J/x. p++1(2) 1
(29)
By [30, Proposition 2.17], we combine (27) and (24) to conclude that each subgradient of the func- tional (29) at 1 is of the form (TI, T2 + ma) a' for some coupling a' P(R x Rd Rd) such that m22
[Id. VV
VU - xdd (x)) Met1 and T13 = Id. 3 PeLLL P4+1-
P++1
But then a' gradient
ld. VV
fgd VU (B(Id=z)ddu-te =))du+1((())))) VOLLL
Mc+1 and (26) has the unique sub- P+tl
Id, VVV Rdd
/
VU 22)d41)))
Pc+1
Me+1-
(30)
Pi+1 #
However, note that now (26) is not differentiable, as it can attain the value +00. We can thus deploy again [30, Proposition 2.17], together with the fact the squared Wasserstein distance differentible at absolutely continuous measures [30, Proposition 2.6], to combine (30) with (21) and conclude that the subgradients of (25) are of the form (TI, T2 + TT) # a, for a P(IR" x IRd Id),, T12 a
Id, VV + .R. VU -z)d (z) + 3 Ppr+1
Mo+1 and (TT x2) # a (TI, - # M4t1 (cf.
Pe++
[30, Definition 2.1]), with 1441 being an optimal transport plan between M441 and μ4z
Necessary condition for optimality, The steps are analogous to Appendices D.3 and D.4. Any minimizer 14++1 of J satisfies Open1 (J(++))). Thus, for pe+1 to be optimal, there must exist 1441 and a, so that (cf. (20))
102 + vzll uall2da da(
d xRd
V(( Jg VU(I4+1 -U(FE1122dd x2)d4+1(2)
2
+ 3 (141(00+1)
+ D2
01,
pe(Ee+1)
V(( + V x)d14+11())
Rd xRd
Jud
V (142(1441)
+ D2
po(0++1)
-
VV(
VU(I++1 x)d1(+1(x)
Rd
Rd
2
+ 3
+ Ft+1 It
driu (2441,
pe(2++1)
Finally, consider " (TT2, Ti) 1441 to conclude.

---

D.6 Proof of Proposition 3.4
We start with an explicit expression for the loss (11) in the case of linearly parametrized functionals:
T
(ol(Zt))) To + (Fe Vib2(1ee11 - 2(41) duc d 02
C(0 0a)
Σ
2 k= xdx
V (144((041)) + +- -x1+1 - 21) dr slt1).
03
(1241)
Since J is convex and quadratic in 0, we can find its minimum by setting the gradient to O. In particular, the derivative of L w.r.t. 01, ((0,, 02 0z), reads
(11(1))1) 0i (fe Voz(zt41 - x(+1) d 02
Σ 1=0 xaxB
(
V Pi+i (1e+1)
1 (11t1 - zi) d2
+ 0g
(1(41
Σ
(or(z))1 (U4t( + 1 (11t1 (01111 zi) (xe( 0441), RddRd
where we used Leibniz integral rule to interchange gradient and integral. Similarly, 1
T-1
(Lee Voa(Te+ -e4a) (1(+1) R2
22 (()
Σ
(C+z((Fe22) To+ 1 - xx)) dre 1441)
(())
T-1
Σ
(142()))11
1 (1+1(1(41)T0 + ( - 2i) dnllas T441) (141))
Rd xRd
1=0
We split integrals and sums to get
0. (())
Σ
V
(
(pperne d
Σ Rdd
-
+
Σ RR (Le
Voze - 2le (d ve(aets)" (pensloons))
8g (()
RR (1pz
[ Vozluete de (1141 (e)ehu((ssseett)
(et+(( (1+1(1141)
ye (1( T (oeen)(
0.. C(0)
jzd
Σ
1 Pe (Iz+1 xe)dda(ae 2mi1)
R' Ptti (x++1)
t=o
We stack these expressions to compactly write the gradient of L
T-1
l ve( (11 T ddptt (00e11(00001)
V VoL(O) =
Σ
Veleet (2,,
Rd xnd
25

---

The expression for follows solving V 0 (and shifting the indices in the first sums from to t' with t' 1).
E
Baselines settings and further comparisons
JKOnet.
We use the default configuration provided in [9]. Specifically:
The potential energy is parametrized and optimized in the same way for JKOnet" (cf. Appen- dices C.4 and C.5).
The optimal transport map is parametrized with an ICNN with two layers of 64 neurons each, with positive weights and Gaussian initialization. The optimization of the transport map is performed with the Adam optimizer [26] with learning rate = 0.01, 81 = 0.5, 32 = 0.9 and ε = 1e 8. The transport map is optimized with 100 iterations (i.e., do not the fixed point approach).
We train the model for 100 epochs in batch sizes (number of particles per snapshot) of 250, without teacher forcing and without any convex regularization. The Sinkhorn algorithm was run with 1.
JKOnet vanilla. The only difference between this method and JKOnet that we replace the ICNN with vanilla MLP with two layers of 64 neurons each as parametrization of the optimal transport map.
JKOnet with Monge Gap. The settings the same as the vanilla JKOnet, with the Monge gap [51] as an additional term in the loss function for the optimal transport map (inner loop). For the computation of the Monge gap used the recommended implementation in [51], namely the one in the OTT-JAX package [13]. The coefficient the Monge gap regularization was set to 1.0. No results are reported on the performance of this method because the computational effort experienced rendered the results not interesting priori (they were on par with JKOnet but the time required per experiment is orders of magnitude more).
Single-cell diffusion dynamics. For the methods in Section 4.4, we report the performances directly from [49, 12]. As mentioned in the main body the paper, the performance of some of the methods not one-to-one comparable with ours, due to slightly different experimental setting. We further clarify the differences here. In our setting and the one used for JKOnet, we use 60% of the data at all time steps for training and 40% for testing. To evaluate performance, between one-step- predictions and ground truth, averaged over the testing data. In the setting of the other algorithms in Section 4.4, instead, all data from all time steps but one is used for training, and the data from the left-out time step is used for testing. Formally, if is the time step used for is used for training. Note that, in all algorithms, the
compute the EMD
testing, then f4 is used for testing and (k,, first and last time steps are always used for training. Since the data includes 5 time steps and one is left out, 80% of the data is used for training and 20% for testing (in practice, these numbers only approximate since the number of particles is not constant over time). In this case, performance evaluated by computing EMD between the prediction and the ground truth at the time step left out, averaged over the left-out time steps. The differences between the two settings demand some caution in comparing and interpreting the performance of the various algorithms. Accordingly, we limit ourselves to observing that JKOnet" achieves state-of-the te--the-n art performance while requiring significantly lower training time (a few minutes compared to hours of the other methods).
F
Functionals
We tested our methods on selection of functionals from standard optimization tests [45]. We report the functionals used with (a projection of) their level set (see Figure 2) for completeness and reproducibility. For v E IRd, we conveniently define
d/2
1
21
d/2
U4
and
22 d
d
dd/2
to extend some two-dimensional functions to multiple dimensions.
26

---

Styblinski-Tang
160, + 5c²)
(31)
Holder table
10 sin (z1) cos (22) exp (1 |11|| )
(32)
Flowers
U1s + 2 * sin(lus().2))
(33)
Oakley-Ohagan
(in((;)) + cos((;) + + vi)
(34)
Watershed
d-1
1
(ws + uf (1i+1 + 4))
(35)
10
Ishigami
1
/
21 + sin(z1)
sin(21) + ssin(zz)2
(36)
10 2
Friedman
110 (2T(21 - ()(22 - -)) + 20 (2(21 7) sin(22 - 7)
1111 1)²
100
2
10 (2(z) 7) cos(22 - 1)² +
1
(z2 7) sin sin - -7))
(37)
10
Sphere
-10
(38)
Bohachevsky
10 (2{ + 22/ - 0.33 - 0.4 cos(4
(39)
Wavy plateau
1 - 30? +1 cos(TU:) + 2
(40)
Zig-zag ridge
(10 Uz41 cos(u:)(1 v441) + of041)
(41)
Double exponential with 20, m 3, and being the d-dimensional vector of ones,
|/ ml 2
|| + m1
exp
+ exp
(42)
27

---

ReLU dimension: any
d Σ
- 50
max(0, v3)
(43)
1=1
Rotational
10ReLU((rrctan((2z + +5 + +5) + +*)
(44)
Flat
dimension: any
0
(45)
G
Failure modes
G.1
On the observability of the energy functionals
When multiple energy functionals are at play or are being considered to describe the evolution of the particles, it not always possible recover the description of the energies. That is, the energy functionals are, in general. not observable. To exemplify this, consider particles initially distributed according to a Gaussian distribution p(to) =0 1) with zero mean and unit variance, diffusing as
dx = ax) + (23dW((),
iee, minimizing the potential energy - and the internal energy. When = 0. at any time the particles distribution (1(()) ( (O, 2,3(1 (o(1-10)) 1o)). When 0, instead, each particle
position at any time is xt)) = x0) and, thus, the particles distribution is paa (t) A(0 e20r). By super-position, the particles distribution at each time t is u(t) = N(0,, e + 2.²t..
Suppose only get two snapshots, (N(0, 1) and (00, of) of the population, at time t 0 and at time / = Ti. For an external observer, there are infinitely many solutions describing the diffusion process: one for each pair (a,, (3 that satisfies
2nTi 2,7 of
While this corner case easily resolved with third observation A (0, aÿ) at time T2, it showcases that it is not obvious necessarily easy to learn the underlying energies. Even when multiple empirical observations are provided, JKOnet" may fail to distinguish the different energy components.
It should be noted, however, that JKOnet" is, to the best of our knowledge, the only method currently capable recovering the distinct energy components. Other methods would associate all the effects to the potential energy, failing to realize there the internal energy is also at play.
In these corner cases, our method can still provide proxy for prediction, but limited insights. We thus believe that understanding under which conditions there is unique selection of energy functionals that describe the diffusion process is fundamental theoretical question.
G.2 When the process is not diffusion
While JKOnet" expands upon the state-of-the-as terms of expressive capabilities, it is still bound to diffusion processes described by (14). To see what happens when this is not the case but one still deploys JKOnet", suppose we are given observations of particles sx()= = (z1(t) x²())) evolving according to
*(t) - f(x(t))) = (x2()) 21 (t) (t(()
(46)
There is interaction between the particles and no noise. Moreover, there is no function V such that fz) - Vz)).. We sample 1000 particles in 1-4.4] 4]2 and observe their evolution over 7 = 5 steps, with т 0.01. We train JKDnet" with the inductive biases 0₂ Ag = 0 for 200 epochs and juxtapose in Figure 11 the resulting predictions with the actual particles trajectory, and the potential vector field with the actual vector field. Not surprisingly, the results are not satisfying. In view of Helmholtz decomposition theorem, we can write f(x) = VV(x) + ((x)), for some (5xx). Thus, a promising direction for future work on the topic the design of methods capable also of learning (x).
28

---

Figure 11: Vector field real (green) and estimated with JKOnet" (blue) of the process (46), which contains solenoidal component. Not surprisingly, when the underlying process is not in the form of (14) the model fails to recover the correct energy terms.
29

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the
relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] "provided proper justification is given (c.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). ). general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: statements in the abstract and introduction are aligned with the main contribution of the paper: to provide method to learn diffusion processes. All claims supported either by rigorous thcoretical results or extensive experimental results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
30

---

Justification: We included paragraph about the limitations of the work in Section 5. We also report all the assumptions made for each statement and practical ways enforce them (e.g., see Remark 3.5).
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiscless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: See Appendix Instead of a proof sketch, we draw the parallel of our results in the euclidean space, which we believe provides the reader with enough intuition to understand the main results the paper.
Guidelines:
The answer NA means that the paper does include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
31

---

Besides providing the code and the data, we discuss in Appendix C the training and model details.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: For anonymity reasons, we have not made our code public and we included in the supplementary material. Upon acceptance, we will release our code on GitHub.
We supplemented the code with instructions, both README files and comments in the code, and running examples. Moreover, when making the code public, we will use tools Read the Docs to complement the open-source code with full documentation.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/ public/guie www.gse//oeeuuumiis..o SooodeeumissiooPoo oolyccy)) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: // http:p//ww.gggeeeeeeo///e/aaaaaa
for more details.
32

---

The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
reproducible, they
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: All the training and model details are specified in Appendix and in appropri- ate configuration files that come along with the code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: When statistics are computed, we report suitable error bars or mean and stan- dard deviations. When statistics are not computed, we report the entirety of the experimental data and make uniform statement (e.g., method performs uniformly better than method Y).
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
33

---

For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix C.7.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: h/w///nuripp..// ctt.pppp..pppppeeee cctppcccc//eeceeecccecc httpehhiiiiuuienneeeeeee...
Question:
Answer: [Yes]
Justification: The research does involve human subjects participants, and it discusses algorithm to train efficiently neural networks on data originating from diffusion processes to understand the source of this data, which has discussed other guidelines does require a discussion of societal applications and potential harmful consequences.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work provides generic algorithm to learn diffusion processes much faster than existing methods. In the guidelines it is reported that algorithms to optimize neural networks do not need any societal impact justification.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
34

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We do not release a trained model or data, and this work mostly focuses on a way to optimize neural networks efficiently learn diffusion processes from data.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: When using code from third party, we report the source of the code directly with the block of code used.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. www.coo//aaaaaa com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
35

---

Answer: [Yes]
Justification: We thoroughly documented the code with instructions in README files and dedicated comments in the code itself and running examples. Moreover, when making the code public, we will use Read the Docs to complement to generate full documentation.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the datasct/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
36


Reviews:
[{"review_id": "p2Q3Rh1Jup", "reviewer": "Reviewer_CXdg", "summary": "This paper studies the problem of learning a diffusion process from samples. It proposes a new scheme based on learning the \"causes mismatch\" of the process, rather than the \"effects mismatch\" as in previous works. The new method is significantly more efficient than the schemes from prior works, and works well in practice.", "questions": "1) Can you provide a more thorough comparison with prior literature, especially the works I have linked above?", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written, and the scheme proposed seems to work well in practice on the examples it was tested on. The loss is intuitive, and resembles the score-matching loss from diffusion models, but is the analogous version for arbitrary diffusion processes. Overall, this seems like a paper that people at NeurIPS would be interested in.", "weaknesses": "I am not familiar enough with the literature, but it seems surprising to me that this scheme has never been proposed before. In particular, the loss is exactly the score-matching in the case of diffusion models, and there are works [1], [2] that have proposed a similar loss for arbitrary diffusion processes. \n\n[1]: https://arxiv.org/abs/2208.09392\n[2]: https://arxiv.org/abs/2209.05442", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We believe the problem of score-matching in diffusion models to be fundamentally different from the one in our paper. \nIn score-matching, one tries to \"reverse\" the time of a known diffusion process, e.g., to recover the uncorrupted state of a corrupted image. \nIn our setting, instead, we use observational population data to learn the energy functional underlying an unknown diffusion process. Our goal is not to \"reverse\" the time of a given diffusion process to reconstruct its initial condition but rather to learn an unknown diffusion process to perform forward-in-time predictions. \n\nFrom a technical perspective, our methodology heavily relies on optimal transport theory and tools from optimization in the probability space. Indeed, our loss can be interpreted as the \"error\" in satisfying a first-order optimality condition in the Wasserstein space. \nTo the best of our knowledge, this approach has not previously appeared in the literature.\nFor instance, the loss function in reference [1] suggested by the reviewer is constructed using tools from stochastic differential equations: Their loss function minimizes the errors in estimating the term $\\nabla_x\\log q_t(x)$ which appears when \"reversing\" the time of a (known) stochastic differential equation. For this reason, their loss cannot be reconciled to ours (in which the stochastic differential equation is unknown).\n\n[1] \"Soft diffusion: Score matching for general corruptions\" by Daras, G., Delbracio, M., Talebi, H., Dimakis, A. G., \\& Milanfar, P. (2022).", "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": "SZp39GZ3R5", "reviewer": "Reviewer_vNRb", "summary": "This paper introduces JKOnet*, a new method for learning diffusion processes from data. It uses first-order optimality conditions of the JKO scheme instead of complex bilevel optimization. JKOnet* can recover potential, interaction, and internal energy components of diffusion processes. The authors provide theoretical analysis and experiments showing JKOnet* outperforms baselines in accuracy, speed, and ability to handle high-dimensional data. They also derive a closed-form solution for linearly parameterized functionals. JKOnet* offers improved computational efficiency and representational capacity compared to existing approaches for modeling diffusion dynamics from population data.", "questions": "- The authors demonstrate JKOnet*'s performance on synthetic datasets. How well does the method perform on real-world diffusion processes? Additional evaluations on empirical data would help understand the method's practical applicability.\n- The paper focuses comparison mainly with JKOnet. How does JKOnet* compare to other recent approaches in learning diffusion processes? \n- In Section 3.4, the authors discuss different parameterizations. How sensitive is JKOnet* to the choice of neural network architecture for the non-linear parameterization case?", "limitations": "The author discusses limitations in section 5", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Develops JKOnet*, a method using first-order optimality conditions of the JKO scheme to learn diffusion processes, avoiding bilevel optimization and improving computational efficiency.\n- Provides theoretical analysis and proofs for JKOnet*, including a closed-form solution for linearly parameterized functionals, backed by comprehensive experiments across various test functions.\n- Demonstrates improved performance in terms of Wasserstein error and computation time compared to existing methods like JKOnet, especially in high-dimensional settings.\n- Enables recovery of potential, interaction, and internal energy components of diffusion processes, expanding the model's applicability to more complex systems and improving interpretability.", "weaknesses": "- The experimental evaluation is limited to synthetic datasets. Real-world data applications would strengthen the practical relevance of the method.\n- While the paper discusses limitations, it does not thoroughly explore potential failure cases or boundary conditions where JKOnet* might underperform.\n- The paper does not provide a comprehensive comparison with other recent approaches in learning diffusion processes beyond JKOnet, which could provide broader context for the method's improvements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Weakness 1\n\nWe thank the reviewer for suggesting us one way to strengthen the presentation of our contributions. \n\nWe deployed our method to learn the diffusion dynamics of embryoid body single-cell RNA sequencing (scRNA-seq) data [1], a popular benchmark in the literature, and compared our results with nine other recent methods in the literature. We discuss the application and the results in the newly added Section 4.4, which will appear in future versions of the paper, and we report the related figure and table in the PDF attached to this response. We briefly summarize our experimental setting and results below. \n\nExperimental setting (briefly): We follow the same data pre-processing as in [2,3]; in particular, we use the same processed artifacts of the embryoid data provided in their work, which contains the first 100 components of the principal components analysis (PCA) of the data and, following [2,3], we focus on the first five. We train on $60\\%$ of the data at each time-step and test $\\mathrm{JKOnet^*}$ to predict the evolution of the left-out data.\n\nResults: $\\mathrm{JKOnet^*}$ outperforms all existing methods in the literature (see the results table in the PDF attached). Importantly, our training takes less than a minute (including the computation of optimal transport plans), whereas the training time of all other existing methods takes hours.\n\nWe visualize the 2 principal components and the interpolations obtained with our method in Figure 5. Note that, for this application, we used a time-varying potential energy, of which we plot the level curves.\n    \n### Weakness 2\n\nWe believe our approach might underperform in the following two cases. \n\nFirst, while diffusion processes include many real-world phenomena, in some real-world applications it might be unknown if the particles are undergoing diffusion. If this is not the case, $\\mathrm{JKOnet^*}$ might underperform. For instance, in the absence of noise and interaction energy, if the vector field is not the gradient of some potential energy function (e.g., it includes a solenoidal component $\\nabla \\times \\psi$ for some function $\\psi: \\mathbb{R}^d \\to \\mathbb{R}^d$), we cannot expect $\\mathrm{JKOnet^*}$ (as well as any other method learning a potential) to infer a reasonable potential. We prepared a dedicated section that we will include in the revision of the paper to discuss this failure mode.\n\nSecond, when learning both potential and interaction energy and noise level, there might be observability issues that prevent the distinction of the different components of the energy functional (e.g., a discrete-time population-level effect might be explained both by a potential energy and a noise level).\n\nWhile we did not experience this issue in our experiments in Section 4.3, we are not aware of rigorous guarantees that ensure observability. In a dedicated section in Appendix G, we provide a small-scale analytical example illustrating this issue. As $\\mathrm{JKOnet^*}$ is the only method capable of simultaneously learning all three energy components, we believe it can serve as the baseline to investigate this observability issue.\n    \n### Weakness 3\n\nWe compared our method with others in our real-world application, discussed in 1) above and in Section 4.4 in the revised version of the paper. \n\n### Question 1, 2\n\nPlease refer to our answer to Weakness 1) above. \n\n### Question 3\n\nIn our experiments, we considered only vanilla parametrizations: two-layers MLP with 64 neurons in each layer, to compare directly with the other works in the literature. What is certainly required is a network that is expressive enough to approximate the energy functional of interest, so standard rationales apply for the choice of activation functions (we use softplus), dimension of the networks, etc. One of the limitations of our works is the data domain (we do not explore, e.g., images). We plan in future work to do so, and in that case more care will be needed to determine the most suitable architecture. Given that the same architecture and hyperparameters worked well across all the experiments (including the real-world experiment), we are confident to state that the learning algorithm itself is not particularly sensitive to the network architecture, which needs of course be chosen to be suitable to represent the energy in the application of choice. We also believe that exciting future work can be done in terms of understanding how different architectures can capture potential, internal, and interaction energies more efficiently: can e.g., transformers be used to learn an interaction energy more efficiently than a vanilla MLP?\n\n[1] \"Visualizing structure and transitions in high-dimensional biological data\" by Kevin R Moon, David van Dijk, Zheng Wang, Scott Gigante, Daniel B Burkhardt, William S Chen, Kristina Yim, Antonia van den Elzen, Matthew J Hirn, Ronald R Coifman, et al. (2019)\n\n[2] \"Improving and generalizing flow-based generative models with minibatch optimal transport\" by Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks,\nKilian Fatras, Guy Wolf, and Yoshua Bengio. (2023)\n\n[3] \"Deep momentum multi-marginal schr\\\"{o}dinger bridge\" by Tianrong Chen, Guan-Horng Liu, Molei Tao, and Evangelos A Theodorou. (2023)", "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": "0i18x4FbZP", "reviewer": "Reviewer_fc3q", "summary": "The authors study diffusion processes from the perspective of Wasserstein gradient flows. Based on the recent fixed-point characterisation for Wasserstein proximal operator methods, they introduce Jordan-Kinderlehrer-Otto (JKO) type methods for learning potential and interaction energies that govern the diffusion process. Such methods are assuming that a sample of the population distribution at each time step is at hand (not necessearily obtained by tracking individual particles) implying important applications across various fields. While theoretical novelties are present (w.r.t. paper [26] that lies in the foundation of this work), the main contribution is the overall methodology for learning diffusion processes.", "questions": "1. In the implementation of the method, a priori computed optimal transport plans are obtained by solving entropy-regularised OT via Sinkhorn-type algorithms or some other methods?\n\n2. What do you think about the applications and/or limitations of the JKOnet$^*$ for the setting of long-trajectories to infer the behaviour in equilibrium, e.g. detection of meta-stable states of Langevin dynamics?", "limitations": "Limitations are addressed adequately.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Paper is, besides minor issues reported bellow, excellently written - very clear, precise and intuitive with well balances technical details between main text and the appendix. Existing ideas are neatly combined to obtain significant improvements of the JKO-type methods and extensive empirical evaluation is presented. The proofs seem correct and well-written.", "weaknesses": "While I do not find important weaknesses, I feel that next several small issues can be addressed to further improve readability:\n\n1. When addressing content presented in the appendix it would be good to refer to the section, e.g. see Figure 6 in Appendix A.\n\n2. It would be good to say what $\\rho_t$ is in Example 2.1\n\n3. While Table 3.1 reports per-epoch complexity for all the methods, it would be important to note that JKOnet$^*$ have additional computational complexity for solving $T$ OT problems of size $N$ in $d$-dimensions. Detailed remark on the initial computational complexity, depending of the algorithm used, should be reported.  \n\n4. In Section 4 it would be helpful to introduce the problems, that is to better explain the task of each experiment and the role of functionals ($V(x)$ ?!)  appearing in Appendix F.  Maybe giving an example on Styblinski-Tang functional appearing in Figures 2, 3 and 4, and then referring to other ones by their names and/or reference equations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Weakness 1\n\nWe thank the reviewer for the suggestion. We prepared a revision of the paper in which we added a reference to the appendix when referencing to content related to the appendix.\n\n### Weakness 2\nGood catch, thank you! We prepared a revision of the paper in which we added the definition of $\\rho$ below the Fokker-Plank equation.\n\n### Weakness 3\nWe thank the reviewer for pointing out this difference, which in fact we believe to be a strength of our method which we did not highlight appropriately. To this end,  we included a dedicated remark, Remark 3.6, for the next revision of the paper. In a nutshell, it is true that $\\mathrm{JKOnet^*}$ requires the construction of the optimal transport couplings beforehand. However, $\\mathrm{JKOnet}$ constructs a new optimal transport plan at each iteration depending on the current estimate of the potential, whereas $\\mathrm{JKOnet^*}$ needs to do so only once, at the beginning. We prepared a revision of the paper in which we added the time required to compute the optimal transport couplings in Section 4.1 ($0.03 \\pm 0.01$s), and we provided an analysis of the comparison between Sinkhorn and plain linear programming for our scope in the newly added Figure 10 in Appendix C.2, which we now expanded (we report Figure 10 in the PDF attached to this rebuttal, and we will include the extended version of the Appendix C.2 for the next revision of the paper). We also better argued why the scaling to high dimensions remains unaffected by the computation of the optimal transport couplings: the dimensionality affects only the construction of the cost matrix in the linear program, and otherwise the computational complexity of the linear program only relates to the number of particles. When the number of particles grows, one can apply the same batching applied in other methods to pre-compute the couplings. We add details regarding these practical considerations in the revised Appendix C.2. Finally, following the comments from reviewer vNRb, we introduced a real-world case study on learning and predicting molecular processes (see Figure 5 and the comparison table in the attached document). This way, we hope we strengthen our contributions not only by achieving state-of-the-art on a real-world application but doing so in under a minute of training (including the computation of the optimal transport couplings) compared to the hours required by the other methods.\n\n### Weakness 4\n\nGood point, thank you for the comment. We have a dedicated section in Appendix B, which we now expanded to discuss the prediction schemes as well, and we have added an introduction to the problems in the experimental section. In particular, we added the data-generation equation, in which the role of the functionals $V(x)$ is apparent: $x_{t+1} = x_t - \\tau \\nabla V(x_t)$. We also added the name and equation of each functional in the figures where it was not listed (e.g., Figure 2).\n\n### Question 1\nThanks for the question. In our implementation, we solved the optimal transport problems via plain linear programming (using the POT python library). We prepared an updated version of Appendix C.2 that we will include in the revision of the paper that contains an analysis of the impact of different ways of solution algorithms on the final outcome in terms of computational time during pre-processing and Wasserstein error (see Figure 10 in the attached PDF). In particular, we conclude that, as long as the couplings are close to the correct one, the algorithm used to compute them does not impact the performance of $\\mathrm{JKOnet^*}$. Since small regularizers slow down the Sinkhorn algorithm, we opted to directly solve the linear program (without regularization). In general, the solver choice can be considered an additional knob that researchers and practitioners can tune when deploying $\\mathrm{JKOnet^*}$. Please see also point 3 above.\n    \n### Question 2\nOnce an energy functional through $\\mathrm{JKOnet^*}$ is learned, the equilibrium state of the system can be inferred in two ways. A simple approach consists of running sufficiently many iterations of the JKO scheme until an equilibrium is reached. Alternatively, the equilibrium state is well-known to be the minimum of the energy functional. Thus, the equilibrium state can be inferred by computing the probability distribution which minimizes the energy functionals, using tools from optimization in the probability space.", "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": "YxSGqQR3XZ", "reviewer": "Reviewer_vFxe", "summary": "This paper considers learning diffusion dynamics from observational data of populations over time, identified as learning the energy functional in Equation 3.  Past research has confronted this inverse problem via complex bilevel optimization, limited to potential energies.  This paper proposes an alternative model JKOnet* that can work with potential, internal, and interaction energies, efficiently minimizes a quadratic loss instead of a complex bilevel optimization, has much lower computational complexity, and out-performs baselines in simulations.  A variant for linearly parameterized functionals has a closed form solution.  The paper's new method reconsiders the JKO scheme using first-order optimality conditions, resulting in decompose the problem into first computing optimal transport plans between adjacent populations and then optimizing a loss for fixed plans.", "questions": "1.  What is JKOnet_l in Table 1?\n\n2.  In Section 4.2, the authors conclude that JKOnet* is well-suited for high-dimensional tasks.  Does this include computing the optimal transport maps?\n\n3.  The discussion in Figure 3 in the text focuses primarily on the speed improvement, yet the performance gains are also quite large, including seemingly between JKOnet* and JKOnet*_l.  Can the authors comment on why the linear parameterization was useful in their experiments?", "limitations": "Limitations are adequately addressed in Section 5", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- Inferring diffusion dynamics from observational data is a difficult and significant problem for which this paper appears to provide a solid contribution.  The paper substantially improves upon JKOnet in terms of multiple directions: better performance (Figure 3), simpler optimization objective (Equation 11), better scalability and efficiency (e.g. Table 1, Section 4.2), and improved generality (Table 1, Section 4.3).  These dimensions are analyzed in experiments across a range of different energy functionals, where the gains are shown in log-scale displaying orders of magnitude improvement.  The paper makes a convincing argument for using JKOnet* over JKOnet.\n- The methodology appears quite strong, well-motivated, and original, with solid intuition given by the authors throughout the paper.", "weaknesses": "Minor weaknesses:\n- While the results are strong, occasionally the language feels too imprecise.  For example, \"runs at lightspeed\" seems inaccurate compared to \"runs very efficiently\".  The authors also mention that they rely upon weeks-old advancements in optimization in the abstract which seems unneeded.\n-  The paper is generally very well-written except for the introduction which could use editing.  It introduces a lot of terminology and details from past research.  Similarly, Figure 1 is referenced multiple times including in the introduction but it was hard to understand until after reading Section 3. \n- The construction of the optimal transport plans does not seem to be included in the computational complexity comparisons.  While this is computed once for JKOnet*, it is additional expense over JKOnet.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### Weakness 1\nWe appreciate the suggestions to improve the exposition and agree with the reviewer's comment. Specifically, we now removed \"phenomenal\", \"runs at lightspeed\" and \"few weeks old advancements\" from the abstract and introduction, rephrasing so to be more factual in the exposition, focusing on the contributions. Thanks!\n\n### Weakness 2 \nWe agree that the introduction blends with the related works rather abruptly. We thus prepared a revised version in which we smoothed it, gradually adding information for the reader less familiar with the literature.\n\nFor instance, we will add the paragraph \"However, the JKO scheme entails an optimization problem in the probability space. Thus,  the problem of finding the energy functional that minimizes a prediction error (w.r.t. observational data) takes the form of a computationally-challenging infinite-dimensional bilevel optimization problem, whereby the upper-level problem is the minimization of the prediction error and the lower-level problem is the JKO scheme.\" to better introduce the reader to what the bilevel optimization problem is about.\n    \nWe also removed the focus from Figure 1 in the introduction.\n\n### Weakness 3\nWe thank the reviewer for pointing out this difference, which in fact is a strength of our method which we did not highlight appropriately and we hope we now did in the Remark 3.6 we prepared for the next revision of the paper. In a nutshell, it is true that $\\mathrm{JKOnet^*}$ requires the construction of the optimal transport couplings beforehand. However, $\\mathrm{JKOnet}$ constructs a new optimal transport plan at each iteration depending on the current estimate of the potential, whereas $\\mathrm{JKOnet^*}$ needs to do so only once, at the beginning. We prepared a revision of the paper in which we added the time required to compute the optimal transport couplings in Section 4.1 ($0.03 \\pm 0.01$s), and we provided an analysis of the comparison between Sinkhorn and plain linear programming for our scope in the newly added Figure 10 in Appendix C.2, which we now expanded (we report Figure 10 in the PDF attached to this rebuttal, and we will include the extended version of the Appendix C.2 for the next revision of the paper). We also better argued why the scaling to high-dimensions remains unaffected by the computation of the optimal transport couplings: the dimensionality affects only the construction of the cost matrix in the linear program, and otherwise the computational complexity of the linear program only relates to the number of particles. When the number of particles grows, one can apply the same batching applied in other methods to pre-compute the couplings. We add details regarding these practical considerations in the revised Appendix C.2. Finally, following the comments from reviewer vNRb, we introduced a real-world case study on learning and predicting molecular processes (see Figure 5 and the comparison table in the attached document). This way, we hope we strengthen our contributions not only by achieving state-of-the-art on a real-world application but doing so in under a minute of training (including the computation of the optimal transport couplings) compared to the hours required by the other methods.\n\n### Question 1\nThanks for pointing out the notation confusion. We defined $\\mathrm{JKOnet^*_l}$ only later on, so now we introduced in the caption to clarify we refer to the linear parametrization.\n\n### Question 2\nThanks for pointing this out. We agree that this is an important point and we will discuss it in the revised version of the paper. In particular, the dimensionality affects only the construction of the cost matrix in the linear program associated with the optimal transport problem, and otherwise the complexity is only related to the number of particles. Specifically, the time required to construct the cost matrix scales linearly with the dimension, and in practice it is minimal and dwarfed by the actual solution of the linear program. When the number of particles grows, one can apply the same batching that is applied in other methods to pre-compute the couplings. We add details regarding these practical considerations in the revision of Appendix C.2. Please see also point 3 above.\n\n### Question 3\nThanks for observing the performance gains and pointing out this aspect. We prepared a revision of Section 4.1 in which we added a paragraph highlighting also the performance gains and discussing linear vs non-linear. In particular, the linear approximation has optimality guarantees, as long as the features are sufficiently rich. In high dimensions, however, the choice of feature is challenging and, thus, we recommend resorting to non-linear parametrizations.", "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:
We thank the reviewer for their time and constructive feedback. Our main changes can be summarized as follows: 

First, we applied our methodology to real data in single-cell diffusion dynamics and compared our results with nine existing methods, as requested by reviewer vNRb. In short, our model, $\mathrm{JKOnet^*}$, outperforms all existing methods both in terms of solution accuracy and training time. Remarkably, our training takes less than a minute, in contrast to the hours of all other methods in the literature. 

Second, we detailed the computation of optimal transport plans/maps, as requested by reviewers fc3q and vFxe. In a nutshell, $\mathrm{JKOnet^*}$ requires computing optimal transport plans only once, before the training. Conversely, $\mathrm{JKOnet}$ demands re-computing optimal transport plans at each training step. Additionally, we clarified that our numerical experiments rely on the linear programming formulation of optimal transport and included an ablation study to compare the linear programming formulation and Sinkhorn algorithm. 

Third, we included a discussion of the failure modes, as requested by vNRb. In short, we envision $\mathrm{JKOnet^*}$ to underperform when the underlying process is not a diffusion process and when observability issues arise (which make the different energy components indistinguishable). 
We included a discussion and examples to illustrate these two corner case phenomena (which, however, we did not experience in our experiments). 

For a detailed response to each reviewer's question and concern, please refer to the responses below.

We believe that these changes both strengthen our contribution and improve the presentation of our results. We thank again the reviewers for the comments that helped us do so!

Meta-review: The paper presents improved methods for learning energy functionals governing a diffusion process from population snapshots observed over time. Existing work uses the JKO scheme, which provides a variational formulation of the probability path, and uses bilevel optimization to learn the energy functionals based on the mismatch between the predicted probability path and the sequence of observed distributions. In this paper, the authors instead use a recent characterization of first-order optimality conditions of the JKO scheme, and find energy functionals to directly satisfy these conditions, leading to significant computational advantages. Reviewers thought the paper was very strong: “excellently written”, “precise”, “intuitive”, addresses a “difficult and significant problem”, “substantially improves upon JKOnet in terms of multiple directions” and so on. Most of the weaknesses were labeled as minor, and were well addressed by the rebuttal. The meta-reviewer appreciates the author’s promise to tone down some language in the introduction (“phenomenal”, “lightspeed”, “few-weeks-old”) and the addition of experiments on a real benchmark dataset with comparisons to more methods during the rebuttal. These will both strengthen the paper.
Decision: Accept (oral)