Accelerated Information Gradient flow

1 INTRODUCTION . Recently , optimization problems on the space of probability and probability models attract increasing attentions from machine learning communities . These problems include variational inference ( Blei et al. , 2017 ) , Bayesian inference ( Liu & Wang , 2016 ) , Generative Adversary Networks ( GAN , Goodfellow et al . ( 2014 ) ) , and policy optimizations ( Zhang et al. , 2018 ) , etc . For instance , variational inference methods approximate a target density by minimizing the Kullback-Leibler ( KL ) divergence as the loss ( objective ) function . Gradient descent methods with sampling efficient properties play essential roles to solve these optimization problems . Here the gradient descent direction often relies on the information metric over the probability space . This direction naturally reflects the change of the loss function with respect to the metric . In literature , two important information metrics , such as the Fisher-Rao metric and the Wasserstein-2 ( in short , Wasserstein ) metric , are of great interests ( Amari , 1998 ; Otto , 2001 ; Lafferty , 1988 ) . For the Fisher-Rao gradient , classical results including Adam ( Kingma & Ba , 2014 ) and K-FAC ( Martens & Grosse , 2015 ) demonstrate its effectiveness in probability models . For the Wasserstein gradient , many classical methods such as Markov chain Monte Carlo ( MCMC ) methods ( Geman & Geman , 1987 ; Neal et al. , 2011 ; Welling & Teh , 2011 ) and particle-based variational inference ( ParVI ) methods ( Liu & Wang , 2016 ; Chen & Zhang , 2017 ; Chen et al. , 2018 ) are based on this framework in the probability space . The strength of using the Wasserstein gradient is also shown in probability models such as GANs . ( Arjovsky et al. , 2017 ; Lin et al. , 2018 ; Li et al. , 2019 ) . The Nesterov ’ s accelerated method ( Nesterov , 1983 ) is widely applied in accelerating the vanilla gradient descent under the Euclidean metric . It corresponds to a damped Hamiltonian flow , known as the accelerated gradient flow ( Su et al. , 2016 ) . A natural question is whether there exists a counterpart of the accelerated gradient flow in the probability space under information metrics . For optimization problems on a Riemannian manifold , the accelerated gradient methods are studied by Liu et al . ( 2017 ) ; Zhang & Sra ( 2018 ) . The probability space embedded with information metric can be viewed as a Riemannian manifold . Several previous works explore accelerated methods in this manifold under the Wasserstein metric . Liu et al . ( 2018 ; 2019 ) propose an acceleration framework of ParVI methods based on manifold optimization . Taghvaei & Mehta ( 2019 ) introduce the accelerated flow from an optimal control perspective . On the other hand , Cheng et al . ( 2017 ) ; Ma et al . ( 2019 ) explore and analyze the acceleration on MCMC , based on the underdamped Langevin dynamics . In this paper , we present a unified framework of accelerated gradient flows in the probability space embedded with information metrics , named Accelerated Information Gradient ( AIG ) flows . From an information-differential-geometry perspective , we derive AIG flows by damping Hamiltonian flows , concerning both the Fisher-Rao metric and the Wasserstein metric . Then we focus on the Wasser- stein metric with the KL divergence loss function . In Gaussian families , we verify the existence of the solution to AIG flows . Here we show that the AIG flow corresponds to a well-posed ODE system in the space of symmetric positive definite matrices . We rigorously prove the convergence rate of AIG flows based on the geodesic convexity of the loss function . Here we note that our proof removes the unnecessary technical assumption in ( Taghvaei & Mehta , 2019 , Theorem 1 ) . Besides , we handle two difficulties in numerical implementations of AIG flows . On the one hand , as pointed out by Taghvaei & Mehta ( 2019 ) ; Liu et al . ( 2019 ) , the logarithm of density term ( Wasserstein gradient of KL divergence ) is hard to approximate in particle formulations . We propose a novel kernel selection method , whose bandwidth is learned by sampling from Brownian motions . We call it the BM method . On the other hand , we notice that the AIG flow can be a numerically stiff system , especially in high-dimensional sample spaces . This is because the solution of AIG flows can be close to the boundary of the probability space . To handle this issue , we propose an adaptive restart technique , which accelerates and stabilizes the AIG algorithm . Numerical results in toy examples , Gaussian measures and Bayesian Logistic regression indicate the validity of the BM method and the acceleration effects of the proposed AIG flow . This paper is organized as follows . Section 2 briefly reviews the information metrics and their corresponding gradient flows and Hamiltonian flows in the probability space . In Section 3 , we formulate various forms of AIG flows and analyze W-AIG flows in Gaussian measures . We theoretically prove the convergence rate of W-AIG flows in Section 4 . Section 5 presents the discrete-time algorithm for W-AIG flows , including the BM method and the adaptive restart technique . Section 6 provides numerical experiments . 2 METRIC AND FLOWS IN THE PROBABILITY SPACE . Suppose that Ω is a region in Rn . Let F ( Ω ) denote the set of smooth functions on Ω . 〈· , ·〉 and ‖ · ‖ are the Euclidean inner product and norm in Rn . ∇ , ∇· and ∆ represent the gradient , divergence and Laplacian operators in Rn . Denote the set of probability density P ( Ω ) = { ρ ∈ F ( Ω ) : ∫ Ω ρdx = 1 , ρ ≥ 0 } . The tangent space at ρ ∈ P ( Ω ) follows Tρ ( Ω ) = { σ ∈ F ( Ω ) : ∫ σdx = 0 . } . The cotangent space at ρ , T ∗ρP ( Ω ) , can be treated as the quotient space F ( Ω ) /R , which are functions in F ( Ω ) defined up to addition of constants . Definition 1 ( Metric in the probability space ) A metric tensor G ( ρ ) : TρP ( Ω ) → T ∗ρP ( Ω ) is an invertible mapping from the tangent space at ρ to the cotangent space at ρ . This metric tensor defines the metric ( inner product ) on the tangent space TρP ( Ω ) . Namely , for σ1 , σ2 ∈ TρP ( Ω ) , we define gρ ( σ1 , σ2 ) = ∫ σ1G ( ρ ) σ2dx = ∫ Φ1G ( ρ ) −1Φ2dx , where Φi is the solution to σi = G ( ρ ) −1Φi , i = 1 , 2 . We present two important examples of metrics on the probability space P ( Ω ) : the Fisher-Rao metric from information geometry and the Wasserstein metric from optimal transport . Example 1 ( Fisher-Rao metric ) The inverse of the Fisher-Rao metric tensor is defined by GF ( ρ ) −1Φ = ρ ( Φ− ∫ Φρdx ) , Φ ∈ T ∗ρP ( Ω ) . The Fisher-Rao metric on the tangent space is given by gFρ ( σ1 , σ2 ) = ∫ Φ1Φ2ρdx− ( ∫ Φ1ρdx ) ( ∫ Φ2ρdx ) , σ1 , σ2 ∈ TρP ( Ω ) , where Φi is the solution to σi = ρ ( Φi − ∫ Φiρdx ) , i = 1 , 2 . Example 2 ( Wasserstein metric ) The inverse of the Wasserstein metric tensor is defined by GW ( ρ ) −1Φ = −∇ · ( ρ∇Φ ) , Φ ∈ T ∗ρP ( Ω ) . The Wasserstein metric on the tangent space is given by gWρ ( σ1 , σ2 ) = ∫ ρ 〈∇Φ1 , ∇Φ2〉 dx , σ1 , σ2 ∈ TρP ( Ω ) , where Φi is the solution to σi = −∇ · ( ρ∇Φi ) , i = 1 , 2 . 2.1 GRADIENT FLOWS . In learning , many problems can be formulated as the optimization problem in the probability space , min ρ∈P ( Ω ) E ( ρ ) . Here E ( ρ ) is a divergence or metric loss functional between ρ and a target density ρ∗ ∈ P ( Ω ) . One typical example of E ( ρ ) is the KL divergence from ρ to ρ∗ , E ( ρ ) = DKL ( ρ‖ρ∗ ) = ∫ log ( ρ ρ∗ ) ρdx . Another example is the maximum mean discrepancy ( MMD , Gretton et al . ( 2012 ) ) , E ( ρ ) = MMD ( ρ , ρ∗ ) = ∫ ∫ ( ρ ( x ) − ρ∗ ( x ) ) K ( x , y ) ( ρ ( y ) − ρ∗ ( y ) ) dxdy , where K ( x , y ) is a given kernel function . The gradient flow for E ( ρ ) in ( P ( Ω ) , gρ ) takes the form ∂tρt = −G ( ρt ) −1 δE ( ρt ) δρt . Here δE ( ρt ) δρt is the L. 2 first variation w.r.t . ρt . We formulate the gradient flow under either the Fisher-Rao metric or the Wasserstein metric . Example 3 ( Fisher-Rao gradient flow ) The Fisher-Rao gradient flow is given by ∂tρt = −ρt ( δE δρt − ∫ δE δρt ρtdx ) . Example 4 ( Wasserstein gradient flow ) The Wasserstein gradient flow writes ∂tρt = ∇ · ( ρt∇ δE δρt ) . 2.2 HAMILTONIAN FLOW . In this subsection , we briefly review the Hamiltonian flow in the probability space . By using the metric gρ in the probability space , we can define a Lagrangian by L ( ρt , ∂tρt ) = 1 2 gρt ( ∂tρt , ∂tρt ) − E ( ρt ) . The Euler-Lagrange equation for the Lagrangian follows ∂t ( δL δ ( ∂tρt ) ) = δL δρt + C ( t ) , ( 1 ) where C ( t ) is a spatially-constant function . Proposition 1 If we let Φt = δL/δ ( ∂tρt ) = G ( ρt ) ∂tρt , the equation ( 1 ) can be formulated as a system of ( ρt , Φt ) , i.e. , ∂tρt −G ( ρt ) −1Φt = 0 , ∂tΦt + 1 2 δ δρt ( ∫ ΦtG ( ρt ) −1Φtdx ) + δE δρt = 0 , ( 2 ) where Φt is up to a spatially-constant function shrift . Here ( 2 ) is the Hamiltonian flow ∂t [ ρt Φt ] − [ 0 1 −1 0 ] [ δ δρt H ( ρt , Φt ) δ δΦt H ( ρt , Φt ) ] = 0 , ( 3 ) with respect to the HamiltonianH ( ρt , Φt ) = 12 ∫ ΦtG ( ρt ) −1Φtdx+ E ( ρt ) . Similarly , we can write the Hamiltonian flow under the Fisher-Rao metric or the Wasserstein metric . Example 5 ( Fisher-Rao Hamiltonian flow ) The Fisher-Rao Hamiltonian flow follows ∂tρt − Φtρt + ∫ Φtρtdx = 0 , ∂tΦt + 1 2 Φ2t − ( ∫ ρtΦtdx ) Φt + δE δρt = 0 . The corresponding Hamiltonian isHF ( ρt , Φt ) = 12 ( ∫ Φ2tρtdx− ( ∫ ρtΦtdx ) 2 ) + E ( ρt ) . Example 6 ( Wasserstein Hamiltonian flow ) The Wasserstein Hamiltonian flow writes ∂tρt +∇ · ( ρt∇Φt ) = 0 , ∂tΦt + 1 2 ‖∇Φt‖2 + δE δρt = 0 . The corresponding Hamiltonian is HW ( ρt , Φt ) = 12 ∫ ‖∇Φt‖2ρtdx + E ( ρt ) . This is identical to the Wasserstein Hamiltonian flow introduced by Chow et al . ( 2019 ) .

Summary: The paper attempts to develop a counterpart of the well-known Nesterov accelerated gradient method for gradient flows on the space of probability measures equipped with an information metric. This is an important problem which is useful for optimization on probability spaces. The accelerated gradient flow is developed by leveraging a damping Hamiltonian flow. The paper focuses mainly on the case with the Wasserstein metric and provides a convergence analysis. Practical considerations such as discretizing the accelerated flow and bandwidth selection are developed for the use of the method in practical problems.