Year: 2021

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.

heading: 1 INTRODUCTION
text: The Wasserstein gradient flow models the gradient dynamics over 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 (Ambrosio et al., 2008; Otto, 2001; Adams et al., 2011; Santambrogio, 2017; Carlier et al., 2017; 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 (Peyré, 2015; Carlier et al., 2017; Li et al., 2020; Carrillo et al., 2021). The computational complexity of these methods scales exponentially as the problem dimension, making them unsuitable for the cases where probability densities over high dimensional spaces are involved.
Our goal is to develop a 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 to a 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 a 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 a variational form of the objective function, which allows the evaluation of the
objective with samples and without density estimation. At the end of the algorithm, a 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 is modeled by the gradient of an input convex neural network, one 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 is applicable to any objective function that has a variational representation. ii). We specialize our algorithm to three important cases where the objective functions are the Kullback-Leibler divergence, the generalized entropy, and the interaction energy . iii). We apply our algorithm to several representative problems including sampling and aggregationdiffusion equation and obtain respectable performance.
Related works: Most existing methods to compute Wasserstein gradient flow are finite difference based (Peyré, 2015; Carlier et al., 2017; Li et al., 2020; Carrillo et al., 2021). These methods require spatial discretization and are thus not scalable to high dimensional settings. Salim et al. (2020) analyze the convergence for a forward-backward 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 et al., 2019b; Frogner & Poggio, 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., 2021; Alvarez-Melis et al., 2021; Yang et al., 2020) were proposed for Wasserstein gradient flow. The first one (Mokrov et al., 2021) focuses on the special case with Kullback-Leibler divergence as objective function. The second one (Alvarez-Melis et al., 2021) uses a density estimation method to evaluate the objective function by back-propagating to the initial distribution, which could become a 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.

heading: 2 BACKGROUND

heading: 2.1 OPTIMAL TRANSPORT AND WASSERSTEIN DISTANCE
text: Given two probability distributions P,Q over the Euclidean space Rn with finite second moments, the optimal transport problem with quadratic cost reads
min T :T]P=Q ∫ Rn ‖x− T (x)‖22dP (x), (1)
where the minimization is over all the feasible transport maps that transport mass from distribution P to distribution Q. The feasibility is characterized by the pushforward operator (Bogachev, 2007) as T]P = Q. When the initial distribution P admits a density, the above optimal transport problem (1) has a unique solution and it is the gradient of a convex function, that is,
T ? = ∇ϕ
for some convex function ϕ(·) : Rn → 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 (1), defines a metric on the space of probability distributions known as the Wasserstein-2 distance (Villani, 2003), denoted by W2(P,Q). The Wasserstein distance has many nice geometrical properties compared with other distances such as L2 distance for probability distributions, making it a popular choice in applications.

heading: 2.2 WASSERSTEIN GRADIENT FLOW
text: Given a functionF(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
∂P
∂t = ∇ ·
( P∇δF
δP
) ,
where δF/δP stands for the gradient of the functionF with respect to the standardL2 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 ) = ∫ Rn P (x) logP (x)dx+ ∫ Rn V (x)P (x)dx, the gradient flow is the Fokker-Planck equation (Jordan et al., 1998) ∂P∂t = ∇ · (P∇V ) + ∆P. When F is the generalized entropy F(P ) = 1m−1 ∫ Rn P
m(x)dx for some positive number m > 1, the gradient flow is the porous medium equation (Otto, 2001; Vázquez, 2007) ∂P∂t = ∆P m.

heading: 3 METHODS AND ALGORITHMS
text: We are interested in solving the optimization problem min P F(P ) (2)
over the space of probability densities P(Rn). In particular, our objective is to develop a particlebased Wasserstein gradient flow algorithm to numerically solve (2).
The objective function F(P ) could exhibit different form depending on the application. In this paper, we present our algorithm for the linear combination of the following three important cases:
Case I The functional is equal to the KL-divergence with respect to a given target distribution Q F(P ) = D(P ||Q) := ∫ log ( P (x)
Q(x)
) P (x)dx. (3)
This is important for the problem of sampling from a target distribution.
Case II The objective functional is equal to the generalized entropy
F(P ) = G(P ) := 1 m− 1
∫ Pm(x)dx.
This case is important for modeling the porous medium.
Case III The objective functional is equal to the interaction energy F(P ) =W(P ) := ∫ ∫ W (x− y)P (x)P (y)dxdy, W : Rn → 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 to 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 to (2) (Santambrogio, 2017). In Section 3.1, 3.2, we first assume F(P ) doesn’t include interaction energy, and introduce JKO/backward scheme to solve (2). We then addW(P ) into consideration and present a forwardbackward scheme in Section 3.3 and close by showing our Algorithm in Section 3.4.

heading: 3.1 JKO SCHEME AND REPARAMETRIZATION
text: To realize the Wasserstein gradient flow, a discretization over time is needed. One such discretization is the famous JKO scheme (Jordan et al., 1998)
Pk+1 = arg min P
1
2a W 22 (P, Pk) + F(P ). (4)
This is essentially a backward Euler discretization or a proximal point method with respect to the Wasserstein metric. The solution to (4) converges to the continuous-time Wassrstein gradient flow when the step size a→ 0. In our method, we reparametrize (4) as an optimization in terms of the transport maps T : Rn → Rn from Pk to P , i.e., by defining P = T]Pk. With this reparametrization, in view of the definition of Wasserstein distance (1), the JKO step (4) becomes
Pk+1 = Tk]Pk, Tk = arg min T
1
2a ∫ Rn ‖x− T (x)‖22dPk(x) + F(T]Pk). (5)
The optimal T is the optimal transport map from Pk to T]Pk and is thus the gradient of a convex function ϕ. Therefore, the JKO scheme can be also expressed as
Pk+1 = ∇ϕk]Pk, ϕk = arg min ϕ∈CVX
1
2a ∫ Rn ‖x−∇ϕ(x)‖22dPk(x) + F(∇ϕ]Pk). (6)
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(P‖Q) AND G(P ) 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 a variational formulation of F . In order to do so, we use the notion of f -divergence between the two distributions P and Q:
Df (P‖Q) = EQ [ f ( P
Q
)] (7)
where f : (0,+∞)→ R is a convex function. The f -divergence admits the variational formulation
Df (P‖Q) = sup h EP [h(X)]− EQ[f∗(h(Y ))]. (8)
where f∗(y) = supx∈R[xy− f(x)] is the convex conjugate of f . The variational form has the special feature that it does not involve the density of P and Q explicitly and can be approximated in terms of samples from P and Q. The functionals D(P ||Q) and G(P ) can both be expressed as f -divergence. With the help of the f -divergence variational formula, when F(P ) = D(P‖Q) or G(P ), the JKO scheme (5) can be equivalently expressed as
Pk+1 = Tk]Pk, Tk = arg min T
{ 1
2a EPk [‖X − T (X)‖2] + sup h V(T, h)
} . (9)
where V(T, h) = EPk [A(T, h)]−EΓ[B(h)], Γ is a user designed distribution which is easy to 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 1. The details appear in Section 3.2.1 and 3.2.2.

heading: 3.2.1 KL DIVERGENCE
text: The KL divergence is the special instance of the f -divergence obtained by replacing f with f1(x) = x log x in (7) Df1(P‖Q) = EQ [ P
Q log
P
Q
] = EP [ log P
Q
] .
Proposition 1. The variational formulation for D(P‖Q) reads
Df1(P‖Q) = 1 + sup h
EP [ log h(X) + log µ(X)
Q(X)
] − Eµ [h(Z)] ,
where µ is a user designed distribution which is easy to sample from. The optimal function h is equal to the ratio between the densities of T]Pk and µ.
The proof for Proposition 1 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 variational form in the JKO scheme (5) yields Pk+1 = Tk]Pk and
Tk = arg min T max h
EPk [ 1
2a ‖X − T (X)‖2 + log h(T (X)) + log µ(T (X)) Q(T (X))
] − Eµ [h(Z)] .(10)
In practice, we choose µ = µk adaptively, where µk is the Gaussian with the same mean and covariance as Pk. We noticed that this choice improves the numerical stability of the the algorithm.

heading: 3.2.2 POROUS MEDIUM EQUATION
text: The generalized entropy can be also represented as f -divergence. In particular, let f2(x) = 1 m−1 (x m − x) 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
Df2(P‖Q) = Ωm−1
m− 1
∫ Pm(x)dx− 1
m− 1 .
Proposition 2. The variational formulation for G(P ) reads
G(P ) = 1 Ωm−1 sup h
( EPk [ m
m− 1 hm−1(X)
] − EQ [hm(Z)] ) . (11)
The optimal function h is equal to the ratio between the densities of T]Pk and Q.
The proof for Proposition 2 is postponed to appendix A. Using this in the JKO scheme yields Pk+1 = Tk]Pk, and
Tk = arg min T max h
1
2a EPk‖X − T (X)‖2 +
1
Ωm−1k
( EPk [ m
m− 1 hm−1(X)
] − EQ [hm(Z)] ) ,
where Ωk is the volume of a set large enough to contain the support of T]Pk for any T that is not too far away from the identity map.
Algorithm 1 Primal-dual gradient flow Input: Objective function F(P ), initial distribution P0, step size a, number of JKO steps K, number of outer loop J1, number of inner loop J2, batch size M . Initialization: Parameterized Tθ and hλ for k = 1, 2, . . . ,K do Pk ← (I − a∇x(W ∗ Pk))]Pk if F(P ) includesW(P ) Tθ ← Tk−1 if k > 1 {// use last iteration Tk−1 as a warm-up} for j1 = 1, 2, . . . , J1 do
Sample Y1, . . . , YM from Pk. Sample Z1, . . . , ZM from Γ. for j2 = 1, 2, . . . , J2 do Apply Adam to λ to maximize 1M ∑M i=1 [A(Tθ, hλ(Yi))− B(hλ(Zi))] end for Apply Adam to θ to minimize 1M ∑M i=1 [ 1 2a‖Yi − Tθ(Yi)‖ 2 +A(Tθ, hλ(Yi)) ]
end for Tk ← Tθ
end for Output: {Tk}Kk=1

heading: 3.3 FORWARD BACKWARD (FB) SCHEME
text: When F(P ) involves the interaction energyW(P ), we add an additional forward step to solve the gradient flow:
Pk+ 12 := (I − a∇x(W ∗ Pk))]Pk (12) Pk+1 := Tk+ 12 ]Pk+ 1 2 , (13)
where I is the identity map, and Tk+ 12 is defined by replacing k by k + 1 2 in (9). In other words, the first gradient descent step (12) is a forward discretization of the gradient flow and the second JKO step (13) is a backward discretization. ∇x(W ∗ P ) can be written as expectation Ey∼P∇x(W (x− y)), thus can also be approximated by samples. Salim 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.
Since W(P ) can be equivalently written as expectation Ex,y∼P [W (x − y)], there exists another non-forward-backward (non-FB) method , i.e., removing the first step and integratingW(P ) into a single JKO step: Pk+1 = Tk]Pk and
Tk = arg min T
{ 1
2a EPk‖X − T (X)‖2 + EX,Y∼Pk [W (T (X)− T (Y ))] + sup h V(T, h)
} .
In practice, we observe the FB scheme is more stable however converge slower than non-FB scheme. The detailed discussion appears in the Appendix B.2, B.3. Remark 1. In principle, one can single out log(Q) term from (10) and perform a similar forward step Pk+ 12 = (I − a(∇xQ)/Q)]Pk (Salim et al., 2020), but we don’t 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 (9) can be both parameterized by neural networks, denoted by Tθ and hλ. With this neural network parametrization, we can then solve the problem by iteratively updating Tθ and hλ. This prime-dual method to solve (2) is depicted in Algorithm 1.
In this work, we implemented two different architectures for the map T . One way is to use a neural network to represent T directly, and another way is to parametrize T as the gradient of a Input convex neural network (ICNN) (Amos et al., 2017) ϕ. The latter has been widely used in optimal transport (Makkuva et al., 2020; Fan et al., 2020; Korotin et al., 2021). In our experiments, we find that the first parameterization gives better result in the sampling application. As we discuss in Section 3.5, when density evaluation is needed, we adopt the ICNN parameterization since we need to compute T−1. Note that ICNN could be modified to be strictly convex and if the function ϕ is strictly convex, the gradient ∇ϕ is invertible.

heading: 3.5 EVALUATION OF THE DENSITY
text: In this section, we assume the solving process doesn’t use forward-backward scheme, i.e. all the probability measures Pk are obtained by performing JKO one by one. Otherwise, the map I−a∇x(W ∗Pk) = I−Ey∼Pk∇x(W (x−y)) includes an expectation term and becomes intractable to push-backward particles to compute density.
If T is invertible, these exists a standard approach to evaluate the density of Pk (Alvarez-Melis et al., 2021; Mokrov et al., 2021) through the change of variables formula. More specifically. we assume T is parameterized by the gradient of an ICNN ϕ that is assumed to be strictly convex. To evaluate the density Pk(xk) at point xk, we back propagate through the sequence of maps Tk = ∇ϕk, . . . , T1 = ∇ϕ1 to get
xi = T −1 i+1 ◦ T −1 i+2 ◦ · · · ◦ T −1 k (xk).
The inverse map T−1j = (∇ϕj)−1 = ∇ϕ∗j can be obtained by solving the convex optimization xj−1 = arg max
x∈Rn 〈x, xj〉 − ϕj(x). (14)
Then, by the change of variables formula, we obtain
log[Pk(xk)] = log[P0(x0)]− k∑ i=1 log ∣∣∇2ϕi(xi−1)∣∣ , (15)
where∇2ϕi(xi−1) is the Hessian of ϕi and |∇2ϕi(xi−1)| is its determinant. By iteratively solving (14) and plugging the resulting xj into (15), we can recover the density Pk(xk) at any point.

heading: 3.6 COMPUTATIONAL COMPLEXITY
text: Per each update k in Algorithm 1, the forward step (12) requires at most O(N2) where N is the total number of particles to push-forward. The backward step (13) requires O(J1kMH) where J1 is 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 Pk requires us to pushforward x0 ∼ P0 through k− 1 maps. However, both Mokrov et al. (2021) and Alvarez-Melis et al. (2021) require O(J1k ( nMH + n3 ) ) which has the cubic dependence on dimension n because they need to query the log det∇2ϕ in each iteration. We refer to Mokrov et al. (2021, Section 5) for the complexity details of calculating the Hessian term. Thus our method has the advantage of independence on the dimension.
We provide training time details in Section 4.2 and Appendix B.4. Other than training and sampling time, the complexity for evaluating the density are the same as 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 have to be normalized. To this end, we consider the Wasserstein gradient flow with objective function F(P ) =∫ P log(P/Q)dx, that is, the KL divergence between distributions P and Q. When this
objective is minimized, P ∝ 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 a = 0.3 and the initial measure is a spherical Gaussian N (0, 2.25Id). Two moons: The two moons distribution is a popular target distribution for sampling task. It is a 2D mixture model composed of 16 Gaussian components; each moon shape consists of 8 Gaussian components. The results are displayed in Figure 1, from which we see that our method is able to generate samples that match the target distribution.
GMM with spherical Gaussians: We also test our 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 n = 8 and n = 13 are depicted in Figure 2. In Figure 2, we not only display the samples as grey dots in the plot, but also the kernel density estimation of generated samples as level sets. As can be seen from the results, both the samples and densities obtained with our algorithm match the target distribution well.

heading: 4.2 ORNSTEIN-UHLENBECK PROCESS
text: We study the performance of our method in modeling the Ornstein-Uhlenbeck Process as dimension grows. The gradient flow is affiliated with the free energy (3), where Q = e(x−b)
TA(x−b)/2 with a positive definite matrix A ∈ Rn × Rn and b ∈ Rn. Given an initial Gaussian distribution N (0, In), the gradient flow at each time t is a Gaussian distribution Pt with mean vector µt = (In − e−At)b and covariance Σt = A−1(In − e−2At) + e−2At (Vatiwutipong & Phewchean, 2019). We calculate Pk with JKO step size a = 0.05 and compare with the Fokker-Planck (FP) JKO (Mokrov et al., 2021). We quantify the error as the SymKL divergence between estimated distribution and the ground truth in Figure 3, where
SymKL(P1, P2) := D(P1‖P2) +D(P2‖P1).
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 to 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 on the dimension.

heading: 4.3 POROUS MEDIA EQUATION
text: We next consider the porous media equation with only diffusion: ∂tP = ∆Pm. This is the Wasserstein gradient flow associated with the energy function F(P ) = 1m−1 ∫ Pm(x)dx. A representative
closed-form solution of the porous media equation is the Barenblatt profile (GI, 1952; Vázquez, 2007) Pgt(t, x) = (t+ t0) −α ( C − β (x− x0)2 (t+ t0)−2α/n ) 1 m−1
+ ,
where α = n
n(m− 1) + 2 , β = (m− 1)α 2mn ,
and t0 > 0 is the starting time and C > 0 is a free parameter.
In principle, our algorithm should match the analytical solution Pgt when the step size a is sufficiently small. When a is not that small, time discretization is inevitable. To account for the time-discretization error of JKO scheme, we 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.3.
In the experiments, we set the stepsize for the JKO scheme to be a = 0.001 and the initial time to be t0 = 0.002. Other parameters are chosen as C = (3/16)1/3,m = 2, n = 1, d = 300. We parametrize the transport map T as the gradient of an ICNN and thus we can evaluate the density following Section 3.5. In Figure 4, we observe that the gap between the density computed using our algorithm and the ground truth density Pgt is dominated by the time-discretization error of the JKO scheme. Our result matches the discrete time solution nicely.

heading: 4.4 AGGREGATION–DIFFUSION EQUATION
text: We finally simulate the evolution of solutions to the following aggregation-diffusion equation:
∂tP = ∇ · (P∇W ∗ P ) + 0.1∆Pm, W (x) = −e−‖x‖ 2 /π.
This corresponds to the energy functionW(P ) + 0.1G(P ). We use the same parameters in Carrillo et al. (2021, Section 4.3.3). The initial distribution is a uniform distribution supported on [−3, 3]× [−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 = 3 on R2 space. With this choice W (x), ∇x(W ∗ Pk) is equal to Ey∼Pk [ 2e−‖x−y‖ 2 /π ] in the gradient descent step (12). And we estimate ∇x(W ∗ Pk) with 104
samples from Pk.
Throughout the process, the aggregation term∇·(P∇W ∗P ) and the diffusion 0.1∆Pm 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).

heading: 5 CONCLUSION
text: 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 of the objective function, we further reformulate the problem in every step as a min-max problem over map T and dual function h respectively. This formulation doesn’t require density estimation using samples and can be optimized using stochastic optimization. It also shows advantages with dimension-independent computation complexity. Our method can also be extended to minimize other objective functions that can be written as f -divergence. Our limitation is the accuracy is not satisfying in sampling tasks with high dimension complex density.

heading: 6 REPRODUCIBILITY
text: The basic setup is distributed in Section 4. We refer to Section 3.4 and Appendix C for all the rest training details. The code of our method is also attached in the supplementary material.

heading: A DETAILS ABOUT VARIATIONAL FORMULA IN SECTION 3.2

heading: A.1 KL DIVERGENCE
text: The KL divergence is the special instance of the f -divergence obtained by replacing f with f1(x) = x log x in (7) Df1(P‖Q) = EQ [ P
Q log
P
Q
] = EP [ log P
Q
] ,
which, according to (8), admits the variational formulation
Df1(P‖Q) = 1 + sup h
EP [h(X)]− EQ [ eh(Y ) ] (16)
where the convex conjugate f∗1 (y) = e y−1 is used.
The variational formulation can be approximated in terms of samples from P and Q. For the case where we have only access to un-normalized density of Q, which is the case for the sampling problem, we use the following change of variable: h→ log(h) + log(µ)− log(Q) where µ is a user designed distribution which is easy to sample from. Under such a change of variable, the variational formulation reads
Df1(P‖Q) = 1 + sup h
EP [ log h(X) + log µ(X)
Q(X)
] − Eµ [h(Z)] .
Note that the optimal function h is equal to the ratio between the densities of T]Pk and µ. Remark 2. The Donsker-Varadhan formula
D(P‖Q) = sup h
EP [h(X)]− logEQ [ eh(Z) ] is another variational representation of KL divergence and it’s a stronger than (16) because it’s a upper bound of (16) for any fixed h. However, we cannot get an unbiased estimation of the objective using samples.

heading: A.2 POROUS MEDIUM EQUATION
text: The generalized entropy can be also represented as f -divergence. In particular, let f2(x) = 1 m−1 (x m − x) 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
Df2(P‖Q) = Ωm−1
m− 1
∫ Pm(x)dx− 1
m− 1 .
As a result, the generalized entropy can be expressed in terms of f -divergence according to
G(P ) = 1 m− 1
∫ Pm(x)dx = 1
Ωm−1 Df2(P‖Q) +
1
Ωm−1(m− 1) .
Upon using the variational representation of the f -divergence with
f∗2 (y) =
( (m− 1)y + 1
m
) m m−1
,
the generalized entropy admits the following variational formulation
G(P ) = sup h
1
Ωm−1
( EP [h(X)]− EQ [( (m− 1)h(Y ) + 1
m
) m m−1 ]) +
1
Ωm−1(m− 1) .
In practice, we find it numerically useful to let h = 1m−1
[ m ( ĥ )m−1 − 1 ] so that
G(P ) = 1 Ωm−1
sup ĥ
( EPk [ m
m− 1 ĥm−1(X)
] − EQ [ ĥm(Z) ]) .
With such a change of variable, the optimal function ĥ = T]Pk/Q.

heading: B ADDITIONAL EXPERIMENT RESULTS AND DISCUSSIONS

heading: B.1 SAMPLING USING ICNN PARAMETERIZATION
text: In Figure 6, we present the sampling results with ∇ϕ parameterized map where ϕ is 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.

heading: B.2 AGGREGATION EQUATION
text: Alvarez-Melis et al. (2021) 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
∂tP = ∇ · (P∇W ∗ P ), W : Rn → R. (17) The forward method is
Pk+1 := (I − a∇x(W ∗ Pk))]Pk. The backward method or JKO is
Pk+1 := Tk]Pk, Tk = arg min T
{ 1
2a EPk [‖X − T (X)‖2] + EX,Y∼Pk [W (T (X)− T (Y ))]
} .
Example 1 We follow the setting in Carrillo et al. (2021, Section 4.3.1 ). The interaction kernel is W (x) = ‖x‖ 4
4 − ‖x‖2
2 , and the initial measure P0 is a GaussianN (0, 0.25I). In this case,∇x(W ∗Pk) becomes Ey∼Pk [ (‖x− y‖2 − 1)(x− y) ] . We use step size a = 0.05 for both methods and show the results in Figure 7.
Example 2 We follow the setting in Carrillo et al. (2021, Section 4.2.3 ). The interaction kernel is W (x) = ‖x‖ 2
2 − ln ‖x‖, and the initial measure P0 is N (0, 1). The unique steady state for this case is
P∞(x) = 1
π
√ (2− x2)+.
The reader can refer to Alvarez-Melis et al. (2021, Section 5.3) for the backward method performance. As for the forward method, ∇x(W ∗ Pk) becomes Ey∼Pk [ x− y − 1x−y ] . Because the kernel W enforces repulsion near the origin and P0 is concentrated around origin,∇x(W ∗ P ) will easily blow up. So the forward method is not suitable for this kind of interaction kernel.
Through the above two examples, if ∇x(W ∗ P ) is smooth, we can notice the backward method converges faster, but is not stable when solving (17). This shed light on the FB and non-FB scheme performance in Section 4.4, B.3. However, if ∇x(W ∗ P ) has bad modality such as Example 2, the forward method loses the competitivity.

heading: B.3 AGGREGATION-DIFFUSION EQUATION WITH NON-FB SCHEME
text: In Figure 8, we show the non-FB solutions to Aggregation-diffusion equation in Section 4.4. FB scheme should be independent with the implementation of JKO, but in the following context, we assume FB and non-FB are both neural network based methods discussed in Section 3. Non-FB scheme reads
Pk+1 = Tk]Pk
Tk = arg min T
{ 1
2a EPk [‖X − T (X)‖2] + EX,Y∼Pk [W (T (X)− T (Y ))] + G(T, h)
} ,
where G(T, h) is represented by the variational formula (11). We use the same step size a = 0.5 and other PDE parameters as in Section 4.4.
Comparing the FB scheme results in Figure 5 and the non-FB scheme results in Figure 8, we observe non-FB converges 1.5× slower than the finite difference method (Carrillo et al., 2021), and FB converges 3× slower than the finite difference method. This may because splitting one JKO step to the forward-backward two steps removes the aggregation term effect in the JKO, and the diffusion term is too weak to make a difference in the loss. Note at the first several k, both Pk and Q are nearly the same uniform distributions, so h is nearly a constant and T (x) exerts little effect in the variational formula of G(P ). Another possible reason is a single forward step for aggregation term converges slower than integrating aggregation in the backward step, as we discuss in Section B.2 and Figure 7.
However, FB generates more regular measures. We can tell the four pulses given by FB are more symmetric. We speculate this is because gradient descent step in FB utilizes the geometric structure of W (x) directly, but integratingW(P ) in neural network based JKO losses the geometric meaning of W (x).

heading: B.4 COMPUTATIONAL TIME
text: Our experiments are conducted on 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 G(P ) and 25 seconds if the energy function contains the KL divergence D(P‖Q).
C IMPLEMENTATION DETAILS
Our code is written in Pytorch-Lightning (Falcon & Cho, 2020). For some parts of plotting in Section 4.1 and 4.2, we adopt the code given by Mokrov et al. (2021).
Without further specification, we use the following parameters: 1) The number of iterations of the outer loop J1 is 600. 2) The number of iterations of the inner loop J2 is 3. 3) The batch size is 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 4 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 use a 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 is always a MLP network with a dropout layer before each layer.

heading: C.1 SAMPLING (SECTION 4.1 AND B.1)
text: Two moons We run K = 10 JKO steps with J2 = 6 inner iterations. h has 6 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 3 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 9 layers and 64 neurons in each layer.

heading: C.2 ORNSTEIN-UHLENBECK PROCESS (SECTION 4.2)
text: For Fokker-Planck JKO, we use the implementation provided by the authors and the default parameters given in Mokrov et al. (2021, Section A.2). We also estimate the SymKL using Monte Carlo according to the author’s instructions.
For our method, we use a linear residual feed-forward NN to work as T , i.e. without activation function. h and T 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 is M = 1000.
However, we estimate SymKL for our algorithm in a different way. Since our map T is a linear transformation, our estimated P̃t is guaranteed to be a Gaussian distribution. We firstly draw 5 · 105 samples from P̃t and calculate the empirical mean µ̃t and covariance Σ̃t. Then we estimateD(P̃t‖Pt) using the following closed form KL divergence between two Gaussians
D(P̃t‖Pt) = 1
2 [ log |Σt| |Σ̃t| − d+ (µ̃t − µt)T Σ̃−1t (µ̃t − µt) + tr(Σ−1t Σ̃t) ] .
D(Pt‖P̃t) is estimated similarly.

heading: C.3 POROUS MEDIA EQUATION (SECTION 4.3)
text: In the experiment, h and T both have 10 neurons in each layer.
To account for the time-discretization error of JKO scheme, we consider the porous media equation in 1D space and use the solution via finite difference method as a reference. More specifically, in the 1D space R, we discretize the density over a fixed grid with grid size d and grid resolution δx. With this discretization, the probability densities become (scaled) probability vectors and the problem (4) can be converted into a convex optimization
min π:δxπT1=P̂k
(δx)2
2a 〈π,M〉+ δx m− 1 (1Tπδx)m1 (18)
where M is the discretized unit transport cost, P̂k ∈ Rd is the probability vector at the previous step, 1 ∈ Rd is the all-ones vector and the optimization variable π ∈ Rd × Rd is the joint distribution between P̂k and P̂k+1. This is a standard convex optimization and can be solved with generic solvers. When an optimal π is obtained, P̂k+1 can be computed as P̂k+1 = π1. We adopt the library CVXOPT1 to solve the convex programming problem (18). In so doing, we arrive at a reference solution P̂0, P̂1, . . . , P̂k, . . . for our algorithm.

heading: C.4 AGGREGATION-DIFFUSION EQUATION (SECTION 4.4 AND B.3)
text: Each JKO step contains J1 = 200 iterations. The batch size is M = 1000.
1http://cvxr.com/cvx/

Review summaries:
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. 

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

This paper proposes a variational formulation of each JKO step for optimizing functionals on measures. Different 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. 

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.