Title: Self-Consistent Velocity Matching of Probability Flows
Venue: NeurIPS 2023
Accepted: yes
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

Abstract

We present a discretization-free scalable framework for solving a large class of mass-conserving partial differential equations (PDEs), including the timedependent 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.

1 Introduction

Mass conservation is a ubiquitous phenomenon in dynamical systems arising from fluid dynamics, electromagnetism, thermodynamics, and stochastic processes. Mathematically, mass conservation is formulated as the continuity equation :

where p t : R d → R is a scalar quantity such that the total mass ∫ p t ( x ) is conserved with respect to t , v t : R d → R d is a velocity field, and T > 0 is total time. We will assume, for all t ∈ [0 , T ] , p t ≥ 0 and ∫ p t ( x ) d x = 1 , i.e., p t is a probability density function. We use µ t to denote the probability measure with density p t . Once a pair ( p t , v t ) satisfies (1), the density p t is coupled with v t in the sense that the evolution of p t in time is characterized by v t (Section 3.1).

We consider the subclass of mass-conserving PDEs that can be written succinctly as

where f t ( · ; µ t ) : R d → R d is a given function depending on µ t , with initial condition µ 0 = µ ∗ 0 for a given initial probability measure µ ∗ 0 with density p ∗ 0 .

Different choices of f t lead to a large class of mass-conserving PDEs. For instance, given a functional F : P 2 ( R d ) → R on the space of probability distributions with finite second moments, if we take where ∇ W 2 F ( µ ) : R d → R d 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 to optimize in the probability measure space. If we take

where b t is a velocity field and D t ( x ) is a positive-semidefinite matrix, then we obtain the timedependent Fokker-Planck equation [Risken and Risken, 1996], which describes the time evolution of the probability flow undergoing drift b t and diffusion with coefficient D t .

A popular strategy to solve (2) is to use an Eulerian representation of the density field p t on a 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-Planck equation and more generally 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 mass-conserving PDEs of the general form (2). Equipped with this formulation, we develop an iterative optimization scheme called self-consistent velocity matching . With the probability flow parameterized as a neural network, at each iteration, we refine the velocity field v t of the current flow to match an estimate of f t 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 to the form of f t , 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-matching objectives for certain PDEs to get rid of computationally expensive quantities such as ∇ log p t 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 a time-varying pushforward map [Biloš et al., 2021] and as a time-varying velocity field [Chen et al., 2018]-both have merits in different scenarios.

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 self-consistency loss efficiently (second column of Figure 2). For PDEs with analytically-known 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-Eijnden, 2023]. Our method faithfully recovers true 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.

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 a 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 high-dimensional problems. Hence, the rest of the section focuses solely on recent neural network-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 is 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 a 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 low-dimensional structure.

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 Wasserstein gradient flow (3). By pushing a reference measure through a chain of neural networks parameterized 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-Leibler (KL) divergence plus a Wasserstein-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 is refined on the fly based on the current positions of the particles. Boffi and Vanden-Eijnden [2023] propose a similar method that applies to a more general class of time-dependent Fokker-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 a velocity field that generates a probability path connecting a 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 FokkerPlanck equation, later extended to McKean-Vlasov type PDEs [Shen and Wang, 2023]. They observe that the velocity field of the flow solution to the Fokker-Planck equation must satisfy a fixed-point equation. They theoretically show that, under certain regularity conditions, a 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., 2018] and the adjoint method. Our work extends the concept of self-consistency to a wider class of PDEs in the form of (2) and circumvents the computationally demanding adjoint method using an 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).

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 ( p t , v t ) t ∈ [0 ,T ] (provided p t is continuous with respect to t and v t is bounded) corresponds to a unique flow map { Φ t ( · ) : R d → R d } t ∈ [0 ,T ] that solves the ordinary differential equations (ODEs) [Ambrosio et al., 2005, Proposition 8.1.8]

and the flow map satisfies µ t = (Φ t ) # µ 0 for all t ∈ [0 , T ] , where (Φ t ) # µ 0 to denote the pushforward measure of µ 0 by Φ t . Moreover, the converse is true: any solution (Φ t , v t ) of (5) with Lipschitz continuous and bounded v t is a solution of (1) with µ t = (Φ t ) # µ 0 [Ambrosio et al., 2005, Lemma 8.1.6]. Thus the Eulerian 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 it automatically satisfies the continuity equation (1).

3.2 Parametrizing probability flows

Our algorithm will be agnostic to the exact parameterization used to represent the probability flow. As such, we need a way to parameterize the flow to access the following quantities for all t ∈ [0 , T ] :

v t : R d → R d , the velocity field of the flow at time t .

Φ t : R d → R d , the flow map, so Φ t ( x ) is the location of a particle at time t if it is at x at time 0 .

µ t ∈ P ( R d ) , the probability measure at time t with sample access and density p t .

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 are finite similar to the uniqueness arguments given in [Chen et al., 2018]. We can now parameterize the flow modeling either the flow map Φ t or the velocity field v t as a neural network.

Time-dependent Invertible Push Forward (TIPF). We first parameterize a probability flow by modeling Φ t : R d → R d as an invertible network for every t . The network architecture is chosen so that Φ t 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 Φ 0 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 E X ∼ µ ∗ 0 ∥ Φ 0 ( X ) -X ∥ 2 . For brevity, we will from now on omit in the text this additional loss term. The velocity field can be recovered via v t ( x ) = ∂ t Φ t (Φ -1 t ( x )) . To recover the density p t of µ t = (Φ t ) # µ 0 , we use the change-of-variable formula log p t ( x ) = log p ∗ 0 (Φ -1 t ( x )) + log det ∣ ∣ J Φ -1 t ( x ) ∣ ∣ (see (1) in Dinh et al. [2016]).

Neural ODE (NODE). We also parameterize a flow by modeling v t : R d → R d as a 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: Φ t ( x ) = x + ∫ t 0 v s (Φ s ( x )) d s and log p t (Φ t ( x )) = log p ∗ 0 ( x ) -∫ t 0 ∇· v s (Φ s ( x )) d s , 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 -v t . With NODE, the initial condition µ 0 = µ ∗ 0 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 to 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 f t ( · ; µ t ) : R d → R d plus the initial condition µ ∗ 0 . If a probability flow µ t with flow map Φ t and velocity field v t satisfies the following self-consistency condition,

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 p t that might be hard to access, we can solve (6) which is a more tractable condition. Shen et al. [2022], Shen and Wang [2023] develop this concept for the Fokker-Planck equation and McKean-Vlasov type PDEs; here we generalize it to a wider class of PDEs of the form (2).

Let θ be the network weights that parameterize the probability flow using TIPF or NODE. The flow's measure, velocity field, and flow map at time t are denoted as µ θ t , v θ t , Φ θ t respectively. One option to solve (6) would be to minimize the self-consistency loss

This formulation is reminiscent of PINNs [Raissi et al., 2019] where a residual of the original PDE is minimized. Direct optimization of (7) is challenging: while the integration over [0 , T ] and µ θ t can be approximated using Monte Carlo, to apply stochastic gradient descent, we must differentiate through µ θ t and f t : 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 a closed-form formula for the time evolution of the quantities within f t , 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 θ k denote the network weights at iteration k . We define iterates

where

Effectively, in each iteration, we minimize (9) by one gradient step where we match the velocity field v θ t to what it should be according to f t based on the network weights θ k from the previous iteration. This scheme can be interpreted as a gradient descent on 7 using the biased gradient estimate ∇ θ F ( θ, θ k ) -see Appendix A for a discussion. We call this iterative algorithm self-consistent velocity matching .

If f t depends on the density of µ t only through the score ∇ log p t (corresponding to a diffusion term in the PDE), then we can apply an integration-by-parts trick [Hyvärinen and Dayan, 2005] to get rid of this density dependency by adding a divergence term of the velocity field. Suppose f t 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 ∈ [0 , T ] , for f t defined in (4) , assume v θ t , D t are bounded and continuously differentiable, and µ θ ′ t is a measure with a continuously differentiable density p θ ′ t that vanishes in infinity and not at finite points. Then we have

We provide the derivation in Appendix B. With Proposition 3.1, we no longer need access to ∇ log p t 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 1 in the appendix. We use Adam optimizer [Kingma and Ba, 2014] to modulate the update (8). For sampling time steps t 1 , . . . , t L in [0 , T ] , we use stratified sampling where t l is uniformly sampled from [ ( l -1) T / L , lT / L ] ; such a 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.

4 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 SCVM-NODE to denote our method with TIPF and NODE parameterization respectively. We use JKO-ICNN to denote the method by Mokrov et 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-Planck equation, and DFE ('discrete forward Euler') to denote the method by Boffi and Vanden-Eijnden [2023]. We implemented all competing methods in JAX-see more details in Appendix C-and we compare quantitatively against these methods when possible.

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], we use the symmetric KullbackLeibler (symmetric KL) divergence, defined as SymKL ( ρ 1 , ρ 2 ) := KL ( ρ 1 ∥ ρ 2 ) + KL ( ρ 2 ∥ ρ 1 ) , where KL ( ρ 1 ∥ ρ 2 ) := E X ∼ ρ 1 [log d ρ 1 ( X ) / d ρ 2 ( X ) ] . Sample estimates of KL can be negative which complicates log-scale plotting, so when this happens, we consider an alternative f -divergence D f ( ρ 1 ∥ ρ 2 ) := E X ∼ ρ 2 [ (log ρ 1 ( X ) -log ρ 2 ( X )) 2 / 2 ] whose sample estimates are always non-negative. We similarly define the symmetric f -divergence Sym D f ( ρ 1 , ρ 2 ) := D f ( ρ 1 ∥ ρ 2 ) + D f ( ρ 2 ∥ ρ 1 ) . 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 ( µ ) = KL ( µ ∥ µ ∗ ) where we have density access to the target measure µ ∗ . To fit into our framework, we set f t ( x ; µ t ) = ∇ log p ∗ ( x ) -∇ log p t ( x ) which matches (4) with b t ( x ) = ∇ log p ∗ ( x ) and D t ( x ) = I d . Following the experimental setup in Mokrov et al. [2021] and Fan et al. [2021], we take µ ∗ to be a mixture of 10 Gaussians with identity covariance and means sampled uniformed in [ -5 , 5] d . The initial measure is µ ∗ 0 = N (0 , 16 I d ) . We solve the corresponding Fokker-Planck PDE for a total time of T = 5 and for d = 10 , . . . , 60 . As TIPF parameterization does not scale to high dimensions, we only consider SCVM-NODE in this experiment.

Figure 1 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.

t =0.04

t=0.32

t = 1.08

t =2.56

t =5.00

Target measure

20

20

20

20

20

20

10

10

10

10

10

10

0

0

0

0

0

0

10

10

10

10

10

10

20

20

0

20

20

20

0

20

20

20

0

20

20

20

0

20

20

20

0

20

20

20

0

20

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 is 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 even if the consistency loss is low-in 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 in dimension 10 , and scaling up to higher dimensions becomes prohibitive.

The rightmost column of Figure 2 shows SCVM-NODE achieves far lower symmetric KL compared to 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 t (Figure 1).

0

2500

5000

7500

10000

Training iteration (

d

= 5

)

10

-

1

10

0

Symmetric KL(

t

=

T

)

SCVM-NODE

ADJ

0

2500

5000

7500

10000

Training iteration (

d

= 5

)

10

0

Consistency

SCVM-NODE

ADJ

0

2500

5000

7500

10000

Training iteration (

d

= 10

)

10

0

10

1

Symmetric KL(

t

=

T

)

SCVM-NODE

ADJ

0

2500

5000

7500

10000

Training iteration (

d

= 10

)

10

0

10

1

Consistency

SCVM-NODE

ADJ

0

2

4

Time

t

(

d

= 30

)

10

0

10

1

10

2

Symmetric KL

SCVM-NODE

JKO-ICNN-PD

JKO-ICNN

0

2

4

Time

t

(

d

= 60

)

10

0

10

1

10

2

Symmetric KL

SCVM-NODE

JKO-ICNN-PD

JKO-ICNN

4.2 Ornstein-Uhlenbeck process

We consider the Wasserstein gradient flow of the KL 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 al. [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 = 5 and d = 10 . However, this is at 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, we 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 in lower probability divergences than SCVMNODE 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 t = 0 . 25 in the left two columns in Figure 3).

For the two JKO methods, they result in much higher errors for t ≤ 0 . 5 compared to later times: this is expected because the dependency of G ( t ) on t is exponential, so convergence to µ ∗ is 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 ∆ t . We choose step size ∆ t = 0 . 01 which empirically gives the best results among { 0 . 1 , 0 . 01 , 0 . 001 , 0 . 0001 } . As DFE achieves far lower symmetric KL divergence or f -divergence compared to alternatives, we only include its Bures-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 et al. [2021], we consider the porous medium equation with only diffusion: ∂ t p t = ∆ p t m with m> 1 . Its solution is the Wasserstein gradient flow of F ( µ ) = ∫ 1 m -1 p ( x ) m d x where p is the density of µ with ∇ W 2 F ( µ )( x ) = ∇ ( m m -1 p m -1 ( x )) -see Appendix D.2 for details. We consider only SCVM-TIPF and JKO methods here.

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

2

.

00

Time

t

(

d

= 5

)

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

Symmetric KL

SCVM-NODE

SCVM-TIPF

ADJ

JKO-ICNN-PD

JKO-ICNN

DFE

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

2

.

00

Time

t

(

d

= 10

)

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

10

1

Symmetric KL

SCVM-NODE

SCVM-TIPF

ADJ

JKO-ICNN-PD

JKO-ICNN

DFE

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

2

.

00

Time

t

(

d

= 5

)

10

-

7

10

-

5

10

-

3

10

-

1

Symmetric

f

-divergence

SCVM-NODE

SCVM-TIPF

ADJ

JKO-ICNN-PD

JKO-ICNN

DFE

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

2

.

00

Time

t

(

d

= 10

)

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

10

1

Symmetric

f

-divergence

SCVM-NODE

SCVM-TIPF

ADJ

JKO-ICNN-PD

JKO-ICNN

DFE

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

2

.

00

Time

t

(

d

= 5

)

10

-

2

6

×

10

-

3

2

×

10

-

2

3

×

10

-

2

4

×

10

-

2

Bures-Wasserstein distance

SCVM-NODE

SCVM-TIPF

ADJ

JKO-ICNN-PD

JKO-ICNN

DFE

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

2

.

00

Time

t

(

d

= 10

)

3

×

10

-

2

4

×

10

-

2

6

×

10

-

2

Bures-Wasserstein distance

SCVM-NODE

SCVM-TIPF

ADJ

JKO-ICNN-PD

JKO-ICNN

DFE

10

20

30

40

50

60

Dimension

d

10

-

3

10

-

2

Symmetric KL

SCVM-NODE

JKO-ICNN-PD

JKO-ICNN

10

20

30

40

50

60

Dimension

d

10

-

3

10

-

2

10

-

1

Symmetric

f

-divergence

SCVM-NODE

JKO-ICNN-PD

JKO-ICNN

10

20

30

40

50

60

Dimension

d

10

-

1

10

0

Bures-Wasserstein distance

SCVM-NODE

JKO-ICNN-PD

JKO-ICNN

DFE

We show the efficiency of SCVM-TIPF compared to JKO-ICNN in dimension d = 1 , 2 , . . . , 6 . We exclude JKO-ICNN-PD because it produces significantly worse results. We visualize the density p t of the solution from SCVM-TIPF and JKO-ICNN on the top of Figure 4 in dimension 1 compared to p ∗ t . Both methods approximate p ∗ t well with SCVM-TIPF being more precise at small t ; this is consistent with the observation in Figure 2 where JKO methods result in bigger errors for small t .

-

1

0

1

t

= 0

.

000

0

1

2

3

4

5

pt

(

x

)

JKO-ICNN

SCVM-TIPF

-

1

0

1

t

= 0

.

004

0

1

2

3

4

5

pt

(

x

)

-

1

0

1

t

= 0

.

006

0

1

2

3

4

5

pt

(

x

)

-

1

0

1

t

= 0

.

015

0

1

2

3

4

5

pt

(

x

)

-

1

0

1

t

= 0

.

025

0

1

2

3

4

5

pt

(

x

)

JKO-ICNN

SCVM-TIPF

2

3

4

5

6

Dimension

d

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

TV distance

JKO-ICNN

t

= 0

.

004

JKO-ICNN

t

= 0

.

025

SCVM-TIPF

t

= 0

.

004

SCVM-TIPF

t

= 0

.

025

2

3

4

5

6

Dimension

d

10

-

6

10

-

5

10

-

4

10

-

3

10

-

2

10

-

1

10

0

Symmetric

f

-divergence

JKO-ICNN

t

= 0

.

004

JKO-ICNN

t

= 0

.

025

SCVM-TIPF

t

= 0

.

004

SCVM-TIPF

t

= 0

.

025

2

3

4

5

6

Dimension

d

10

-

3

10

-

2

Wassserstein-2 distance

JKO-ICNN

t

= 0

.

004

JKO-ICNN

t

= 0

.

025

SCVM-TIPF

t

= 0

.

004

SCVM-TIPF

t

= 0

.

025

10

10

-

10

-

0

1

2

2

SCVM-TIPF

SCVM-NODE

DFE

ADJ

SDE-EM

4

6

Time

t

(

d

= 3

)

On the bottom row of Figure 4, we plot the total variation (TV) distance, the symmetric f -divergence, and the Wasserstein-2 distance (details on the TV distance are given in Appendix C.6) between the recovered solution p t and p ∗ t for both methods at t = 0 . 004 and t = 0 . 025 . Note that the values of all metrics are very low implying that the solution from either method is 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 6 , training JKO-ICNN took 102 minutes compared to 21 minutes for SCVM-TIPF.

4.4 Time-Dependent Fokker-Planck equation

In this section, we qualitatively evaluate our method for solving a PDE that is not a Wasserstein gradient flow. In this case, JKO-based methods cannot be applied. Consider the OU process from Section 4.2 when the mean β and the covariance matrix Γ become time-dependent as β t and Γ t . The resulting PDE is a time-dependent Fokker-Planck equation of the form (4) with

In this configuration, when the initial measure p 0 is Gaussian, the solution µ t can again be shown to be Gaussian with mean and covariance following an ODE-see Appendix D.3 for more details. We consider, in dimension 2 and 3 , time-dependent attraction towards a harmonic mean β t = a (sin( πωt ) , cos( πωt )) using the expression of β t from Boffi and Vanden-Eijnden [2023], augmented to β t = a (sin( πωt ) , cos( πωt ) , t ) in dimension 3 .

We apply both SCVM-TIPF and SCVM-NODE to this problem and compare our 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 a 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 a non-Fokker-Planck PDE that can be readily solved by our method.

2

4

6

8

10

Time

t

(

d

= 2

)

10

-

3

10

-

2

10

-

1

Symmetric

f

-divergence

SCVM-TIPF

SCVM-NODE

DFE

ADJ

SDE-EM

2

4

6

8

10

Time

t

(

d

= 3

)

10

-

2

10

-

1

10

0

Symmetric

f

-divergence

SCVM-TIPF

SCVM-NODE

DFE

ADJ

SDE-EM

2

4

6

8

10

Time

t

(

d

= 2

)

10

-

2

2

×

10

-

2

3

×

10

-

2

4

×

10

-

2

Wassserstein-2 distance

SCVM-TIPF

SCVM-NODE

DFE

ADJ

SDE-EM

2

4

6

8

10

Time

t

(

d

= 3

)

10

-

1

4

×

10

-

2

6

×

10

-

2

Wassserstein-2 distance

SCVM-TIPF

SCVM-NODE

DFE

ADJ

SDE-EM

8

10

4.5 Additional qualitative low-dimensional dynamics

To demonstrate the flexibility of our method, we apply our algorithm to model more general massconserving dynamics than the ones considered in the previous sections.

Flow splashing against obstacles. We model the phenomenon of a 2-dimensional flow splashing against obstacles using a Fokker-Planck equation (4) where b t encodes the configuration of three

4

3

2

×

×

×

10

10

10

-

-

-

10

-

2

2

2

2

Wassserstein-2 distance

2

4

Time

t

(

d

6

= 2

)

SCVM-TI

SCVM-N

DFE

ADJ

SDE-EM

8

obstacles that repel the flow (See Appendix D.5 for details). We solve this PDE using SCVM-NODE for T = 5 and visualize the recovered flow in (6). When solving the same PDE using SDE-EM, the flow incorrectly crosses the bottom right obstacle due to a finite time step size (Figure 14). When using DFE, 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 be trained without infinite loss in this example.

t =0.05

t = 0.2

t =0.5

t = 1.0

t = 2.0

t = 4.0

Smooth interpolation of measures. To illustrate the flexibility of our method, we demonstrate two ways to formulate the problem of smoothly interpolating a list of measures. First, we model the interpolation as a time-dependent Fokker-Planck equation and use it to interpolate MNIST digits 1, 2, and 3, starting from a Gaussian (Figure 7). Next, we adopt an optimal transport formulation and use it to generate an animation sequence deforming a 3D hand model to a different pose and then to a 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.

t = 0.37

t = 0.40

t = 0.41

t = 0.44

t = 2.10

t = 2.40

t = 2.71

t = 4.21

t = 4.45

t = 5.41

00°0 =

t =0.60

t = 1.20

t = 1.80

t = 2.40

t = 3.00

t = 3.60

t = 4.20

t = 4.80

t = 5.40

5 Conclusion

By extending the concept of self-consistency from Shen et al. [2022], we present an iterative optimization method for solving a wide class of mass-conserving PDEs without temporal or spatial discretization. In all experiments considered, our method achieves strong quantitative results with significantly less training time than JKO-based methods and the adjoint method in high dimensions.

Below we highlight a few future directions. First, as discussed, the two ways to parameterize a probability flow, TIPF, and NODE, both have their specific limitations. Finding a new parameterization that combines the advantages of both TIPF and NODE is an important next step. Secondly, we hope to extend our approach to incorporate more complicated boundary conditions. Finally, given that the proposed algorithm is highly effective empirically, it 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 W911NF2010168 and W911NF2110293, of Air Force Office of Scientific Research award FA9550-19-1-031, of National Science Foundation grant CHS-1955697, 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.

SH acknowledges the financial support from the University of Bordeaux (UBGR grant) and the French Research Agency (PostProdLEAP).

References

Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. arXiv preprint arXiv:2209.15571 , 2022.

David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. 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 J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning , pages 146-155. PMLR, 2017.

Marin Biloš, Johanna Sommer, Syama Sundar Rangapuram, Tim Januschowski, and Stephan Günnemann. Neural flows: Efficient alternative to neural odes. Advances in Neural Information Processing Systems , 34:21325-21337, 2021.

Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. arXiv preprint arXiv:2105.15183 , 2021.

Nicholas M Boffi and Eric Vanden-Eijnden. Probability flow solution of the fokker-planck equation. Machine Learning: Science and Technology , 4(3):035012, jul 2023. doi: 10.1088/2632-2153/ ace2aa. URL https://dx.doi.org/10.1088/2632-2153/ace2aa .

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://github.com/google/jax .

Martin Burger, José A Carrillo, and Marie-Therese 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 K 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-stratonovitch equation. Probability Theory and Related Fields , 115(4):549-578, 1999.

Marco Cuturi, Laetitia Meng-Papaxanthos, 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 arXiv:1605.08803 , 2016.

Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted 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.02424 , 2021.

Jean Feydy, Thibault Séjourné, François-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 divergences. 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.com/google/flax .

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-planck 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:1412.6980 , 2014.

Alexey Kroshnin, Vladimir Spokoiny, and Alexandra Suvorikova. Statistical inference for bureswasserstein barycenters. The Annals of Applied Probability , 31(3):1264-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 to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 , 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:15243-15256, 2021.

Gabriel Peyré. Entropic approximation of wasserstein gradient flows. SIAM Journal on Imaging Sciences , 8(4):2323-2351, 2015.

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-vlasov type pdes. arXiv preprint arXiv:2303.11205 , 2023.

Zebang Shen, Zhenfu Wang, Satyen Kale, Alejandro Ribeiro, Aim Karbasi, and Hamed Hassani. Self-consistency of the fokker-planck equation. arXiv preprint arXiv:2206.00860 , 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 , 44(1):133-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.

Appendix

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.

Algorithm 1 Self-consistent velocity matching

end for

Output:

optimized θ .

A Biased Gradient Interpretation

Assume µ θ t has density p θ t . Using f θ t ( x ) to denote f t ( x ; µ θ t ) from (7), the self-consistency loss can be written as,

Assuming all terms involving θ are differentiable with respect to θ , the gradient of L ( θ ) with respect to the neural network parameters θ can be written as:

Here we use J θ to denote the Jacobian with respect to θ . On the other hand, the gradient used in the updates (8) is ∇ θ F ( θ, θ ′ ) at θ ′ = θ :

We see that (15) is exactly the middle term (13). Hence our formulation can be interpreted as doing gradient descent with a biased gradient estimator. It remains a future work direction to

theoretically analyze the amount of bias in (15) and the condition under which the dot product ⟨∇ L ( θ ) , ∇ θ F ( θ, θ ′ ) | θ ′ = θ ⟩ ≥ 0 . The central challenge would be to relate J θ v θ t and J θ f θ t ; this depends on the neural network architecture and the type of the PDE.

B Integration-by-Parts Trick

This is a common trick used in score-matching literature [Hyvärinen and Dayan, 2005].

Proof of Proposition 3.1. Fix t > 0 . The form of f in (4) is

where we use integration-by-parts to get the second last equation and the assumption that v θ t , D t are bounded and p θ t ( x ) → 0 as ∥ x ∥ → ∞ .

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 a low-rank linear skip connection x ↦→ A ( t ) x + b ( t ) where A ( t ) = L ( t ) L ⊤ ( t ) and L ( t ) is a d × 20 matrix to 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 a sinusoidal embedding for the time input t plus two fully connected layers of 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-parts 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 N train = 10 5 or 2 × 10 5 , B = 1000 , L = 10 or 20 . We use Adam [Kingma and Ba, 2014] with

a cosine decay learning rate scheduler, with initial learning rate 10 -3 , the number of decay steps same as N train, and α = 0 . 01 (so the final learning rate is 10 -5 ). Since we are effectively performing gradient descent using a biased gradient, we set b 2 = 0 . 9 in Adam (instead of the default b 2 = 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://github.com/facebookresearch/w2ot . Compared to the original ICNN implementation by Amos et al. [2017], we add an 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 is set to 20. The activation layer is taken to be CELU.

To implement the method by Fan et al. [2021], we model the dual potential as 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), we 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 at 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 of 10 -3 to obtain finer inverses, but it takes 6x longer compared to 10 -2 ).

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 steps-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 a 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-Eijnden [2023]. We use 5000 particles. For score estimation, we use the same network architecture as 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 ∆ t . For the OU process experiment in dimension 60, when ∆ t = 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 ∆ t can result in bigger errors. We choose ∆ t = 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 x -axis in the log-scale plots in Figure 3.

To compute Bures-Wasserstein distance [Kroshnin et al., 2021], we first fit a 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 4 and Figure 8 to compare the estimated and ground-truth solutions. It is approximated by the L 1 distance between the densities calculated over 50000 samples uniformly distributed on the compact

D Additional Experimental Details

D.1 Ornstein-Uhlenbeck process

The OU process is the Wasserstein gradient flow of the KL divergence with respect to a Gaussian µ ∗ = N ( β, Γ -1 ) where β ∈ R d and Γ is a d × d positive-definite matrix. When the initial distribution is µ ∗ 0 = N (0 , I d ) , the gradient flow at time t is known to be a Gaussian distribution G ( t ) with mean ( I d -e -t Γ ) β and covariance Γ -1 ( I d -e -2 t Γ ) + e -2 t Γ . 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:

where t 0 > 0 is the starting time, α = m d ( m -1)+2 , β = ( m -1) α 2 dm , and C > 0 is a free constant. Similar to Fan et al. [2021], we choose m = 2 and total time T = 0 . 025 . The initial measure follows a Barenblatt distribution supported in [ -0 . 25 , 0 . 25] d ( C is chosen accordingly) with t 0 = 10 -3 . We use Metropolis-Hastings to sample from µ 0 .

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

Time

t

(

d

= 3

)

10

-

3

10

-

2

TV distance

JKO-ICNN

SCVM-TIPF

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

Time

t

(

d

= 3

)

10

-

4

10

-

3

10

-

2

10

-

1

10

0

Symmetric

f

-divergence

JKO-ICNN

SCVM-TIPF

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

Time

t

(

d

= 3

)

10

-

2

3

×

10

-

3

4

×

10

-

3

6

×

10

-

3

Wassserstein-2 distance

JKO-ICNN

SCVM-TIPF

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

Time

t

(

d

= 6

)

10

-

6

10

-

5

10

-

4

TV distance

JKO-ICNN

SCVM-TIPF

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

Time

t

(

d

= 6

)

10

-

5

10

-

3

10

-

1

10

1

Symmetric

f

-divergence

JKO-ICNN

SCVM-TIPF

0

.

005

0

.

010

0

.

015

0

.

020

0

.

025

Time

t

(

d

= 6

)

3

×

10

-

2

4

×

10

-

2

Wassserstein-2 distance

JKO-ICNN

SCVM-TIPF

D.3 Time-Dependant Fokker-Planck equation

We consider a time-dependent Fokker-Planck equation of the form (4) with the velocity field

When the initial measure p 0 is Gaussian, the solution µ t can again be shown to be Gaussian with mean m t and covariance Σ t solutions of the differential equations:

In practice, we experiment with constant Γ t = diag (1 , 3) and D t = σ 2 I d . We also experience in dimension 3 by considering and Γ t = diag (1 , 3 , 1) . We set a = 3 , ω = 1 , σ = √ 0 . 25 and pick as initial distribution p 0 a Gaussian with mean b 0 and covariance σ 2 I d . We set the total time to T = 10 .

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 SCVMNODE gives visually indistinguishable trajectories. We also plot in Figure 10 the evolution of particles simulated by Euler-Maruyama (EM-SDE) discretization of the Fokker-Planck equation. Corresponding animated GIFs be found at this link.

t

= 0

.

0

t

= 0

.

5

t

= 1

.

0

t

= 1

.

5

t

= 2

.

0

t

= 2

.

5

t

= 3

.

0

t

= 3

.

5

t

= 4

.

0

t

= 4

.

5

-1

-1

-1

-2

2

Z-

2

Z-

3

3

3

3

E-

4

0

D

-1

-1

-1

-1

-2

2

2

-2

-3

3

3

Z-

-2

Z-

-2

0

-2

0

2

t

= 0

.

0

t

= 0

.

5

t

= 1

.

0

t

= 1

.

5

t

= 2

.

0

t

= 2

.

5

t

= 3

.

0

t

= 3

.

5

t

= 4

.

0

t

= 4

.

5

T-

Z-

Z-

-3

-3

-3

-3

I-

1

2

-2

-2

-2

-3

E-

-3

-2

Z-

D.4 Flock of birds

We model the dynamics of a flock of birds by augmenting the time-dependent Fokker-Planck equation (11) with an interaction term:

This is similar to the harmonically interacting particles experiment in Boffi and Vanden-Eijnden [2023], but we use a population expectation E [ µ t ] instead of an empirical one in modeling the repulsion from the mean. Since f t needs to access E [ µ t ] , 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 [ µ t ] using Monte Carlo samples from µ t . We use a similar setup as in Section 4.4, except we now use an 'infinity sign" attraction β t = a (cos(2 πωt ) , 0 . 5 sin(2 πωt )) along with a sinusoidal α t = 2sin( πwt ) . Depending on the sign of α t , particles are periodically attracted towards or repulsed from their mean. Both SCVM-TIPF and SCVM-NODE produce similar visual results as shown in Figure 11 and Figure 12.

We use a constant Γ t = I d and a constant diffusion matrix D = σ 2 I d . We set a = 3 , ω = 0 . 5 , and σ = √ 0 . 25 . We pick as initial distribution p 0 a Gaussian with mean (0 , 0) and covariance σ 2 I d . We set the total time to T = 10 .

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 at this link.

t

= 0

.

0

t

= 0

.

5

t

= 1

.

0

t

= 1

.

5

t

= 2

.

0

t

= 2

.

5

t

= 3

.

0

t

= 3

.

5

t

= 4

.

0

t

= 4

.

5

1

2

-2

Z-

7-

t

= 0

.

0

t

= 0

.

5

t

= 1

.

0

t

= 1

.

5

t

= 2

.

0

t

= 2

.

5

t

= 3

.

0

t

= 3

.

5

t

= 4

.

0

t

= 4

.

5

0

-1

-3

3

3

7-

-2

Z-

D

Z-

D.5 Flow splashing against obstacles

We use the following formulation for modeling the flow. Each obstacle is modeled as a line segment. The endpoints of the three obstacles are:

We model the dynamics as a Fokker-Planck equation where f t of the form (4) is defined as

where q sink = (4 , 0) , and π O i ( x ) is the projection of x onto obstacle i represented as a line segment, and p N (0 , 0 . 04) is the density of an 1-dimensional Gaussian with variance 0 . 04 .

The initial distribution is chosen to be N (0 , 0 . 25 I 2 ) . We train SCVM-NODE for 10 4 with an 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.

t =0.05

t = 0.2

t = 0.5

t = 1.0

t = 2.0

t = 4.0

D.6 Smooth interpolation of measures

Suppose we are to smoothly interpolate M measures ν 1 , . . . , ν M with densities q 1 , . . . , q M , and we want the flow to approximate ν i at time r i . To achieve this goal, we present two formulations using different choices of f t . We use SCVM-NODE in this section.

Measure interpolation using time-dependent Fokker-Planck equations. We model the dynamics as a Fokker-Planck equation where f t of the form (4) is taken to be

where ϕ ( t ) is defined as the continuous bump function

for bandwidth h = 1 . 0 .

Below we provide details for the MNIST interpolation result in the top row of Figure 7. We use the first three images of 1 , 2 , 3 from the MNIST dataset. To construct ν i from a digit image, we use a mixture of Gaussians where we put one equally-weighted Gaussian with covariance 0 . 02 2 I 2 on the pixels with values greater than 0.5 (images are first normalized to have values in [0 , 1] ). The initial distribution is N ((0 . 5 , 0 . 5) , 0 . 04 I 2 ) . To train SCVM-NODE, we use an initial learning rate of 10 -4 with cosine decay for a total of 10 5 iterations. This takes roughly 1 hour to train.

Measure interpolation using optimal transport. An alternative way to interpolate measures using our framework is to use optimal transport to define f t . Recall µ t denotes the probability flow at time t . We then define

where W 2 2 is the squared Wasserstein-2 distance and ∇ W 2 W 2 2 is its Wasserstein gradient. In practice, we compute ∇ W 2 W 2 2 using sample access and we employ the debiased Sinkhorn divergence [Genevay et al., 2018, Feydy et al., 2019] implemented in the JAX OTT library [Cuturi et al., 2022]. This formulation differs from the one in Zhang et al. [2022] in that here we prescribe the precise PDE based on f t , whereas in Zhang et al. [2022] an optimal transport loss is used to fit the keyframes along with many regularizers on the velocity field v t to promote the smoothness and other desirable properties. In contrast, we do not use any regularizer on v t .

To train SCVM-NODE to produce the hand-hand-ball animation sequence in the bottom row of Figure 7, we first sample 20000 points from the interior of the three meshes (a hand mesh, a hand mesh in a different pose, and a ball mesh), and we set ν i to be the empirical measure of the corresponding point cloud. Note that different from the first formulation using Fokker-Planck equations, in the optimal transport formulation, throughout we only require sample access from each ν i . We use an initial learning rate of 10 -4 with cosine decay for a total of 5 × 10 4 iterations. This takes 4.5 hours, which is significantly longer than the training time in Zhang et al. [2022] (reported to be 15 minutes). We leave further improvement of our method to interpolate measures faster as future work.

To render the animation, we sample 20000 points and render the point cloud at 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 in Zhang et al. [2022] which makes sure the key measures v i 's are fit exactly in the resulting animation. We also did not use unbalanced optimal transport, which as reported in Zhang et al. [2022] can make the fingers of the hand more separated, but requires careful parameter tuning.

Peer reviews:
comments: 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.

STRENGTHS:
The class of PDEs being solved is important, including the celebrated Fokker-Planck and Wasserstein gradient flow. 

The computational complexity of this method is much lower than the previous ones, while preserving good empirical results.

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

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

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

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

STRENGTHS:
- The paper is well-written. The motivation and the idea is clear. 

- The description of the literature is comprehensive.

- Two sets of the proposed methods in very different characteristics, NODEs and TIPFs, are interesting.

- Empirical evidence is strong.

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.

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

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nSome 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?\n\nLIMITATIONS:\nNot concerned.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\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.\n\nWEAKNESSES:\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.\n\nQUESTIONS:\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\"?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nl. 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.\n\nLIMITATIONS:\nsee above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-23"}