Learning diffusion at lightspeed

Authors: antonio  terpin
nicolas  lanzetti
martín  gadea
florian  dörfler
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:01:01

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.

Main paper: 
 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 a potential field, the interaction with other particles, and a stochastic term. If 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 of particle trajectories is the Jordan-Kinderlehrer-Otto (JKO) scheme [24] which predicates that the particles as a whole move to decrease an aggregate energy, 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 computationallychallenging 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. 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.effects mismatch = W 2 (µ 1 , μ1 ) 2 causes mismatch : i ∇V θ (x t+1 i ) + 1 τ (x t+1 i -x t i ) 2 µ t+1 = argmin µ∈P(R d ) J θ (µ) + 1 2τ W 2 (µ, µ t ) 2 µ 0 µ 1 µ 2 µ T μ1 μ2μT Figure 1 : Given a sequence of snapshots (µ 0 , . . . , µ T ) of a population of particles undergoing diffusion, we want to find the parameters θ of the parametrized energy function J θ that best explains the particles evolution. Given θ, the effects mismatch is the Wasserstein distance between the observed trajectory and the predicted trajectory obtained iteratively solving the JKO step with J θ . The first-order optimality condition in [30] applied to the JKO step suggests that the "gradient" of J θ with respect to each μt vanishes at optimality, i.e., for μt = µ t . For J θ (µ) = R d V θ (x)dµ(x), this condition is depicted on the right. The gradient (dashed blue arrows) of the true V (level curves in dashed blue) at each observed particle x t+1 i (blue circles) in the next snapshot µ t+1 opposes the displacement (dotted red arrows) from a particle x t i (red triangles) in the previous snapshot µ t . Instead, the gradient (solid green arrows) of the estimated V θ (level curves in solid green) at each observed particle x t+1 i (square) does not oppose the displacement from a particle x t i in the previous snapshot µ t . This mismatch in the causes of the diffusion process is what JKOnet * minimizes. Furthermore, to be practical, it is limited to 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 a method to learn the causes. In [23, 42] , the authors try to learn a geometry that explains the observed transport maps. Unfortunately, the cost between two configurations along a cost-minimizing trajectory is often not a 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 of the infinite-dimensional bilevel optimization problem. We provide a closed-form solution in the case of linearly parametrized energy functionals and a simple, interpretable, and efficient algorithm for non-linear parametrizations. Via exhaustive numerical experiments, we show that, in the case of potential energies only, JKOnet * outperforms the state-of-the-art 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 a fraction of the computational cost. Figure 1 shows an overview of our method, detailed in Section 3. 
diffusion processes via optimal transport 2.1 preliminaries: 
 The gradient of ρ : R d → R is ∇ρ ∈ R d and the Jacobian of ϕ : R d → R n is ∇ϕ ∈ R n×d . We say that f : R d → R has bounded Hessian if ∇ 2 f (x) ≤ C for some C > 0 (and some matrix norm ∥•∥). The divergence of F : R d → R n is ∇ • F and its laplacian is ∇ 2 F . The identity function is Id : R d → R d , Id(x) = x. We denote by P(R d ) the space of (Borel) probability measures over R d with finite second moment. For µ ∈ P(R d ), supp(µ) is its support. The Dirac's delta measure at x ∈ R d , is δ x . All the functions are assumed to be Borel, and for f : R d → R, R d f (x)dµ(x) is the (Lebesgue) integral of f w.r.t. µ. If µ is absolutely continuous w.r.t. the Lebesgue measure, µ ≪ dx, then it admits a density ρ : R d → R ≥0 , and the integral becomes R d f (x)ρ(x)dx. The pushforward of µ via a (Borel) map f : R d → R d is the probability measure f # µ defined by (f # µ)(B) = µ(f -1 (B)); when µ is empirical withN , µ = 1 N N i=1 δ xi , then f # µ = 1 N N i=1 δ f (xi). Given µ, ν ∈ P(R d ), we say that a probability measure γ ∈ P(R d × R d ) is a transport plan (or coupling) between µ and ν if its marginals are µ and ν. We denote the set of transport plans between µ and ν by Γ(µ, ν). The Wasserstein distance between µ and ν isEQUATION.(EQUATIONWhen µ and ν are discrete, (1) is a linear program. If, additionally, they have the same number of particles, the optimal transport plan is γ = (Id, T ) # µ for some (transport) map T : R d → R d [39] . When µ is absolutely continuous, γ = (Id, ∇ψ) # µ for some convex function ψ [7] .Many continuous-time diffusion processes can be modeled by partial differential equations (PDEs) or stochastic differential equations (SDEs): Example 2.1 (Fokker-Planck). The Fokker-Planck equation, ∂ρ(t, x) ∂t = ∇ • (∇V (x)ρ(t, x)) + β∇ 2 ρ(t, x),describes the time evolution of the distribution ρ of a set of particles undergoing drift and diffusion, dX(t) = -∇V (X(t))dt + 2βdW (t),where X(t) is the state of the particle, V (x) the driving potential, and W (t) the Wiener process.The pioneering work of Jordan, Kinderlehrer, and Otto [24] , related diffusion processes to energyminimizing trajectories in the Wasserstein space (i.e., probability space endowed with the Wasserstein distance), providing a discrete-time counterpart of the diffusion process, the JKO scheme,EQUATIONwhere J : P(R d ) → R ∪ {+∞} is an energy functional and τ > 0 is the time discretization. Example 2.2 (Fokker-Plank as a Wasserstein gradient flow). The Fokker-Plank equation (2) results from the continuous-time limit (i.e., τ → 0) of the JKO scheme (3) for the energy functionalJ(µ) = R d V (x)dµ(x) + β R d ρ(x) log(ρ(x))dx with dµ(x) = ρ(x)dx.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 as follows: We have access to sample populations µ 0 , µ 1 , . . . , µ T , and we want to learn the energy functional governing their dynamics. A direct approach to tackle the inverse problem is a bilevel optimization, used, among others, for the model JKOnet in [9] . This approach bases on the following two facts. First, by Brenier's theorem, the solution of (3), µ t+1 , can be expressed 1 as the pushforward of µ t via the gradient of a convex function ψ t : R d → R and, thus,W 2 (µ t , µ t+1 ) 2 = R d ∥x -∇ψ t (x)∥ 2 dµ t (x).Second, the optimization problem (3) is equivalently written asargmin ψt∈C J(∇ψ t# µ t ) + 1 2τ R d ∥x -∇ψ t (x)∥ 2 dµ t (x),where C is the class of continuously differentiable convex functions from R d to R. Therefore, the learning task can be cast into the following bilevel optimization problem, which minimizes the discrepancy between the observations (µ t ) and the predictions of the model (μ t ):EQUATIONA practical implementation of the above requires a parametrization of both the transport map and the energy functional. The former problem has been tackled via input convex neural network (ICNN) parametrizing ψ t [2, 8] or via the "Monge gap" [51] . The second problem is only addressed for energy functional of the form J(µ) = R V θ (x)dµ(x), without interaction and internal energies, where V θ 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 challenging. Second, most energy functionals are not potential energies but include interactions and internal energy terms as well. Although it is tempting to include other terms in the energy functional J (e.g., parametrizing interaction and internal energies), the complexity of the bilevel optimization problem renders such an avenue viable only in principle. 
learning diffusion at lightspeed: 
 Our methodology consists of replacing the optimization problem (3) with its first-order necessary conditions for optimality. This way, we bypass its computational complexity which, ultimately, leads to the bilevel optimization problem (4). Perhaps interestingly, our methodology for learning diffusion is based on first principles: whereas e.g. [9] minimizes an error on the effects (the predictions), we minimize an error on the causes (the energy functionals driving the diffusion process); see Figure 1 . As we detail in Section 4, the resulting learning algorithms are significantly faster and more effective.To start, we illustrate our idea in the Euclidean case (i.e., R d ) and later generalize it to the probability space (i.e., P(R d )). Consider the problem of learning the energy functional J : R d → R ∪ {+∞} of the analog of the JKO scheme in the Euclidean space, the proximal operatorEQUATION)Under sufficient regularity, we can replace (5) by its first-order optimality conditionEQUATIONGiven a dataset (x 0 , x 1 , . . . , x T ), we seek the energy functional that best fits the optimality condition:EQUATIONIn 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 in the probability space [30] , allows us to formulate the learning task as a single-level optimization problem.Consider initially the case where the energy functional is a potential energy, for V : R d → R,EQUATIONThe following proposition is the counterpart of (6) in P(R d ): Proposition 3.1 (Potential energy). Assume V is continuously differentiable, lower bounded, and has a bounded Hessian. Then, the JKO scheme (3) has an optimal solution µ t+1 and, if µ t+1 is optimal for (3), then there is an optimal transport plan γ t between µ t and µ t+1 such thatR d ×R d ∇V (x t+1 ) + 1 τ (x t+1 -x t ) 2 dγ t (x t , x t+1 ) = 0.Proposition 3.1 is by all means the analog of ( 6), since for the integral to be zero, ∇V (x t+1 ) + 1 τ (x t+1x t ) = 0 must hold for all (x t , x t+1 ) ∈ supp(γ t ). Since the collected population data µ 0 , µ 1 , . . . , µ T are not optimization variables in the learning task, the optimal transport plan γ t can be computed beforehand. Thus, we can learn the energy functional representation by minimizing over a class of continuously differentiable potential energy functions the loss functionEQUATIONConsider now the general case where the energy functional consists of a potential energy (with the potential function V : R d → R), interaction energy (with interaction kernel U : R d → R), and internal energy (expressed as the entropy weighted by β ∈ R ≥0 ):EQUATIONThe first-order necessary optimality condition for the JKO scheme then reads as follows. Proposition 3.2 (General case). Assume V and U are continuously differentiable, lower bounded, and have a bounded Hessian. Then, the JKO scheme (3) has an optimal solution µ t+1 which is absolutely continuous with density ρ t+1 and, if dµ t+1 (x) = ρ t+1 (x)dx is optimal for (3), then there is an optimal transport plan γ t between µ t and µ t+1 such that0 = R d ×R d ∇V (x t+1 ) + R d ∇U (x t+1 -x ′ t+1 )dµ t+1 (x ′ t+1 ) + β ∇ρ t+1 (x t+1 ) ρ t+1 (x t+1 ) + 1 τ (x t+1 -x t ) 2 dγ t (x t , x t+1 ).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 β ∈ R ≥0 the loss functionEQUATIONRemark 3.3. We generalize the setting to time-varying energies in Section 4.4 and in Appendix B.For our model JKOnet * , we parametrize the energy functional at a measure dµ = ρ(x)dx as follows:J θ (µ) = R d V θ1 (x)dµ(x) + R d ×R d U θ2 (x -y)d(µ × µ)(x, y) + θ 3 R d ρ(x) log(ρ(x))dx,where θ 1 , θ 2 ∈ R n , θ 3 ∈ R, and we set ). 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 is the number of 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] .θ = θ ⊤ 1 , θ ⊤ 2 , θ ⊤ 3 ⊤ ∈ R 2n+1 . Model FLOPS per epoch Seq. op. per particle Generality V (x) β U (x) JKOnet O T DN d + N 2 log(N ) ε 2 O T DN d + N 2 log(N ) ε 2 ✓ × × w/o TF JKOnet O T DN d + N 2 log(N ) ε 2 O DN d + N 2 log(N ) ε 2 ✓ × × w/ TF JKOnet O T D N d + N 2 log(N ) ε 2 O T D N d + N 2 log(N ) ε 2 ✓ × × w/ MG w/o TF JKOnet O T D N d + N 2 log(N ) ε 2 O D N d + N 2 log(N ) ε 2 ✓ × × w/ MG, TF JKOnet * w/o U (x) O (T N d) O (d) ✓ ✓ × JKOnet * w/ U (x) O T N 2 d O (N d) ✓ ✓ ✓ JKOnet * l w/o U (x) O T N dn + n 3 O T N dn + n 3 ✓ ✓ × JKOnet * l w/ U (x) O T N 2 dn + n 3 O T N 2 dn + n 3 ✓ ✓ ✓Linear parametrizations. When the parametrizations are linear, i.e. V θ1 (x) = θ ⊤ 1 ϕ(x), U θ2 (xy) = θ ⊤ 2 ϕ(x-y) for feature maps ϕ 1 , ϕ 2 : R d → R n , the optimal θ * can be computed in closed-form: Proposition 3.4. Assume that the features ϕ 1,i and ϕ 2,i are continuously differentiable, bounded, and have bounded Hessian. Define the matrix y t : R d → R (2n+1)×d byy t (x t ) := ∇ϕ 1 (x t ) ⊤ , R d ∇ϕ 2 (x t -x ′ t ) ⊤ dµ t (x ′ t ), ∇ρt(xt) ρt(xt)⊤ and suppose that the data is sufficiently exciting so that T t=1 R d y t (x t )y t (x t ) ⊤ dµ t (x t ) is invertible. Then, the optimal solution of (11) isEQUATION)Remark 3.5. The excitation assumption can be enforced via regularization terms λ i ∥θ i ∥ 2 , with λ i > 0, in the loss (11) . Another practical alternative is to use pseudoinverse or solve the leastsquares 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 of the parameters θ 1 , θ 2 , θ 3 to zero, the corresponding energy term is 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 is deterministic and driven by an external potential, one can set θ 2 = θ 3 = 0. Similarly, if it can be assumed that the interaction between the particles is negligible, we can set θ 2 = 0.Here, we motivate the theoretical benefits of JKOnet * over JKOnet using the desiderata: Styblinski-Tang (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). 2. Per-particle computational complexity (i.e., the cost to process a single particle when maximally parallelizing the algorithm, prior to 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 a reformulation of the learning problem that decouples prediction of the population evolution and learning of the dynamics (such coupling is the crux of ( 4)). As a 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 * l generally needs more computation per epoch (primarily related to the number of features) but requires a single epoch. In Table 1 we also report the computational complexity of the variants of JKOnet using a vanilla multi-layer perceptron (MLP) 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 use 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 a new optimal transport plan at each iteration depending on the current estimate of the potential, whereas JKOnet * needs to do so only once, at the beginning. Moreover, as discussed in Section 4.1, in Section 4.2, in the application to single-cell diffusion dynamics in Section 4.4, and in the ablations in Appendix C.2, this additional cost is minimal.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 D. 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 of the paper.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experiments: 
 The code for the experiments is available at https://github.com/antonioterpin/ jkonet-star. We include the training and architectural details for the JKOnet * models family in Appendix C. The settings of the baselines considered are the one provided by the corresponding papers, 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 * V introduces the inductive bias θ 2 = θ 3 = 0. Similarly, we refer to the linear parametrizations by JKOnet * l,V and JKOnet * l .Metrics. To evaluate the prediction capabilities we use the one-step-ahead earth-mover distance (EMD), min γ∈Γ(µt,μt) R d ×R d ∥x -y∥dγ(x, y), where µ t and μt are the observed and predicted populations. In particular, we consider the average and standard deviation over a trajectory.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 a synthetic dataset (see Appendix B) consisting of particles subject to a non-linear drift, x t+1 = x tτ ∇V (x t ), with τ = 0.01, T = 5, and the potential functions V (x) (31)- (45) in Appendix F, shown in Figure 2 and in Figure 6 in Appendix A. Potential V (x), from left to right: (31)- (45) in Appendix F EMD (norm.)JKOnet * JKOnet * V JKOnet * l JKOnet * l,V JKOnet JKOnet-vanilla 0 500 1,000 0 0.5 1 Epochs Convergence ratio 10 -2 10 -1 10 0 10 1 JKOnet-vanilla JKOnet JKOnet * V JKOnet * JKOnet * l,V JKOnet * l Time per epoch [s]Figure 3 : Numerical results of Section 4.1. The scatter plot displays points (x i , y i ) where x i indexes the potentials in Appendix F and y i 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 0 and 1 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 statistics 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 a theoretically guided loss may provide strong computational benefits on par with sophisticated model architectures. Our linearly parametrized models, JKOnet * l and JKOnet * l,V , require a computational time per epoch comparable to the JKOnet family, but they only need one epoch to solve the problem optimally. Our non-linear models, JKOnet * and JKOnet * V , instead both require significantly lower time per epoch and converge faster than the JKOnet 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 in Figure 2 and Figure 6 in Appendix A. Compared to JKOnet, our model also requires a 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 a vanilla MLP deprives the method of the theoretical connections with optimal transport, which, in our experiments, appears to be associated with stability (NaN in Figure 3 ).The results suggest orders of magnitude of improvement also in terms of accuracy of the predictions. These performance gains can be observed also between the linear and non-linear parametrization of JKOnet * . 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 a 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 in Appendix G.Experimental setting. We assess the performance of JKOnet *V to recover the correct potential energy given N ∈ {1000, 2500, 5000, 7500, 10000} particles in dimension d ∈ {10, 20, 30, 40, 50}, generated as in Section 4.1. 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) to highlight the time-dependency.PC 2 PC 1 PC 1 PC 1 PC 1 PC 1 PC 1 PC 1 PC 1 PC 1 PC 1Results. We summarize our 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 is related to the Euclidean norm, it is expected to 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 in Appendix C.2. We thus conclude that JKOnet * is well suited for high-dimensional tasks.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 β ∈ {0.0, 0.1, 0.2}. To our knowledge, this is the first model to recover all three energy terms.Results. We summarize our findings on the right. Compared to the setting in Section 4.1, there are two 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 * to recover the energy components that best explain the observed populations. RNA sequencing (scRNA-seq) data [35] describing the differentiation of human embryonic stem cells over a period of 27 days. We follow the data pre-processing 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 t ∈ R and minimize the lossT -1 t=0 R d ×R d ∇V (x t+1 , t + 1) + 1 τ (x t+1 -x t ) 2 dγ t (x t , x t+1 ).To predict the evolution of the particles, then, we use the implicit scheme (see Appendix B)x t+1 = x t -τ ∇V (x t+1 , t + 1).We train the time-varying extension of JKOnet * V , JKOnet and JKOnet-vanilla for 100 epochs on 60% of the data at each time and we compute the EMD between the observed µ t (40% remaining data) and one-step ahead prediction μt 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 interpolated predictions). As indicated by the table on the right, JKOnet * outperforms JKOnet. We also compare JKOnet * with recent work in the literature which focuses on the slightly different setting, namely the inference of µ t from the evolution at all other time steps µ k , k ̸ = t, without train/test split of the data (the numerical values are taken directly from [12, 49] and our statistics are computed over the timesteps). Since the experimental 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 a few minutes, while the methods listed take hours to run. We further discuss these results in Appendix E.EMD JKOnet [9] 1.363 ± 0.214 JKOnet-vanilla [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 * V (time-varying) 0.624 ± 0.007Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: Besides providing the code and the data, we discuss in Appendix C the training and model details. 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 a 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 it 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 to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , 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.In the case of closed-source models, it 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 to have some path to reproducing or verifying the results. 
conclusion and limitations: 
 Contributions. We introduced JKOnet * , a 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 in 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 a minute, compared to the hours of all existing methods.Limitations. Our work did not address a few important challenges, which we believe to be exciting open questions. On the practical side, JKOnet * owns its performances to a loss function motivated by deep theoretical results. However, its architecture is 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.Outlook. 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] .Question: Does the paper provide open access to the data and code, with sufficient instructions to 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. Justification: All the training and model details are specified in Appendix C and in appropriate 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 a 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, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: When statistics are computed, we report suitable error bars or mean and standard deviations. When statistics are not computed, we report the entirety of the experimental data and make uniform statement (e.g., method X 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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are 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). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: The research does involve human subjects or participants, and it discusses an algorithm to train efficiently neural networks on data originating from diffusion processes to understand the source of this data, which has discussed in other guidelines does not 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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
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 a 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 is 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 or 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 a 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 disinformation. 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.• 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 or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
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 to learn diffusion processes from data.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: When using code from a 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 a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?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 not release new assets.• Researchers should communicate the details of the dataset/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 is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix: We collect the level curves of the ground-truth 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-Tang (31) Holder table (32) Flowers (33) Oakley-Ohagan (34) Watershed ( 35) Ishigami (36) Friedman ( 37) Sphere ( 38) Bohachevski (39) Wavy plateau (40) Zig-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 in Appendix F, from top-left to bottom-right, row-by-row.The prediction scheme for each particle,EQUATIONwhere n t is sampled from the d-dimensional Gaussian distribution at each t, follows from the Euler-Maruyama discretization (see [27] ) of the diffusion processEQUATIONIn particular, we sample 2N particles (in Section 4.1, N = 1000; in Section 4.2, N is indicated in the experiment) uniformly in [-4, 4] d and update their state according to (13) for 5 timesteps, for a total of 6 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. Each heat-map corresponds to a functional in Appendix F, from top-left to bottom-right, row-by-row. The x-axis corresponds to the dimension and the y-axis corresponds to the number of particles. The colors represent the EMD error. Thus, a method that scales well to high-dimensional settings should display a relatively stable color along the rows: the error is related to the norm and, thus, is linear in the dimension d; here, the growth is sublinear. 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, one can adopt the implicit scheme x t+1 = x tτ ∇V (x t+1 ) which is also suggested by the loss in Proposition 3.1. In particular, we adopt this scheme for the time-varying potential in Section 4.4:EQUATION)We refer to this scheme as implicit, and to (13) as explicit, since in (15) x t+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, we work on R (instead of the space of probability measures) and consider the analysis in Section 3.1 for the time-varying functionV (x, t) = 0 if 0.2 ≤ t ≤ 0.3 or 0.7 ≤ t ≤ 0.8, -0.75 • x 2 otherwise.Suppose one adapts the loss in (7) so that the value of the potential at the previous timestep is used, i.e.,EQUATIONThis explicit adaptation, in our experiments, requires more observations to correctly recover the energy potential; see the first two figures on the left in Figure 8 . Since in Section 4.4 we are given only 5 timesteps, we explored the use of an implicit adaptation, i.e.,EQUATIONHowever, 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 in time. With the implicit prediction scheme in (15) , instead, we are able to recover the correct potential with only 10 timesteps; see the right-most figure in Figure 8 . We rely on these considerations to learn the time-varying energy potential of the scRNA-seq data in Section 4.4.This section describes the training details for our model JKOnet * .Compute Couplings Fit Densities {(xEQUATIONt+1 , γ t (x(i) t , x (j) t+1 ), ρ t+1 )} t=T -1 i,j,t=0T Figure 9 : Data pipeline for JKOnet * .In principle, a 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 ).The raw population data µ 0 , µ 1 , . . . , µ T are typically empirical. Since they are not optimization variables in the learning task, the optimal couplings γ 0 , γ 1 , . . . , γ T can be computed beforehand.Differently from [9] , we do not need to track the gradients within the Sinkhorn algorithm to compute the back-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, we compute the couplings in batches of 1000. In our experiments, we solve 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 ε. 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 JKOnet * . 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 * .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 = N i=1 µ(x i )δ xi and µ 2 = M j=1 µ(y j )δ yj can be cast as the linear programmin γij ≥0 ij c ij γ ij s.t. N i=1 γ ij = µ 2 (y j ) ∀j ∈ {1, . . . , M } M j=1 γ ij = µ 1 (x i ) ∀i ∈ {1, . . . , N },where γ ij is the weight of the coupling between x i and y j , γ(x i , y j ) = γ ij andc ij = ∥x i -x j ∥ 2 isthe ij th entry of the cost matrix C. Therefore, the dimensionality d of the data impacts the solving time only in 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 at different timesteps can be computed in parallel. As a result, only the number of particles (i.e., the size of the ) or with Sinkhorn-type algorithms with regularization ε. In the plot on the left, the left y-axis (in red, crosses) represents the EMD error (normalized so 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 in 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 a 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, we highlight that this operation is done only once and independently from, e.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.The case with θ 3 non-zero imposes the estimation of the density ρ t and its gradient ∇ρ t from the empirical probability measures µ t . To the extent of estimating ρ t , there are many viable options (e.g., see the SciPy package [53] ). We use a mixture of 10 gaussians for all the experiments in this paper.To compute ∇ρ t , we rely on the autograd feature of JAX [6] . The complexity related to ρ t and ∇ρ t , instead, is bypassed if there is no internal energy in (10) (i.e., θ 3 = 0).We use the Adam optimizer [26] with the parameters β 1 = 0.9, β 2 = 0.999, ε = 1e-8, and constant learning rate lr = 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.The neural networks of potential and interaction energies are multi-layer perceptrons with 2 hidden layers of size 64 with softplus activation functions and a one-dimensional output layer (cf. [9] ). Future work shall investigate different architectures for the interaction energy, for instance using (Set-)Transformers [31, 52] or Deep Sets [55] .The features used in this work are polynomials up to degree 4 and radial basis exponential functions exp -∥v-c∥ 2 σ with σ = 0.5 and c in the discretized (10 points per dimension) grid [-4, 4] d .Finally, we use a regularizer λ = 0.01 on the square norm of the parameters θ= θ ⊤ 1 , θ ⊤ 2 , θ ⊤ 3 ⊤ .The empirical data was collected entirely on an Ubuntu 22.04 machine equipped with an AMD Ryzen Threadripper PRO 5995WX processor and a Nvidia RTX 4090 GPU. Interestingly, both JKOnet * and JKOnet * V demonstrate comparable computational times when executed on a GPU, highlighting the method's strong parallelizability. Among the methods compared in Section 4.1, JKOnet * is the only one that significantly benefits from GPU parallelization in our experiments, while all others exhibit similar computational times when run on a CPU. When parallelizing over CPU cores, we used on parallel [46] .We briefly collect and summarizes the notation, definitions, and results of [30] used in the proofs of this work. For more intuition and details, we refer the reader to [30] . A sequence(µ n ) n∈N ⊆ P(R d ) converges narrowly to µ if R d ϕ(x)dµ n (x) → R d ϕ(x)dµ(x)for all bounded continuous functions ϕ : R d → R. We say that the convergence is "in Wasserstein" if W 2 (µ n , µ) → 0 or, equivalently,R d ϕ(x)dµ n (x) → R d ϕ(x)dµ(x)for all continuous functions ϕ : R d → R with at most quadratic growth (i.e., ϕ(x) ≤ A(1 + ∥x∥2 ) for some A > 0). For i, m ∈ N, 1 ≤ i ≤ m, we denote by π i : (R d ) m → R d the projection map on the i th component, i.e., π i (x 1 , . . . , x m ) = x i . We compose projections as π ij (x) = (π i (x), π j (x)).Wasserstein calculus. We use two notions of subdifferentability in the Wasserstein space.• Regular Wasserstein subdifferentiability. A transport plan ξ ∈ P(R d × R d ) belongs to the regular subdifferential of a functional J : P(R d ) → R at μ ∈ P(R d ), denoted ∂J(μ), if for all µ ∈ P(R d ) and ξ ∈ (π 1 , π 2π 1 ) # Γ 0 (μ, µ) we haveJ(µ) -J(μ) ≥ max α∈P(R d ×R d ×R d ) π12 # α= ξ,π13 # α=ξ R d ×R d ×R d v⊤ v dα(x, v, v) + o(W 2 (μ, µ)).Here, α can be interpreted as a coupling of the tangent vectors ξ (the subgradient) and ξ (the tangent vector µμ).• (General) Wasserstein subdifferentiability. A transport plan ξ ∈ P(R d × R d ) belongs to the (general) subdifferential of a functional J :P(R d ) → R at μ, denoted ∂J(μ), if there are sequences (µ n ) n∈N ⊂ P(R d ) and (ξ n ) n∈N ⊂ P(R d × R d ) so that (i) ξ n ∈ ∂J(µ n ), (ii) ξ n converges in Wasserstein to ξ.From these definitions, the analogous ones for the supergradient follows: The regular supergradients of J at μ are the elements of -∂(-J)(μ) and the supergradients as the elements of -∂(-J)(μ). We then call the gradient of J the unique element, if it exists, of -∂(-J)(μ) ∩ ∂J(μ), and we say that J is differentiable.Necessary Conditions for Optimality. For J proper and lower semi-continuous w.r.t. convergence in Wasserstein (i.e., if µ n converges in Wasserstein to µ, then lim inf n→∞ J(µ n ) ≥ J(µ)), consider the optimization problem infEQUATIONThen, if µ * ∈ P(R d ) is optimal for (18), i.e., J(µ * ) = inf µ∈P(R d ) J(µ), it must satisfy the "Fermat's rule in Wasserstein space" (see [30, Theorem 3.3] ):EQUATIONIn particular, this implies that for some ξ ∈ ∂J(µ * ) it must holdEQUATION)Outline of the appendix and proofs 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 to familiarize with the techniques used in this work. Since the internal energy (perhaps surprisingly) simplifies the setting by restricting to absolutely continuous measures (i.e., {µ ∈ P(R d ) : µ ≪ dL}), 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 proofs follow the same recipe. First, we prove existence of a solution. Second, we characterize the Wasserstein subgradients for the functional considered. Finally, we conclude deploying [ In particular, here we consider µ ≪ dx since the internal energy is otherwise +∞ by definition and µ is certainly not a minimum.The JKO step in Proposition 3.1 is the optimization problem, resembling (18) ,inf µ∈P(R d ) J(µ) := R d V (x)dµ(x) + 1 2τ W 2 (µ, µ t ) 2 .Since V is lower bounded, up to replacing V by V -min x∈R d V (x), we can without loss of generality assume that V is non-negative. We now proceed in three steps.Existence of a solution.• J has closed level sets w.r.t. narrow convergence. As an intermediate step to prove compactness of the level sets we prove their closedness, which is equivalent to lower semi-continuity of J. Since V is continuous and lower bounded, the functional µ → R d V (x)dµ(x) is lower semicontinuous [29] . The Wasserstein distance is well-known to be lower semi-continuous [30] . Thus, their sum J is also lower semi-continuous.• J has compact (w.r.t. narrow convergence) level sets. Without loss generality, assume λ is large enough so that the level set is not empty. Otherwise, the statement is trivial. Closedness of the level sets follows directly from lower semi-continuity, thus it 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 ,{µ ∈ P(R d ) : J(µ) ≤ λ} ⊂ µ ∈ P(R d ) : 1 2τ W 2 (µ, µ t ) 2 ≤ λ = {µ ∈ P(R d ) : W 2 (µ, µ t ) 2 ≤ 2τ λ}.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 α := inf µ∈P(R d ) J(µ) and let (µ n ) n∈N be a minimizing sequence so that J(µ n ) → α. Then, by definition of convergence, µ n ∈ {µ ∈ P(R d ) : J(µ) ≤ 2α} for all n sufficiently large. By compactness of the level sets, µ n converges narrowly (up to subsequences) to some µ ∈ {µ ∈ P(R d ) : J(µ) ≤ λ}. 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 n→∞ J(µ n ) = α. Thus, µ is a minimizer of J.Wasserstein subgradients. By [30, Proposition 2.17], we combine ( 21) and ( 22) to conclude that any subgradient of J at µ t+1 must be of the form (π 1 , π 2 + π 3 ) # α, for some "sum" couplingα ∈ P(R d × R d × R d ), π 12# α = (Id, ∇V ) # µ t+1 and π 13# α = π 1 , π1-π2 τ # γ ′ t+1 (cf. [30, Definition 2.1]), with γ ′t+1 being an optimal transport plan between µ t+1 and µ t .Necessary condition for optimality. Any minimizer µ t+1 of J satisfies 0 µt+1 ∈ ∂J(µ t+1 ). Thus, for µ t+1 to be optimal, there must exist γ ′ t+1 and α, so that (cf. ( 20))0 = R d ×R d ×R d ∥v 2 + v 3 ∥ 2 dα(x t+1 , v 1 , v 2 ) = R d ×R d ×R d ∥∇V (x t+1 ) + v 2 ∥ 2 dα(x t+1 , v 1 , v 2 ) = R d ×R d ∥∇V (x t+1 ) + v 2 ∥ 2 d((π 13 ) # α)(x t+1 , v 2 ) = R d ×R d ∥∇V (x t+1 ) + v 2 ∥ 2 d π 1 , π 1 -π 2 τ # γ ′ t+1 (x t+1 , v 2 ) = R d ×R d ∇V (x t+1 ) + x t+1 -x t τ 2 dγ ′ t+1 (x t+1 , x t ).Finally, consider γ t = (π 2 , π 1 ) # γ ′ t+1 to conclude.As a preliminary step to Proposition 3.2, consider the case β = 0. The JKO step is the optimization problem, resembling (18) ,EQUATION)Since V and U are lower bounded, up to replacing J with Jinf µ∈P(R d ) J(µ), we can without loss of generality assume that J is non-negative.Existence of a solution. Since also µ →R d ×R d U (x -y)d(µ × µ)(x, y) is lower semi-continuous w.r.t. narrow convergence [29] , the proof of existence is analogous to the one in Proposition 3.1.Wasserstein subgradients. By [30, Proposition 2.17], we combine ( 22) and ( 23) to conclude that each subgradient of the functional 26) has the unique subgradientµ → R d V (x)dµ(x) + R d R d U (x -y)dµ(y)dµ(x) (26) at µ t+1 is of the form (π 1 , π 2 + π 3 ) # α ′ for some coupling α ′ ∈ P(R d × R d × R d ) such that π 12# α ′ = (Id, ∇V ) # µ t+1 and π 13# α ′ = Id, R d ∇U (Id -x)dµ t+1 (x) # µ t+1 . But then α ′ = Id, ∇V, R d ∇U (Id -x)dµ t+1 (x) # µ t+1 and (EQUATIONSimilarly, we conclude that Id, -∇V + R d ∇U (Id -x))dµ(x) # µ t+1 is the unique supergradient of ( 26) at µ and, thus, it is differentiable there. We can thus deploy again [30, Proposition 2.17] to combine (27) with (21) and conclude that the subgradients of ( 25) are of the form (π1 , π 2 + π 3 ) # α, for α ∈ P(R d × R d × R d ), π 12# α = Id, ∇V + R d ∇U (Id -x)dµ t+1 (x) # µ t+1 and (π 1 , π 2 ) # α = π 1 , π1-π2 τ # γ ′ t+1 (cf. [30, Definition 2.1]), with γ ′ t+1 being an optimal transport plan between µ t+1 and µ t .Necessary condition for optimality. The steps are analogous to Appendix D.3. Any minimizer µ t+1 of J satisfies 0 µt+1 ∈ ∂J(µ t+1 ). Thus, for µ t+1 to be optimal, there must exist γ ′ t+1 and α, so that (cf. ( 20))0 = R d ×R d ×R d ∥v 2 + v 3 ∥ 2 dα(x t+1 , v 1 , v 2 ) = R d ×R d ×R d ∇V (x t+1 ) + R d ∇U (x t+1 -x)dµ t+1 (x) + v 2 2 dα(x t+1 , v 1 , v 2 ) = R d ×R d ∇V (x t+1 ) + R d ∇U (x t+1 -x)dµ t+1 (x) + v 2 2 d((π 13 ) # α)(x t+1 , v 2 ) = R d ×R d ∇V (x t+1 ) + R d ∇U (x t+1 -x)dµ t+1 (x) + x t+1 -x t τ 2 dγ ′ t+1 (x t+1 , x t ).Finally, consider γ t = (π 2 , π 1 ) # γ ′ t+1 to conclude. Importantly, for β = 0 we do not need to restrict to µ ≪ dx in Proposition 3.2, so that the statement holds regardless of β (see also Appendix D.5).For β = 0, see Appendix D.4. Let β > 0 here. The JKO step in Proposition 3.1 is the optimization problem, resembling (18) ,EQUATION)Since V, U , and R d ρ(x) log(x)dx are lower bounded, up to replacing J with Jinf µ∈P(R d ) J(µ), we can without loss of generality assume that J is non-negative.Existence of a solution. Since also µ →R d ×R d U (x -y)d(µ × µ)(x, y) and µ → R d ρ(x) log(x)dx if µ ≪ L, +∞ else.are lower semi-continuous w.r.t. narrow convergence [29] , the proof of existence is analogous to the one in Proposition 3.1.Wasserstein subgradients. When β > 0, only µ t+1 ≪ dL can be optimal and we are thus interested in characterizing the subgradients of (28) only at these probability measures. Let ρ t+1 be the density of µ t+1 . Consider first the functionalµ t+1 → R d V (x)dµ t+1 (x)+ R d R d U (x-y)dµ t+1 (y)dµ t+1 (x)+β R d ρ t+1 (x) log(ρ t+1 (x))dx.(29) By [30, Proposition 2.17], we combine ( 27) and (24) to conclude that each subgradient of the functional (29) at µ t+1 is of the form (π 1 , π 2 + π 3 ) # α ′ for some coupling α ′ ∈ P(R d × R d × R d ) such that π 12# α ′ = Id, ∇V + R d ∇U (Id -x)dµ t+1 (x) # µ t+1 and π 13# α ′ = Id, β ∇ρt+1 ρt+1 # µ t+1 . But then α ′ = Id, ∇V + R d ∇U (Id -x)dµ t+1 (x), β ∇ρt+1 ρt+1 #µ t+1 and ( 26) has the unique subgradientEQUATIONHowever, note that now ( 26) is not differentiable, as it can attain the value +∞. We can thus deploy again [30, Proposition 2.17], together with the fact the squared Wasserstein distance is 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 (π1 , π 2 + π 3 ) # α, for α ∈ P(R d × R d × R d ), π 12# α = Id, ∇V + R d ∇U (Id -x)dµ t+1 (x) + β ∇ρt+1 ρt+1 # µ t+1 and (π 1 , π 2 ) # α = π 1 , π1-π2 τ # γ ′ t+1 (cf. [30, Definition 2.1]), with γ ′t+1 being an optimal transport plan between µ t+1 and µ t .Necessary condition for optimality. The steps are analogous to Appendices D.3 and D.4. Any minimizer µ t+1 of J satisfies 0 µt+1 ∈ ∂J(µ t+1 ). Thus, for µ t+1 to be optimal, there must exist γ ′ t+1 and α, so that (cf. ( 20))0 = R d ×R d ×R d ∥v 2 + v 3 ∥ 2 dα(x t+1 , v 1 , v 2 ) = R d ×R d ×R d ∇V (x t+1 ) + R d ∇U (x t+1 -x)dµ t+1 (x) + β ∇ρ t+1 (x t+1 ) ρ t (x t+1 ) + v 2 2 dα(x t+1 , v 1 , v 2 ) = R d ×R d ∇V (x t+1 ) + R d ∇U (x t+1 -x)dµ t+1 (x) + β ∇ρ t+1 (x t+1 ) ρ t (x t+1 ) + v 2 2 d((π 13 ) # α)(x t+1 , v 2 ) = R d ×R d ∇V (x t+1 ) + R d ∇U (x t+1 -x)dµ t+1 (x) + β ∇ρ t+1 (x t+1 ) ρ t (x t+1 ) + x t+1 -x t τ 2 dγ ′ t+1 (x t+1 , x t ).Finally, consider γ t = (π 2 , π 1 ) # γ ′ t+1 to conclude.We start with an explicit expression for the loss (11) in the case of linearly parametrized functionals:L(θ 1 , θ 2 , θ 3 ) := 1 2 T -1 t=0 R d ×R d ∇ϕ 1 (x t+1 ) ⊤ θ 1 + R d ∇ϕ 2 (x t+1 -x ′ t+1 ) ⊤ dµ t+1 (x ′ t+1 ) θ 2 + θ 3 ∇ρ t+1 (x t+1 ) ρ t+1 (x t+1 ) + 1 τ (x t+1 -x t ) 2 dγ t (x t , x t+1 ).Since J is convex and quadratic in θ, we can find its minimum by setting the gradient to 0. In particular, the derivative of L w.r.t.θ 1 , ∇ θ1 L(θ 1 , θ 2 , θ 3 ), reads T -1 t=0 R d ×R d ∇ϕ 1 (x t+1 ) ∇ϕ 1 (x t+1 ) ⊤ θ 1 + R d ∇ϕ 2 (x t+1 -x ′ t+1 ) ⊤ dµ t+1 (x ′ t+1 ) θ 2 + θ 3 ∇ρ t+1 (x t+1 ) ρ t+1 (x t+1 ) + 1 τ (x t+1 -x t ) dγ t (x t , x t+1 ). = T -1 t=0 R d ×R d ∇ϕ 1 (x t+1 ) y t+1 (x t+1 ) ⊤ θ + 1 τ (x t+1 -x t ) dγ t (x t , x t+1 ),where we used Leibniz integral rule to interchange gradient and integral. Similarly,∇ θ2 L(θ) = T -1 t=0 R d ×R d R d ∇ϕ 2 (x t+1 -x ′ t+1 )dµ t+1 (x ′ t+1 ) • y t+1 (x t+1 ) ⊤ θ + 1 τ (x t+1 -x t ) dγ t (x t , x t+1 ) ∇ θ3 L(θ) = T -1 t=0 R d ×R d ∇ρ t+1 (x t+1 ) ⊤ ρ t+1 (x t+1 ) y t+1 (x t+1 ) ⊤ θ + 1 τ (x t+1 -x t ) dγ t (x t , x t+1 ).We split integrals and sums to get∇ θ1 L(θ) = T -1 t=0 R d ∇ϕ 1 (x t+1 )y t (x t+1 ) ⊤ dµ t+1 (x t+1 ) θ + T -1 t=0 R d ×R d 1 τ ∇ϕ 1 (x t+1 )(x t+1 -x t )dγ t (x t , x t+1 ) ∇ θ2 L(θ) = T -1 t=0 R d R d ∇ϕ 2 (x t+1 -x ′ t+1 )dµ t+1 (x ′ t+1 ) y t (x t+1 ) ⊤ dµ t+1 (x t+1 ) θ + T -1 t=0 R d ×R d 1 τ R d ∇ϕ 2 (x t+1 -x ′ t+1 )dµ t+1 (x ′ t+1 ) (x t+1 -x t )dγ t (x t , x t+1 ) ∇ θ3 L(θ) = T -1 t=0 R d ∇ρ t+1 (x t+1 ) ⊤ ρ t+1 (x t+1 ) y t (x t+1 ) ⊤ dµ t+1 (x t+1 ) θ + T -1 t=0 R d ×R d 1 τ ∇ρ t+1 (x t+1 ) ⊤ ρ t+1 (x t+1 ) (x t+1 -x t )dγ t (x t , x t+1 ) .We stack these expressions to compactly write the gradient of L∇ θ L(θ) = T -1 t=0 R d y t (x t+1 )y t (x t+1 ) ⊤ dµ t+1 (x t+1 ) θ + 1 τ T -1 t=0 R d ×R d y t (x t+1 )(x t+1 -x t )dγ t (x t , x t+1 ).The expression for θ follows solving ∇ θ L(θ) = 0 (and shifting the indices in the first sums from t to t ′ with t ′ = t + 1).JKOnet. We use the default configuration provided in [9] . Specifically:• The potential energy is parametrized and optimized in the same way as for JKOnet * (cf. Appendices 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 lr = 0.01, β 1 = 0.5, β 2 = 0.9 and ε = 1e -8.The transport map is optimized with 100 iterations (i.e., we do not use 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 is that we replace the ICNN with a vanilla MLP with two layers of 64 neurons each as parametrization of the optimal transport map.JKOnet with Monge Gap. The settings are 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 we used the recommended implementation in [51] , namely the one in the OTT-JAX package [13] . The coefficient of 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 a 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 of the paper, the performance of some of the methods are not one-to-one comparable with ours, due to a 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, we compute the EMD between one-step-ahead 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 t is the time step used for testing, then µ t is used for testing and µ k , k ̸ = t is used for training. Note that, in all algorithms, the 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 are only approximate since the number of particles is not constant over time). In this case, performance is 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-art performance while requiring significantly lower training time (a few minutes compared to hours of the other methods).We tested our methods on a 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 ∈ R d , we conveniently definez 1 := 1 ⌊d/2⌋ ⌊d/2⌋ i=1 v i and z 2 := 1 d -⌊d/2⌋ d i=⌈d/2⌉ v ito extend some two-dimensional functions to multiple dimensions.EQUATIONHolder tableEQUATIONEQUATIONOakley-OhaganEQUATIONEQUATIONIshigami sin(z 1 ) + 7 sin(z 2 ) 2 + 1 10EQUATIONFriedmanEQUATIONEQUATIONWavy plateauEQUATIONZig-zag ridgeEQUATIONDouble exponential with σ = 20, m = 3, and 1 being the d-dimensional vector of ones,EQUATIONReLU dimension: anyEQUATIONRotational 10ReLU(arctan2(z 2 + 5,EQUATIONFlat dimension: any 0G Failure modes G.1 On the observability of the energy functionalsWhen multiple energy functionals are at play or are being considered to describe the evolution of the particles, it is not always possible to recover the true 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 µ(t 0 ) = N (0, 1) with zero mean and unit variance, diffusing asdX(t) = αX(t)dt + 2βdW (t),i.e., minimizing the potential energy V (x) = -α 2 x 2 and the internal energy. When α = 0, at any time t the particles distribution is µ β (t) = N (0, 1 + 2β(tt 0 )). When β = 0, instead, each particle position at any time t is x(t) = e αt x(0) and, thus, the particles distribution is µ α (t) = N (0, e 2αt ). By super-position, the particles distribution at each time t is µ(t) = N (0, e 2αt + 2βt).Suppose we only get two snapshots, N (0, 1) and N (0, σ 2 1 ) of the population, at time t = 0 and at time t = T 1 . For an external observer, there are infinitely many solutions describing the diffusion process: one for each pair (α, β) that satisfies e 2αT1 + 2βT 1 = σ 2 1 . While this corner case is easily resolved with a third observation N (0, σ 2 2 ) at time T 2 , it showcases that it is not obvious or 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 of 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 a proxy for prediction, but limited insights. We thus believe that understanding under which conditions there is a unique selection of energy functionals that describe the diffusion process is a fundamental theoretical question.While JKOnet * expands upon the state-of-the-art in 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 particlesx(t) = [x 1 (t) x 2 (t)]⊤ evolving according toEQUATION)There is no interaction between the particles and no noise. Moreover, there is no function V such that f (x) = -∇V (x). We sample 1000 particles in [-4, 4] 2 and observe their evolution over T = 5 steps, with τ = 0.01. We train JKOnet * with the inductive biases θ 2 = θ 3 = 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) = ∇V (x) + ∇ × ψ(x), for some ψ(x). Thus, a promising direction for future work on the topic is the design of methods capable also of learning ψ(x).Figure 11 : Vector field real (green) and estimated with JKOnet * (blue) of the process in (46) , which contains a solenoidal component. Not surprisingly, when the underlying process is not in the form of (14) the model fails to recover the correct energy terms.Justification: We included a paragraph about the limitations of the work in Section 5. We also report all the assumptions made for each statement and practical ways to 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 a 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, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on 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 a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.Under an absolute continuity assumption.We exclude the functional(32) from the interaction energies due to numerical issues in the data generation. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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. 

[1]: https://arxiv.org/abs/2208.09392
[2]: https://arxiv.org/abs/2209.05442
questions: questions: 1) Can you provide a more thorough comparison with prior literature, especially the works I have linked above?
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: - The experimental evaluation is limited to synthetic datasets. Real-world data applications would strengthen the practical relevance of the method.
- While the paper discusses limitations, it does not thoroughly explore potential failure cases or boundary conditions where JKOnet* might underperform.
- 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.
questions: 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.
- The paper focuses comparison mainly with JKOnet. How does JKOnet* compare to other recent approaches in learning diffusion processes? 
- 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: limitations: The author discusses limitations in section 5
paper_summary: paper_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.
strength_weakness: strength_weakness: While I do not find important weaknesses, I feel that next several small issues can be addressed to further improve readability:

1. When addressing content presented in the appendix it would be good to refer to the section, e.g. see Figure 6 in Appendix A.

2. It would be good to say what $\rho_t$ is in Example 2.1

3. 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.  

4. 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.
questions: 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?

2. 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: Limitations are addressed adequately.
paper_summary: paper_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.
strength_weakness: strength_weakness: Minor weaknesses:
- 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.
-  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. 
- 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.
questions: questions: 1.  What is JKOnet_l in Table 1?

2.  In Section 4.2, the authors conclude that JKOnet* is well-suited for high-dimensional tasks.  Does this include computing the optimal transport maps?

3.  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: Limitations are adequately addressed in Section 5

Comments: title: Acknowledgement of the rebuttal
comment: Thank you for your response. These addressed my questions. I have raised my score.
I appreciate the detailed response and additional experiments in the rebuttal, and continue to recommend paper acceptance.
I thank the authors for their rebuttal. I remain confident of the quality of their paper, suggest the acceptance and keep my score.