Variational Wasserstein gradient flow
Authors: Jiaojiao Fan
Amirhossein Taghvaei
Yongxin Chen
Venue: ICLR 2022 Submitted (2022)
Keywords: Wasserstein gradient flow
JKO
f-divergence

Abstract: The gradient flow of a function over the space of probability densities with respect to the Wasserstein metric often exhibits nice properties and has been utilized in several machine learning applications. The standard approach to compute the Wasserstein gradient flow is the finite difference which discretizes the underlying space over a grid, and is not scalable. In this work, we propose a scalable proximal gradient type algorithm for Wasserstein gradient flow. The key of our method is a variational formulation of the objective function, which makes it possible to realize the JKO proximal map through a primal-dual optimization. This primal-dual problem can be efficiently solved by alternatively updating the parameters in the inner and outer loops. Our framework covers all the classical Wasserstein gradient flows including the heat equation and the porous medium equation. We demonstrate the performance and scalability of our algorithm with several numerical examples.

Under review as a conference paper at ICLR 2022
VARIATIONAL WASSERSTEIN GRADIENT FLOW
Anonymous authors
Paper under double-blind review
ABSTRACT
The gradient flow of function over the space of probability densities with respect to the Wasserstein metric often exhibits nice properties and has been utilized in several machine learning applications. The standard approach to compute the Wasserstein gradient flow is the finite difference which discretizes the underlying space over a grid, and is not scalable. In this work, we propose a scalable proximal gradient type algorithm for Wasserstein gradient flow. The key of our method is a variational formulation of the objective function, which makes it possible to realize the JKO proximal map through primal-dual optimization. This primal-dual problem can be efficiently solved by alternatively updating the parameters in the inner and outer loops. Our framework covers all the classical Wasserstein gradient flows including the heat equation and the porous medium equation. We demonstrate the performance and scalability of our algorithm with several numerical examples.
INTRODUCTION
The Wasserstein gradient flow models the gradient dynamies the space of probability densities with respect to the Wasserstein metric. It was first discovered by Jordan, Kinderlehrer, and Otto (JKO) in their seminal work (Jordan et al., 1998). They pointed out that the Fokker-Planck equation is in fact the Wasserstein gradient flow of the free energy, bringing tremendous physical insights to this type of partial differential equations (PDEs). Since then, the Wasserstein gradient flow has played an important role in optimal transport, PDEs, physics, machine learning, and many other areas Santambrogio 2017 Carlier et al. 2017
(Ambrosio al. 2008 Otto 2001: Adams et al.
Frogner & Poggio 2020
Despite the abundant theoretical results on the Wasserstein gradient flow established over the past decades (Ambrosio et al., 2008 Santambrogio 2017), the computation of it remains a challenge.
Most existing methods are either based on finite difference of the underlying PDEs or based on finite dimensional optimization: both require discretization of the underlying space (Peyre) 2015 Carlier et al. 20177; Li et al. 2020, Carrillo et al., 202T). The computational complexity of these methods scales exponentially the problem dimension, making them unsuitable for the cases where probability densities over high dimensional spaces are involved.
Our goal is to develop scalable method to compute the Wasserstein gradient flow without discretizing the underlying space. One target application we are specifically interested in is optimization over the space of probability densities. Many problems such as variational inference can be viewed as special cases of such optimization. We aim to establish a method for this type of optimization that is applicable large class of objective functions.
Our algorithm is based on the JKO scheme (Jordan et al. 1998), which is essentially a backward Euler time discretization method for the continuous time Wasserstein gradient flow. In each step of JKO scheme, one needs to find a probability density that minimizes a weighted sum of the Wasserstein distance (square) to the probability density at the previous step and the objective function. We reparametrize this problem in each step so that the optimization variable becomes the optimal transport map from the probability density at the previous step and the one we want to optimize, recasting the problem into stochastic optimization framework. This transport map can either be modeled by a standard feedback forward network or the gradient of an input convex neural network. The latter is justified by the fact that the optimal transport map for the optimal transport problem with quadratic cost with any marginals is the gradient of a convex function. Another crucial ingredient of our algorithm is variational form of the objective function, which allows the evaluation of the

---

Under review as a conference paper at ICLR 2022
objective with samples and without density estimation. At the end of the algorithm, sequence of transport maps connecting an initial distribution and the target distribution are obtained. One can then sample from the target distribution by sampling from the initial distribution (often Gaussian) and then propagating these particles through the sequence of transport maps. When the transport map modeled by the gradient of input convex neural network, can evaluate the target density at every point.
Our contributions can be summarized as follows.
i). We develop a neural network based algorithm to compute Wasserstein gradient flow without spatial discretization. Our algorithm applicable any objective function that has variational representation.
ii). We specialize our algorithm to three important cases where the objective functions are the Kullback-L Libler divergence, the generalized entropy, and the interaction energy. iii). We apply our algorithm several representative problems including sampling and aggregation- diffusion equation and obtain respectable performance.
Related works: Most existing methods to compute Wasserstein gradient flow are finite difference
based (Peyré) 20155] Carlier et al.) 2017: Li et al. 2020 Carrillo et al. (02T). These methods require spatial discretization and are thus not scalable to high dimensional settings. [Salim et al. (2020) analyze the convergence for
scheme but leave the implementation of JKO an open question. There is a line of research that uses particle-based method to estimate the Wasserstein gradient flow (Carrillo al., 2019b Frogner & 2020. In these algorithms, the current density value is often estimated using kernel method whose complexity scales at least quadratically with the number of particles. More recently, three interesting neural network based methods (Mokrov et al.)
202T; Alvarez-Melis et al., 2021 Yang et al. (0200) were proposed for Wasserstein gradient flow.
The first one (Mokrov et al.) 2021) focuses on the special case with Kullback-Leible divergence as
objective function. The second one (Alvarez- Melis et al. 2021) uses density estimation method to
evaluate the objective function by back-propagating to the initial distribution, which could become computational burden when the number of time discretization is large. The third one (Yang et al.) 2020) is based on a forward Euler time discretization of the Wasserstein gradient flow and is more sensitive to time stepsize. Over the past few years, many neural network based algorithms have been proposed to compute optimal transport map or Wasserstein barycenter (Makkuva et al. 2020 Korotin et al. 2019: Fan et al.) [2020] Korotin et al., (2021). These can be viewed as special cases of Wasserstein gradient flows or optimizations over the space of probability densities.
2 BACKGROUND
2.1 OPTIMAL TRANSPORT AND WASSERSTEIN DISTANCE
Given two probability distributions P. P.Q over the Euclidean space R" with finite second moments, the optimal transport problem with quadratic cost reads
min T T-T:P=Q / T(x)
where the minimization is over all the feasible transport maps that transport mass from distribution P distribution Q. The feasibility is characterized by the pushforward operator (Bogachev) (007) as TiP = Q. When the initial distribution P admits density, the above optimal transport problem has unique solution and it the gradient of a convex function, that is,
Vip
for some convex function sp.) : R² + R. In this paper, we assume probability measures admit densities and use the notation for the measure and the density interchangeably.
The square-root of the minimum transport cost, namely, the minimum of TT. defines a metric on the space of probability distributions known the Wasserstein-2 distance (Villani) 2003), denoted by (P,, Q). The Wasserstein distance has many nice geometrical properties compared with other distances such T.2 distance for probability distributions, making it popular choice in applications.
2

---

Under review as a conference paper at ICLR 2022
2.2 WASSERSTEIN GRADIENT FLOW
Given function F(P) over the space of probability densities, the Wasserstein gradient flow describes the dynamics of the probability density when it follows the steepest descent direction of the function F(P) with respect to the Wasserstein metric W2. The Wasserstein gradient flow can be explicitly represented by the PDE
OP 01
6F PV
-
8P
where / stands for the gradient of the function with respect to the standard L2 metric (Villani 2003 Ch. 8)
Many important PDEs are the Wasserstein gradient flow for minimizing certain objective functions
F(P). For instance, when F is the free energy F(P) f
P(x) log P(x)dr + JR.. V(x) ((x))(
the gradient flow is the Fokker-Plane -lankkk- equation (Jordan et al.
1998)
OP V (PVV) AP. When
F is the generalized entropy F(P) m-1 fg pm dr for some positive number 1, the
gradient flow is the porous medium equation Otto
2001
Vázquez 2007 aP APm
3 METHODS AND ALGORITHMS
We are interested in solving the optimization problem
F(P)
(2)
the space of probability densities P(RR). In particular, our objective is develop particle- based Wasserstein gradient flow algorithm to numerically solve (2).
The objective function F(P) could exhibit different form depending on the application. this paper, we present our algorithm for the linear combination of the following three important cases:
Case The functional is equal the KL-divergence with respect to a given target distribution Q
F(P) D(PIIQ) : : log
P(x)
P(r)dr.
(3)
(x))
This is important for the problem of sampling from a target distribution.
Case II The objective functional is equal the generalized entropy
F(P)= = CP) = m - pm (#)dz
This case is important for modeling the porous medium.
Case III The objective functional equal to the interaction energy
F(P) = (PP)=((P)--) = f / ((x - P()) W : R" + R.
This case is important for modeling the aggregation equation.
These functionals have been widely studied in the Wasserstein gradient flow literature (Carlier et al. 2017: Santambrogio 2017; Ambrosio et al., (2008) due their desirable properties. It can be shown that if F(P) is composed by the above functionals, under proper assumptions, Wasserstein gradient flow associated with F(P) converges to the unique solution (2) (Santambrogio 2017).
we first assume F(P) doesn't 't include interaction energy, and introduce JKO/backward scheme to solve (2). We then add W(P) into consideration and present a forward-
In Section 3.1
3.2
backward scheme in Section 3.3 and close by showing our Algorithm in Section 3.4.
3.1 JKO SCHEME AND REPARAMETRIZATION
To realize the Wasserstein gradient flow, a discretization over time is needed. One such discretization
is the famous JKO scheme (Jordan et al. 1998)
P Pk+1 arg arg min p 2a
W
(4)
3

---

Under review as a conference paper at ICLR 2022
This is essentially backward Euler discretization proximal point method with respect to the Wasserstein metric. The solution to (4) converges to the continuous-timg time Wassrstein gradient flow when the step size + O.
In our method, we reparametrize as an optimization terms of the transport maps T: R + IR
from Ps to P, i.e., by defining P TiP. With this reparametrization, in view of the definition of Wasserstein distance [ the JKO step 4 becomes
1 2a
1
/ J
Pk+1 Trip. Th arg min
|| (T)))
(5)
T
The optimal T is the optimal transport map from Ps: TiP; and is thus the gradient of convex function sp Therefore, the JKO scheme can be also expressed as
Pa+1 VypuPka PP arg min 1
| R
(c(())[BdP<())) F(V()Pk).
PECVX
2a
where CVX stands for the space of convex functions. We use the preceding two schemes (5) and (6) in our numerical method depending on the application.
3.2 D(PIIQ) AND (PP) REFORMULATION WITH VARIATIONAL. FORMULA
The main challenge in implementing the JKO scheme is to evaluate the functional F(P) in terms of samples from P. We achieve this goal by using variational formulation of F. In order to do so, use the notion of f-divergence between the two distributions P and Q:
(y()) = Ee f (o)
where (0, + R is a convex function. The f-divergence admits the variational formulation
Dy(1)(Q)= sup -B (h(Y))
where f*(() sup bzy f(x)] the convex conjugate of f. The variational form has the special feature that does not involve the density of P and Q explicitly and can be approximated in terms of samples from P and Q The functionals ((Pl)) and ((P) can both be expressed as f-divergence.
With the help of the f-divergence variational formula, when F(P) D(PIIQ) or ((P), the JKO scheme (5) can be equivalently expressed as
Pa+1 = -THPk Te arg min
[ V(T,))
T(X)
T
2a
where V(T, Ep [A(T,))) Er [(h),, is user designed distribution which is easy sample from, and A and B are functionals whose form depends on the functional F. The form of these two functionals for the KL divergence and the generalized entropy appears in Table The details appear in Section 3.2. and B.2.2
Table I: Variational formula for D(PI(Q) and ((P)
Energy function
A(T
B8hh)
S P Plog(P D//)de
log ((T + log (T)) log Q(T)
Gaussian distribution μ
pm do
(h(T))
m-1
hm
Uniform distribution Q
Qm-T
I
3.2.1 KL DIVERGENCE
The KL divergence is the special instance of the f-divergence obtained by replacing of with fi(z) =
log in 7
P
P
D = Eg
log
= Ep
log
Q
Q
4

---

Under review as a conference paper at ICLR 2022
Proposition 1. The variational formulation f for D(PIIQ) reads
((X)
(a(P))Q) 1+ sup
-up Ep log (x)) log
Ex [k(2)],
Q(X)
where is a user designed distribution which is easy to sample from. The optimal function h is equal the ratio between the densities of THr and p.
The proof for Proposition D can be found in appendix A It becomes practical when we have only access to un-normalized density of Q, which is the case for the sampling problem. Using this T and
variational form in the JKO scheme (5) yields Pk+1
1 X -
(T())) log log + log Q(T(X)) LE,, (10)
TB arg min max Ex
T
2a
In practice, we choose l = l4 adaptively, where A is the Gaussian with the same mean and covariance Pk. We noticed that this choice improves the numerical stability of the the algorithm.
3.2.2 POROUS MEDIUM EQUATION
The generalized entropy can be also represented f-divergence. In particular, let f2(x) m- (x I) and let Q be the uniform distribution on a set which is the superset of the support of density P(x) and has volume Ω. Then
mm (pl))) -
/ (r)dd -
Proposition The variational formulation for (P)) reads
G(P) Qm h
sup m
E.r.
1m-1- hmm (x)
RQ 11T (h("(()) (2))
(11)
The optimal function h is equal to the ratio between the densities of TtPP and Q.
The proof for Proposition 2 is postponed to appendix
A
Using this in the JKO scheme yields
Pk+1 TatPks and
Tg arg min max
Ep IIX 2(X)1
Ep.
hm
(x)
Eq [h (mm(2))
2a
2
m
where Qs the volume of set large enough to contain the support of TiPx for any T that not too far away from the identity map.
Algorithm 1 Primal-dual gradient flow
Input: Objective function F(P), initial distribution Po, step size a, number of JKO steps K, number of outer loop J1, number of inner loop J2, batch size M.
Initialization: Parameterized To and hA
for k = 1.2..... KK do
Ps < (I <V Pk)) if F(P) includes W(P) (// use last iteration T--1 as a warm-up)
To T2-1ff
for j1 = 1,,, J1 do
Sample Y YM from Ps. Sample Zi,..., . ZM from r.
for j2 = .,,,.... J2 do
Apply Adam to À to maximize end for
M ΣM 11 [A(Tg, (X()) - B(hx
Apply Adam to to minimize
V - 2 + (xx)))
Za
end for Th + To end for
Output: (Tx K
5

---

Under review as a conference paper at ICLR 2022
3.3 FORWARD BACKWARD (FB) SCHEME
When F(P) involves the interaction energy W(P), we add an additional forward step to solve the gradient flow:
Pk+] (I a (W W Pa))
(12)
Pk+1 == T + +
(13)
where / the identity map, and T defined by replacing k by le + in 9 other words, the first gradient descent step (12) is forward discretization of the gradient flow and the second JKO step v)), thus can also be approximated by samples. (alim et al. (2020) firstly propose this method to solve Wasserstein gradient flow and provide the theoretical convergence analysis. We make this scheme practical by giving a scalable implementation of JKO.
(13) is backward discretization. V.(W P) can be written as expectation V
Since W(P) be equivalently written expectation Ez [W(x y)]. there exists another con-forrward-backwardd (non-FB) method, i.e., removing the first step and integrating W(P) into a single JKO step: PE+1 TetPs and
F. | T(X)II + ExyY (W((())))((()----- (TV)) sup V(T, V(T,h) m 2a
Th arg min
T
In practice, observe the FB scheme is stable however converge slower than non-FB scheme.
The detailed discussion appears in the Appendix B.2 B33
Remark 1. In principle, one single oud log Q) term from (T)) and perform similar forward
step Pk+ = (T al (VVV)//Q)) P/ (Salim et al. 2020), but we don't I observe improved performance
of doing this in sampling task.
3.4 PRIMAL-DUAL. ALGORITHM AND PARAMETRIZATION OF T AND h
The two optimization variables T' and h in our minimax formulation ( can be both parameterized by neural networks, denoted by To and h2- With this neural network parametrization, we can then solve the problem by iteratively updating To and h2 This prime-dual method to solve (2) is depicted in Algorithm
In this work, we implemented two different architectures for the map T. One way is to use a neural network to represent 7 directly, and another way is to parametrize T as the gradient of a Input convex neural network (ICNN) (Amos et al. (017) yo. The latter has been widely used in optimal transport (Makkuva et al. 2020 Fan et al. 2020 Korotin et al., 2021). In experiments, we find that the first parametcrization gives better result in the sampling application. As we discuss in Section [3. when density evaluation is needed, adopt the ICNN parameterization since need to compute T-1 Note that ICNN could modified to be strictly convex and if the function is strictly convex, the gradient Vip is invertible.
3.5 EVALUATION OF THE DENSITY
In this section, we assume the solving process doesn't 't use forward-backward scheme, i.e. all the probability measures Ps. are obtained by performing JKO one by one. Otherwise, the map I- a- Px) I Ey (W((xx))) includes an expectation term and becomes intractable to push- backwad particles to compute density.
If T' is invertible, these exists a standard approach to evaluate the density of Pa Alvarez-Melin al. 2021 Mokrov al. 202T) through the change of variables formula. More specifically. we assume Th, is parameterized by the gradient of an ICNN sp that is assumed to be strictly convex. To evaluate the density Ps (xx) at point Ik, we back propagate through the sequence of maps Th = Ve Vipi to get
T T
T (œx)-
The inverse map T (Vy,) -1
V can be obtained by solving the convex optimization
Tj- arg max (x, 2.)) (x).
(14)
zER

---

Under review as a conference paper at ICLR 2022
Then, by the change of variables formula, obtain
log[P((xx)] = [og[[[((0]]
log V ((t-1)))
(15)
where (3 is the Hessian of (o; and
is its determinant. By iteratively solving and plugging the resulting æy into (15). we can recover the density Ps (x(xx) at any point.
114
3.6
COMPUTATIONAL COMPLEXITY
Per each update k in Algorithm 1 the forward step 12) requires at most ( where N is the total
number particles to push-forward. The backward step (3) requires O kM where J1 the number of iterations per each JKO step, M is the batch size, and H is the size of the network. k shows up in the bound because sampling Px requires us to pushforward zy ~ Po through
1 maps.
However, both Mokrov et al. (2021) and Alvarez-Melis et al. (2021) require O(J1k (nMH + n²))
which has the cubic dependence dimension because they need to query the log det V in each iteration. We refer to Mokrov et al. (2021) Section 5) for the complexity details of calculating the Hessian term. Thus method has the advantage of independence on the dimension.
We provide training time details Section and Appendix [8.4] Other than training and sampling time, the complexity for evaluating the density are the same the above two methods due to the standard density evaluation process (see Section 3.5)
4 NUMERICAL EXAMPLES
4.1 SAMPLING
We first consider the sampling problem to sample from a target distribution Q. Note that Q doesn't 't have to be normalized. To this end, we consider the Wasserstein gra- dient flow with objective function F(P)= P Ploo((P/))ce that is, the KL divergence between distributions P and Q. When this
(a) ground truth
(b) ours
The left figure shows samples from the tar- get 16- 6--GMM distribution and the right figure shows samples obtained by our method. Each plot contains 4000 points.
Figure 1:
objective is minimized, P x Q. In our experiments, we consider two types of target distribution: the two moons distribution and the Gaussian mixture model (GMM) with spherical Gaussian components. In this set of experiments, the step size is set to be 0.3 and the initial measure is a spherical Gaussian (0, 2.25
Two moons: The two moons distribution is popular target distribution for sampling task. It is 2D mixture model composed of 16 Gaussian components; each moon shape consists of 8 Gaussian components. The results are displayed in Figure generate samples that match the target distribution.
I from which we see that our method is able to
Stationary measure
Fitted measure lours)
Stationary measure
Fitted messure (ours)
(a) Dimension n.
(b) Dimension n 13
Figure 2: Comparison between the target GMM and fitted measure of generated samples by method. Samples are projected onto 2D plane by performing PCA. We refer the reader to Mokrov et al. 2021) for the performance of another algorithm in similar setup.
7

---

Under review as a conference paper at ICLR 2022
GMM with spherical Gaussians: We also test algorithm in sampling from GMM in higher dimensional space. The target GMM has 9 Gaussian components with equal weights and the same covariances. The results with dimension we not only display the samples as grey dots the plot, but also the kernel density estimation of generated samples level sets. As can be seen from the results, both the samples and densities obtained with our algorithm match the target distribution well.
8 and 13 depicted in Figure 2 In Figure 2
4.2 ORNSTEIN-UILENBERR PROCESS
We study the performance of our method in modeling the Orastein-Uhlenbeck Process as dimension grows. The gradient flow is affiliated with the free energy (3), where Q = 3 A(=-)((=)))))) with positive definite matrix the gradient flow at each time is Gaussian distribution P. with mean vector 14; (7n A
R R and E R" Given an initial Gaussian distribution (O, I,).
and covariance Σ, - (I, e -2At) + -2Ar (Vatiwutipong Phewchean 2019). We calculate P& with JKO step size = 0.05 and compare with the Fokker-Plance (FP) JKO (Mokrov et al. 2021).
We quantify the error as the SymKL divergence between estimated distribution and the ground truth in Figure where
SymKL(PI, P2) = D(P)IIB) - Pt)-
Fokkeer Planck
Fokiler Flankk
1.5
1.5
-2.0
1og, Symkt
SynScL
-3.0
-3.5
-4.
15 D. dimenssio
25 D. dimeension
(a) Time t = 0.5
(b) Time t = 0.9
Figure 3: We repeat the experiments for 15 times dimensions d = 2, 12, 16,32.
We also compare the training time per every two JKO steps with FP JKO. The computation time for FP JKO is around 20s when d = 2 and increases 100s when d = 32. Our method's training time remains at 20s + 2s for all the dimensions d = 2 ~ 32. This is due to we fix the neural network size for both methods and our method's computation complexity doesn't depend the dimension.
4.3 POROUS MEDIA EQUATION
1.0
180
(a) 0.002
(b) 0.004
(c) f 0.006
(d) 0.008
Figure 4: Comparison among exact density, finite difference method solution given by CVXOPT, and the density given by our method. To better visualize the distributed particles from each distribution, we also plot the histograms of our method as the blue shadow.
We next consider the porous media equation with only diffusion: 0,P APm, This is the Wasser- stein gradient flow associated with the energy function F(P) n pm (x:)dr. representative
8

---

Under review as a conference paper at ICLR 2022
closed-form solution of the porous media equation is the Barenblatt profile (Gl 1952 V Vázquez 2007
2)) 10) (c
3 (x xo)² (o)
20//
n
(m a)a
where a n(m +)
B
2mn
and to > is the starting time and > 0 a free parameter.
In principle, algorithm should match the analytical solution Pye when the step size a is sufficiently small. When a is not that small, time discretization is inevitable. To account for the time-iiscrtizzation error of JKO scheme, consider the porous media equation in 1D space and use the solution via finite difference method as a reference. The details appear in appendix C.33
In the experiments, we set the stepsize for the JKO scheme to be 0.001 and the initial time to be to = 0.002. Other parameters are chosen C' (3/16) m 2, 1, 300. We parametrize the transport map T as the gradient an ICNN and thus we can evaluate the density following Section 3.5 In Figure 4 observe that the gap between the density computed using algorithm and the ground truth density Pgt is dominated by the time-discretizatio error of the JKO scheme. Our result matches the discrete time solution nicely.
4.4 EGGREGATION-DIIFUSSE EQUATION
(a) k = 24
(b) l: = 36
(c) k = 60
(d) k 84
(e) l: = 92
Figure Histogram for simulated measures P& by FB scheme different
We finally simulate the evolution of solutions to the following aggregation-diffuee diffuusion equation:
0,P = V (PVW * P) + 0.1AP" W(x) - -----l 1=1l? I=lll//T. /T.
This corresponds to the energy function W(P) + 0.1)(P). We use the same parameters in Carrillo uniform distribution supported on |-3, x 3,3) and the JKO step size a = 0.5. In Figure 5, we utilize FB scheme to simulate the gradient flow for this equation with m on TR.2 space. With this choice Wwx), V. Pk) is equal to
et al. (2021, Section 4.3.3). The initial distribution
Ey-Pa 2e | /T]] in the gradient descent step
And estimate V (W Pk) with 10*
samples from Pk-
Throughout the process, the aggregation term (PVPW PP) and the diffusion 0.1AP" adversarially exert their effects and cause the probability measure split to four pulses and converge to a single pulse in the end (Carrillo et al., 2019a).
CONCLUSION
In this paper we presented a novel neural network based algorithm to compute the Wasserstein gradient flow. Our algorithm follows the JKO time discretization scheme. We reparametrize the problem so that the optimization variable becomes the transport map T between the consecutive steps. By utilizing a variational formula the objective function, we further reformulate the problem in every step as min-max problem over map T' and dual function / respectively. This formulation doesn't 't require density estimation using samples and can be optimized using stochastic optimization. It also shows advantages with be extended to minimize other objective functions that can be written as f-divergence. Our limitation the accuracy is not satisfying in sampling tasks with high dimension complex density.
computation complexity. Our method also

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY
The basic setup distributed Section 4 We refer to Section a.d Appendix C for all the rest training details. The code of our method is also attached in the supplementary material.
10

---

Under review as a conference paper at ICLR 2022
REFERENCES
Stefan Adams, Nicolas Dir, Mark Peletier, and Johannes Zimmer. From a large-deviations principle the Wasserstein gradient flow: a new micro-macro passage. Communications in Mathematical Physics, 307(3):711-815, 2011.
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, 2008.
Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning, pp. 146-155, PMLR, 2017.
Vladimir Bogachev. Measure theory, volume 1. Springer Science & Business Media, 2007.
Guillaume Carlier, Vincent Duval, Gabriel Peyré, and Bernhard Schmitzer. Convergence of entropic schemes for optimal transport and gradient flows, SIAM Journal on Mathematical Analysis, 49(2): 1385-1418, 2017.
Jose A Carrillo, Sabine Hittmeir, Bruno Volzone, and Yao Yao. Nonlinear aggregation-diffusil equations: radial symmetry and long time asymptotics. Inventiones mathematicae, 218(3):889-977, 2019a.
Jose A Carrillo, Katy Craig, Li Wang, and Chaozhen Wei. Primal dual methods for Wasserstein gradient flows. Foundations of Computational Mathematics, pp. 1-55, 2021.
José Antonio Carrillo, Katy Craig, and Francesco S Patacchini. A blob method for diffusion. Calculus of Variations and Partial Differential Equations, 58(2):1-53, 2019b.
William Falcon and Kyunghyun Cho. framework for contrastive self-supervised learning and designing a new approach. arXiv preprint arXiv:2009.001/ 2020.
Jiaojiao Fan, Amirhossein Taghvaei, and Yongxin Chen. Scalable computations of Wasserstein barycenter via input convex neural networks. arXiv preprint arXiv:2007.04462, 2020.
Charlie Frogner and Tomaso Poggio. Approximate inference with Wasserstein gradient flows. In International Conference on Artificial Intelligence and Statistics, pp. 2581-2590. PMLR, 2020.
Barenblatt GI. On some unsteady motions of a liquid and gas in porous medium. Prikl. Mat. Mekh., 16:67-78, 1952.
Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the Fokker-Planck equation. SIAM journal on mathematical analysis, 2911:1111, 1998.
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. arXiv preprint 1rXiv: 13082, 2019.
Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous Wasserstein-2 barycenter estimation without minimax optimization. arXiv preprint 11X0v:2102.01752, 2021.
Wuchen Li, Jianfeng Lu, and Li Wang. Fisher information regularization schemes for Wasserstein gradient flows. Journal of Computational Physics, 416:109449; 2020.
Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672-6681. PMLR, 2020.
Petr Mokrov, Alexander Korotin, Lingxiao Li, Aude Genevay, Justin Solomon, and Evgeny Burnaev. Large-scale Wasserstein gradient flows. arXiv preprint 1rXiv=2106.00736, 2021.
Felix Otto. The geometry of dissipative evolution equations: the porous medium equation. 2001.
Gabriel Peyré. Entropic approximation of Wasserstein gradient flows. SIAM Journal on Imaging Sciences, (44:2333-2355) 2015.
11

---

Under review as a conference paper at ICLR 2022
Adil Salim, Anna Korba, and Giulia Luise. The Wasserstein proximal gradient algorithm. arXiv preprint arXiv:2002.030:5,, 2020.
Filippo Santambrogio. Euclidean, metric, and wasserstein gradient flows: an overview. Bulletin of Mathematical Sciences, 7(1):87-154, 2017.
Vatiwutipong and N Phewchean. Altemntive way derive the distribution of the multivariate ornstein-uhlentee process. Advances in Difference Equations, 2019(1):1-7, 2019.
Juan Luis Vázquez. The porous medium equation: mathematical theory. Oxford University Press on Demand, 2007.
Cédric Villani. Topics in optimal transportation. Number 58. American Mathematical Soc., 2003.
Zhuoran Yang. Yufeng Zhang, Yongxin Chen, and Zhaoran Wang. Variational transport: A convergent particle-base algorithm for distributional optimization. arXiv preprint arXiv:2012. 2020.
12

---

Under review as a conference paper at ICLR 2022
A DETAILS ABOUT VARIATIONAL FORMULA IN SECTION 3.2
A.1 KL DIVERGENCE
The KL divergence is the special instance of the f-divergence obtained by replacing with fi(x)
x log x in
P
P
P Q
D₁, (PI()) EQq
log Ep log
Q
Q
which, according to ( admits the variational formulation
(pl()) = 1 sup [h(x)] - BQ (
(16)
h
where the convex conjugate (i e is used.
The variational formulation can be approximated in terms of samples from P and Q. For the case where have only access un-normalized density of Q, which is the case for the sampling problem, we use the following change of variable: h - log(h) + log(p) - log(Q) where is user designed distribution which easy sample from. Under such change of variable, the variational formulation reads
(PI()) sup Ep log h(X) log (xX)
EE [I(2))
Q(x)
Note that the optimal function ht is equal the ratio between the densities of TFPss and μ.
Remark 2. The Donsker- Varadhan formula
D(PII)) = sup Ep [M(X)] - log Ee (h(22)
is another variational representation of KL divergence and 's stronger than (16) because it's upper bound of (16) for any fixed h. However, we cannot get an unbiased estimation of the objective using samples.
A.2 POROUS MEDIUM EQUATION
The generalized entropy can be also represented as f-divergence. particular, let f2(x) m (x I) and let Q be the uniform distribution on set which is the superset of the support of density P(x)) and has volume Ω. Then
m
/ pm (x)dæ
D₂,
m
m
As result, the generalized entropy can be expressed in terms of diivergencce according
/
1
G(P)
pm (x)dz
Q 1/2 + fTT 1(m 1)
Upon using the variational representation of the f-divergence with
(z = (m 1)y +
the generalized entropy admits the following variational formulation
Eq (
(m 1)A(Y) +
)
G(P) sup h
Er[h(X)] -E
+
I
Q Qm-1 (m
1)
m-
In practice, find it numerically useful let h
m
h
so that
m
G(P) = Qm
sup h
E.
m
7 m 1x))
- EEQ [h. (2)))
m
With such change of variable, the optimal function h THPP//
13

---

Under review as a conference paper at ICLR 2022
B ADDITIONAL EXPERIMENT RESULTS AND DISCUSSIONS
B.1 SAMPLING USING ICNN PARAMETERIZATION
Stationary meesure
Fitted measure lours)
Stationary measure
Fitted messure (ours)
(a) Dimension n
(b) Dimension n 13
Figure 6: Sampling Gaussian mixture models by parameterizing the map by V..
In Figure 6 we present the sampling results with Vi parameterized map where op a ICNN neural network. The experiment setting is the same as Section 4.1 and we can observe a MLP network map gives better fitted measures.
B.2 AGGREGATION EQUATION
Alvarez-Melis et al. 2 proposes using the neural network based JKO, i.e. the backward method, to solve (17). They parameterize T as the gradient of the ICNN. In this section, we use two cases to compare the forward method and backward when F(P) W(P). This could help explain the FB
and non-FB scheme performance difference later in Section B.3
We study the gradient flow associated with the aggregation equation
(PVW P), W"
(17)
The forward method is
Pe+ = (I aW PR))
The backward method JKO is
E -(()) Ex (((()) - -7Y))) 2a
Pa+1 = TatPe
TA arg min T
Example We follow the setting in Carrillo et 2022 Section 4.3.1 The interaction kernel is W(x)
|=||" ||] and the initial measure Po is a Gaussian (0, 0.254). In this case, (. PP))
becomes Ey-P, [[]| y/ - 1) v)] We use step size 0.05 for both methods and show
the results Figure
Example 2 We follow the setting in Carrillo et al. 2202 Section 4.2.3 The interaction kernel is W(x)
| In ||=||, and the initial measure Po is (( 1). The unique steady state for this case
2
(o( (22 (2 - x²)
The reader can refer to Alvareeee-sMlss al. 12021 Section 5.3) for the backward method performance. As for the forward method, (W Px) becomes Ey-rr, y z-9 Because the kernel W enforces repulsion near the origin and Po is concentrated around origin, V. (W PP will easily blow up. So the forward method is not suitable for this kind interaction kernel.
Through the above two examples, if V (W * P)) is smooth, we can notice the backward method converges faster, but is not stable when solving (17). This shed light the FB and non-FB scheme performance in Section 4.4 B.3 However, (W * P) has bad modality such as Example 2, the forward method loses the competitivity.
14

---

Under review as a conference paper at ICLR 2022
(a) Forward method k 23, t 1.15
(b) Forward method & 200, / = 10
(c) Backward method 23, 1.15
(d) Backward method k 40,
Figure 7: The steady state is supported ring of radius 0.5. Backward converges faster to the steady rate but unstable. As goes large, it cannot keep the regular ring shape and will collapse after h >> 50.
B.3
EQUATION WITH NON-FB SCHEME
In Figure we show the non-FB solutions to Aggregation-diin -iififuoin-ifiionnnn equation in Section 4.4 FB scheme should be independent with the implementation of JKO, but in the following context, assume FB and non-FB are both neural network based methods discussed in Section Non-FB scheme reads
P&+1 TrPPk
1 2a
T3 arg min
Ex.-a,,
where G(T,h) h) is represented by the variational formula other PDE parameters as in Section 4.4
We use the same step size a = 0.5 and
(a) l 18
(b) k = 24
(c) = 30
(d) k = 42
Figure 8: Histograms for simulated measures Ps by non-FB scheme at different k.
Comparing the FB scheme results in Figure 5 and the non-FB scheme results in Figure we observe non-FB converges 1.5x slower than the finite difference method (Carrillo et al.) 2021), and FB converges 3x slower than the finite difference method. This may because splitting one JKO step to the forward-backwand term is too weak to make a difference in the loss. Note at the first several l. both Ps. and Q are nearly the same uniform distributions, so h is nearly constant and T(x) exerts little effect in the variational formula of ((P). Another possible reason a single forward step for aggregation term converges slower than integrating aggregation in the backward step, as discuss in Section B.2 and Figure
two steps removes the aggregation term effect in the JKO, and the diffusion
15

---

Under review as a conference paper at ICLR 2022
However, FB generates regular measures. We tell the four pulses given by FB more symmetric. We speculate this is because gradient descent step in FB utilizes the geometric structure of W (a) directly, but integrating W(P) in neural network based JKO losses the geometric meaning of W (a))
B.4 COMPUTATIONAL TIME
Our experiments conducted GeForce RTX 3090. The forward step (12 takes about 14 seconds to pushforward one million points.
Assume each JKO step involves 500 iterations and the number of inner iteration J1 =3, then each JKO step (13) takes 100 seconds if the energy function contains the generalized energy ((P) and 25 seconds the energy function contains the KL divergence D(PI))).
C IMPLEMENTATION DETAILS
Our code is written in Pytorch-Lightn Falcon & Cho 2020)
For some parts of plotting in Section
we adopt the code given by Mokrov al. (2021).
Without further specification, we use the following parameters:
1) The number iterations of the outer loop J/1 600.
2) The number of iterations of the inner loop J/ is 3.
3) The batch size fixed to be M 100.
4) The learning rate is fixed to be 0.001.
5) All the activation functions are set to be PReLu.
6) h has layers and 16 neurons in each layer.
7) T has 5 layers and 16 neurons in each layer.
The transport map T can be parametrized in different ways. We residual MLP network for it in Section 4.1 4.2 B.2 and the gradient of a strongly convex ICNN in Section 4.3 4.4 B.1
B.3
The
dual test function h always MLP network with dropout layer before each layer.
C.I
SAMPLING (SECTION 4. AND B.
Two moons We run = 10 JKO steps with J2 = inner iterations. h has layers. T has 5 layers.
GMM 8D example trains for K 50 JKO steps. h has 6 layers and 64 neurons in each layer. T has layers and 128 neurons in each layer.
13D example trains for K 20 JKO steps. h has 8 layers and 64 neurons in each layer. T has layers and 64 neurons in each layer.
C.2 ORNSTEIN-UHLENBERR PROCESS (SECTION 4.2
For Fokker-Plaan Planck JKO, use the implementation provided by the authors and the default parameters given in Mokrov et al. (201 Section A.2). We also estimate the SymKL using Monte Carlo according to the author's 's instructions.
For our method, we use a linear residual feed-forward NN to work as T. i.e without activation function. and 7 both have 3 layers and 64 hidden neurons per layer for all dimensions. We also train them for J1 = 500 iterations per each JKO with learning rate 0.005. The batch size M = 1000.
However, estimate SymKL. for algorithm in different way. Since map T is linear transformation, our estimated P, is guaranteed to be Gaussian distribution. We firstly draw 10 samples from Pt and calculate the empirical mean fi and covariance Σ, Then we estimate D (P)))P.) using the following closed form KL divergence between two Gaussians
Σx
P(P))P)) = log
d + (T4 - -4) Σ (T4. 14) + tr( 55x)
ΣS
D D(P..IP)) is estimated similarly.
16

---

Under review as a conference paper at ICLR 2022
C.3 POROUS MEDIA EQUATION (SECTION 4.3
In the experiment, and 7' both have 10 neurons in each layer.
To account for the time-discretization error of JKO scheme, we consider the porous media equation in ID space and use the solution via finite difference method as a reference. More specifically, in the ID space R, we discretize the density over fixed grid with grid size d and grid resolution dx. With this discretization, the probability densities become (scaled) probability vectors and the problem can be converted into convex optimization
(6x)² 2a
dæ m
min I=P.
M)
(1' "
(18)
where M the discretized unit transport cost, Pr. Rd the probability vector the previous step, 1 RC is the all-ones vector and the optimization variable R IR is the joint distribution between Pa. and Pk+1- This is standard convex optimization and be solved with generic solvers. When an optimal π is obtained, P can be computed Px+1 m1 We adopt the library CVXOPT PrroPTTto to solve the convex programming problem (8). In so doing, we arrive reference solution Po: P Pr....,,........
.k....
for for our algorithm.
C.4 AGGREGATION-
EQUATION (SECTION 4.4 AND B.3
Each JKO step contains J1 = 200 iterations. The batch size is M = 1000.
http://ovxr.com//rn
17


Reviews:
[{"review_id": "zYGuvjjK9mb", "reviewer": "Reviewer_S3mH", "summary": "", "questions": "", "limitations": "", "rating": 3, "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 paper identifies a crucial challenge in the existing works of emulating JKO steps, namely the expensive computation of the densities of the form $P_k(x)$ for each $k$. The solution the paper suggests is reasonable, but it comes at the cost of adding another inner maximization which makes the optimization a lot more difficult (e.g. unstable due to high variance). Overall I think the amount of contribution is very limited. The variational formulation of the objective functionals considered is all well known, and putting it together with the JKO step is fairly straightforward. The experiments are not convincing enough to demonstrate the practical advantages of the proposed method compared to the alternatives.\n\nDetailed comments:\n- In Eq (9), a distribution $\\Gamma$ is introduced. What is the point of introducing this measure? Is it essentially an importance sampling of $Q$ for which we only know the density? Why is it enough to just choose Gaussians, which could be very different from $Q$?\n- Many sections are very similar to [Mokrov 2021]. For example, there is nothing new in Sec 3.5, and the experiment's setup in 4.1 and 4.2 are exactly the same. Yet the more challenging experiments from [Mokrov 2021] are not reproduced here, such as posterior inference and non-linear filtering. \n- In Sec 3.6, it is common to use the Hutchinson trace estimator to approximate the gradient of $\\log \\det$ (in addition to a linear solve), which could speed up the competing methods. It might be good to include a comparison to that.\n- The results in Figure 2 are visibly worse than those of [Mokrov 2021]. Moreover, here only up to dimension 13 is included, whereas [Mokrov 2021] contains dimension 32.\n- The results in Figure 3 of the proposed method are better than those of [Mokrov 2021]. I'm wondering why this is the case since in [Mokrov 2021] the KL divergence is calculated exactly, whereas in the proposed method additional bias could be introduced due to the failure of maximizing $h$.\n", "summary_of_the_paper": "This paper proposes a variational formulation of each JKO step for optimizing functionals on measures.\nDifferent from existing recent works on emulating JKO steps by training pushforward neural networks (either directly or as gradients of convex functions), the variational formulation involves another inner maximization of a function, without needing density access that typically requires cubic time complexity due to computing the log determinants of the pushforwards. Experiments are done to demonstrate the practicality of the algorithm.\n", "summary_of_the_review": "The paper studies an important challenge in JKO steps encountered by recent works, but the contributions are incremental without demonstrating convincing practical advantages.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Wl5ufYQpO6k", "reviewer": "Reviewer_CWde", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "**Benefits**\n1) The variational approach might be a solution to issues of direct computation such as cubically growing complexity as a function of the dimension;\n\n**Drawbacks**\n\n2) The experiments of the paper are weak and do not sufficiently support the main claims;\n3) The relation to the prior work is not fully disclosed;\n4) The scope of the paper (WGFs) is narrow.\n\n*Detailed comments are below.*\n\n**Relation to prior work.** A large part of the algorithm proposed by the authors matches the previously known art. More precisely, sections 3.0 and 3.1 exploit the reformulation of JKO via functions that have already been proposed in [1] (not cited in the paper!) and extensively used in [2,3]. I think this is not very clear from the text and might add extra (inexistent) value to the current work.\n\nIf I correctly understand, the actual difference w.r.t. the prior work, e.g., [2], is the way that the f-divergence is optimized. The authors use a variational approximation in contrast to direct computation. Here I have two remarks.\n\nFirst, the key claim of the paper is that direct computation scales cubically and is not feasible in high dimensions. I agree, but what about fast approximations? In [3] the authors explicitly state they use a fast estimator based on the Hutchinson method. In [2] the authors state that to speed up the computation fast approximations can be used. It is unfair to ignore this and compare in the experimental section only with [2] using and only by using the direct computation.\n\nSecond, the variational approximation of the f-divergence is not novel, see, for example, f-GANs [5]. I wonder if such approximations have already been used in Bayesian machine learning (BML)? I encourage the authors to include a detailed discussion of this. Why is this important? I am not an expert in BML, but it seems to me that this particular variational approximation can be applied to a dozen of tasks involving KL divergence. The current paper demonstrates that it outperforms direct computation. If this is indeed true, why hasn’t this approach been used for other BML tasks? If it has already been, the authors should include relevant references to support the current experimental findings.\n\nOverall, the paper should more carefully acknowledge the prior work and make it transparent what is new and what is already well known. Paper [4] also seems relevant.\n\n**Experiments.** The experimental section is very weak both in terms of quality and quantity. In terms of quantity, the paper looks poor compared to predecessors, e.g., [2], a paper to which they compare their method. In particular, the results provided in Figure 2 visually suffer from notable artifacts, which is a contrast to that reported by [2] in a similar setup (Figure 2 of [2]). Besides, the dimensions considered in the current paper are lower, which is suspicious.\n\nImportantly, the quantitative results provided in Figure 3 raise questions. How is it possible that on such a toy example (evolving Gaussian distributions -- linear pushforward maps grad psi), the variational methods so drastically outperform the direct computation (up to 10 times)? This is quite unbelievable, in particular, in small dimensions. Did the authors use the same network architectures and other shared hyperparameters for comparison? A discussion here is needed.\n\n**Scope.** If I correctly understand, the authors do not provide any high-dimensional applications of WGFs. Due to this, I currently tend to think the scope of the paper is narrow and the usefulness of the proposed approach for the community of ICLR is questionable. Adding an application would definitely benefit the paper. \n\n**Correctness of the method.** Overall, the method is correct. However, in one of the experiments, the authors approximate the pushforward of the JKO step directly as the neural network, not as the gradient of the input-convex network. In this case, the push forward distribution might not have density damaging the entire JKO scheme.\n\n**Clarity.** What for the FB scheme is introduced (section 3.3)? This is not clear from the text.\n\n**References**\n\n[1] Benamou, J. D., Carlier, G., Mérigot, Q., & Oudet, E. (2016). Discretization of functionals involving the Monge–Ampère operator. Numerische mathematik, 134(3), 611-636.\n\n[2] Mokrov, P., Korotin, A., Li, L., Genevay, A., Solomon, J., & Burnaev, E. (2021). Large-Scale Wasserstein Gradient Flows. arXiv preprint arXiv:2106.00736.\n\n[3] Alvarez-Melis, D., Schiff, Y., & Mroueh, Y. (2021). Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks. arXiv preprint arXiv:2106.00774.\n\n[4] Bunne, C., Meng-Papaxanthos, L., Krause, A., & Cuturi, M. (2021). JKOnet: Proximal Optimal Transport Modeling of Population Dynamics. arXiv preprint arXiv:2106.06345.\n\n[5] Nowozin, S., Cseke, B., & Tomioka, R. (2016, December). f-gan: Training generative neural samplers using variational divergence minimization. In Proceedings of the 30th International Conference on Neural Information Processing Systems (pp. 271-279).", "summary_of_the_paper": "The paper proposes a method to compute Wasserstein Gradient Flows (WGFs) via neural networks and the JKO scheme. In contrast to prior works, to compute WGFs of functionals involving f-divergences, the authors use variational approximations rather than direct computations. It is claimed to work faster and perform better.", "summary_of_the_review": "My overall impression of the paper is that it is unfinished. While the idea of variational approximation is reasonable, I suppose this paper requires a major revision with a dozen text improvements and experiments. Therefore, I vote to reject this paper in its current form.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "IfbHE9O91_W", "reviewer": "Reviewer_pvCT", "summary": "", "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": "Strengths:\n* An ingenious use of f-divergence duality (aka variational formulation) to re-write density-depending functional flow objectives, such as the generalized entropy, as (optimization of) expectations, allowing for computation via finite-sample approximation \n* The resulting method seems to be empirically valid in low- and mid-dimensional settings \n\nWeaknesses:\n* Novelty, the main (and perhaps, only) novelty of this approach compared to (Mokrov et al., 2021; Alvarez-Melis et al., 2021; Bunne et al., 2021) is the use of the variational formulation of f-divergences, something that is itself well known. Note that this is only relevant for functional objectives that depend on the density of the measure itself, e.g., those that cannot be expressed as an expectation over the measure. All other objectives, including the interaction energy considered here, can be tackled with the exact same approach of the three papers above. \n* The novel contributions should be more clearly separated from non-novel ones. E.g., the reformulation of WGF as optimization over convex functions might seem like a contribution of this paper to an uninitiated reader, although it has been explored extensively in other work, most prominently in Benamou et al (2014), which is puzzling not cited here, and more recently in the 3 JKO + ICNN papers cited above.  \n* Some statements are not clear, not correct or too vague. These should all be clarified or corrected. For example:\n    * \"that is not too far away from the identity map\" in Pg 5. \n    * In section 2.2, the $\\delta F / \\delta P$ is not formally a gradient, but the first-variation of a functional\n    * In Section 3.6, it is stated that prior works requires explicit computation of log determinants of Hessians at cubic complexity, but this is not the case. At least one of those methods uses cuadrtic-cost (matrix-vector-product) stochastic estimators of the Hess logdet.\n* The motivation for using the FB scheme for the interaction energy is not clear. Some of the intuition/motivation discussed in sections B.2 and B.3 should probably be moved to the main text. \n* Limited experimental validation, especially with regards to high-dimensional settings, which is arguably the main promise of this work. In addition, all experiments are synthetic. A more compelling evaluation framework should include at least one high-dimensional, realistic dataset. \n* The evaluation is mostly qualitative. The paper purposefuly chooses PDEs with known solutions, but then provides mostly qualitative/visual results. A more compelling evaluation framwork would include quantitative comparison against these known solutions in sections 4.3 and 4.4.\n* While the computational complexity summary provided in the paper is useful, there are some many hidden constants in those bounds that they are hardly useful. These should be complemented with a thorough empirical runtime analysis comparing against exact and inexact methods, such as those cited here as related work. \n\nMinor comments:\n* While the motivation is different, there are deep and unexplored connections between the dualizatin of the f-divergence objectives used here and the dual of the kantorovich OT problem that has been recently explored extensively for learning Monge maps between distributions. Many of these rely on convex conjugacy to reformulate a sup objective as a sup-inf one (see Korotin et al. 2021b for an excellent survey on these). It would have been great (though not obligatory!) to see a discussion on these connections \n* I would suggest moving Table 1 after Proposition 1, so that all objects in the definition of $\\mathcal{A}(T,h)$ have been already introduced. I spent some time trying to figure out what $\\mu(T)$ was before finding it further down.\n\nMissing related work:\n\n* Benamou et al, \"Discretization of functionals involving the Monge-Ampère operator\". 2014.\n* Huang et al, \"Convex Potential Flows: Universar Probability Distributions with Optimal Transport and Convex Optimization\", ICLR 2021.\n* Bunne et al., \"JKOnet: Proximal Optimal Transport Modeling of Population Dynamics\", 2021.", "summary_of_the_paper": "This paper proposes a method to solve Wasserstein gradient flows based on the JKO scheme using variational formulations of functional objectives, such as the KL divergence or the generalized entropy (non-linear diffusion). Relying on known reformulations of the JKO scheme as optimization over convex functions, the paper departs from recent related methods in expressing certain objectives as f-divergences, and in turn using the dual formulation of these divergences to circumvent the need to do explicit density computation in these. The resulting method involves parametrizing two types of operators as neural networks (one of them as an input-convex neural network), and solving a mini-max objective. The paper presents experiments on simple PDEs (mostly in 1D or 2D) with known solutions. ", "summary_of_the_review": "The paper provides an interesting variation on recent JKO-based methods for computational Wasserstein gradient flows, but its limited novelty and empirical evaluation diminish its contribution, and make it a borderline paper in my view. That being said, if the issues I raise in my review are properly addressed, I would be willing to increase my score. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "7JcBbeRRQe2", "reviewer": "Reviewer_72fj", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "Strengths:\n\nThe paper is well written and high level ideas on the theory are explained. The approach is clear and reasonable. The simulations show promising results, they do not seem to have been cherry picked. They show what is claimed by the authors and there is no surprise given the simplicity of the approach.\n\nWeaknesses:\n\nThe theoretical consistence of the proposed methods is ignored. Are the proposed scheme consistent time discretizations of the WGF? Are the assumptions satisfied? Is the alternative maximization/minimization strategy (Algo 1) consistent? Do all measures considered admit density wrt Lebesgue? The paper mainly provides intuition for Algo 1 without solid theoretical foundations. \n\nMoreover, the technical contribution is rather limited (see the summary of the paper above). One could argue that it is an \"easy paper\" in the sense that only the simulations seem new. The novelty is mainly to have parametrized the functions in the objective of the JKO by neural networks. ", "summary_of_the_paper": "This paper studies the implementation of some Wasserstein Gradient Flows (WGF) in discrete time but without discretizing the space. The methods proposed are based on the JKO operator to discretize WGF in time. The implementation of the JKO can be challenging. The strategy of the authors is to first reparametrize the JKO as a minimization over a space of functions (instead of measures) via pushforward. Then, when the objective function is a f-divergence, the objective inside the JKO admit a variational representation and can be expressed as a sup. Conclusion: each JKO is written as a min max over a space of functions. To solve it, they parametrize the functions by neural networks and alternatively maximize and minimize the problem using Adam. An important feature is that the objective in the min max can be approximated with samples of the current distribution (its density doesn't appear, only integrals wrt to the current distribution). ", "summary_of_the_review": "The approach is reasonable, the simulation promising but I do not see a significant technical contribution (They essentially reparametrized a min max problem over a function space using neural nets and running Adam to alternatively maximize and minimize). ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes  a min/max reformulation for JKO gradient flows appealing the variational formulation of f-divergences. This would alleviate the need of an explicit density. All reviewers pointed out the limited novelty in the work and the limited experimentation. 

We encourage authors to add a theoretical analysis to their work and further strengthening of the experimental section with high dimensional experiments, and to resubmit the work on an upcoming venue.
Decision: Reject