Large-Scale Wasserstein Gradient Flows
Authors: Petr Mokrov
Alexander Korotin
Lingxiao Li
Aude Genevay
Justin Solomon
Evgeny Burnaev
Venue: NeurIPS 2021 Poster (2021)
Keywords: wasserstein-2 gradient flows
JKO stepping
optimal transport
input-convex neural networks
diffusion processes

Abstract: Wasserstein gradient flows provide a powerful means of understanding and solving many diffusion equations. Specifically, Fokker-Planck equations, which model the diffusion of probability measures, can be understood as gradient descent over entropy functionals in Wasserstein space. This equivalence, introduced by Jordan, Kinderlehrer and Otto, inspired the so-called JKO scheme to approximate these diffusion processes via an implicit discretization of the gradient flow in Wasserstein space. Solving the optimization problem associated with each JKO step, however, presents serious computational challenges. We introduce a scalable method to approximate Wasserstein gradient flows, targeted to machine learning applications. Our approach relies on input-convex neural networks (ICNNs) to discretize the JKO steps, which can be optimized by stochastic gradient descent. Contrarily to previous work, our method does not require domain discretization or particle simulation.  As a result, we can sample from the measure at each time step of the diffusion and compute its probability density. We demonstrate the performance of our algorithm by computing diffusions following the Fokker-Planck equation and apply it to unnormalized density sampling as well as nonlinear filtering. 

Large-Scale Wasserstein Gradient Flows
Petr Mokrov"
Alexander Korotin*
Skolkovo Institute of Science and Technology Moscow Institute of Physics and Technology
Skolkovo Institute Science and Technology
Artificial Intelligence Research Institute
Moscow, Russia
Moscow, Russia koootin@skoltech..i.
petr www.ooowwwwwoleeoooooo
Lingxiao Li Massachusetts Institute of Technology Cambridge, Massachuseuss, USA lingxiao0mit.cc edu
Aude Genevay
Massachusetts Institute of Technology
Cambridge, Massachusetls, USA aude. aeaeeegeneva@@gmal.ccmm
Justin Solomon
Evgeny Burnaev
Massachusetts Institute of Technology Cambridge, Massachusetts, USA jsolomon@mit. edu
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute Moscow, Russia burnaev@skoltech.c ru
Abstract
Wasserstein gradient flows provide powerful means of understanding and solving many diffusion equations. Specifically, Pokker-Planck, equations, which model the diffusion of probability measures, can be understood as gradient descent over entropy functionals in Wasserstein space. This equivalence, introduced by Jordan, Kinderlehrer and Otto, inspired the so-called JKO scheme to approximate these diffusion processes via an implicit discretization of the gradient flow in Wasserstein space. Solving the optimization problem associated each JKO step, however, presents serious computational challenges. We introduce a scalable method to approximate Wasserstein gradient flows, targeted machine learning applications. Our approach relies on input-convex neural networks (ICNNs) to discretize the JKO steps, which can be optimized by stochastic gradient descent. Unlike previous work, our method does not require domain discretization or particle simulation. As result, we can sample from the measure each time step of the diffusion and compute its probability density. We demonstrate our algorithm's is performance by computing diffusions following the Fokker-Planck equation and apply it unnormalized density sampling as well as nonlinear filtering.
Introduction
Stochastic differential equations (SDEs) are used model the evolution of random diffusion processes across time, with applications in physics (63), finance (22 [52), and population dynamics [35]. In machine learning, diffusion processes also arise in applications filtering [34] 21 and unnormalized
posterior sampling via discretization of the Langevin diffusion [70].
The time-evolving probability density Pt of these diffusion processes is governed by the Fokker- Planck equation. Jordan, Kinderlehrer, and Otto 32 showed that the Bokker-Plance equation is
Equal contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

equivalent to following the gradient flow of an entropy functional in Wasserstein space, i.e., the space of probability measures with finite second order moment endowed with the Wasserstein distance. This inspired a simple minimization scheme called JKO scheme, which consists an implicit Fuler discretization of the Wasserstein gradient flow. However, each step of the JKO scheme is costly as requires solving minimization problem involving the Wasserstein distance.
One way to compute the diffusion to use fixed discretization of the domain and apply standard 40 to get Pt. For example, (50) proposes a method to approximate the diffusion based on JKO stepping and antropy-regularing optimal transport. However, these methods are limited to small dimensions since the discretization of space grows exponentially,
numerical integration methods 18 49 15
An alternative to domain discretization is particle simulation. It involves drawing random samples (particles) from the initial distribution and simulating their evolution via standard methods such as Fuler-Maruyama scheme [36] 89.2]. After convergence, the particles approximately distributed according to the stationary distribution, but no density estimate readily available.
Another way to avoid discretization is to parameterize the density of Pr. Most methods approximate only the first and second moments Pt. e.g., via Gaussian approximation. Kalman filtering approaches
can then compute the dynamics [4 39 33 61]. More advanced Gaussian mixture approximations
[65 or more general parametric families have also been studied [64 669. In 488 variational methods are used to minimize the divergence between the predictive and the true density.
Recently, 24 introduced a parametric method to compute JKO steps via entropy-regularize optimal transport. The authors regularize the Wasserstein distance in the JKO step ensure strict convexity and solve the unconstrained dual problem via stochastic program on a finite linear subset of basis functions. The method yields unnormalized probability density without direct sample access.
Recent works propose scalable continuous optimal transport solvers, parametrizing the solutions by reproducing kernels [10, fully-connected neural networks [22]. or Input Convex Neural Networks (ICNNs)
In particular, ICNNs gained attention for Wasserstein-2 transport since their R. can represent OT maps for the quadratic cost. These continuous solvers scale better to high dimension without discretizing the input measures, but they are too computationally expensive to be applied directly to JKO steps.
37
44 38
gradients Vua R
Contributions. We propose scalable parametric method to approximate Wasserstein gradient flows via JKO stepping using input-convex neural networks (ICNNs) (6. Specifically, we leverage Brenier's theorem bypass the costly computation of the Wasserstein distance, and parametrize the optimal transport map as the gradient of an ICNN. Given sample access to the initial measure po. we use stochastic gradient descent (SGD) sequentially learn time-discretized JKO dynamics of pt. The trained model can sample from continuous approximation of Pt and compute its density de: (z). We compute gradient flows for the Fokker-Plance free energy functional FFP given by (5). but our method generalizes other cases. We demonstrate performance by computing diffusion following the Fokker-Planck equation and applying it to unnormalized density sampling as well as nonlinear filtering.
Notation, P2 PR.) denotes the set of Borel probability measures on IR with finite second moment. Pa denotes its subset of probability measures absolutely continuous with respect Lebesgue measure. For p E (T) we denote by de (x)) its density with respect to the Lebesgue measure. de
I(p, v) denotes the set of probability measures on R.,, x R.P with marginals and v. For measurable T:R + RR,, denote by T1l the associated push-forward operator between measures.
2 Background on Wasserstein Gradient Flows
We consider gradient flows in Wasserstein space (P2(RD), W2), the space of probability measures with finite second moment on DD endowed with the Wasserstein-2 metrie Wig.
Wasserstein-2 distance. The (squared) Wasserstein-2 metric W2 between μ4 v E (2(RP) is
Wa(H,)) def min =ET((a.) RR
|| 21 dr v),
where the minimum is over measures on DRD x RD with marginals μ and respectively [68].

---

For je E P2.ac (R (R)), there exists y-unique map Vij*: : : + R. that is the gradient of a convex
function v² : _RD + R J {oo} satisfying d = y [46). From Brenier's theorem [13.. it follows that = i Vub"itu is the unique minimizer of i.e.
{ JRO
v)
|| V (x)ll
Wasserstein Gradient Flows. In the Euclidean case, gradient flows along function f IR follow the steepest descent direction and are defined through the ODE du. f Discretiza- tion of this flow leads to the gradient descent minimization algorithm. When functionals are defined over the space of measures equipped with the Wasserstein-2 metric, the equivalent flow is called the Wasserstein gradient flow. The idea is similar: the flow follows the steepest descent direction, but this time the notion of gradient is more complex. We refer the reader to [4] for exposition of gradient flows in metric spaces, or [59 Chapter 8] for an accessible introduction.
A curve of measures (PETERR following the Wasserstein gradient flow of functional F solves the continuity equation
Pp: 0t
V
s.t. P0 pp;
(2)
where F.). is the first variation of F 4 Theorem 8.3.1]. The term on the right can be understood as the gradient of F in Wasserstein space, a vector field perturbatively rearranging the mass in Pr to yield the steepest possible local change of F.
Wasserstein gradient flows are used various applied tasks. For example, gradient flows are applied in training 8
43 25 or refinement 7 of implicit generative models. In reinforcement learning,
gradient flows facilitate policy optimization [55] 72). Other tasks include crowd motion modelling
45
58 50
dataset optimization 12] and in-between animation [26].
Many applications come from the connection between Wasserstein gradient flows and SDEs. Consider an IR D -vauedd stochastic process {Xt)eR: governed by the following Itō SDE:
dxx, = (X, 22, s.t. Xo pd
where 4 RD R is the potential function. W, is the standard Wiener process, and 3 0 is the magnitude. The solution of (3) is called an inrection-diffusian process. The marginal measure Pt of X2 at each time satisfies the mokker-Pllnnee equation with fixed diffusion coefficient:
Op: div (Z)()))) + 3- 3''Apt Ot
s.t. Po = pp
(4)
Equation (4) is the Wasserstein gradient flow 2 for JF given by the Fokker-Plance Planckk free energy functional 332
Fir(p) u(p) 8-1E(p)
(5)
where u(p) = SRD (x))d((z) is the potential energy and E(p) - S log da (z)dp(x) is the entropy. As the result, to solve the SDE B. one may compute the Wasserstein gradient flow of the Fokker-Planck equation with the free-energy functional Frp given by (5)
JKO Scheme. Computing Wasserstein gradient flows is challenging. The closed form solution is typically unknown, necessitating numerical approximation techniques. Jordan, Kinderlehrer, and Otto proposed method later abbreviated as JKO integration to approximate the dynamics of pe in 332. It consists of time-discretization update of the continuous flow given by:
(k) +- arg min
W2 (2(k F(p) + 2h
(6)
(EP)
where ((0) = p0 is the initial condition and > is the time-discretization step size, The discrete time gradient flow converges to the continuous one as h - O, i.c., (k)) Pkb. The method was further developed in [4] 600, but performing JKO iterations remains challenging thanks to the minimization with respect to Wz.
A common approach to perform JKO steps is to discretize the spatial domain. For support size 2 1 (6) can be solved by standard optimal transport algorithms [5I]. dimensions > 3, discrete supports can hardly approximate continuous distributions and hence the dynamics of gradient flows.
To tackle this issue, (24] propose a stochastic parametric method to approximate the density of pt- Their method uses entropy-ree regulai-ee optimal transport (OT), which is biased.
3

---

Computing Wasserstein Gradient Flows with ICNNS
We now describe our approach compute Wasserstein gradient flows via JKO stepping with ICNNs.
3.1 JKO Reformulation via Optimal Push-forwards Maps
Our key idea replace the optimization (6 over probability measures by an optimization over convex functions, an idea inspired by M.. Thanks to Brenier's theorem, for any E P2,ac there exists unique p(k-1)) -eeasurbble gradient V RDD _ TR of convex function xé satisfying Vitt(k--1, We set Veitp(k-1) and rewrite (6) as an optimization over convex s:
(b(()) arg min Convex
F(Velpp + 2h W (otk--1) Febdp(k-1))
To proceed to the next step of JKO scheme, we define p()) def
V Vell (p(k-1)
Since is the pushforward of p(k-1) by the gradient of a convex function Vil. the W term in 7 can be evaluated explicitly, simplifyyng the Wasserstein-2 distance
7
g6(4)) arg min
F(Vytol-
1
/ || Vo(z))
2h RR
Convex
This formulation avoids the difficulty of computing Wasserstein-2 distances. An additional advantage is that we can sample from g(*), Since p! = [Vbbll) ....o (eb(")))pp) one may sample vn p()), and then (bb(k) o..... .... V (xr("))))) gives sample from p(*), Moreover, if functions 15() are strictly convex, then gradients Vib(C) are invertible. In this case, the density tpp is computable by the change of variables formula (assuming 2b(-) are twice differentiable)
do
di of p Vis(k) ....o
(k( (zx) [det 22 (xx-1)]-1... [det V
(
dx (Ho),
dx
where r; Veili) for i 1, I,...,k. l and ' dx the density of (())
3.2 Stochastic Optimization for JKO via ICNNs
In general, the solution v((k) functions. To tackle this issue, [1 discretizes the space of convex function. The approach also requires discretization of measures (8) limiting this method to small dimensions.
is intractable since it requires optimization over all convex
We propose to parametrize the search space using input convex neural networks (ICNNs) 6 satisfying a universal approximation property among convex functions [20). ICNNs are parametric models of the form the D IR with the convex w.r.t the input. ICNNs are constructed from neural network layers, with restrictions on the weights and activation functions to preserve the input-convexity. see 6 $3.1] (37 SB.2.. The parameters are optimized via deep learning optimization techniques such as SGD.
The JKO step then becomes finding the optimal parameters 6* for thg
g + min min
F( (n + 2h JRD
||æ Vvo( 1
(10)
0
If the functional F can be estimated stochastically using random batches from p(k-1) then SGD can be used to optimize 0. FFP given by (5) is an example of such functional:
Theorem (Estimator of FFe). Let p E (2 and T RD DB be diffeomorphism. For
random batch 21...., EN ~ p. the expression (1 - 13-------8 AE , (NN) where
def (T(xx)) and
N n=l
AE (e,.,....,
xx)
log det VT(xx)),
N
n=1
is an estimator of (TP(()t)) up constant (want. T) shift given by 8 - E(p).

---

Proof. Ur is straightforward unbiased estimator (TTp)) Let p and pr be the densities of and Táp. Since T is diffeomorphism, we have pr(y) = p(x) det (T where x (y)
Using the change of variables formula, we write
E E(TTp)
(v()) log pT (v)dy
R2o
{ p(x) det VT(x)|-1 p(x) det VT(x) dddet (T))))
RR
= Sco p(x) log + J p(x) log det VT(z)
E(p) / RRD
log det (T(x))dr,
AET(p) del E(TTp) E(p) - Jzp log det VT(x))dp(z)
which explains that ET is an unbiased estimator of (E)(p). As the result, Ür 18-1--88 is an estimator for FFp(Ttp) (Tppp) - 3 (TTE(T up to a shift of B (())
To apply Theorem to case, take T' +- Vih and p p ((-1)) obtain stochastic estimator for Fip( V Vrott fixed, so the offset of the estimator plays no role the optimization w.r.t. 0.
(k- 1) in [TO.. Here, (-1E(p(--1) is (-independent and constant since p(k-1) is
Algorithm details our stochastic JKO method for FFP. The training is done solely based on random samples from the initial measure p": its density is not needed.
Algorithm 1: Fokker-Planck JKO via ICNNS
Input Initial measure pd accessible by samples;
JKO discretization step h >> O, number of JKO steps K >> 0; target potential (x), diffusion process temperature / B=1; batch size N;
Otranee :trained ICNN models (1)(K) |= representing JKO steps for k = 1,,,..,,..... 1,. K do
rhg basic ICNN model;
for ; = 1,2,... do
Sample batch Z p of size N:
X + (k
V (Z);
W - Σ
V (x(x) x z113:
xEX
û N
Σ
V
(z)))
zEX
Σ
log det (z));
xEX
-- 2h
W
+ -
22. Perform gradient step over by using 70
(k)
This algorithm assumes the Fokker- Planck diffusion energy functional. However, our method admits straightforward generalization to any FF that can be stochastically estimated; studying such functionals promising avenue for future work.
3.3 Computing the Density of the Diffusion Process
Our algorithm provides a computable density for p(k), As discussed 3.1 it is possible to sample from (f5) while simultaneously computing the density of the samples. However, this approach does
not provide direct way to evaluate dp(k) (Ix) for arbitrary 2b D. We resolve this issue below.
5

---

If convex function is strongly convex, then its gradient is bijective on R By the change of variables formula for zy E IR it holds de de (Ix) = dolt (Ia-1) (det
where
da
Ik = Vvilk) (xx (xx-1). To compute -1 one needs to solve the convex optimization problem:
Fk (x==)(
-k-1 arg max (z, (z)).
(11)
zERD
If we know the density of p", to compute the density of (()) at zy we solve &s convex problems
2k-1 arg max (z, z&) HERD
2b)()) (x)]
20 arg arg max (2, e) (x)]
ZER
to obtain -g-1 To and then evaluate the density as
dpk dx
dp dx
(mr)
(mo)
[II det (2:100((3)1)))
Note the steps above provide general method for tracing back the position of a particle along the flow, and density computation is simply a byproduct.
Experiments
In this section, we evaluate our method on toy and real-world applications. Our code is written in PyTorch and is publicly available
The experiments are conducted on a GTX 1080Ti. most cases, we performed several random restarts to obtain mean and variation of the considered metric. As the result, experiments require about 100-150 hours of computation. The details are given in Appendix A
Neural network architectures. In all experiments, we use the DenselCNN architecture for the Algorithm [ with SoftPlus activations. The network th twice differentiable w.r.t. the input æ and has bijective gradient Vidg RRD - R with positive semi-definite Hessian V2 (x+
37 Appendix B.2]
at each x. We use automatic differentiation to compute Vaga and V2
Metric. To qualitatively compare measures, we use the symmetric Kullback Leibler divergence
SymKL(p1, P2) def (L (p2) KL(p2llp1),
(12)
dei (x dpp (x) is the Kullback -Lebllr divergence. For particle-bri based where KL(pillpe) def J&p log dpz methods, we obtain an approximation of the distribution by kernel density estimation.
4.1 Convergence Stationary Solution
Starting from an arbitrary initial measure p", an ddvection-diffusiin process unique stationary solution p" [56) with density
4 converges to the
dp" dz
Z
(13)
where Z = fg cxp(-80(x)) dr is the normalization constant. This property makes it possible to compute the symmetric KI. between the distribution to which our method converges and the ground truth, provided Z is known.
We use (0, 16/p) as the initial measure p and a random Gaussian mixture as the stationary measure p². In our method, we perform K = 40 JKO steps with step size h 0.1. We compare with parti- cle simulation method (with 103 10 10 particles) based the Euler-Mann Maruyama EM] approximation [36, 89.2]. We repeat the experiment 5 times and report the averaged results in Figure
0.5
0.0
10.5
[EM]
[EM] 10k
(EM) 50K
Ours
10
dimension
In Figure we present qualitative results of our Figure 1: SymKL between the computed and method converging to the ground truth in D = 13, 32. the stationary measure in D 2, 2..... 12

---

Statienary measure
Fitted measure (ours)
Stationary measure
Fitted measure (ours)
10
-5
-
-10
-0
10
10
-5
(a) Dimension D 13
(b) Dimension D 32
Figure 2: Projections to 2 first PCA components of the true stationary measure and the measure approximated our method in dimensions D 13 (on the left) and D 32 (on the right).
4.2 Modeling Orastein-Uhlenbeck Processes
processes are advection-diffusion processes (4 with 4(x) -2 A(= Axx=> for symmetric positive definite A E R and b E RR. They are among the few examples where we know Pr for any E R++ in closed form, when the initial measure p Gaussian [67]. This allows to quantitatively evaluate the computed dynamics of the process, not just the stationary measure.
We choose A, 6 random and set pd be the standard Gaussian measure (0, ID). We approximate the dynamics of the process by our method with JKO step / 0.05 and compute SymKL between the true Pt and the approximate one time 0.5 and 0.9. We repeat the experiment 15 times in dimensions D 1.2.... 12 and report the performance at in Figure 3 The baselines are [EM] with 10 10 5 x 10t particles, EM particle simulation endowed with the Proximal Recursion operator EM PR] with 10* particles [16), and the parametric dual inference method [24) for JKO
steps Duall JKO]. The detailed comparison for times 0.1, 0.2. given Appendix C
2
[Dual [KO]
[EM]
[EM] 1K
(EM PR] 10K
[EM] 10K
Ours
logroSymkt
legtoSymki
-1
[uual
[EM] 50K
[EM] 1K
[EM PR] 10K
-3
[EM]
D. dimension
D. dimension
(a) Time t = 0.5
(b) Time f 0.9
Figure 3: SymKL values between the computed measure and the true measure P: at 0.5 (on left) and ( = 0.9 (on the right) in dimensions D = 1,2...., 12. Best viewed in color.
4.3 Unnormalized Posterior Sampling in Bayesian Logistic Regression
An important task in Bayesian machine learning to which our algorithm can be applied sampling from unnormalized posterior distribution. Given the model parameters IR with the prior M distribution po(x) well as the conditional density p(S]z) I p(sm (x) {si...., SM the posterior distribution is given by
of the data S
P(z|S) p(((x))o(z) ((S)=)po(z)) po(x) II P(S)
M
p(smm/r))
Computing the normalization constant P(S) is in general intractable, underscoring the need for estimation methods that sample from p(S ()) given the density only up a normalizing constant.

---

In our context, sampling from p(r S) can be solved similarly to the task in $4.1 From (13), it follows that the advection-diffusee cess with temperature ,3 > and (xx) = -3 log [po(z)- p[S]=)] has (x) p(r|S) as the stationary distribution. Thus, we can use our method to approximate the diffusion process and obtain sampler for p(x]S) as result.
Accuracy
Log-I Log-Liielihood
Dataset
Ours
[SVGD]
Ours
[SVGD]
pro-
covtype
0.75
0,75
-0.515
-0.515
german
0.67
0.65
-0.6
-0.6
diabetis
0.775
0.78
-0.45
-0.46
twonorm
0.98
0.98
-0.059
-0.062
ringnorm
0.74
0.74
-0.5
-0.5
banana
0.55
0.54
-0.69
-0.69
splice
0.845
0.85
-0.36
-0.355
waveform
0.78
0.765
-0.485
-0,465
The potential energy (p)) Jgo ((z)dp(z) can be estimated efficiently by using a trick sim- ilar to the ones in stochastic gradient Langevin dynamics (70), which consists in resampling samples in S uniformly. For evaluation, con-
image 0.82
0.815
-0.43
-0.44
Table 1: Comparison of our method with [SVGD] 42 for Bayesian logistic regression.
sider the Bayesian linear regression setup of 42 We use the 8 datasets from [47]. The number of features ranges from 60 and the dataset size from 700 to 7400 data points. We also use the Covertype datase with 500K data points and 54 features. The prior on regression weights w is given by Po(w (o) NNwl0, a with po(a) Gamma(a|1. L, 0.01), so the prior on parameters [w, al of the model is given by Po(x)= = Po(w, a) = po (w] ( ) po()). We randomly split each dataset into train Synin and test Swe ones with ratio 4:1 and apply the inference on the posterior ((anS In Table we report accuracy and log- ligliikli of the predictive distribution on Swst- As the baseline, we use particle-based Stein Variational Gradient Descent 42
We use the author's
implementation with the default hyper-parameters.
4.4 Nonlinear Filtering
We demonstrate the application of our method to filtering a nonlinear diffusion. In this task, we consider a diffusion process X, governed by the Fokker-Planck equation 4 At times ( < l2 < tk we obtain noisy observations of the process Y₁ = Xt2, + Uk with Uk ~ ~0(,o)). The goal is to compute the predictive distribution Pt,x (x(Yxxx)) for t tx given observations YI.K-
For each & and 2 & predictive distribution Pt,x (2) Yysk) follows the diffusion process on time interval [4x, t] with initial distribution Pt,,x (v]) If t = t then
Xtx. (x) (1a) x P(Y = x) x (x:VVx11-)
(14)
For k = 1,..,, K, we sequentially obtain the predictive distribution Pt..x (ar Vi.k) by using the previous predictive distribution -x [1). First, given access to Pt.- x (2) (i we approximate the diffusion on interval (4,1, 4u] with initial distribution Pt. x (x) Y 1) by our
Algorithm| I to get access to Pt... (x])\ix=1)) Next, we use
(14) to get unnormalized density and
Metropolis Heatopol algorithm [57] to sample from Pt (x) Y1:k) We give details in Appendix
For evaluation, we consider the experimental setup of [24 S6.3]. We assume that the 1-dimensional diffusion process X, has potential function (zx) highly nonlinear. We simulate nonlinear filtering on the time interval Iatat
sin (2n2) + 4 which makes the process
sec., 4in 5 sec. and take the noise observations each 0.5 sec. The noise variance o² = 1 and p( (Xo) N (X. 1).
We predict the conditional density Piam. x (x) Y1.9) and compare the prediction with ground truth obtained with numerical integration method by Chang and Cooper [19], who use a fine discrete grid, As the baselines, we use Dual JKO] 24 as well as the Bayesian Bootstrap filter [BBF] 27, which combines particle simulation with bootstrap resampling at observation times.
We repeat the experiment 15 times. In Figure 44l we report the SymKL between predicted density and true cX,, Y1.9). We visually compare the litted and true conditional distributions in Figure 4b
5 Discussion
Complexity of training and sampling. Let T be the number of operations required to evaluate ICNN (z). and assume that the evaluation of 4 (x) in the potential energy 24 takes O(1) time.
htttps://www.cc
cww..csseeeett.. edu. http://ggggg//tttoo//t////t/tt//

---

Discrepancy comparision sec.
Diffusion pdfs comparison sec.
0.5
Ours
-0.5
logroSymki
0.4
[Dual [KO]
[0BF] 50000
0.3
ground
-1.5
-2.0
0.1
-2.5
0.0
Ours [Dual [KO]
[BBF]
[BBF]
[8BF]
[B6F]
(a) SymKI, values.
(b) Visualized probability density functions.
Figure 4: Comparison of the predicted conditional density and true p( tx. Yı:)).
Recall that computing the gradient is small constant factor harder than comput- ing the function itself [41] Thus, evalu- TRD requires O(T) operations and evaluating the Hessian (z)o(x) : RD I DDD takes O(DT)) time. To compute log det (z)) we need O(D*) extra operations. Sampling from p(k-1) tpo involves pushing To ~ p forward by a sequence of ICNNS gi()) of length k - 1, requiring O((k- - 1)T) operations. The forward pass to evaluate the JKO step objective Z in Algorithm requires
Operation
Time Complexity
Eval. tho, V th.. Vi0bo Eval. log det V2ebo
T, O(DT+D*)
ation of Viho())
IRD
Sample x ~ p(k)
(((k-))T)
Eval. Cn a ~ p(k)
O(DT+D*)
Eval. aā 50 on on p(k)
O(DT+D*)
Vis(t-1) o
Sample ~ -p((t) and Eval. dp!4)
((I--))TD++
de (z)
Table Complexity of operations in our method for computing JKO steps via ICNNS.
O(DT + D operations, as does the backward pass to compute the gradient až 39 w.r.t. 0.
The memory complexity more difficult characterize, since it depends on the autodiff implementa- tion. It does not exceed the time complexity and is linear in the number of JKO steps k.
Wall-clock times. All particle-based methods considered in $ and |uall JKO require from several seconds to several minutes CPU computation time. Our method requires from several minutes to few hours on GPU, the time is explained by the necessity train a new network at each step.
Advantages. Due to using continuous approximation, our method scales well to high dimensions, as we show in $44.1 1 and $4.2 After training, we can produce infinitely many samples 24 p(0) together with their trajectories 1k-1: Zk-2, .,...,.......................... To along the gradient flow. Morcover, the densities of samples in the flow (h(h)
dx (ar)),
do
di (II-1), 1).... dp 0) (zo) can be evaluated immediately.
In contrast, particle-based and domain discretization methods do not scale well with the dimension (Figure 3) and provide no density. Interestingly, despite its parametric approximation. [Dual JKO performs comparably to particle simulation and worse than ours (see additionally [24, Figure 3J).
Linitations. To train 1: JKO steps, our method requires time proportional to 222 due the increased complexity of sampling addition, for very high dimensions D, exact evaluation of log det (z)) is time-consuming.
(kk), This may be disadvantageous for training long diffusions. In
Future work. To reduce the computational complexity of sampling from pkk) at step one may regress an invertible network H .R. - R
p 31 to satisfy (xv)) = Vab(R) ..... V Vi ()) (zp)
and use Hipn (k) simplify sampling. An alternative is use variational inference 12 54
to approximate (A)) To mitigate the computational complexity of computing log det (v)( fast approximation can be used [66] 28). More broadly, developing ICNNs with casily-computable exact Hessians is critical avenue for further research ICNNs continue to gain attention in machine
learning 37 38
30
23
3.
Potential impact. Diffusion processes appear in numerous scientific and industrial applications, including machine learning, finances, physics, and population dynamics. Our method will improve models in these areas, providing better scalability. Performance, however, might depend on the expressiveness of the ICNNs, pointing to theoretical convergence analysis as a key topic for future study to reinforce confidence in our model.

---

In summary, we develop an efficient method to model diffusion processes arising in many practical tasks. We apply our method common Bayesian tasks such as unnormalized posterior sampling (44.3) and nonlinear filtering (94.4). Below we mention several other potential applications:
Population dynamics. In this task, one needs to recover the potential energy (x)) included in the Fokker-Planck free energy functional Fip based on samples from the diffusion obtained at timesteps f1,..... E,n see [29]. This setting can be found in computational biology, see $6.3 of [29]. A recent paper [14] utilizes ICNN-powered JKO to model population dynamics.
Reinforcement learning. Wasserstein gradient flows provide optimize an agent policy in reinforcement learning, see [55] 72]. The idea of the method is to maximize the expected total reward (see (10) in (72) using the gradient flow associated with the Fokker-Planee functional (see (12) in (72). The authors of the original paper proposed discrete particle approximation method to solve the underlying JKO scheme. Substituting their approach with our ICNN based JKO can potentially improve the results.
eeoretically-groundee way to
Refining Generative Adversarial Networks. In the GAN setting, given trained generator G and discriminator D, one can improve the samples from by D via considering gradient flow w.r.t. entropy- regularizee f-divergence between real and generated data distribution (see [7]. in particular, formula (4) for reference). Using KL-divergence makes the gradient flow consistent with our method: the functional F defining the flow has only entropic and potential energy terms. The usage our method instead of particle simulation may improve the generator model.
Molecular Discovery. In 33. in parallel to our work the JKO-ICNN scheme is proposed. The authors consider the molecular discovery as an application. The task is to increase the drug-likenen, of given distribution of molecules while staying close to the original distribution Pu- The task reduces to optimizing the functional F(p) (x=+ DDp, Pu) for certain potential & (V
the notation of (3 and discrepancy D. The authors applied the JKO-ICNN method to minimize
on MOSES 53 molecular dataset and obtained promising results.
BCKNOWLEDGEMENTS Skoltech acknowledges the support of the Ministry of Science and Higher Education grant No. generous support of Army Research Office grants
The MIT Geometric Data Processing group acknowledges the and W911NF21 10293, of Air Force Office of Scientific Research award FA9550-19-1-00 19-1- of National Science Foundation grants I 1 and CHS-1955697, from the CSAIL Systems that Learn program, from the MIT-IBM Watson AI Laboratory, from the Toyota-CSAIL Joint Research Center, from gift from Adobe Systems, from an MIn. Immersion Lab/NCSOFT Gaming Program seed grant, and from the Skoltech-MIT Next Generation Program.
References
[1] Juha Ala-Luhtala, Simo Särkkä, and Robert Piché. Gaussian filtering and variational approx- imations for Bayesian smoothing in dontintous-discrete stochastic dynamic systems. Signal Processing, 111:124-133 136, 2015.
[2] David Alvarez-Melis and Nicolò Fusi. Gradient flows in dataset space. arXiv preprint 1xXi::2010.1176000 2020.
[3] David Alvarez-Melis, Yair Schiff, and Youssef Mroueh. Optimizing functionals on the space of probabilities with input convex neural networks. arXiv preprint arXiv:2106.000 00774, 2021.
[4] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient flows: in metric spaces and in the space of probability measures. Springer Science Business Media, 2008.
[5] Brandon Amos and J Zico Kolter. Optnet:
Differentiable optimization as layer in neural networks. In International Conference on Machine Learning, pages 136-145. PMLR, 2017.
[6] Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pages 146-155, PMLR, 2017.
[7] Abdul Fatir Ansari, Ming Liang Ang, and Harold Soh. Refining deep generative models via Wasserstein gradient flows. arXiv preprint arXiv:2012.00780, 2020.
10

---

[8] Michael Arbel, Anna Korba, Adil Salim, and Arthur Gretton. Maximum mean discrepancy gradient flow. arXiv preprint arXiv: 04370, 2019.
[9] Lynton Ardizzone, Jakob Kruse, Sebastian Wirkert, Daniel Rahner, Eric W Pellegrini, Ralf Klessen, Lena Maier-Hein, Carsten Rother, and Ullrich Köthe. Analyzing inverse problems with invertible neural networks. arXiv preprini pooo::00880555558 2018.
[10] Genevay Aude, Marco Cuturi, Gabriel Peyré, and Francis Bach, Stochastic optimization for large-scale optimal transport. arXiv preprint arXiv: 2016.
[11] Jean-David Benamou, Guillaume Carlier, Quentin Mérigot, and Edouard Oudet. Discretization of functionals involving the Monge-Ampère operator. Numerische mathematik, 134())611-633,
2016.
[12] David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):559-877, 2017.
[13] Yann Brenier. Polar factorization and monotone rearrangement of vector-valued functions. Communications on pure and applied mathematics, 44(4):375-417, 1991.
[14] Charlotte Bunne, Laetitia Mleg--Pppaxanthos Andreas Krause, and Marco Cuturi. Jkonet: Proximal optimal transport modeling of population dynamics. arXiv preprint drXiv:2106.06335, 2021.
[15] 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.
[16] Kenneth F. Caluya and Abhishek Halder. Proximal recursion for solving the fokker-plang equation, 2019.
[17] José A Carrillo, Alina Chertock, and Yanghong Huang. A finite-volumt method for nonlinear nonlocal equations with gradient flow structure. Communications in Computational Physics, 17(1):233-258 2015.
[18] IS Chang and G Cooper. A practical difference scheme for fokker-planck equations. Journal of Computational Physics, 6(1):1-16, 1970.
[19] J.S Chang and G Cooper. A practical difference scheme for fokker-planck equations. Journal of Computational Physics, 6(1):1-16, 1970.
[20] Yize Chen, Yuanyuan Shi, and Baosen Zhang. Optimal control via neural networks: convex approach. arXiv preprint 188.v:1805.11833 2018.
[21] Arnaud Doucet and Adam M Johansen. A tutorial particle filtering and smoothing: Fifteen years later. Handbook of nonlinear filtering, 12(656-704):3, 2009.
[22] Nicole El Karoui, Shige Peng, and Marie Claire Quenez. Backward stochastic differential equations in finance. Mathematical finance, 7(1):1-71, 1997.
[23] Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Scalable computations of wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.04462 2020.
[24] Charlie Frogner and Tomaso Poggio. Approximate inference with Wasserstein gradient flows. In International Conference on Artificial Intelligence and Statistics, pages 2581-2590. PMLR, 2020.
[25] Yuan Gao, Yuling Jiao, Yang Wang, Yao Wang, Can Yang, and Shunkang Zhang, Deep generative learning via variational gradient flow. In International Conference on Machine Learning, pages 2093-2100 PMLR, 2019.
[26] Yuan Gao, Guangzhen Jin, and Jian-Guo Liu. Inbetweening auto-animation via fokker-planck dynamics and thresholding. arXiv preprint arXiv:2005.0 08858, 2020.
[27] N.
Gordon, D. Salmond, and A. Smith. Novel approach to CatneearrnoonGaussiie Bayesian state estimation. 1993.

---

[28] Insu Han, Dmitry Malioutov, and Jinwoo Shin. Large-scale log log-determinant computation through stochastic chebyshev expansions. In International Conference on Machine Learning. pages 908-917. PMLR, 2015.
[29] Tatsunori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-level diffusions with generative rnns, In International Conference on Machine Learning, pages 2417-2426, PMLR, 2016.
[30] Chin-Wei Huang, Ricky TQ Chen, Christos Tsirigotis, and Aaron Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. arXiv preprint arXiv:2012..5592,, 2020.
[31] Jöm-Henrik Jacobsen, Arnold Smeulders, and Edouard Oyallon. i-revnet: Deep invertible networks. arXiv preprint 18X2:1880200088. 2018.
[32] 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.
[33] Simon J Julier, Jeffrey K Uhlmann, and Hugh F Durrant- Whyte. A new approach for filtering nonlinear systems. In Proceedings of 1995 American Control Conference-ACC"95, volume pages 1628 1632. IEEE, 1995.
[34] Rudolph Kalman and Richard S Bucy. New results in linear filtering and prediction theory. 1961.
[35] Søren Klim, Stig Bousgaard Mortensen, Niels Rode Kristensen, Rune Viig Overgaard, and Henrik Madsen. Population stochastic modelling (psm)-an package for mixed-effects models based stochastic differential equations. Computer methods and programs in biomedicine, 94(3)277-2299, 2009.
[36] Peter E. Kloeden. Numerical solution of stochastic differential equations/ Peter E. Kloeden, Eckhard Platen. Applications of mathematics; v. 23. Springer, Berlin, 1992.
[37] Alexander Korotin, Vage Egiazarian, Arip Asadulaev. Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations,
2021.
[38] Alexander Korotin, Lingxiao Li. Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-
2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021.
[39] Harold Kushner. Approximations to optimal nonlinear filters. IEEE Transactions on Automatic Control, 12(5):546-5555 1967.
[40] Hugo Lavenant, Sebastian Claici, Edward Chien, and Justin Solomon. Dynamical optimal transport on discrete surfaces. ACM Transactions on Graphics (TOG), 37(6):1-16, 2018.
[41] Seppo Linnainmaa. The representation of the cumulative rounding error of an algorithm as taylor expansion of the local rounding errors. Master's 's Thesis (in Finnish), Univ. Helsinki, pages 6-7, 1970.
[42] Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. arXiv preprint 1rXhv:1608.04477 2016.
[43] Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian- Robert Stöter. Sliced-Wassen Wasserstein flows: Nonparametric generative modeling via optimal transport and diffu- sions. In International Conference Machine Learning, pages 4104 4113. PMLR, 2019.
[44] Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks, In International Conference on Machine Learning, pages 6672-6681. PMLR, 2020.
[45] Bertrand Maury, Aude Boudneff-Chupple and Filippo Santambrogio. A macroscopic crowd motion model of gradient flow type. Mathematical Models and Methods in Applied Sciences, 178701117771181,, 2010.
12

---

[46] Robert McCann et al. Existence and uniqueness of monotone measure- measure-presente preserving maps. Duke Mathematical Journal, 80(2):309-3222 324, 1995.
[47] Sebastian Mika, Gunnar Ratsch, Jason Weston, Bernhard Scholkopf, and Klaus-Robert Mullers. Fisher discriminant analysis with kernels. In Neural networks for signal processing IX: Proceed- ings of the 1999 IEFE signal processing society workshop (cat. no. 98th8468), pages 41-48. Icce, 1999.
[48] Manfred Opper. Variational inference for stochastic differential equations. Annalen der Physik, (3)::3300233,, 2019.
[49] Lorenzo Pareschi and Mattia Zanella. Structure preserving schemes for nonlinear tokker-plande equations and applications. Journal of Scientific Computing, 74(3):1555-1600, 2018.
[50] Gabriel Peyré. Entropic approximation of Wasserstein gradient flows. SIAM Journal Imaging Sciences, 8(4):2323-2351, 2015.
[51] Gabriel Peyré, Marco Cuturi, al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 11(5-6)3355607,, 2019.
[52] Eckhard Platen and Nicola Bruti-Liberati. Numerical solution of stochastic differential equations with jumps in finance, volume 64. Springer Science & Business Media, 2010.
[53] Daniil Polykovskiy, Alexander Zhebrak. Benjamin Sanchez-Lengeling, Sergey Golovanov, Oktai Tatanov, Stanislav Belyaev, Rauf Kurbanov, Aleksey Artamonov, Vladimir Aladinskiy, Mark Veselov, Artur Kadurin, Simon Johansson, Hongming Chen, Sergey Nikolenko, Alán Aspuru-Guzik, and Alex Zhavoronkov. Molecular sets (MOSES): A benchmarking platform for molecular generation models. Frontiers in Pharmacology, 11:1931, 2020.
[54] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning, pages 1530-1538. PMLR, 2015.
[55] Pierre H Richemond and Brendan Maginnis. On Wasserstein reinforcement learning and the Fokker-Planch equation. arXiv preprint arXiv:1712.07/88: 2017.
[56] Hannes. Risken. The Fokker-Planck Equation: Methods of Solution and Applications (Springer Series in Synergetics). Springer,, 1996.
[57] Christian P Robert and George Casella. The Metropolis Hastings algorithm. In Monte Carlo Statistical Methods, pages 231-283. Springer, 1999.
[58] Filippo Santambrogio. Gradient flows in Wasserstein spaces and applications crowd move- ment. Séminaire Equations aux dérivées partielles (Polytechnique), pages 1-16, 2010.
[59] Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58- 63):94, 2015.
[60] Filippo Santambrogio. Euclidean, Metric, and Wasserstein gradient flows: an overview, 2016.
[61] Simo Sarkka. On unscented kalman filtering for state estimation of continuous-time nonlinear systems. IEEE Transactions on automatic control, 52(9): 2007.
[62] Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv:1711.02283. 2017.
[63] Kazimierz Sobczyk. Stochastic differential equations: with applications to physics and engi- neering, volume 40. Springer Science Business Media, 2013.
[64] Tobias Sutter, Arnab Ganguly, and Heinz Koeppl. A variational approach to path estimation and parameter inference of hidden diffusion processes. The Journal of Machine Learning Research, 17(1):654466580 2016.
13

---

[65] Gabriel Terejanu, Puneet Singla, Tarunraj Singh, and Peter D Scott. A novel gaussian sum filter method for accurate solution the nonlinear filtering problem. In 2008 11th International Conference Information Fusion, pages 1-8. IEEF, 2008.
[66] Shashanka Ubaru, Jie Chen, and Yousef Saad. Fast estimation of tr(((a)) via stochastic lanczos quadrature. SIAM Journal on Matrix Analysis and Applications, 38(4):1075-1099 2017.
[67] Vatiwutipong and N Phewchean. Alternative way to derive the distribution the multivariate ornstein uubeenbnk process. Advances in Difference Equations, 2019(1)11-7,001 2019.
[68] Cédric Villani. Optimal transport: old and new, volume 338. Springer Science Business Media, 2008.
[69] Michail D Vrettas, Manfred Opper, and Dan Cornford. Variational mean-field algorithm for efficient inference large systems of stochastic differential equations. Physical Review E, 91(1):012148, 2015.
[70] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681-688, Citeseer, 2011.
[71] Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational inference. IEEE transactions on pattern analysis and machine intelligence, 41(8):2008-202) 2018.
[72] Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization as Wasserstein gradient flows. In International Conference on Machine Learning, pages 5737- 5746. PMLR, 2018.
14


Reviews:
[{"review_id": "wwKSfg58on2", "reviewer": "Reviewer_FY52", "summary": "The Fokker-Planck equation can be treated as gradient descent over entropy functionals in Wasserstein space. To solve this problem, Jordan Kinderlehrer and Otto proposed the JKO scheme. Later, through the connection between the Brenier theorem and this scheme, [10] treat it as an optimization problem which aims at finding the convex Brenier functional $\\psi$ in equation (7). However, the solution given by [10] is usually intractable in high dimensional space. Thus, this paper proposes to use the input convex neural networks (ICNNs) proposed by [5] to replace $\\psi$. And then use SGD to optimize the model parameter. This forms the core contribution of this work. Additionally, the experiments are interesting.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- My main concern about this paper is its lack of novelty. As summarized above, the core contribution is to replace the convex functional $\\psi$ in equation (7) proposed by [10] with the input convex neural networks (ICNNs) proposed by [5]. It is a common trick to replace the convex Brenier potential with the ICNN network in optimal transport community. The only difference is the background of the applications.\n- In theorem 1, how to compute the high dimensional determinant $det \\nabla T(x_n)$? \n- In Sec. 4.1, since $W_2(\\rho_{k-1}, \\rho_k)$ for different $k$ has been computed during the optimization, it should be reported. The plot can help illustrate the performance of the proposed method.\n- The comparison between the proposed method and [10] in low dimensional space in terms of memory usage/training time/accuray is necessary, because it is a direct generalization of [10].", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "wFZICnCC_Tm", "reviewer": "Reviewer_yWvM", "summary": "This paper proposes an input-convex neural networks (ICNNs) based method to approximate the JKO scheme for discretizing Wasserstein gradient flows. Following a previous work [1], the key idea is to reformulate the distribution optimization problem in each JKO iteration as optimization over convex functions due to the well-known Brenner’s polar factorization theorem [2]. ICNNs are used to parameterize these convex functions from computational considerations. The proposed method is applied to simulate Wasserstein gradient flows for sampling from unnormalized density and nonlinear filtering.\n[1] Discretization of functionals involving the Monge-Ampere operator.\n[2] Polar factorization and monotone rearrangement of vector-valued functions.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "It is a novel idea to introduce ICNNs for addressing the computational challenges of the JKO scheme. ICNNs may be a promising tool to make the JKO scheme feasible for large-scale computations.\nQuality: The work is technically sound.\nPros: The methodology part is well-organized, and the advance of the proposed method is clear.\nThe potential applications are interesting, such as sampling from unnormalized density, nonlinear filtering.\nCons: (1) Maybe the title is a little ambiguous. It would be better to include input-convex neural networks (ICNNs) in the title since the proposed method is entirely based on ICNNs.\n(2) As the title indicates, the focus of this work is a large-scale setting. However, it is hard to conclude that the large-scale setting is well considered in the experimental part. For example, the considered data dimension is at most 32 in Section 4.1, 12 in Section 4.2, 60 in Section 4.3, 1 in Section 4.4.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "tUTHFIYodXZ", "reviewer": "Reviewer_tTtz", "summary": "This paper is on the computation of Wasserstein gradient flow, a type of gradient flow based on the optimal transport theory. The proposed algorithm is built on the famous JKO scheme. To implement the JKO scheme, input-convex neural networks (ICNNs) are used to model the optimal transport map between the distributions at consecutive steps. ", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality:\nComputing Wasserstein gradient flow is a classical problem in optimal transport. Most existing algorithms requires discretization of the underlying space and are thus not scalable. The present work starts from the JKO scheme, an implicit scheme, for Wasserstein gradient flow, and utilizes ICNNs to realize a JKO step. The objective free energy is also estimated efficiently using samples as in Theorem 1. The combination of these two techniques is new. A related work on particle-based Fokker-planck is [1].\n\nQuality:\nThe main results on the Wasserstein gradient flow of the Fokker-planck free energy are technically sound. \n1. Since this gradient flow is mainly for sampling purpose, the authors can include more comprehensive comparison with more advanced sampling algorithms other than overdamped Langevin (EM). \n2. There is only one example (4.2) to illustrate the temporal behavior of the gradient flow and only two time points (t=0.5, 0.9) are evaluated. More examples are needed.\n3. The claim ``our method admits straightforward generalization to any F that ….,’’ should either be removed or more properly explained. The proposed algorithm does use some special property of the Fokker-planck free energy that cannot be generalized.\n4. The results of dual JKO in figure 4 seems to be in consistent with [2].\n5. The algorithm seems to have larger variance than EM, especially for low dimensional problems (Figure 1, 3, 4).\n\nClarity:\nThe paper is over-all well-written. \n1. The experiments can be improved by adding more details. For instance, how is Accuracy in Table 1 defined?\n2. The notation T is used in Theorem 1 and Section 5 for different purposes.\n3. It should be argmax in (11) and the following equations\n\nSignificance:\nThe computation of Wasserstein gradient flow is an important problem. The proposed algorithm provides a possible way to calculate Wasserstein gradient flow in high dimensional setting. One limitation of this algorithm is that it is only applicable to the Fokker-planck free energy.\n\n[1] K.F. Caluya et al, Proximal recursion for solving the Fokker-Planck equation\n[2] C. Frogner et al, Approximate inference with Wasserstein gradient flows", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations and societal impact are properly discussed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "OhQZ5mDm9rP", "reviewer": "Reviewer_JdcD", "summary": "In this work the authors introduce an innovative discretisation of the Jordan Kinderleher Otto scheme (which describes the evolution of the gradient flow of relative entropy in the Wasserstein metric).  To do so the authors iteratively construct a transformation map as a composition of gradients of input-convex neural networks, which satisfy a universal approximation theory among convex functions.  Using Brenier's theorem, an iterative variational formulation for the density at discrete time-steps is obtained, from which one can generate samples and or evaluate an unnormalised density.\n\nA variety of applications are presented, including finding the stationary distribution, as well as nonlinear filtering.   ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This submission brings two important ideas together -- the idea of using input-convex neural networks with Brenier's theorem to formulate a variational problem for the density.   While straightforward, it is a nice idea which provides a means of evaluating the density of an overdamped Langevin dynamics which overcomes the curse of dimensionality associated with spatial discretisation approaches. \n\nWhile I find the idea interesting and the submission is clearly written, I would argue that it provides a very elegant solution to a quite narrow problem.   Evaluating the density of an SDE beyond the usual overdamped langevin SDE is not supported (granted one can adjust the metric tensor to accomodate some cases), in particular, non-ergodic diffusions cannot be approached via this formulation.\n\nSecondly, for problems which involve stationary sampling as presented in the examples, we are severely limited by the various costs, mainly the quadratic cost in time-step which forbids long time -simulations.\n\nFinally, it's worth remembering that this is still a discretisation in time, and it is not clear from this manuscript whether discretisation error induced by finite h has really been considered or studied.  I feel that this should have been addressed more carefully in the examples.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "this has been adequately addressed", "review_text": "", "extra_scores": {"ethical_concerns": "none", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work focuses on a new method to approximate Wasserstein gradient flows, allowing to sample from and compute the density at each time step. The use of ICNN to play the role of the convex potential allows to discretize the time dimension efficiently. Based on the overall evaluation by the reviewers - which is mostly positive - and my own, I recommend to accept this work.

There was a serious and fruitful discussion between reviewers and authors, the latter providing interesting perspective on the motivation of their work, that should be included in a camera-ready version. 
Decision: Accept (Poster)