Self-Consistent Velocity Matching of Probability Flows
Authors: Lingxiao Li
Samuel Hurault
Justin Solomon
Venue: NeurIPS 2023 poster (2023)
Keywords: JKO
mass-conservation
PDE
Fokker-Planck
scalable
discretization-free
neural ODE

Abstract: We present a discretization-free scalable framework for solving a large class of mass-conserving partial differential equations (PDEs), including the time-dependent Fokker-Planck equation and the Wasserstein gradient flow. The main observation is that the time-varying velocity field of the PDE solution needs to be self-consistent: it must satisfy a fixed-point equation involving the probability flow characterized by the same velocity field. Instead of directly minimizing the residual of the fixed-point equation with neural parameterization, we use an iterative formulation with a biased gradient estimator that bypasses significant computational obstacles with strong empirical performance. Compared to existing approaches, our method does not suffer from temporal or spatial discretization, covers a wider range of PDEs, and scales to high dimensions. Experimentally, our method recovers analytical solutions accurately when they are available and achieves superior performance in high dimensions with less training time compared to alternatives.

Self-Consistent Velocity Matching of Probability Flows
Lingxiao Li MIT CSAIL
Samuel Hurault
Justin Solomon MIT CSAIL jsol weolomon@mii..o edu
Univ. Bordeaux, Bordeaux INP, CNRS, IMB samuel. huraul twmuth.mthhh..mm www.lltttttt.uu.otooaaoeae www.oreaax....... fr
1 ingxiao@mit..ee
Abstract
We present a discretization-free scalable framework for solving a large class of mass-conserving partial differential equations (PDEs), including the time- dependent Fokker-Planck equation and the Wasserstein gradient flow. The main observation is that the time-varying velocity field of the PDE solution needs be self-consistent: it must satisfy a fixed-point equation involving the probabil- velocity field. Instead of directly minimizing
ity flow characterized by the the residual of the fixed-point equation with neural parameterization, we use an iterative formulation with a biased gradient estimator that bypasses significant computational obstacles with strong empirical performance. Compared to existing approaches, our method does not suffer from temporal or spatial discretization, covers wider range of PDEs, and scales to high dimensions. Experimentally. our method recovers analytical solutions accurately when they are available and achieves superior performance in high dimensions with less training time compared alternatives.
Introduction
Mass conservation is ubiquitous phenomenon in dynamical systems arising from fluid dynamics, clectromagnetism,, thermodynamics, and stochastic processes. Mathematically, mass conservation is formulated the continuity equation:
(.p((x) V (u.P.), Vz, [0,
(1)
where pt Rd - R is scalar quantity such that the total mass Pr is conserved with respect Ur Rd Rd is velocity field, and T' is total time. We will assume, for all E [0, 7'). Pt > and p:(x)) dr = 1, i.e., pr is a probability density function. We use to denote the probability measure with density Pt- Once pair (p), satisfies (I), the density Pt is coupled with v, in the sense that the evolution of pt in time is characterized by (Section 3.1).
We consider the subclass of mass monse-onnsrriig PDEs that can be written succinetly (fr(:; p.)p.) Vx,t t [0,T].
(pp(())
(2)
where fi(: (4) given initial probability measure μz with density p5-
Rd - Rd is given function depending on pe, with initial condition (40 1g for
Different choices of Se lead to large class of mass-conserving PDEs. For instance, given functional F (P(R) - on the space of probability distributions with finite second moments, if we take
[e(:: = - ((r)((
where F(n)(()) Rd Rd is the Wasserstein gradient of F. then the solution to (2) is the Wasserstein gradient flow of F [Santambrogio, 2015, Chapter 8]. Thus, solving (2) efficiently allows us optimize in the probability measure space. If we take
f((;;)) b, (z) loges)))
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

where by is a velocity field and Dz(x) is dependent Fokker-Plande equation [Risken and Risken, 1996]. which describes the time evolution of the probability flow undergoing drift b, and diffusion with coefficient D..
matrix, then we obtain the time-
A popular strategy to solve (2) is to discretized mesh or as a neural network [Raissi et al., 2019]. However, these approaches do not fully exploit the mass-conservation principle and are usually limited to low dimensions. Shen et al. [2022], Shen and Wang [2023] recently introduced the notion of self-consistency for the Fokker-Plande equation and more generally McKean- McKean-Vlasov type PDEs. This notion is a Lagrangian formulation of (2). They apply the adjoint method to optimize self-consistency. In this work, we extend their notion of self-consistency to maass-cnnsrrving PDEs of the general form (2). Equipped with this formulation, we develop an iterative optimization scheme called self- celf-connistent velocity matching.
Eulerian representation of the density field pt on
With the probability flow parameterized as neural network, at each iteration, we refine the velocity field 2y of the current flow to match an estimate of ft evaluated using the network weights from the previous iteration. Effectively, we minimize the self-consistency loss with a biased but more tractable gradient estimator.
This simple scheme has many benefits. First, the algorithm is agnostic the form of ft, thus covering a wider range of PDEs compared to past methods, Second, we no longer need to differentiate through differential equations using the adjoint method as in Shen and Wang [2023], which is orders of magnitude slower than our method with worse performance in high dimensions. Third, this iterative formulation allows us to rewrite the velocity-matchine computationally expensive quantities such as V log Pt in the Fokker-Planck equation (Proposition 3.1). Lastly, our method is flexible with probability flow parameterization: we have empirically found that the two popular ways of parameterizing the flow-as as a time-varying pushforward map [Biloš et al., 2021] ]] and as a time-varying velocity field [Chen et al., 2018]-both both have merits in different scenarios.
objectives for certain PDEs get rid of
Our method tackles mass-conserving PDEs of the form (2) in a unified manner without temporal or spatial discretization. Despite using a biased gradient estimator, in practice, our method decreases the celficonsi loss efficiently (second column of Figure 2). For PDEs with sellyticallykknoowe solutions, we quantitatively compare with the recent neural JKO-based methods [Mokrov et al., 2021, Fan et al., 2021, Alvarez-Melis et al., 2021], the adjoint method [Shen and Wang, 2023], and the particle-based method [Boffi and Vanden- solutions with quality on par with the best previous methods in low dimensions and with superior quality in high dimensions. Our method is also significantly faster than competing methods, especially in high dimensions, at the same time without discretization. We further demonstrate the flexibility of our method on two challenging experiments for modeling flows splashing against obstacles and smooth interpolation of measures where the comparing methods are either not applicable or have noticeable artifacts.
2023]. Our method faithfully recovers true
2
Related Works
Classical PDE solvers for mass-conserving PDEs such as the Fokker-Planck equation and the Wasserstein gradient flow either use an Eulerian representation of the density and discretize space as a grid or mesh [Burger et al., 2010, Carrillo et al., 2015, Peyré, 2015] or use Lagrangian representation, which discretizes the flow as a collection of interacting particles simulated forward in time [Crisan and Lyons, 1999, Westdickenberg and Wilkening, 2010). Due to spatial discretization, these methods struggle with -ige--imensional problems. Hence, the rest of the section focuses solely on recent neural network-basee based methods.
Physics-informed neural networks. Physics-informed neural networks (PINNs) are prominent methods that solve PDEs using deep learning [Raissi et al., 2019, Karniadakis et al., 2021]. The main idea to minimize the residual of the PDE along with loss terms to enforce the boundary conditions and to match observed data. Our notion of self-consistency is Lagrangian analog of the residual in PINN. Our velocity matching only occurs along the flow of the current solution where interesting dynamics happen, while in PINNs the residual is evaluated on collocation points that occupy the entire domain. Hence our method is particularly suitable for high-dimensional problems where the dynamics have a dow-dimensions structure.
2

---

Neural JKO methods. Recent works [Mokrov et al., 2021, Alvarez-Melis et al.. 2021, Fan et al., 2021] apply deep learning to the time-discretized JKO scheme [Jordan et al., 1998] to solve Wasser- stein gradient flow (3). By pushing a reference measure through a chain of neural networks parameter- ized as input-convex neural networks (ICNNs) [Amos et al., 2017], these methods avoid discretizing the space. Mokrov et al. [2021] optimize one ICNN to minimize Kullback-Leible (KL) divergence plus a Wasserstein-2 2 distance term at each JKO step. This method is extended to other functionals by Alvarez-Melis et al. [2021]. Fan et al. [2021] use the variational formulation of f-divergence to obtain a faster primal-dual approach.
An often overlooked problem of neural JKO methods is that the total training time scales quadratically with the number of JKO steps: to draw samples for the current step, initial samples from the reference measure must be passed through a long chain of neural networks, along with expensive quantities like densities. However, using too few JKO steps results in large temporal discretization errors. Moreover, the optimization at each step might not have fully converged before the next step begins, resulting in an unpredictable accumulation of errors. In contrast, our method does not suffer from temporal discretization and can be trained end-to-end. It outperforms these neural JKO methods with less training time in experiments we considered.
Velocity matching. A few recent papers employ the idea of velocity matching to construct a flow that follows a learned velocity field. di Langosco et al. [2021] simulate the Wasserstein gradient flow of the KL divergence by learning a velocity field that drives a set of particles forward in time for Bayesian posterior inference. The velocity field refined on the fly based on the current positions of the particles. Boffi and Sanden-Fijnder [2023] propose a similar method that applies to a more general class of dime-dependent Fokker-Plan Planck equations. These two methods approximate probability measures using finite particles which might not capture high-dimensional distributions
well. Liu et al. [2022], Lipman et al. [2022], Albergo and Vanden-Eijnden [2022] use flow matching for generative modeling by learning velocity field that generates probability path connecting reference distribution to the data distribution. Yet these methods are not designed for solving PDEs.
Most relevant to our work, Shen et al. [2022] propose the concept of self-consistency for the Fokker- Planck equation, later extended to McKean- that the velocity field of the flow solution to the Bokker-Planck equation must satisfy fixed-pon -point equation. They theoretically show that, under certain regularity conditions, form of probability divergence between the current solution and the true solution is bounded by the self-consistency loss that measures the violation of the fixed-point equation. Their algorithm minimizes such violation using neural ODE parameterization [Chen et al., 201 8] and the adjoint method. Our work extends the concept of self-consisenncy to a wider class of PDEs in the form of (2) and circumvents the computationally demanding adjoint method using iterative formulation. We empirically verify that our method is significantly faster and reduces the self-consistency loss more effectively in moderate dimensions than that of Shen and Wang [2023] (Figure 2).
type PDEs [Shen and Wang. 2023]. They observe
3 Self-Consistent Velocity Matching
3.1 Probability flow of the continuity equation
A key property of the continuity equation (1) is that any solution (pt. ve)ee(O.T) (provided is continuous with respect to and v is bounded) corresponds to unique flow map (-.(-)= R -
R hejo.T that solves the ordinary differential equations (ODEs) [Ambrosio et al., 2005, Proposition 8.1.8]
d oo(x)= =x dt =.(())=(((((()))) Vz, E [0,7],
(5)
and the flow map satisfies 14 ( for all t E [0, Tl, where (+. ) Mo to denote the pushforward measure of μi by 4,. Moreover, the converse is true: any solution (Ø,, v2) of (5) with Lipschitz continuous and bounded v, a solution of (1) with μ4 = # [Ambrosio et al., 2005, Lemma 8.1.6]. Thus the Fulerian viewpoint of (1) is equivalent to the Lagrangian viewpoint of (5). We next exploit this equivalence by modeling the probability flow using the Lagrangian viewpoint so that automatically satisfies the continuity equation (1).
3

---

3.2 Parametrizing probability flows
Our algorithm will be agnostic to the exact parameterization used to represent the probability flow.
As such, we need way to parameterize the flow to access the following quantities for all t E [0, T]:
+ Rd + Rd, the flow map, so (x(xr) is the location of a particle at time if is at x at time O.
Ut Rd + Rd, the velocity field of the flow at time
Me E P(R"), the probability measure at time with sample access and density pt.
We will assume all these quantities are sufficiently continuous and bounded to ensure the Eulerian and Lagrangian viewpoints in Section 3.1 are equivalent. This can be achieved by using continuously differentiable activation functions in the network architectures and assuming the network weights finite similar to the uniqueness arguments given in [Chen et al., 2018) We can now parameterize the flow modeling either the flow map 9 or the velocity field By as neural network.
Time-dependent Invertible Push Forward (TIPF), We first parameterize a probability flow by modeling 0, R - Rd as invertible network for every t. The network architecture is chosen so that 4, has an analytical inverse with a tractable Jacobian determinant, similar to [Biloš et al., 2021]. We augment RealNVP from Dinh et al. [2016] so that the network for predicting scale and translation takes t as an additional input. To enforce the initial condition, we need Py to be the identity map. This condition can be baked into the network architecture [Biloš et al., 2021] or enforced by adding an additional loss term Ex~ || [x)) X x For brevity, we will from now on omit in the text this additional loss term. The velocity field can be recovered via vy = 0 ( (x)). #
To recover the density Pt of je log P (0, (x)) + log det J (x)) (see (1) in Dinh et al. [2016)).
the change-of variiiaell formula log (c))
Neural ODE (NODE). We also parameterize flow by modeling U2 R - R neural network; this is used in Neural ODE [Chen et al., 2018]. The network only needs to satisfy the minimum requirement of being continuous. The flow map and the density can be recovered via numerical integration: (, = x + Sg u, (z)) ds and log Pr (. (x)) = log (5 Io' v Ug ( (x)) ds, a direct consequence of (1) also known as the instantaneous change-of-variable formula [Chen et al., 2018]. To obtain the inverse of the flow map, we integrate along -t. With NODE, the initial condition 140 = M is obtained for free.
We summarize the advantages and disadvantages of TIPF and NODE as follows. While the use of invertible coupling layers in TIPF allows exact access to samples and densities, TIPF becomes less effective in higher dimensions as many couple layers are needed retain good expressive power, due to the invertibility requirement. In contrast, NODE puts little constraints on the network architecture, but numerical integration can have errors. Handling the initial condition is trivial for NODE while an additional loss term or special architecture is needed for TIPF. As we will show in the experiments, both strategies have merits.
3.3 Formulation
We now describe our algorithm for solving mass-conserving PDEs (2). A PDE of this form is determined by fe(s)))= : Rd Rd plus the initial condition vūj If probability flow μt with flow map o and velocity field satisfies the following self-consistency condition,
(6) then the continuity equation of this flow implies the corresponding PDE (2) is solved. Conversely, the velocity field of any solution of (2) will satisfy (6). Hence, instead of solving (2) which is a condition on the density Pt that might be hard access. we can solve (6) which is a more tractable condition. Shen et al. [2022], Shen and Wang [2023] develop this concept for the McKean- Vlasov type PDEs; here we generalize to a wider class of PDEs of the form (2).
ve(x) = f(: pc), Vz in the support of Pr:
equation and
Let 0 be the network weights that parameterize the probability flow using TIPF or NODE. The flow's measure, velocity field, and flow map at time are denoted nig of. & respectively. One option to solve (6) would be to minimize the self-consistency loss
Ex [lx)(()) f((x; ((xx)111))
min jo

---

This formulation is reminiscent of PINNs [Raissi et al., 2019] where residual of the original PDE is minimized. Direct optimization of (7) challenging: while the integration over [o, 2' and μul can be approximated using Monte Carlo, to apply stochastic gradient descent, we must differentiate through i4 and ft: this can be either expensive or intractable depending on the network parameterization. The algorithm by Shen and Wang [2023] minimizes (7) with the adjoint method specialized to Fokker- Planck equations and McKean-Vlasov type PDEs; extending their approach to more general PDEs requires closed-form formula for the time evolution of the quantities within fo, which at best can only be obtained on a case-by-case basis.
Instead, we propose the following iterative optimization algorithm to solve (7). Let fg denote the network weights at iteration k. We define iterates
0k+1 == 0g 11 VVolo=0
(8)
where
F(O,0a) == /
_ Ex - f fx(x:
dt.
(9)
Effectively, in each iteration, we minimize (9) by one gradient step where we match the velocity field og to what it should be according to Sa based on the network weights 0g from the previous iteration. This scheme can be interpreted as a gradient descent on 7 using the biased gradient estimate VoF(0, -x)-sce matching.
see Appendix A for discussion. We call this iterative algorithm self-consistent velocity
If fa depends on the density of 14 only through the score V log Pt (corresponding to a diffusion term in the PDE), then we can apply an integration-by by-paarts trick [Hyvärinen and Dayan, 2005] to get rid of this density dependency by adding a divergence term of the velocity field. Suppose ft is from the Fokker-Planck equation (4). Then the cross term in (9) after expanding the squared norm has the following alternative expression.
Proposition 3.1. For every t E [0, for ft defined in (4), assume vi Dr are bounded and continuously differentiable, and ph is measure with continuously differentiable density ve that vanishes in infinity and not at finite points. Then we have
E Ex=el' f((;;u()) = Ex=ap' (x)
(1 b.) + V (DT (
(10)
We provide the derivation in Appendix B. With Proposition 3.1, we no longer need access to V log p when computing F. This is useful for NODE parameterization since obtaining the score would otherwise require additional numerical integration.
Our algorithm is summarized in Algorithm I in the appendix. We use Adam optimizer [Kingma and Ba, 2014] to modulate the update (8). For sampling time steps .i,....tt tt in [0, TJ. we use stratified sampling where tą is uniformly sampled from [0-1)T/I, IT/L; such sampling strategy results in more stable training in our experiments. We implemented our method using JAX [Bradbury et al., 2018] and FLAX [Heek et al., 2020]. See Appendix C for the implementation details.
Experiments
We show the efficiency and accuracy of our method on several PDEs of the form (2). We start with three Wasserstein gradient flow experiments (Section 4.1, Section 4.2, Section 4.3). Next, we consider the time-dependent Fokker-Planck equation that simulates attraction towards a harmonic mean in Section 4.4. Finally, in Section 4.5, we apply our framework to generate complicated low-dimensional dynamics including flows splashing against obstacles and smooth interpolation of measures. We will use SCVM-TIPF and SCCM-NODE to denote our method with TIPF and NODE parameterization respectively. We use JKO-ICNN to denote the method by Mokrov al. [2021]] JKO-ICNN-PD to denote the method by Fan et al. [2021] (PD for "primal-dual"), ADJ to denote the adjoint method by Shen and Wang [2023], SDE-EM to denote the Euler-Maruyama method for solving the SDE associated with the Fokker-Plande equation, and DFE ("discrete forward Fuler") to denote the method by Boffi and Vanden-Fijnden [2023]. We implemented all competing methods JAX- see more details in Appendix C-and we compare quantitatively against these methods when possible.
5

---

In Table 1, we compare the time complexity of training the described methods, where we show that SCVM-TIPF and SCVM-NODE have low computational complexity among all methods.
Evaluation metrics. For quantitative evaluation, we use the following metrics. To compare measures with density access, following Mokrov et al. [2021].
the symmetric Kullback- Leibler (symmetric KL) divergence, defined as SymKL(p1, P2) := KL(p1 P2) + KL(p2 || Pt), where KL(p1 p2) := Ex-p, [log dp: (x)(x)))))X)) Sample estimates of KL can be negative which complicates log-scale plotting, so when this happens, we consider an alternative f-divergence
Dy(p1 || p2) := Ex~p2 [ [1) px))) whose sample estimates are always non-negative We similarly define the symmetric f-divergence SymD, (p1,
(PP..P) P2) == (p || P2) + ( || p1).
For particle-based methods. we use kernel density estimation (with Scott's rule) to obtain the density function before computing symmetric KL or f-divergence. We also consider the Wasserstein-2 distance [Bonneel et al., 2011] and the Bures- Wasserstein distance [Kroshnin et al., 2021]; these two measures only require sample access. All metrics are computed using i.i.d. samples. See Appendix C.6 for more details.
4.1 Sampling from mixtures of Gaussians
We consider computing the Wasserstein gradient flow of the KL divergence F(H) = KL(P || p*) where we have density access to the target measure μ To fit into framework, set fa (z; (b))
V log p" (x) V log pe()) which matches (4) with (x) = V loogp" (x) and D₂(x) Ia. Following the experimental setup in Mokrov al. [2021] and Fan et al. [2021].
take y be mixture of 10 Gaussians with identity covariance and means sampled uniformed in [-5, 5, 5) The initial measure μ (O, 16la). We solve the corresponding Fokker-Planck PDE for total time o 5 and for d = 10, 60. As TIPF parameterization does not scale to high dimensions, we only consider SCVM-NODE in this experiment.
Figure shows the probability flow produced by SCVM-NODE in dimension 60 at different time steps: as we can see, the flow quickly converges to the target distribution.
Targel measure
-10
Figure 1: Probability flow produced by SCVM-NODE for a 60-dimensional mixture of Gaussians at irregular time steps. Samples projected onto the first two PCA components and kernel density estimation is used to generate the contours.
In Figure 2, we quantitatively compare our method with Mokrov et al. [2021]. Fan et al. [2021], and Shen and Wang [2023]. Training time is reported for all methods.
In the left two columns of Figure 2, we find that even though the adjoint method ADJ [Shen and Wang, 2023] minimizes the self-consistency loss (7) directly, the decay of self-consistency can be much slower than that of SCVM-NODE as the dimension increases. We suspect this due to the amount of error accumulated in the adjoint method which involves two numerical integration passes to obtain the gradient. Moreover, ADJ requires up to third-order spatial derivatives of the parameterized neural velocity field which can be inaccurate comparison SCVM-NODE, only requires one integration pass and the first-order spatial derivative of the network. Despite the bias of the gradient used in SCVM-NODE. it finds more efficient gradient trajectories than ADJ. Additionally, ADJ takes 80 times longer to train than SCVM-NODE dimension 10, and scaling up to higher dimensions becomes prohibitive.
if the consistency loss is low-in
The rightmost column of Figure shows SCVM-NODE achieves far lower symmetric KL compared the JKO methods. The gradient flow computed by JKO methods does not decrease KL divergence monotonically, likely because the optimization at each JKO step has yet to reach the minimum even though we use 2000 gradient updates for each step. For both JKO methods, the running time for each JKO step increases linearly because samples (and for JKO-ICNN also log det terms) need to be pushed through a growing chain of ICNNs; as a result, the total running time scales quadratically with the number of JKO steps. JKO methods also take about 40 times as long evaluation time as

---

SCVM-NODE in dimension 60 since density access requires solving an optimization problem for every JKO step. On top of the computational advantage and better results, our method also does not have temporal discretization: after being trained, the flow can be accessed at any time / (Figure 1).
Symmanic Kaft 1]
SCVM-NODE
SCVM NODE
10P
Consatorcy
NDU
lof-
Symmericc
2.P
SCVIA-NODE
JKG-ICNN-PD
2500
3000
7000
10000
2540
5000
7500
10000
Training Stnation (d 02
Traininng Renation (s)
Time (0)
Synmeticc Kuf -T
10
Consiacrry
KL
LDF
SCVM NODE
SCVM NODE
E0.
SCVM-NODE
10t
JKO ICNN PD IKO-ICNR
IDP
2600 3000
7500 10)
10000
2500:
5000
75000
100.00
2 Time (
Training Deration
Training Iteration
10)
Figure 2: Quantitative comparison for the mixture of Gaussians experiment. The left two columns plot the symmetric KL (at t = T compared against the target measure) and consistency (7) versus the training iterations for SCVM-NODE (ours) and ADJ [Shen and Wang, 2023]. The rightmost column plots the symmetric KL across time t (compared against the target measure) for SCVM-NODS and the JKO methods in high dimensions. Training time: for = 10, SCVM-NODE takes 7.37 minutes, ADJ takes 585.2 minutes; for d 60, SCVM-NODE takes 23.9 minutes, JKO-ICNN takes 375.2 minutes, and JKO-ICNN-PD takes 24.4 minutes.
4.2
process
We consider the Wasserstein gradient flow of the KI divergence with respect to a Gaussian with the initial distribution being a Gaussian, following the same experimental setup in Mokrov et al. [2021], ] Fan et [2021]. In this case, the gradient flow at time t is a Gaussian G(t) with a known mean and covariance; see Appendix D.1 for details. We quantitatively compare all methods in Figure 3:
ADJ achieves the best results in dimensions d and d = 10. However, this is the cost of high training time: in dimension 10, ADJ takes 341 minutes to train, while SCVM-TIPF and SCVM-NODE take 23 and 9 minutes respectively for the same 20k training iterations. As such, omit ADJ in higher-dimensional comparisons.
Both our SCVM-TIPF and SCVM-NODE achieve good results second only to ADJ in dimensions 5 and 10. In low dimensions, SCVM-TIPF results lower probability divergences than SCVM- NODE likely due to having exact density access. Although not shown, SCVM-TIPF also satisfies the initial condition well (numbers at t = 0 are comparable to those at f = 0.25 in the left two columns in Figure 3).
For the two JKO methods, they result in much higher errors for < 0.5 compared to later times: this expected because the dependency of Gtt) on t exponential, so convergence to μ² faster in the beginning, yet a constant step size is used for JKO methods.
For DFE, the result is highly sensitive to the forward Euler step size At. We choose step size AL = 0.01 which empirically gives the best results among 10.1, 0.01, 0.001, 0.0001). As DFE achieves far lower symmetric KI. divergence or f-divergence compared to alternatives, we only include its Bures-Wasserrs Wasserstein distance in high dimensions where its number is slightly worse than alternatives (bottom right plot of Figure 3).
4.3 Porous medium equation
Following Fan al. [2021], we consider the porous medium equation with only diffusion: d.p: App" with m > 1. Its solution is the Wasserstein gradient flow of F(u)= = -=TP(x)"" da where p is the density of μ with V (()(())) = V -= pm (1 (x)) see Appendix D.2 for details. We
consider only SCVM-TIPF and JKO methods here.
7

---

SCOR WODE
SCOU-TIP
SCHSSTIP
1074
ACJ
ACJ
JKG4CWN
JKG4CW
HODE
DF2
DFF
PROVOWI
307
2073
07'0
8.0
110
1.50
075
068
1.55
Time 66
Time (22 311
SCHASTIP ACJ
JKS-CVN-RD
JKG4CWN
DFF2
KO
DF7
0%
ala
120
1.75
075
008
110
1.55
Time 66
Time 200
SCAMNNDDE SCWALT PC ADJ
33.30
KGS-CWN-DD
KG--CW
DF2
107
KOJEANEO
KOITNN
6:
0.5
0.5
110
1.5
175
Times 64
Time E 311
Figure 3: Quantitative results for the OU process experiment. The left two columns show the metrics (symmetric KL, symmetric f f-divergence, and Bures- Wasssesssnn distance) versus time of various methods computed against the closed formed solution G(t) in dimension 5, 10. The right column shows the metrics averaged across / versus dimension in higher dimensions.
We show the efficiency of SCVM-TIPF compared to JKO-ICNN in dimension d = 1,2,,..,6 6. We exclude JKO-ICNN-PD because it produces significantly worse results. We visualize the density pr of the solution from SCVM-TIPF and JKO-ICNN on the top of Figure 4 in dimension 1 compared to p². Both methods approximate of well with SCVM-TIPE being more precise small t; this is consistent with the observation in Figure 2 where JKO methods result in bigger errors for small t.
3000
2000
2
JKD-CM
SCVM
1000
2- 104
1- 1000
1- 1003
SCIN TET 2000 SCAN-TT
SciM 899 2010
SCWM 18V 210)
Figure Top: visualization the densities of P and p. for the porous medium equation in dimension 1 at varying time steps i for SCVM-TIPF and JKO-ICNN. Bottom: total variation distance, symmetric f-divergence, and Wasserstein-2 distances across dimensions at 0.004 and 0.025 between Pe and pł for solving the porous medium equation.

---

On the bottom row of Figure 4, we plot the total variation (TV) distance, the symmetric f-divergence, and the W distance (details on the TV distance are given Appendix C.6) between the recovered solution Pt and pi for both methods at / = 0.004 and f. = 0.025. Note that the values of all metrics are very low implying that the solution from either method very accurate, with SCVM-TIPF more precise in TV distance and symmetric f-divergence, especially for d > 3. Like with the experiments in previous sections, JKO-ICNN is much slower to train: in dimension training JKO-ICNN took 102 minutes compared to 21 minutes for SCVM-TIPE.
4.4 TimeeDeeendeet Fokker-Plante equation
In this section, we qualitatively evaluate our method for solving a PDE that is not a Wasserstein gradient flow. In this case, IKO-based methods cannot be applied. Consider the OU process from Section 4.2 when the mean and the covariance matrix I' become time-dependent as 8r and T.. The resulting PDE is time-dependent Fokker- Planck equation of the form (4) with
p4) = (B X) DV log (X
(11)
In this configuration, when the initial measure Po is Gaussian, the solution μ4 can again be shown be Gaussian with mean and covariance following an ODE see Appendix D.3 for more details. We consider, in dimension 2 and 3, deeedeeendeett attraction towards harmonic mean Bg = a(sin (Twt), cos(cooto) (Tw))) using the expression of 8, from Boffi and Eijnden [2023], augmented to Br = a(sin (ssin(wwl)) cos(Twl)), 1) in dimension 3.
We apply both SCVM-TIPF and SCVM-NODE to this problem and compare results with alternatives. Similar to Figure 3, as shown in Figure 5, both SCVM-TIPF and SCVM-NODE achieve results on par with ADJ, with both SCVM methods being 30 times faster than ADJ in dimension 10. DFE results in good Wasserstein-2 metrics but worse divergences. Visualization of the evolution of few sampled particles are given in Figure 9 and Figure 10.
In Appendix D.4, we augment (11) with an interaction term to simulate a flock of (infinitely many) birds, resulting in non- Fokker-Planck PDE that can be readily solved by method.
SCYMITIPE SEVM-KOOE
10-
Synwere (dvegges
ADJ SDE.EM
Synwer
SCVM-TFF
DIF
013
10
Time 24 (2)
T me (3
DRF
ADJ
SCVV TIPF
SOVM-NDDE ADJ SDEEM
Rx
10-2
Time (4))))
Time
Figure 5: Symmetric KI. divergence and Wasserstein-2 distances across time for d = 2,3 between the recovered flows and the ground truth for the time-dependent Bokker-Planch equation.
4.5 Additional qualitative low-dimensional dynamics
To demonstrate the flexibility of our method, we apply our algorithm to model more general mass- conserving dynamics than the considered in the previous sections.
Flow splashing against obstacles. We model the phenomenon against obstacles using Fokker-Plann equation (4) where b, encodes the configuration of three
2-dimensional flow splashing

---

obstacles that repel the flow (See Appendix D.5 for details). We solve this PDE using SCVM-NODE for T' and visualize the recovered flow (6). When solving the same PDE using SDE-EM, the flow incorrectly crosses the bottom right obstacle due to finite time step size (Figure 14). When using DFF, the path of initial samples appears jagged (right of Figure 13); our method has no such issue and results in continuous sample paths (left of Figure 13). Method ADJ suffers from numerical instability and cannot trained without infinite loss in this example.
1-010
Figure 6: A flow splashing against three obstacles (in purple) produced by SCVM-NODE. Particles are colored based on the initial coordinates.
Smooth interpolation of measures. To illustrate the flexibility of our method, demonstrate two ways to formulate the problem of smoothly interpolating list of measures. First, we model the interpolation time-dependent Tokker-Plandee equation and use interpolate MNIST digits and 3, starting from a Gaussian (Figure 7). Next, we adopt an optimal transport formulation and use generate an animation sequence deforming 3D hand model a different pose and then to ball, similar to the setup in Zhang et al. [2022]. Note that the optimal transport formulation is not solvable using competing methods. See Appendix D.6 for more details.
F-
LD
1-80
Figure 7: Smooth interpolation of measures. Top: interpolating MNIST digits to 3. Bottom: interpolating hand from the initial pose to different pose and then to ball.
5 Conclusion
By extending the concept of self-consistency from Shen et al. [2022], we present an iterative optimization method for solving wide class of mass-conserving PDEs without temporal spatial discretization. In all experiments considered, our method achieves strong quantitative results with significantly less training time than JKO-base methods and the adjoint method in high dimensions.
Below we highlight a few future directions. First, as discussed, the two ways to parameterize a proba- bility flow. TIPF, and NODE, both have their specific limitations. Finding new parameterization that combines the advantages of both TIPF and NODE is important next step. Secondly, we hope to extend our approach to incorporate more complicated boundary conditions. Finally, given that the proposed algorithm highly effective empirically, would be an interesting theoretical step to explore its convergence properties.
Acknowledgements The MIT Geometric Data Processing group acknowledges the generous support of Army Research Office grants W91 and W91 10293, of Air Force Office of Scientific Research award 115500-19-1-333 of National Science Foundation grant 195-6195,,77 from the CSAIL Systems that Learn program, from the MIT-IBM Watson AI Laboratory. from the Toyota-CSAIL Joint Research Center, from a gift from Adobe Systems, and from a Google Research Scholar award.
acknowledges the financial support from the University of Bordeaux (UBGR grant) and the French Research Agency (PostProdLEAP).
10

---

References
Michael Albergo and Eric Vanden- Eijnden.. Building normalizing flows with stochastic interpolants. arXiv preprint 2rX1v:2209/5557, 2022.
David Alvarez-Melis, Yair Schiff, and Youssef Mrouch. Optimizing functionals on the space of probabilities with input convex neural networks. arXiv preprint arXiv:2106.00774, 2021.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. International Conference on Machine Learning, pages 146-155. PMLR, 2017.
Marin Biloš, Johanna Sommer, Syama Sundar Rangapuram, Tim Januschowski, and Stephan Gün- nemann. Neural flows: Efficient alternative to neural odes. Advances in Neural Information Processing Systems, 34:21325 22333377 2021.
Mathicu Blondel, Quentin Berthet, Marco Cuturi. Roy Frostig, Stephan Hoyer, Felipe Llinares-Lóper Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. arXiv preprint 1.X1v:2005.11183, 2021.
Nicholas M Boffi and Eric Vanden-Fijnde Probability flow solution of the fokker-planck equation. Machine Learning: Science and Technology, 4(3):035012, jul 2023, doi: 10. 10.1088/2632-2153/ acc2aa. URL https: / http//xxxxoo/// /ttpdxxxioooooooooooooooooooooot org//////// 10. 1088 1088/2632-2222 211/2/26322222222222 21532/11222
Nicolas Bonneel, Michiel Van De Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using lagrangian mass transport. In Proceedings of the 2011 SIGGRAPH Asia conference, pages 1-12, 2011.
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. URL
http:// http:///ittuut github. com/ foom/google/jpi
Martin Burger, José A Carrillo, and Marie-Therest Wolfram. A mixed finite element method for nonlinear diffusion equations. Kinetic & Related Models, 3(1):59, 2010.
José A Carrillo, Alina Chertock, and Yanghong Huang. A finite-volume method for nonlinear nonlocal equations with a gradient flow structure. Communications in Computational Physics, 17 (1):233-258, 2015.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Dan Crisan and Terry Lyons. A particle approximation of the solution of the kushner-sttattnnovichh equation. Probability Theory and Related Fields, 115(4):549-5 578, 1999.
Marco Cuturi, Lactitia Mengg-apaxanntoor Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201.12324. 2022.
Lauro Langosco di Langosco, Vincent Fortuin, and Heiko Strathmann. Neural variational gradient descent. arXiv preprint arXiv:2107.10731 2021.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint /60vv: 08803, 2016.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighter linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3-11, 2018.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Variational wasserstein gradient flow. arXiv preprint arXiv:2112.022 2021.

---

Jean Feydy, Thibault Séjourné, François-Xavier Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2681-2690. PMLR, 2019.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn diver- gences. In International Conference on Artificial Intelligence and Statistics, pages 1608-1617. PMLR, 2018.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Marc van Zee. Flax: A neural network library and ecosystem for JAX, 2020. URL http://github..cm www.uuu..coo/ooee////ee.e
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker-plande equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998.
George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3(6):422 440, 2021.
Patrick Kidger. On Neural Differential Equations. PhD thesis, University of Oxford, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 6980, 2014.
Alexey Kroshnin, Vladimir Spokoiny, and Alexandra Suvorikova. Statistical inference for bures- wasserstein barycenters. The Annals of Applied Probability, 3163):1244-1298, 2021.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.08. 000003, 2022.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay. Justin M Solomon, and Evgeny Burnaev. Large-scale wasserstein gradient flows. Advances in Neural Information Processing Systems, 34: 1523--1555,, 2021.
Gabriel Peyré. Entropic approximation of wasserstein gradient flows. SIAM Journal on Imaging Sciences, 8(4))2323323555
Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Hannes Risken and Hannes Risken. Fokker-planck equation. Springer, 1996.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94 2015.
Zebang Shen and Zhenfu Wang. Entropy-dissipation informed neural network for mckean- type pdes. arXiv preprint arXiv:2303. 1205, 2023.
Zebang Shen, Zhenfu Wang, Satyen Kale, Alejandro Ribeiro, Aim Karbasi, and Hamed Hassani. Self-consistency of the fokker-planee equation. arXiv preprint arXiv:2206.00S8: 2022.
Juan Luis Vázquez. The porous medium equation: mathematical theory. Oxford University Press on Demand, 2007.
Michael Westdickenberg and Jon Wilkening. Variational particle schemes for the porous medium equation and for the system of isentropic euler equations. ESAIM: Mathematical Modelling and Numerical Analysis, 1331)1133-166, 2010.
Paul Zhang, Dmitriy Smirnov, and Justin Solomon. Wassersplines for neural vector field-controlled animation. In Computer Graphics Forum, volume 41, pages 31-41. Wiley Online Library, 2022.
12

---

L....,L. SCVM-TIPF SCVM-NODE ADJ
JKO-IONN JKO-ICNN-PD
O(STd")
D(STNose O
d (
Table 1: Time complexity of training for iterations for the methods considered, in terms of dimension d. We assume for simplicity that the same batch size is used, the training is done for T time steps, and any network forward pass takes ((d) time. For SCVM-NODE and ADJ. N denotes the number of ODE integration steps. For SCVM-TIPF, RealNVP is used to build the coupling layer, and we use d coupling layers, hence the extra multiple of d. For ADJ, the de term comes from the third-order spatial derivatives. For JKO-ICNN, the d' term is due computing the log-determinant term. For both JKO methods, the quadratic dependency on I is due to maintaining growing chain of neural networks of size T as described in the related works section.
Appendix
In this appendix, we provide details and further justification proposed method. In Appendix A, provide an interpretation of the update (9) gradient descent step with biased gradient. In Appendix B, we explain the aategration-by-palt trick used to prove (3.1). In Appendix C, we provide implementation details of all considered methods. In Appendix D, provide additional experimental details and results.
Algorithm 1 Self-consistent velocity matching
Input: fc(:s),, p. T. Nuain B, L. Initialize network weights 0.
for L..... Muain do 8' + 0.
Sample 21..... ........ FB ~ Mi. .1....,ty [0, Ib - 4% (It), vb = -b...,.,,
- 0 nf Vo 0 Bt. Evs (2(()).)) fulva, end for
Output: optimized
A Biased Gradient Interpretation
Assume 2 has density pf. Using ff (x) denote ft(:; p44) from (7), the self-consistency loss can be written as,
L(0) = J (z) (x)]]
dr.
Assuming all terms involving are differentiable with respect to 6 the gradient of L(0) with respect the neural network parameters can be written as:
VL(0) / (a
52 de
(12)
+2 ''' I (nlCe)oool (x)) dæ
(13)
2 I"" (ROON (vf(2) ff (z)) dr dt.
(14)
Here we use Jo to denote the Jacobian with respect to 0. On the other hand, the gradient used in the updates (8) VoF 0²) at 0' 0:
VoF(0, 0-=0
-2 =2 [ (1((()))))
(x)) dx dt.
(15)
We see that (15) is exactly the middle term (13). Hence our formulation can interpreted as doing gradient descent with a biased gradient estimator. It remains a future work direction to
13

---

theoretically analyze the amount of bias in (15) and the condition under which the dot product (VL(Ø), VoF 0')))= The central challenge would be relate Jae and Jg f this depends on the neural network architecture and the type of the PDE.
B Integration-by-Part Trick
This is common trick used in score-matching literature [Hyvärinen and Dayan, 2005].
Proof of Proposition 3.1. Fix t > 0. The form of ft in (4) is
f(( ((()) Dx
Hence
E Ex-uy [o (x)( f((:: = Ex (x)) E Ex (.(X))D((XVllll (x)]
The second term can be written as
E Exmep' [xT(X) Dz
(xx)x (x)] ((()))DD((VVVog))) (x) (z)
/ ( (z)/p (x) pt (z)
/ (l)x)) T Di(x) (z)
/ v (D. (z))) (x)
Ex X of V (Dx(X) (X)),,
where integration-by-part to get the second last equation and the assumption that of, Dr bounded and p. (z) - 0 as || + co.
C
Implementation Details
C.1 Network architectures for SCVM.
For TIPF, our implementation follows Dinh et al. [2016]. Each coupled layer uses 3-layer fully connected networks with layer sizes 64, 128, 128 for both scale and translation prediction. We use twice as many coupling layers as the dimension of the problem while each coupling layer updates one coordinate; we found using fewer layers with random masking gives much worse results.
For NODE, we use a 3-layer fully connected network for modeling the velocity field with layer size 256. We additionally add low-rank linear skip connection i A(()x + b(l) where A(I) L (t) and L(t)) is d x 20 matrix make A(t) low-rank.
We use SILU activation [Elfwing et al., 2018] which is smooth for all layers for both TIPF and NODE. For NODE, we apply layer normalization before applying activation. We also add sinusoidal embedding for the time input plus two fully connected layers size 64 before concatenating it with the spatial input.
The numerical integration for NODE is done using Diffrax library [Kidger, 2021] with a relative and absolute tolerance of 10 4; we did not find considerable improvement when using a lower tolerance.
We use the integration- by--prtt trick for SCVM-NODE whenever possible. Since TIPF has tractable log density, we do not use such a trick and optimize (9) directly for SCVM-TIPF which we found to produce better results.
C.2 Hyperparameters.
Unless mentioned otherwise, we choose the following hyperparameters for Algorithm 1. We set Nuain 10s 2 103, B 1000, L 10 20. We use Adam [Kingma and Ba, 2014] with
14

---

cosine decay learning rate scheduler, with initial learning rate 10 10m3, the number of decay steps same as Minin and a 0.01 (so the final learning rate 10 5). Since we are effectively performing gradient descent using a biased gradient, we set b2 = 0.9 in Adam (instead of the default b2 0.999), so that the statistics in Adam can be updated more quickly; we found this tweak improves the results noticeably.
C.3 Implementation of JKO methods.
We base our JAX implementation of ICNN on the codebase by the original ICNN author: https:/ bestps://githen /tt.bb/..tu.....// thub. com/facebookreseerrd.o w2ot. Compared to the original ICNN implemen- tation by Amos al. [2017], we add additional convex quadratic skip connections used by Mokrov et al. [2021]. which we found to be crucial for the OU process experiment. For ICNNs, we use hidden layer sizes 64, 128, 128, 64. The quadratic rank for the convex quadratic skip connections set to 20. The activation layer is taken to be CELU.
To implement the method by Fan et al. [2021], model the dual potential a 4-layer fully connected network with layer size 128, with CELU activation. For the gradient flow of KL divergence and generalized entropy (used in Section 4.3), follow closely the variational formulation and the necessary change of variables detailed in Fan et al. [2021, Corollary 3.3, Corollary 3.4).
In order to compute the log density any JKO step, following Mokrov et al. [2021], we need to solve a convex optimization to find the inverse of the gradient of an ICNN, We use the LBFGS algorithm from JAXopt [Blondel et al., 2021] to solve the optimization with tolerance 10-2 (except for Section 4.3 we use a tolerance 10 -0 to obtain finer inverses, but it takes 6x longer compared to 10 10.
We always use 40 JKO steps, consistent with past works. For each JKO step, we perform 1000 stochastic gradient descent using Adam optimizer with a learning rate of 10-3, except for the mixture of Gaussians experiment, we use 2000 deps-using fewer steps will result in worse results. We have tested with the learning rate schedules used in Fan et al. [2021], Mokrov et al. [2021] and did not notice any improvement.
C.4 Implementation of ADJ method.
We implement the adjoint method carefully following the formulation in Shen and Wang [2023]. The neural network for parameterizing the velocity field is identical to the one used in SCVM-NODE. The ODE integration also uses the same hyperparameters as that of SCVM-NODE. This way we can compare ADJ with SCVM-NODE in fair manner since they only differ in the gradient estimation.
C,5 Implementation of DFE method.
We implement DFE following the algorithm outlined in Boffi and Vanden-Fijnden [2023]. We 5000 particles. For score estimation, use the same network architecture in NODE. At each time step, we optimize the score network 100 steps. We found the result of DFE depends tremendously on the time step size At. For the OU process experiment in dimension 60, when At = 0.1, 0.01, 0.001, 0.0001, the resulting Bures- Wasserstein distance at the final time to the target measure is 28. 11, 0.31, 0.46, 9.21 respectively. Surprisingly, a smaller At result in bigger errors. We choose At = 0.01 since it gives the best results.
C.6 Evaluation metrics
For all our experiments, calculations of all metrics are repeated 20 times on 1000 samples from each distribution. Our plots show both the average and the standard deviation calculated over these 20 repetitions.
When estimating symmetric KL divergence using samples, due to the finite sample size and the numerical error in estimating the log density, the estimated divergence can be very close to zero or even negative (when this occurs we take absolute values). This explains why the standard deviation regions touch the z-axis in the log-scale plots in Figure 3.
15

---

To compute Bures-Wer Wasserstein distance [Kroshnin et al., 2021], we first fit Gaussian to the samples of either distribution and then compute the closed-form Wasserstein-2 distance between the two Gaussians.
For the porous medium equation (Section 4.3), the total variation distance is used in Figure and Figure 8 to compare the estimated and ground-trute truuh solutions. It is approximated by the L distance between the densities calculated over 50000 samples uniformly distributed on the compact
11.22 m.x 1.25xmaxl with Fmax
C/ (8t to) to) - being the bound of the support of p.-
D Additional Experimental Details
D.1
Ornstein-Uhlenber process
The OU process is the Wasserstein gradient flow of the KI. divergence with respect to Gaussian "" = N(8,1'1 where 3 E R and I' is a x d positive-definite matrix. When the initial distribution pD (0, Ia). the gradient flow at time t is known to be Gaussian distribution G(t) with mean (La e 8 and covariance 1'=1 (la - e 211" e 221" We set the total time T = 2.
D.2 Porous medium equation
This flow has as closed-form solution given by the Barenblatt profile Vázquez [2007] when initialized accordingly:
pt (1) (c |||0|| (t (o)
(16)
where to > 0 is the starting time, = d(m- + 5
(m-1)a and C > 0 is free constant. Similar
2dm
to Fan et al. [2021], choose = 2 and total time T 0.025. The initial measure follows
Barenblatt distribution supported in |-0.55, 0.25] (C is chosen accordingly) with lo = 10-3 3, We
Cetropolis-Hasttnng to sample from /A0-
J000 1m
Figure 8: Metrics (TV, Symmetric f-divergence and Wasserstein-2 distance) across time for dimen- sions 3 and 6 between the estimated μ4 and the ground truth 4r when solving the Porous Medium Equation.
D.3 Time-Dependant Fokker- Planck equation
We consider time-dependent Fokker-Planck equation of the form (4) with the velocity field
fe(X, pr) = T+(X - -3) - D, (X().
(17)
When the initial measure Po is Gaussian, the solution f4 again be shown be Gaussian with mean my and covariance Σ, solutions of the differential equations:
m [m, 33) Σ = -r, Σ,T. + 2D.
(18)
In practice, we experiment with constant r, d1,g(1.3) and D, I² We also experience in
dimension 3 by considering and r, = diag(1 1, 3, 1). We set a 3, w 1, = v0.25 and pick as initial distribution po Gaussian with mean by and covariance o² 22 We set the total time to T 10.
16

---

We plot in Figure 9, for dimension 2, snapshots at different time steps of particles following the flow given by our method with TIPF parametrization. We only show SCVM-TIPF because SCVM- NODE gives visually indistinguishable trajectories. We also plot in Figure 10 the evolution of particles simulated by Fuler-Maruyane (EM-SDE) discretization of the Fokker-Planck equation. Corresponding animated GIFs be found at this link.
Figure 9: Evolution of particles (in blue) following the flow learned with SCVM-TIPP for the time-dependent OU process (Section 4.4). In red is the moving attraction trap.
2005
Figure 10: Evolution of particles (in blue) obtained by SDE-EM discretization for the time-dependent OU process (Section 4.4). In red is the moving attraction trap.
D.4 Flock of birds
We model the dynamies flock of birds by augmenting the time-dependent Fokker-Planck equation (1I) with an interaction term:
fi(X, 14) [6 - X) ax E(p,)) DV log (X).
This is similar to the harmonically interacting particles experiment Boffi and Vanden-Eijnden [2023], but we use population expectation E(pt) instead of an empirical one in modeling the repul- sion from the mean. Since fo needs to access Eluu]. the resulting PDE is not a Fokker-Planck equation (4) and hence not solvable using the method in Boffi and Vanden-Eijnden [2023] but can be solved with our method by estimating E[pr] using Monte Carlo samples from μ4c- We use a similar setup in Section 4.4, except now use an "infinity sign" attraction Br a(cos(2mwt)) 0.5 sin (2twt)) along with a sinusoidal a = sin( (w)t))). Depending on the sign au, particles are periodically attracted towards or repulsed from their mean. Both SCVM-TIPE and SCVM-NODE produce similar visual results as shown in Figure and Figure 12.
We use constant r, Ta and a constant diffusion matrix D = v0.25. We pick as initial distribution Po a Gaussian with mean (0, 0) and covariance T We set the total time to T 10.
o2 We set = 3, === 0.., and
17

---

We respectively show in Figure 11 and Figure 12 simulations of particles following the flow learned with SCVM-TIPF and SCVM-NODE. Corresponding animated GIFs be found this link.
Figure 11: Evolution of particles following the flow trained with TIPF parametrization on the flock of birds PDE (Section 4.5). In red shows the moving attraction mean.
Figure 12: Evolution of particles following the flow trained with NODE parametrization on the flock of birds PDE (Section 4.5). In red shows the moving attraction mean.
D.5 Flow splashing against obstacles
We use the following formulation for modeling the flow. Each obstacle modeled line segment. The endpoints of the three obstacles are:
(00, 3), (3, 0.5)), ((1, 0), (1.5, 0)), (6,, (6.
We model the dynamics as a Fokker- Planck equation where f. of the form (4) is defined as
Σ x (x) 1=1
bz(z) (Asink x)) 20
|| To, (x)) PN1000.0))( To, (x)1l),
De(x)) = =12
where (sink (4, 0). and To, (x) is the projection of x onto obstacle i represented as a line segment, and PA((0.0.04) is the density of 1-dimensional Gaussian with variance 0.04.
The initial distribution is chosen to be (0, 0.2572). We train SCVM-NODE for 1 with initial learning rate of 10-4 Training takes 5.4 minutes. The time step size for SDE-EM used to produce Figure 14 is 0.005. Corresponding animated GIFs be found at this link.
18

---

Figure 13: Trajectory of 200 random particles across time using the same setup as in Figure 6. Left are sample paths obtained by our method, and right are sample paths obtained by DFE [Boffi and Vanden-Eijndens 2023].
Figure 14: Same setup in Figure but with SDE-EM. We see the paths of the particles continuous. Moreover, the particles spill over the obstacle on the bottom right due to finite time step size. comparison, SCVM-NODE does not have such a problem.
D.6 Smooth interpolation of measures
Suppose we are to smoothly interpolate M measures P1s.... PM with densities Q1. c...... QMM and want the flow to approximate D, at time To achieve this goal, we present two formulations using different choices of ft- We use SCVM-NODE in this section.
Measure interpolation using time-dependent Nokker-Planck equations. Tokker-Plande equation where fi of the form (4) is taken be
We model the dynamics
be(x) Σ plt r:) (ogg(<)) log):(x))
D.(x)= I2,
where ((t) is defined as the continuous bump function
0.5h < 0.6/
<(t)
(0.6h (4)((0.1))
0.0
otherwise,
for bandwidth h 1.0.
Below we provide details for MNIST interpolation result the top row Figure 7. use the first three images of 1, from the MNIST dataset. To construct from digit image, we use mixture of Gaussians where we put one equally- pixels with values greater than 0.5 (images are first normalized to have values in [0, 1)). The initial distribution is (O.5, 0.5), 0.04[2). To train SCVM-NODE, we use an initial learning rate of -4 with cosine decay for total of 10 iterations. This takes roughly hour to train.
Gaussian with covariance 0.02² In on the
Measure interpolation using optimal transport. An alternative way to interpolate measures using our framework is to optimal transport to define ft. Recall μr denotes the probability flow at time We then define
M
f((z) Σ olt rs V.. (2r,
19

---

where W is the squared Wasserstein-2 distance and Ww, W is its Wasserstein gradient. In practice, we compute Ww, W using sample access and we employ the debiased Sinkhorn divergence [Genevay et al., 2018, 8, Feydy et al., 2019] implemented in the JAX OTT library [Cuturi et al., 2022]. This formulation differs from the one in Zhang et [2022] in that here we prescribe the precise PDE based on ft. whereas in Zhang et al. [2022] an optimal transport loss is used to fit the keyframes along with many regularizers on the velocity field it to promote the smoothness and other desirable properties. In contrast, we do not use any regularizer Ut-
To train SCVM-NODE Figure 7, we first sample 20000 points from the interior of the three meshes (a hand mesh, hand mesh different pose, and ball mesh), and we set V; to the empirical measure of the corresponding point cloud. Note that different from the first formulation using Fokker-Planche equations, in the optimal transport formulation, throughout we only require sample access from each 14, We use an initial learning rate of 10-4 with cosine decay for total of 5 x 10 iterations. This takes 4.5 hours, which is significantly longer than the training time Zhang al. [2022] (reported to be 15 minutes). We leave further improvement of our method interpolate measures faster as future work.
produce the hand-hand mald-ball animation sequence in the bottom row of
To render the animation, we sample 20000 points and render the point cloud each time step using metaballs along with smoothing, similar to the procedure in Zhang et al. [2022]. We did not use the barycentric interpolation postprocessing step Zhang et [2022] which makes sure the key measures ey 's are fit exactly in the resulting animation. We also did not use unbalanced optimal transport, which reported in Zhang al. [2022] can make the fingers of the hand more separated, but requires careful parameter tuning.
20


Reviews:
[{"review_id": "Oon6JviECf", "reviewer": "Reviewer_AfQf", "summary": "This paper is concerned with solving mass-conserving PDEs using neural network parameterizations. The proposed methodology relies on the minimization self-consistency loss presented in [1, 2]. While the work of [2] is specialized to FP equation and McKean-Vlasov type PDEs, their approach cannot be generalized to more general cases. The author propose a simple yet effective way of minimizing the said loss. The proposed method is agnostic to the form of $f_t$ and furthermore, it does not require differentiating through $f_t$ and $\\mu_\\theta$, which is expensive in practice. \n\n[1] Shen, Zebang, et al. \"Self-consistency of the fokker planck equation.\" Conference on Learning Theory. PMLR, 2022.  \n[2] Shen, Z., & Wang, Z. (2023). Entropy-dissipation Informed Neural Network for McKean-Vlasov Type PDEs. arXiv preprint arXiv:2303.11205.\n\n", "questions": "l. 243, it is mentionned that TIPF, which involves normalizing flows, does not scale well with the dimension. This is surprising to me since normalizing flows do scale relatively well with the dimension. For $d = 60$ i do not see what the bottleneck is. ", "limitations": "see above. ", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- I found the paper to be clearly written, well detailed and with an extensive discussion of the related works.\n- The proposed method seems to have a significant advantage compared with JKO schemes and the adjoint method of [2]. First, it is simple to implement yet effective and fast. Next, for the FP equation or for the Wasserstein gradient flow, which are of practical interest, it is not required to differentiate through $\\nabla \\log p_t$ (which can be expensive). \n- Overall, the numerical experiments are interesting and extensive. ", "weaknesses": "I think that it may have been better to have a more detailed discussion/comment of the biased gradient estimate of the self-consistency loss $\\nabla_\\theta F(\\theta, \\theta_k)$. The resulting optimization procedure is reminiscent of a minimization-minimization procedure and I believe that there might be some connection that I haven't had time to dig in further. I guess that the paper would be a bit more interesting if the authors are able to justify that eq. (9) is a principled criterion that, when properly minimized, ensures the decrease of some statistical distance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review. Please find below our response to your specific questions.\n\n**A detailed discussion of the biased gradient estimate and connection to the minimization-minimization procedure**\n\nPlease refer to (a) in the common response to all reviewers.\n\n**Why doesn’t TIPF scale to $d=60$?**\n\nWe use [RealNVP](https://arxiv.org/pdf/1605.08803.pdf) as the backbone for TIPF. The building blocks are coupling layers that update one coordinate while masking out the others. The bottleneck comes from the fact that we cannot afford too many coupling layers, since the number is proportional to the dimension. We have tested the option of updating multiple coordinates in one coupling layer but got significantly inferior results.\n\nIn the original paper, RealNVP has been applied successfully to image space which is high-dimensional. For image space, it is possible to update many coordinates in a structured way (e.g. using a checkerboard pattern; see Figure 3 of the RealNVP paper) and still achieve good performance. However, we are dealing with unstructured vector data so similar tricks cannot be applied easily.\n\nIt is possible that there exists a better choice of normalizing flow other than RealNVP, and TIPF could be improved with a better architecture. If you happen to know such an architecture, please let us know."}, {"review_id": "QKgdAU5pwA", "reviewer": "Reviewer_3qHN", "summary": "The authors propose a method for\nsolving mass-conserving partial differential\nequations (PDEs) within an iterative optimisation scheme\n without resorting\nto either spatial or temporal discretisation.\n\nInspired from previous work [1] on imposing \nself-consistency for efficiently solving\nFokker-Planck equations, they\n impose a self-consistency constrain to\nobtain the time-dependent velocity field that \nis associated with mass conserving\nPDE. Thus instead of solving a PDE in terms of a probability density, they solve more efficiently for obtaining the associated\nvelocity field. \nThey represent the probability flow of the mass-conserving PDE with a neural network, and within an iterative scheme they optimise\nthe weights of the network to satisfy the self-consistency constraint evaluating the relevant terms using the network weights of the previous iteration.\n\nThey test the performance of their approach\non several numerical experiments and compare\nwith analytic solutions (when available). The explored examples study diverse model systems starting from simple analytically tractable systems (Fokker-Planck equation with stationary density comprising a mixture of Gaussians and  Ornstein-Uhlenbeck processes of dimension $d=10;30;60$) , the porous medium and time-dependent Fokker-Planck equations, as well as some systems with many-body interacting dynamics, and smooth interpolation between non-trivial probability measures. Although the method of [1] seems to perform better in higher dimensnions, the current proposed approach claims to be considerably faster than [1].\n", "questions": "\n\n\n- Do you have any insight what happens in Figure 2 in the d=30;60 experiments and the JKO-PD method has such a non monotonic decrease in sym KL with time?\n\n- Could you probably perform a systematic study for the computational costs of the current approach for increasing system dimension and compare it with the other competing methods?\n\n- Do you have any empirical insights on when the biased optimisation of the self-consistency constrain will fail?\n\nMinor:\n\n- There is a verb missing in Line 218\n- Line 64 \"comparing\" seems incorrect here. Do you mean \"competing\"?", "limitations": "- No guarantee that the optimisation will converge on the correct solution given random weight initializations.\n\nReferences:\n\n[1] Shen, Zebang, and Zhenfu Wang. \"Entropy-dissipation Informed Neural Network for McKean-Vlasov Type PDEs.\" arXiv preprint arXiv:2303.11205 (2023).\n\n[2] Shen, Zebang, et al. \"Self-consistency of the fokker planck equation.\" Conference on Learning Theory. PMLR, 2022.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "\n- Impressive that the proposed approach does not require either time or space discretisation.\n\n- well written and presented manuscript.\n\n- applicable both for representing probability flows in terms of neural ODEs and as a time-dependent push forward, that making it possible to relate the current approach with several existing frameworks for solving PDEs/transforming prob. densities.\n\n- the authors present extensive numerical experiments in multiple settings employing systems of relatively high dimension.", "weaknesses": "\n- Relatively small contribution when considering [1] and [2], but nevertheless impressive results.\n- Some of the plots are extremely difficult to read [eg fig 4].\n\n- There are no convergence guarantees of the framework and the solution relies on the convergence of the biased optimisation of the self consistency equation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review. Please find below our response to your specific questions.\n\n**Some of the plots are extremely difficult to read.**\n\nThank you for pointing this out. We will make the figures bigger in the revision.\n\n**Although the method of [1] seems to perform better in higher dimensions, the current proposed approach claims to be considerably faster than [1].**\n\nWe’d like to point out that the ADJ method of [1] *does not perform better* in higher dimensions (it becomes significantly worse than SCVM-NODE at $d=10$): this is shown in Figure 2 second row. As written in the first bullet point in Sec 4.1, we suspect this is due to the amount of error accumulated in the two integration passes in ADJ and the usage of up to third-order spatial derivatives of the neural velocity field. Although the formulation of ADJ has no bias, the numerical procedure can introduce more errors.\n\n**Do you have any insight into what happens in Figure 2 in the $d=30;60$ experiments and the JKO-PD method has such a non-monotonic decrease in sym KL with time?**\n\nJKO-PD is a primal-dual method that involves solving a min-max optimization problem. It is only equivalent to the JKO flow (for which the objective must be monotonic) if the inner maximization is solved perfectly. This is not the case in practice and is hence a disadvantage of JKO-PD. It is possible the result will be better with more inner steps. We use 2 outer minimization steps and 3 inner maximization steps, which is the default recommendation from the JKO-PD paper.\n\n**Could you probably perform a systematic study of the computational costs of the current approach for increasing system dimension and compare it with the other competing methods?**\n\nPlease refer to (b) in the common response to all reviewers.\n\n**Do you have any empirical insights on when the biased optimization of the self-consistency constraint will fail?**\n\nFor all experiments presented in the paper, the biased optimization is surprisingly effective without needing to tune the parameters. For SCVM-NODE, we did notice a failure case when the network parameters are initialized from a Gaussian with a very large variance (corresponding to a rapidly changing velocity field with high magnitude). However, the default weight initialization from FLAX suffices in all experiments. We also note other methods rely on proper weight initialization as well.\n\nPlease also refer to (a) in the common response to all reviewers."}, {"review_id": "GobgGs5gNH", "reviewer": "Reviewer_v66q", "summary": "This article deals with the problem of solving partial differential equations (PDEs) using neural networks. More specifically, the authors focus on mass-conserving PDEs:\n$$ \\forall x,\\quad \\forall t\\in[0,T],\\quad \\partial_t p_t(x) = -\\nabla\\cdot( f_t(x;\\mu_t)p_t ) . (Equation~2) $$\n\nThe authors have developed an iterative optimization scheme called \"self-consistent velocity matching\" (Algorithm 1 in the supplementary material). This method is based on self-consistency, _i.e._ by denoting $\\Phi_t\\colon\\mathbb{R}^d\\to\\mathbb{R}^d$ the flow map at time $t$ associated to Equation (2) and $v_t\\colon\\mathbb{R}^d\\to\\mathbb{R}^d$ its velocity at time $t$, asking for\n$$ \\forall x,\\quad v_t(x) = -f_t(x;\\mu_t) . (Equation~6) $$\n\nAssuming that the flow map $\\Phi_t$ (or the velocity field $v_t$) is a neural network (TIPF for $\\Phi_t$, or NODE for $v_t$) parametrized by a parameter $\\theta$, the \"self-consistent velocity matching\" algorithm iterates three steps:\n1. Define a space-time grid where to solve the equation: $\\forall b=1,\\ldots,B$, $\\forall\\ell=1,\\ldots,L$, $~$ $x_b\\sim\\mu^\\ast$, $t_\\ell\\sim\\mathcal{U}([0,T])$,\n2. Let particles $x_b$ drift along the flow $\\Phi_t$: $\\forall b$, $\\forall\\ell$, $~$ $y_{b,\\ell}=\\Phi_{t_\\ell}^{\\theta^k}(x_b)$,\n3. Update the parameter $\\theta^k$ according to a _biased_ gradient descent step: $\\theta^{k+1}=\\theta^k-\\eta\\frac1{BL} \\nabla_\\theta \\sum_{b,\\ell}\\lVert v_{t_\\ell}^{\\theta}(y_{b,\\ell})-f_{t_\\ell}(y_{b,\\ell};\\mu^{\\theta^k}_{t_\\ell})\\rVert^2$.\n\nThe authors conduct several numerical experiments to compare themselves with the state of the art. The results obtained are comparable or better in larger dimensions. In particular, the proposed method seems more in line with the laws of physics (_cf._ flow against obstacles experiment, Figure 13 in the supplementary material).\n\nThe theoretical convergence of the proposed method is not investigated.", "questions": "Some comments:\n- Perhaps emphasize on lines 142-143 that the flow *or* the velocity field will be parameterized. \n- Line 59: Which figure in Figure 2?\n- Line 147: realNVP not defined\n- The way references are written down is not consistent.\n- When reference is made to figures in the supplementary material (Figure 8 and following), specify that they are in the supplementary material; the same applies to Algorithm 1.\n- Algorithm 1: I think it should be $t_\\ell$ instead of $t$ in the $\\theta$ update (last line before *end for*).\n- Equations (13), (14), and (15): Should there not be a 2?", "limitations": "Not concerned.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- A self-sufficient article with precise references to the supplementary material that makes it easy to read. \n- The proposed algorithm is precisely defined, step by step, and its limitations are discussed.\n- In the supplementary material, all the details of the numerical implementation are given, making the experiments reproducible.\n", "weaknesses": "The introduction lacks context. Similarly, the supplementary material would benefit from being preceded by an opening paragraph with a title and summary of its contents.\n\nThe proposed method is based on the use of TIPF or NODE networks, each with its own advantages and disadvantages. It is not clear to me when to use one or the other.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review. Please find below our response to your specific questions.\n\n**The introduction lacks context. Similarly, the supplementary material would benefit from being preceded by an opening paragraph with a title and summary of its contents.**\n\nThank you for pointing this out. We have revised the text according to your comments and added an opening paragraph in the appendix as follows:\n\n“In this appendix, we provide details and further justification of the proposed method. In Appendix A, we provide an interpretation of the update (9) as a gradient descent step with a biased gradient. In Appendix B, we explain the integration-by-parts trick used to prove (3.1). In Appendix C, we provide implementation details of all considered methods. In Appendix D, we provide additional experimental details and results.”\n\n**The proposed method is based on the use of TIPF or NODE networks, each with its own advantages and disadvantages. It is not clear to me when to use one or the other.**\n\nPlease refer to (c) in the common response to all reviewers."}, {"review_id": "YGt6u7rQwW", "reviewer": "Reviewer_5yDs", "summary": "The authors present a framework for solving mass-conserving partial differential equations (PDE) from Lagrangian viewpoint. In this framework the flow map or velocity field of the problem is parameterized by two neural network approaches: Time-dependent-Invertible-Push-Forward (TIPF) and Neural ODE (NODE). Then the self-consistency condition is enforced on the velocity field parameterized by the neural network methods mentioned above through an iterative optimization process. In this iterative procedure a biased estimate of the  gradient of the self-consistency condition is used. The proposed approach reduces the computational cost while providing high performance in terms of error metrics such as KL-divergence, total variation (TV) distance. The authors evaluate their framework by considering a wide range of PDEs and comparing against the baseline models developed in recent years. ", "questions": "- In Section 3.2 it is mentioned that to enforce initial condition, an extra loss term in the form of L2 loss is included to the loss function for TIPF method. This does not guarantee the satisfaction of the constraint 100% and it is just a soft constraint enforcement that has been used by PINNs model as well. After the model is trained, did you check the model error for predicting the initial condition? Or after training process is finished, how does the initial condition loss trend looks like and how much is it at the end of training?\n- In addition to initial condition, how do you handle other constraints such as boundary conditions? have you checked whether the predicted solution satisfies the boundary constraints, and if not, how much is the discrepancy?\n- In page 4, I did not follow how the solution density $p_t$ can be recovered using the flow map. In particular, where did that equation come from: $log p_t(x) = log p_0^*(\\phi_t^{-1}(x)) + log \\text{det} | J \\phi_t^{-1} (x)|$. It is important for authors to put themselves in readers' shoes and provide enough information when describing their methodology.\n- Regarding the gradient estimate, could you provide some discussion about the missing terms in calculating the gradient? The proposed methods result in more accurate predictions compared to other baselines, does it imply that for these specific canonical problems the effect of those terms are negligible? I guess that may not be true. Is it useful to compare the predictions against accurate solutions from a numerical solver with fine time and space resolution for a stiff problem, in order to better capture the effect of missing terms?", "limitations": "- Authors describe a few future directions to investigate some existing limitations in the current manuscript. I would encourage them to first provide some numerical evidences about the potential limitations they want to explore in the future works. For example, as asked above, using their current framework, they can numerically investigate how much is the discrepancy in satisfying initial condition, and boundary conditions. \n- Finding more accurate estimate of gradient in their iterative framework is another important aspect to be explored. For the current paper, I still hope to see that they provide some insight about the effect of missing terms.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- It is nice to see an extensive array of empirical evaluations against other baseline models.\n- Overall the paper is well-written and easy to follow (despite few confusing statements addressed below)\n", "weaknesses": "- I wish there were some preliminary analysis on the effect of the biased gradient estimate. As authors clearly described in Appendix A, they eliminate some terms in the gradient of self-consistency loss (to alleviate some computational cost), and just use the term (13). To me, it seems like that term (14) can be computed in the same way as term (13) and I was wondering if there is any reason that it is eliminated. It is unclear how other terms could impact the model performance (accuracy) with the cost of increasing computational complexity of the model and how would this trade-off look like. \n- I did not find any description of what boundary conditions were used for some of the problems such as Fokker-Planck equation and how much the model is capable of satisfying that. If they are not satisfied, how much is the effect on accuracy of the predicted solution.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review. Please find below our response to your specific questions.\n\n**To me, it seems like that term (14) can be computed in the same way as term (13) and I was wondering if there is any reason that it is eliminated.**\n\nTerm (14) involves the term $J_\\theta f_t^\\theta(x) = J_\\theta f_t(x; \\mu_t^\\theta)$. This term is challenging to compute because 1) the derivative of $f_t(x;\\mu_t)$ w.r.t. $\\mu_t$ can be complicated (e.g. in optimal transport formulation in Sec D.6), and 2) the derivative of $\\mu_t^\\theta$ w.r.t. $\\theta$ can also be complicated: for instance, for NODE, numerical integration is needed to obtain $\\mu_t^\\theta$, so differentiating it w.r.t. $\\theta$ must differentiate through each integration step which is costly; we verified this in our early experiments.\n\nIn comparison, the term (13) is easy to compute because for both NODE and TIPF we have closed-form expressions for $v_t^\\theta$.\n\nPlease also refer to (a) in the common response to all reviewers.\n\n**Are boundary conditions used?**\n\nAside from the initial condition, we did not consider any other boundary condition. This is in line with competing methods such as [Shen and Wang 2023, Mokrov et al. 2021, Fan et al. 2021]. As acknowledged in the conclusion, finding new ways to handle boundary conditions is an important future step.\n\n**After training SCVM-TIPF, what is the error in predicting the initial condition?**\n\nFor OU process experiments, we have computed the same metrics from Figure 3 for the recovered flow at $t = 0$ — we did not include these metrics at $t=0$ because the plot is in log-scale and all other methods have exactly zero discrepancies at $t=0$. Since metrics are computed against the ground truth closed-form solution, the numbers at $t=0$ will reveal how well the initial condition is satisfied. For TIPF, the numbers for all metrics at $t=0$ are very small and are close to the numbers at $t=0.25$. For instance, for $d=10$, at $t=0$ the symmetric KL is $10^{-3.26}$ and the Bures-Wasserstein distance is $10^{-1.44}$. We have added a comment about this in Sec 4.2.\n\nFor the porous media equation experiment, the first plot in Figure 4 shows visually the initial condition is well satisfied.\n\n**How is the density $p_t$ recovered from the flow map?**\n\nThe formula $\\log p_t(x) = \\log p_0^*(\\Phi_t^{-1}(x)) + \\log \\det |J\\Phi_t^{-1}(x)|$ is a direct consequence of the change-of-variable formula. See (1) in [Dinh et al. 2016]. We have added this reference to the paper.\n\n**Further discussion about the missing terms in calculating the gradient?**\n\nPlease refer to (a) in the common response to all reviewers.\n"}, {"review_id": "m4iaLxED0z", "reviewer": "Reviewer_wUYN", "summary": "The paper proposes a method for solving mass-conserving partial differential equations (PDEs). The method leverages the property of self-consistency to devise an iterative optimization method for velocity matching. The proposed method was tested on several important PDEs.", "questions": "As elaborated in the manuscript, NODEs and TIPFs have their own weaknesses: namely, training time for NODEs and scalability for TIPFs. Although some numbers, e.g., training time, are provided in some benchmark problems, it is hard to assess the pros and cons of the two approaches. Are there some systematic comparisons on this aspect? In-depth comparisons between SCVM-NODEs and SCVM-TIPFs? ", "limitations": "The authors address the limitation to some extent, described in the form of future directions. One of those limitations is being asked in Questions. ", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-written. The motivation and the idea is clear. \n\n- The description of the literature is comprehensive.\n\n- Two sets of the proposed methods in very different characteristics, NODEs and TIPFs, are interesting.\n\n- Empirical evidence is strong.\n", "weaknesses": "- As opposed to the very well-written texts, some of the figures are in poor quality. Figures 3, 4, 5 are particularly hard to read. It needs some work to make them more legible, e.g., different markers, etc.\n\n- In-depth comparisons between SCVM-NODEs and SCVM-TIPFs would give more information/intuitions to readers on which method to use given a specific PDE problem. Please see Questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review. Please find below our response to your specific questions.\n\n**Figures are of poor quality.**\n\nThank you for pointing this out. We agree the figures are too small, partially due to the tight page limit. We will make the figures easier to read in the new revision.\n\n**In-depth comparison between NODE and TIPF would give more information/intuition.**\n\nPlease see (c) in the common response to all reviewers.\n"}, {"review_id": "wDKHzZYOOg", "reviewer": "Reviewer_egp4", "summary": "This paper proposes to solve a class of mass preserving PDEs by minimizing a loss defined with a self-consistency condition. Also, instead of minimizing this loss directly, an iterative formulation with a biased gradient estimator is proposed. Proposed method is a extension of two previous papers, extending their self-consistency to mass preserving PDEs, and getting rid of the usage of the expensive adjoint method.", "questions": "What is the computational complexity of the method, with regard to the dimension of the PDEs?", "limitations": "Proposed method only works for mass-preserving PDEs.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The class of PDEs being solved is important, including the celebrated Fokker-Planck and Wasserstein gradient flow. \n\nThe computational complexity of this method is much lower than the previous ones, while preserving good empirical results.\n\nProposed method does not require spatial and temporal discretization, making it possible to use in high-dimensional case.", "weaknesses": "Although the mass preserving PDE (2) will always induce the self consistency condition (6), the other way around does not work, and it is not argued in much details when this will fail. \n\nThe proposed method does not require discretization, and it is motivated that the proposed method works better than the previous methods on high-dimensional cases, but very few experiments were performed in the high-dimensional environment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable review. Please find below our response to your specific questions.\n\n**Although the mass-preserving PDE (2) will always induce the self-consistency condition (6), the other way around does not work, and it is not argued in much detail when this will fail.**\n\nWe would like to remind the reviewer that the self-consistency condition (6) does imply (2) under fairly general conditions (if the velocity field is Lipschitz continuous and bounded) — this is discussed in Sec. 3.1.\n\n**Very few experiments were performed in the high-dimensional environment.**\n\nOne challenge is that there is no closed-form ground truth solution for high-dimensional PDEs. We are only aware of OU processes which we tested extensively on. If the reviewer has suggestions on other high-dimensional PDEs that we could test our method on, please let us know.\n\n**What is the computational complexity of the method, with regard to the dimension of the PDEs?**\n\nPlease see (c) in the common response to all reviewers.\n"}]

Author rebuttal:
We thank all reviewers for their constructive feedback. We appreciate the recognition of the simplicity and empirical effectiveness of our method. Please find below our common response to the initial reviews.

**(a) A more detailed discussion of the biased gradient estimate and other interpretations.**

Aside from Appendix A, it remains future work to quantify precisely the amount of bias in the biased gradient. As noted in Appendix A, a central challenge is to relate the Jacobians $J_\theta v_t^\theta$ and $J_\theta f_t^\theta$ in (13) and (14) for a given $\theta$. This is difficult because the derivatives are with respect to the neural network parameters, and hence the analysis must involve the specific network architecture.  Moreover, the dependency of $f_t^\theta$ on $\theta$ is via integrating along a velocity field parameterized by $\theta$ (if using NODE), which can be convoluted and hard to compute/analyze.

We have attempted a few alternative interpretations of the biased-gradient descent but they don’t offer any new insight. Reviewer AfQf suggested “minimization-minimization”, but we cannot find any reference to that; we suppose it is similar to coordinate descent. However, our iterates do not fit the framework of coordinate descent. To illustrate this difficulty, suppose we want to minimize $F(\theta, \theta’)$ defined in (9) using coordinate descent over two (block) coordinates $(\theta, \theta’)$. Then we cannot avoid computing the gradient of $F$ with respect to $\theta’$, which is very costly as argued in the previous paragraph. 

**(b) The computational complexity of all methods with respect to the dimension.**

As with evaluating any optimization method, comparing computational cost is only meaningful when we control different methods to attain the same accuracy since increasing the computational cost (e.g. more iterations) would increase the accuracy. Since most presented methods in this work involve complicated training schemes and neural network parameterizations, it is hard to give a precise formula for the time complexity in terms of accuracy.

We can, however, comment on the time complexity of training for $S$ iterations for methods considered, in terms of dimension $d$. Assuming for simplicity that the same batch size is used, the training is done for $T$ time steps, and any network evaluation takes O(d) time,

- SCVM-NODE takes $O(ST N_{ODE}  d)$ time, where $N_{ODE}$ is the number of ODE integration steps.
- SCVM-TIPF (with RealNVP) takes $O(ST d^2)$ time since there will be $d$ coupling layers.
- ADJ takes $O(ST N_{ODE} d^3)$ time, where $d^3$ comes from the third-order spatial derivatives.
- JKO-ICNN takes $O(ST^2 d^3)$ time where $d^3$ is due to the log-determinant term.
- JKO-ICNN-PD takes $O(ST^2 d)$ time.

For JKO methods, the quadratic dependency on $T$ is due to the fact that, to draw samples from the current step, initial samples from the reference measure must be passed through all previous time steps. We also highlight that the inference time for JKO methods is significantly higher because a convex optimization needs to be solved to invert the ICNN flow for a total of $T$ times.

We will add a table about the time complexity of training for each method in the revised paper.


(c) A more in-depth comparison between SCVM-NODE and SCVM-TIPF.

Below is a summary of the advantages and disadvantages of TIPF and NODE. We will include this comparison in a new section in the appendix for the revised paper.

TIPS:
* Advantages:
  - TIPF works well in low dimensions (e.g. OU process in Sec 4.2, Porous medium equation in Sec 4.3). 
  - It has exact and cheap density access, whereas NODE requires numerical integration to obtain the density. 
* Disadvantages:
  - It does not scale up to high dimensions due to requiring many coupling layers, due to using RealNVP architecture.
  - The invertibility requirement of TIPF can limit its expressivity.
  - We need a regularizer to enforce the initial condition.

It is an interesting future direction to find alternatives to RealNVP that work better in higher dimensions.

NODE:
* Advantages:
  - NODE works well in high dimensions.
  - The requirement for the network parameterization is minimal (just needs to be a continuous vector field).
  - The Initial condition is automatically satisfied.
* Disadvantages:
  - Sampling and density queries require numerical integration, which can be slow and can incur numerical errors.

Our goal is to present two very different ways of parameterizing the probability flow to demonstrate the flexibility of the proposed generic framework (8-9). Finding new effective ways to parameterize the probability flow is an active research area, so we expect there will be better parameterizations in the future that can be plugged into our framework.


Meta-review: The authors present a new approach to solving mass-conserving partial differential equations (PDEs) designing an iterative optimization scheme which avoids the need for spatial or temporal discretization. The main idea revolves around considering self-consistency constraints satisfied by the time-dependent velocity field associated with mass-conserving PDEs. This method then consists in learning an approximate velocity field based on these constraints, rather than directly solving the PDE in terms of a probability density. 
The proposed approach is tested through a range of numerical experiments and is compared against analytic solutions whenever feasible. 

The reviewers and I concur that the paper is an interesting contributions and should be accepted. 

Please take into consideration reviewers' comments before submitting the final version of your work.
Decision: Accept (poster)