Accelerated Gradient Flow for Probability Distributions

This paper presents a methodology and numerical algorithms for constructing accelerated gradient flows on the space of probability distributions . In particular , we extend the recent variational formulation of accelerated gradient methods in Wibisono et al . ( 2016 ) from vector valued variables to probability distributions . The variational problem is modeled as a mean-field optimal control problem . The maximum principle of optimal control theory is used to derive Hamilton ’ s equations for the optimal gradient flow . The Hamilton ’ s equation are shown to achieve the accelerated form of density transport from any initial probability distribution to a target probability distribution . A quantitative estimate on the asymptotic convergence rate is provided based on a Lyapunov function construction , when the objective functional is displacement convex . Two numerical approximations are presented to implement the Hamilton ’ s equations as a system of N interacting particles . The continuous limit of the Nesterov ’ s algorithm is shown to be a special case with N = 1 . The algorithm is illustrated with numerical examples . 1 INTRODUCTION . Optimization on the space of probability distributions is important to a number of machine learning models including variational inference ( Blei et al. , 2017 ) , generative models ( Goodfellow et al. , 2014 ; Arjovsky et al. , 2017 ) , and policy optimization in reinforcement learning ( Sutton et al. , 2000 ) . A number of recent studies have considered solution approaches to these problems based upon a construction of gradient flow on the space of probability distributions ( Zhang et al. , 2018 ; Liu & Wang , 2016 ; Frogner & Poggio , 2018 ; Chizat & Bach , 2018 ; Richemond & Maginnis , 2017 ; Chen et al. , 2018 ) . Such constructions are useful for convergence analysis as well as development of numerical algorithms . In this paper , we propose a methodology and numerical algorithms that achieve accelerated gradient flows on the space of probability distributions . The proposed numerical algorithms are related to yet distinct from the accelerated stochastic gradient descent ( Jain et al. , 2017 ) and Hamiltonian Markov chain Monte-Carlo ( MCMC ) algorithms ( Neal et al. , 2011 ; Cheng et al. , 2017 ) . The proposed methodology extends the variational formulation of ( Wibisono et al. , 2016 ) from vector valued variables to probability distributions . The original formulation of Wibisono et al . ( 2016 ) was used to derive and analyze the convergence properties of a large class of accelerated optimization algorithms , most significant of which is the continuous-time limit of the Nesterov ’ s algorithm ( Su et al. , 2014 ) . In this paper , the limit is referred to as the Nesterov ’ s ordinary differential equation ( ODE ) . The extension proposed in our work is based upon a generalization of the formula for the Lagrangian in Wibisono et al . ( 2016 ) : ( i ) the kinetic energy term is replaced with the expected value of kinetic energy ; and ( ii ) the potential energy term is replaced with a suitably defined functional on the space of probability distributions . The variational problem is to obtain a trajectory in the space of probability distributions that minimizes the action integral of the Lagrangian . The variational problem is modeled as a mean-field optimal problem . The maximum principle of the optimal control theory is used to derive the Hamilton ’ s equations which represent the first order optimality conditions . The Hamilton ’ s equations provide a generalization of the Nesterov ’ s ODE to the space of probability distributions . A candidate Lyapunov function is proposed for the convergence analysis of the solution of the Hamilton ’ s equations . In this way , quantitative estimates on convergence rate are obtained for the case when the objective functional is displacement convex ( McCann , 1997 ) . Table 1 provides a summary of the relationship between the original variational formulation in Wibisono et al . ( 2016 ) and the extension proposed in this paper . We also consider the important special case when the objective functional is the relative entropy functional D ( ρ|ρ∞ ) defined with respect to a target probability distribution ρ∞ . In this case , the accelerated gradient flow is shown to be related to the continuous limit of the Hamiltonian MonteCarlo algorithm ( Cheng et al. , 2017 ) ( Remark 1 ) . The Hamilton ’ s equations are finite-dimensional for the special case when the initial and the target probability distributions are both Gaussian . In this case , the mean evolves according to the Nesterov ’ s ODE . For the general case , the Lyapunov function-based convergence analysis applies when the target distribution is log-concave . As a final contribution , the proposed methodology is used to obtain a numerical algorithm . The algorithm is an interacting particle system that empirically approximates the distribution with a finite but large number of N particles . The difficult part of this construction is the approximation of the interaction term between particles . For this purpose , two types of approximations are described : ( i ) Gaussian approximation which is asymptotically ( as N →∞ ) exact in Gaussian settings ; and ( ii ) Diffusion map approximation which is computationally more demanding but asymptotically exact for a general class of distributions . The outline of the remainder of this paper is as follows : Sec . 2 provides a brief review of the variational formulation in Wibisono et al . ( 2016 ) . The proposed extension to the space of probability distribution appears in Sec . 3 where the main result is also described . The numerical algorithm along with the results of numerical experiments appear in Sec . 4 . Comparisons with MCMC and Hamiltonian MCMC are also described . The conclusions appear in Sec . 5 . Notation : The gradient and divergence operators are denoted as∇ and∇· respectively . With multiple variables , ∇z denotes the gradient with respect to the variable z . Therefore , the divergence of the vector field U is ∇ · U ( x ) = ∑d n=1∇xnUn ( x ) . The space of absolutely continuous probability measures on Rd with finite second moments is denoted by Pac,2 ( Rd ) . The Wasserstein gradient and the Gâteaux derivative of a functional F is denoted as∇ρF ( ρ ) and ∂F∂ρ ( ρ ) respectively ( see Appendix C for definition ) . The probability distribution of a random variable Z is denoted as Law ( Z ) . 2 REVIEW OF THE VARIATIONAL FORMULATION OF WIBISONO ET AL . ( 2016 ) . The basic problem is to minimize a C1 smooth convex function f on Rd . The standard form of the gradient descent algorithm for this problem is an ODE : dXt dt = −∇f ( Xt ) , t ≥ 0 ( 1 ) Accelerated forms of this algorithm are obtained based on a variational formulation due to Wibisono et al . ( 2016 ) . The formulation is briefly reviewed here using an optimal control formalism . The Lagrangian L : R+ × Rd × Rd → R is defined as L ( t , x , u ) : = eαt+γt ( 1 2 |e−αtu|2︸ ︷︷ ︸ kinetic energy − eβtf ( x ) ︸ ︷︷ ︸ potential energy ) ( 2 ) where t ≥ 0 is the time , x ∈ Rd is the state , u ∈ Rd is the velocity or control input , and the time-varying parameters αt , βt , γt satisfy the following scaling conditions : αt = log p − log t , βt = p log t+ logC , and γt = p log t where p ≥ 2 and C > 0 are constants . The variational problem is Minimize u J ( u ) = ∫ ∞ 0 L ( t , Xt , ut ) dt Subject to dXt dt = ut , X0 = x0 ( 3 ) The Hamiltonian function is H ( t , x , y , u ) = y · u− L ( t , x , u ) ( 4 ) where y ∈ Rd is dual variable and y · u denotes the dot product between vectors y and u . According to the Pontryagin ’ s Maximum Principle , the optimal control u∗t = arg max v H ( t , Xt , Yt , v ) = e αt−γtYt . The resulting Hamilton ’ s equations are dXt dt = +∇yH ( t , Xt , Yt , ut ) = eαt−γtYt , X0 = x0 ( 5a ) dYt dt = −∇xH ( t , Xt , Yt , ut ) = −eαt+βt+γt∇f ( Xt ) , Y0 = y0 ( 5b ) The system ( 5 ) is an example of accelerated gradient descent algorithm . Specifically , if the parameters αt , βt , γt are defined using p = 2 , one obtains the continuous-time limit of the Nesterov ’ s accelerated algorithm . It is referred to as the Nesterov ’ s ODE in this paper . For this system , a Lyapunov function is as follows : V ( t , x , y ) = 1 2 ∣∣x+ e−γty − x̄∣∣2 + eβt ( f ( x ) − f ( x̄ ) ) ( 6 ) where x̄ ∈ arg minx f ( x ) . It is shown in Wibisono et al . ( 2016 ) that upon differentiating along the solution trajectory , ddtV ( t , Xt , Yt ) ≤ 0 . This yields the following convergence rate : f ( Xt ) − f ( x̄ ) ≤ O ( e−βt ) , ∀t ≥ 0 ( 7 ) 3 VARIATIONAL FORMULATION FOR PROBABILITY DISTRIBUTIONS . 3.1 MOTIVATION AND BACKGROUND . Let F : Pac,2 ( Rd ) → R be a functional on the space of probability distributions . Consider the problem of minimizing F ( ρ ) . The ( Wasserstein ) gradient flow with respect to F ( ρ ) is ∂ρt ∂t = ∇ · ( ρt∇ρF ( ρt ) ) ( 8 ) where ∇ρF ( ρ ) is the Wasserstein gradient of F. An important example is the relative entropy functional where F ( ρ ) = D ( ρ|ρ∞ ) : =∫ Rd log ( ρ ( x ) ρ∞ ( x ) ) ρ ( x ) dx where ρ∞ ∈ Pac,2 ( Rd ) is referred to as the target distribution . The gradient of relative entropy is given by∇ρF ( ρ ) = ∇ log ( ρρ∞ ) . The gradient flow ∂ρt ∂t = −∇ · ( ρt∇ log ( ρ∞ ) ) + ∆ρt ( 9 ) is the Fokker-Planck equation ( Jordan et al. , 1998 ) . The gradient flow achieves the density transport from an initial probability distribution ρ0 to the target ( here , also equilibrium ) probability distribution ρ∞ ; and underlies the construction and the analysis of Markov chain Monte-Carlo ( MCMC ) algorithms . The simplest MCMC algorithm is the Langevin stochastic differential equation ( SDE ) : dXt = −∇f ( Xt ) dt+ √ 2 dBt , X0 ∼ ρ0 where Bt is the standard Brownian motion in Rd . The main problem of this paper is to construct an accelerated form of the gradient flow ( 8 ) . The proposed solution is based upon a variational formulation . As tabulated in Table 1 , the solution represents a generalization of Wibisono et al . ( 2016 ) from its original deterministic finite-dimensional to now probabilistic infinite-dimensional settings . The variational problem can be expressed in two equivalent forms : ( i ) The probabilistic form is described next in the main body of the paper ; and ( ii ) The partial differential equation ( PDE ) form appears in the Appendix . The probabilistic form is stressed here because it represents a direct generalization of the Nesterov ’ s ODE and because it is closer to the numerical algorithm .

Summary: This paper introduces a functional extension of the Bregman Lagrangian framework of Wibisono et al. 2016. The basic idea is to define accelerated gradient flows on the space of probability distribution. Because the defined flows include a term depending on the current distribution of the system, which is difficult to compute in general, the authors introduce an interacting particle approximation as a practical numerical approximation. The experiments are a proof-of-concept on simple illustrative toy examples.