Large-Scale Wasserstein Gradient Flows

1 Introduction . Stochastic differential equations ( SDEs ) are used to model the evolution of random diffusion processes across time , with applications in physics [ 63 ] , finance [ 22 , 52 ] , and population dynamics [ 35 ] . In machine learning , diffusion processes also arise in applications filtering [ 34 , 21 ] and unnormalized posterior sampling via a discretization of the Langevin diffusion [ 70 ] . The time-evolving probability density ρt of these diffusion processes is governed by the FokkerPlanck equation . Jordan , Kinderlehrer , and Otto [ 32 ] showed that the Fokker-Planck equation is ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . equivalent to following the gradient flow of an entropy functional in Wasserstein space , i.e. , the space of probability measures with finite second order moment endowed with the Wasserstein distance . This inspired a simple minimization scheme called JKO scheme , which consists an implicit Euler discretization of the Wasserstein gradient flow . However , each step of the JKO scheme is costly as it requires solving a minimization problem involving the Wasserstein distance . One way to compute the diffusion is to use a fixed discretization of the domain and apply standard numerical integration methods [ 18 , 49 , 15 , 17 , 40 ] to get ρt . For example , [ 50 ] proposes a method to approximate the diffusion based on JKO stepping and entropy-regularized optimal transport . However , these methods are limited to small dimensions since the discretization of space grows exponentially . An alternative to domain discretization is particle simulation . It involves drawing random samples ( particles ) from the initial distribution and simulating their evolution via standard methods such as Euler-Maruyama scheme [ 36 , M9.2 ] . After convergence , the particles are approximately distributed according to the stationary distribution , but no density estimate is readily available . Another way to avoid discretization is to parameterize the density of ρt . Most methods approximate only the first and second moments ρt , e.g. , via Gaussian approximation . Kalman filtering approaches can then compute the dynamics [ 34 , 39 , 33 , 61 ] . More advanced Gaussian mixture approximations [ 65 , 1 ] or more general parametric families have also been studied [ 64 , 69 ] . In [ 48 ] , variational methods are used to minimize the divergence between the predictive and the true density . Recently , [ 24 ] introduced a parametric method to compute JKO steps via entropy-regularized optimal transport . The authors regularize the Wasserstein distance in the JKO step to ensure strict convexity and solve the unconstrained dual problem via stochastic program on a finite linear subset of basis functions . The method yields unnormalized probability density without direct sample access . Recent works propose scalable continuous optimal transport solvers , parametrizing the solutions by reproducing kernels [ 10 ] , fully-connected neural networks [ 62 ] , or Input Convex Neural Networks ( ICNNs ) [ 37 , 44 , 38 ] . In particular , ICNNs gained attention for Wasserstein-2 transport since their gradients ∇ψθ : RD → RD can represent OT maps for the quadratic cost . These continuous solvers scale better to high dimension without discretizing the input measures , but they are too computationally expensive to be applied directly to JKO steps . Contributions . We propose a scalable parametric method to approximate Wasserstein gradient flows via JKO stepping using input-convex neural networks ( ICNNs ) [ 6 ] . Specifically , we leverage Brenier ’ s theorem to bypass the costly computation of the Wasserstein distance , and parametrize the optimal transport map as the gradient of an ICNN . Given sample access to the initial measure ρ0 , we use stochastic gradient descent ( SGD ) to sequentially learn time-discretized JKO dynamics of ρt . The trained model can sample from a continuous approximation of ρt and compute its density dρtdx ( x ) . We compute gradient flows for the Fokker-Planck free energy functional FFP given by ( 5 ) , but our method generalizes to other cases . We demonstrate performance by computing diffusion following the Fokker-Planck equation and applying it to unnormalized density sampling as well as nonlinear filtering . Notation . P2 ( RD ) denotes the set of Borel probability measures on RD with finite second moment . P2 , ac ( RD ) denotes its subset of probability measures absolutely continuous with respect to Lebesgue measure . For ρ ∈ P2 , ac ( RD ) , we denote by dρdx ( x ) its density with respect to the Lebesgue measure . Π ( µ , ν ) denotes the set of probability measures on RD×RD with marginals µ and ν . For measurable T : RD → RD , we denote by T ] the associated push-forward operator between measures . 2 Background on Wasserstein Gradient Flows . We consider gradient flows in Wasserstein space ( P2 ( RD ) , W2 ) , the space of probability measures with finite second moment on RD endowed with the Wasserstein-2 metricW2 . Wasserstein-2 distance . The ( squared ) Wasserstein-2 metricW2 between µ , ν ∈ P2 ( RD ) is W22 ( µ , ν ) def = min π∈Π ( µ , ν ) ∫ RD×RD ‖x− y‖22 dπ ( x , y ) , ( 1 ) where the minimum is over measures π on RD × RD with marginals µ and ν respectively [ 68 ] . For µ ∈ P2 , ac ( RD ) , there exists a µ-unique map ∇ψ∗ : RD → RD that is the gradient of a convex function ψ∗ : RD → R t { ∞ } satisfying ∇ψ∗ ] µ = ν [ 46 ] . From Brenier ’ s theorem [ 13 ] , it follows that π∗ = [ idRD , ∇ψ∗ ] ] µ is the unique minimizer of ( 1 ) , i.e. , W22 ( µ , ν ) = ∫ RD ‖x−∇ψ∗ ( x ) ‖22 dµ ( x ) . Wasserstein Gradient Flows . In the Euclidean case , gradient flows along a function f : R → R follow the steepest descent direction and are defined through the ODE dxtdt = −∇f ( xt ) . Discretization of this flow leads to the gradient descent minimization algorithm . When functionals are defined over the space of measures equipped with the Wasserstein-2 metric , the equivalent flow is called the Wasserstein gradient flow . The idea is similar : the flow follows the steepest descent direction , but this time the notion of gradient is more complex . We refer the reader to [ 4 ] for exposition of gradient flows in metric spaces , or [ 59 , Chapter 8 ] for an accessible introduction . A curve of measures { ρt } t∈R+ following the Wasserstein gradient flow of a functional F solves the continuity equation ∂ρt ∂t = div ( ρt∇xF ′ ( ρt ) ) , s.t . ρ0 = ρ0 , ( 2 ) where F ′ ( · ) is the first variation of F [ 4 , Theorem 8.3.1 ] . The term on the right can be understood as the gradient of F in Wasserstein space , a vector field perturbatively rearranging the mass in ρt to yield the steepest possible local change of F . Wasserstein gradient flows are used in various applied tasks . For example , gradient flows are applied in training [ 8 , 43 , 25 ] or refinement [ 7 ] of implicit generative models . In reinforcement learning , gradient flows facilitate policy optimization [ 55 , 72 ] . Other tasks include crowd motion modelling [ 45 , 58 , 50 ] , dataset optimization [ 2 ] , and in-between animation [ 26 ] . Many applications come from the connection between Wasserstein gradient flows and SDEs . Consider an RD-valued stochastic process { Xt } t∈R+ governed by the following Itô SDE : dXt = −∇Φ ( Xt ) dt+ √ 2β−1dWt , s.t . X0 ∼ ρ0 ( 3 ) where Φ : RD → R is the potential function , Wt is the standard Wiener process , and β > 0 is the magnitude . The solution of ( 3 ) is called an advection-diffusion process . The marginal measure ρt of Xt at each time satisfies the Fokker-Planck equation with fixed diffusion coefficient : ∂ρt ∂t = div ( ∇Φ ( x ) ρt ) + β−1∆ρt , s.t . ρ0 = ρ0 . ( 4 ) Equation ( 4 ) is the Wasserstein gradient flow ( 2 ) for F given by the Fokker-Planck free energy functional [ 32 ] FFP ( ρ ) = U ( ρ ) − β−1E ( ρ ) , ( 5 ) where U ( ρ ) = ∫ RD Φ ( x ) dρ ( x ) is the potential energy and E ( ρ ) = − ∫ RD log dρ dx ( x ) dρ ( x ) is the entropy . As the result , to solve the SDE ( 3 ) , one may compute the Wasserstein gradient flow of the Fokker-Planck equation with the free-energy functional FFP given by ( 5 ) . JKO Scheme . Computing Wasserstein gradient flows is challenging . The closed form solution is typically unknown , necessitating numerical approximation techniques . Jordan , Kinderlehrer , and Otto proposed a method—later abbreviated as JKO integration—to approximate the dynamics of ρt in ( 2 ) [ 32 ] . It consists of a time-discretization update of the continuous flow given by : ρ ( k ) ← arg min ρ∈P2 ( Rn ) [ F ( ρ ) + 1 2h W22 ( ρ ( k−1 ) , ρ ) ] ( 6 ) where ρ ( 0 ) = ρ0 is the initial condition and h > 0 is the time-discretization step size . The discrete time gradient flow converges to the continuous one as h → 0 , i.e. , ρ ( k ) ≈ ρkh . The method was further developed in [ 4 , 60 ] , but performing JKO iterations remains challenging thanks to the minimization with respect toW2 . A common approach to perform JKO steps is to discretize the spatial domain . For support size / 106 , ( 6 ) can be solved by standard optimal transport algorithms [ 51 ] . In dimensions D ≥ 3 , discrete supports can hardly approximate continuous distributions and hence the dynamics of gradient flows . To tackle this issue , [ 24 ] propose a stochastic parametric method to approximate the density of ρt . Their method uses entropy-regularized optimal transport ( OT ) , which is biased . 3 Computing Wasserstein Gradient Flows with ICNNs . We now describe our approach to compute Wasserstein gradient flows via JKO stepping with ICNNs . 3.1 JKO Reformulation via Optimal Push-forwards Maps . Our key idea is to replace the optimization ( 6 ) over probability measures by an optimization over convex functions , an idea inspired by [ 11 ] . Thanks to Brenier ’ s theorem , for any ρ ∈ P2 , ac there exists a unique ρ ( k−1 ) -measurable gradient ∇ψ : RD → RD of a convex function ψ satisfying ρ = ∇ψ ] ρ ( k−1 ) . We set ρ = ∇ψ ] ρ ( k−1 ) and rewrite ( 6 ) as an optimization over convex ψ : ψ ( k ) ← arg min Convex ψ [ F ( ∇ψ ] ρ ( k−1 ) ) + 1 2h W22 ( ρ ( k−1 ) , ∇ψ ] ρ ( k−1 ) ) ] . ( 7 ) To proceed to the next step of JKO scheme , we define ρ ( k ) def= ∇ψ ( k ) ] ρ ( k−1 ) . Since ρ is the pushforward of ρ ( k−1 ) by the gradient of a convex function ∇ψ , theW22 term in ( 7 ) can be evaluated explicitly , simplifying the Wasserstein-2 distance term in ( 7 ) : ψ ( k ) ← arg min Convex ψ [ F ( ∇ψ ] ρ ( k−1 ) ) + 1 2h ∫ RD ‖x−∇ψ ( x ) ‖22dρ ( k−1 ) ( x ) ] . ( 8 ) This formulation avoids the difficulty of computing Wasserstein-2 distances . An additional advantage is that we can sample from ρ ( k ) . Since ρ ( k ) = [ ∇ψ ( k ) ◦ · · · ◦ ∇ψ ( 1 ) ] ] ρ0 , one may sample x0 ∼ ρ ( 0 ) , and then∇ψ ( k ) ◦ · · · ◦ ∇ψ ( 1 ) ( x0 ) gives a sample from ρ ( k ) . Moreover , if functions ψ ( · ) are strictly convex , then gradients ∇ψ ( · ) are invertible . In this case , the density dρ ( k ) dx of ρ ( k ) = ∇ψ ( k ) ◦ · · · ◦ ∇ψ ( 1 ) ] ρ0 is computable by the change of variables formula ( assuming ψ ( · ) are twice differentiable ) dρ ( k ) dx ( xk ) = [ det∇2ψ ( k ) ( xk−1 ) ] −1 · · · [ det∇2ψ ( 1 ) ( x0 ) ] −1 · dρ ( 0 ) dx ( x0 ) , ( 9 ) where xi = ∇ψ ( i ) ( xi−1 ) for i = 1 , . . . , k and dρ ( 0 ) dx is the density of ρ ( 0 ) .

Summary: This paper is on the computation of Wasserstein gradient flow, a type of gradient flow based on the optimal transport theory. The proposed algorithm is built on the famous JKO scheme. To implement the JKO scheme, input-convex neural networks (ICNNs) are used to model the optimal transport map between the distributions at consecutive steps.