Variational Wasserstein gradient flow

1 INTRODUCTION . 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 . 2 BACKGROUND . 2.1 OPTIMAL TRANSPORT AND WASSERSTEIN DISTANCE . 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 . 2.2 WASSERSTEIN GRADIENT FLOW . 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 . 3 METHODS AND ALGORITHMS . 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 . 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 ) 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 .

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