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: In this work the authors introduce an innovative discretisation of the Jordan Kinderleher Otto scheme (which describes the evolution of the gradient flow of relative entropy in the Wasserstein metric).  To do so the authors iteratively construct a transformation map as a composition of gradients of input-convex neural networks, which satisfy a universal approximation theory among convex functions.  Using Brenier's theorem, an iterative variational formulation for the density at discrete time-steps is obtained, from which one can generate samples and or evaluate an unnormalised density.  A variety of applications are presented, including finding the stationary distribution, as well as nonlinear filtering.