Kernelized Wasserstein Natural Gradient

1 INTRODUCTION . The success of machine learning algorithms relies on the quality of an underlying optimization method . Many of the current state-of-the-art methods rely on variants of Stochastic Gradient Descent ( SGD ) such as AdaGrad ( Duchi et al. , 2011 ) , RMSProp ( Hinton et al. , 2012 ) , and Adam ( Kingma and Ba , 2014 ) . While generally effective , the performance of such methods remains sensitive to the curvature of the optimization objective . When the Hessian matrix of the objective at the optimum has a large condition number , the problem is said to have a pathological curvature ( Martens , 2010 ; Sutskever et al. , 2013 ) . In this case , the first-order optimization methods tend to have poor performance . Using adaptive step sizes can help when the principal directions of curvature are aligned with the coordinates of the vector parameters . Otherwise , an additional rotation of the basis is needed to achieve this alignment . One strategy is to find an alternative parametrization of the same model that has a better-behaved curvature and is thus easier to optimize with standard first-order optimization methods . Designing good network architectures ( Simonyan and Zisserman , 2014 ; He et al. , 2015 ) along with normalization techniques ( LeCun et al. , 2012 ; Ioffe and Szegedy , 2015 ; Salimans and Kingma , 2016 ) is often critical for the success of such optimization methods . The natural gradient method ( Amari , 1998 ) takes a related but different perspective . Rather than reparametrizing the model , the natural gradient method tries to make the optimizer itself invariant to reparameterizations by directly operating on the manifold of probability distributions . This requires endowing the parameter space with a suitable notion of proximity formalized by a metric . An important metric in this context is the Fisher information metric ( Fisher and Russell , 1922 ; Rao , 1992 ) , which induces the Fisher-Rao natural gradient ( Amari , 1985 ) . Another important metric in probability space is the Wasserstein metric ( Villani , 2009 ; Otto , 2001 ) , which induces the Wasserstein natural gradient ( Li and Montufar , 2018a ; b ; Li , 2018 ) ; see similar formulations in Gaussian families ( Malagò et al. , 2018 ; Modin , 2017 ) . In spite of their numerous theoretical advantages , applying natural gradient methods is challenging in practice . Indeed , each parameter update requires inverting the metric tensor . This becomes infeasible for current deep learning models , which typically have millions of parameters . This has motivated research into finding efficient algorithms to estimate the natural gradient ( Martens and Grosse , 2015 ; Grosse and Martens , 2016 ; George et al. , 2018 ; Heskes , 2000 ; Bernacchia et al. , 2018 ) . Such algorithms often address the case of the Fisher metric and either exploit a particular structure of the parametric family or rely on a low rank decomposition of the information matrix . Recently , Li et al . ( 2019 ) proposed to estimate the metric based on a dual formulation and used this estimate in a proximal method . While this avoids explicitly computing the natural gradient , the proximal method also introduces an additional optimization problem to be solved at each update of the model ’ s parameters . The quality of the solver will thus depend on the accuracy of this additional optimization . In this paper , we use the dual formulation of the metric to directly obtain a closed form expression of the natural gradient as a solution to a convex functional optimization problem . We focus on the Wasserstein metric as it has the advantage of being well defined even when the model doesn ’ t admit a density . The expression remains valid for general metrics including the Fisher-Rao metric . We leverage recent work on Kernel methods ( Sriperumbudur et al. , 2017 ; Arbel and Gretton , 2017 ; Sutherland et al. , 2017 ; Mroueh et al. , 2019 ) to compute an estimate of the natural gradient by restricting the functional space appearing in the dual formulation to a Reproducing Kernel Hilbert Space . We demonstrate empirically the accuracy of our estimator on toy examples , and show how it can be effectively used to approximate the trajectory of the natural gradient descent algorithm . We also analyze the effect of the dimensionality of the model on the accuracy of the proposed estimator . Finally , we illustrate the benefits of our proposed estimator for solving classification problems when the model has an ill-conditioned parametrization . The paper is organized as follows . In Section 2 , after a brief description of natural gradients , we discuss Legendre duality of metrics , and provide details on the Wasserstein natural gradient . In Section 3 , we present our kernel estimator of the natural gradient . In Section 4 we present experiments to evaluate the accuracy of the proposed estimator and demonstrate its effectiveness in supervised learning tasks . 2 NATURAL GRADIENT DESCENT . We first briefly recall the natural gradient descent method in Section 2.1 , and its relation to metrics on probability distribution spaces in Section 2.2 . We next present Legendre dual formulations for metrics in Section 2.3 where we highlight the Fisher-Rao and Wasserstein metrics as important examples . 2.1 GENERAL FORMULATION . It is often possible to formulate learning problems as the minimization of some cost functional ρ 7→F ( ρ ) over probability distributions ρ from a parametric model PΘ . The set PΘ contains probability distributions defined on an open sample space Ω⊂Rd and parametrized by some vector θ∈Θ , where Θ is an open subset of Rq . The learning problem can thus be formalized as finding an optimal value θ∗ that locally minimizes a loss function L ( θ ) : =F ( ρθ ) defined over the parameter space Θ . One convenient way to solve this problem approximately is by gradient descent , which uses the Euclidean gradient of L w.r.t . the parameter vector θ to produce a sequence of updates θt according to the following rule : θt+1 =θt−γt∇L ( θt ) . Here the step-size γt is a positive real number . The Euclidean gradient can be viewed as the direction in parameter space that leads to the highest decrease of some linear modelMt of the cost function L per unit of change of the parameter . More precisely , the Euclidean gradient is obtained as the solution of the optimization problem : ∇L ( θt ) =−argmin u∈Rq Mt ( u ) + 1 2 ‖u‖2 . ( 1 ) The linear modelMt is an approximation of the cost function L in the neighborhood of θt and is simply obtained by a first order expansion : Mt ( u ) =L ( θt ) +∇L ( θt ) > u . The quadratic term ‖u‖2 penalizes the change in the parameter and ensures that the solution remains in the neighborhood where the linear model is still a good approximation of the cost function . This particular choice of quadratic term is what defines the Euclidean gradient descent algorithm , which can often be efficiently implemented for neural network models using back-propagation . The performance of this algorithm is highly dependent on the parametrization of the model PΘ , however ( Martens , 2010 ; Sutskever et al. , 2013 ) . To obtain an algorithm that is robust to parametrization , one can take advantage of the structure of the cost function L ( θ ) which is obtained as the composition of the functional F and the model θ 7→ρθ and define a generalized natural gradient ( Amari and Cichocki , 2010 ) . We first provide a conceptual description of the general approach to obtain such gradient . The starting point is to choose a divergenceD between probability distributions and use it as a new penalization term : −argmin u∈Rq Mt ( u ) + 1 2 D ( ρθt , ρθt+u ) . ( 2 ) Here , changes in the model are penalized directly in probability space rather than parameter space as in ( 1 ) . In the limit of small u , the penalization term can be replaced by a quadratic term u > GD ( θ ) uwhere GD ( θ ) contains second order information about the model as measured byD . This leads to the following expression for the generalized natural gradient∇DL ( θt ) where the dependence inD is made explicit : ∇DL ( θt ) : =−argmin u∈Rq Mt ( u ) + 1 2 u > GD ( θt ) u . ( 3 ) From ( 3 ) , it is possible to express the generalized natural gradient by means of the Euclidean gradient : ∇DL ( θt ) =GD ( θt ) −1∇L ( θt ) . The parameter updates are then obtained by the new update rule : θt+1 =θt−γtGD ( θt ) −1∇L ( θt ) . ( 4 ) Equation ( 4 ) leads to a descent algorithm which is invariant to parametrization in the continuous-time limit : Proposition 1 . Let Ψ be an invertible and smoothly differentiable re-parametrization ψ = Ψ ( θ ) and denote by L̄ ( ψ ) : =L ( Ψ−1 ( ψ ) ) . Consider the continuous-time natural gradient flows : θ̇s=−∇Dθ L ( θs ) , ψ̇s=−∇Dψ L̄ ( ψs ) , ψ0 =Ψ ( θ0 ) Then ψs and θs are related by the equation ψs=Ψ ( θs ) at all times s≥0 . This result implies that an ill-conditioned parametrization of the model has little effect on the optimization when ( 4 ) is used . It is a consequence of the transformation properties of the natural gradient by change of parametrization : ∇Dψ L̄ ( ψ ) = ∇θΨ ( θ ) ∇Dθ L ( θ ) which holds in general for any covariant gradient . We provide a proof of Proposition 1 in Appendix C.1 in the particular the case when D is either Kullback-Leibler divergence F , or the squared Wasserstein-2 distanceW using notions introduced later in Section 2.3 and refer to Ollivier et al . ( 2011 ) for a detailed discussion . The approach based on ( 2 ) for defining the generalized natural gradient is purely conceptual and can be formalized using the notion of metric tensor from differential geometry which allows for more generality . In Section 2.2 , we provide such formal definition in the case when D is either the Kullback-Leibler divergence F , or the squared Wasserstein-2 distanceW . 2.2 INFORMATION MATRIX VIA DIFFERENTIAL GEOMETRY . WhenD is the Kullback-Leibler divergence or relative entropy F , then ( 3 ) defines the Fisher-Rao natural gradient∇FL ( θ ) ( Amari , 1985 ) andGF ( θ ) is called the Fisher information matrix . GF ( θ ) is well defined when the probability distributions inPΘ all have positive densities , and when some additional differentiability and integrability assumptions on ρθ are satisfied . In fact , it has an interpretation in Riemannian geometry as the pull-back of a metric tensor gF defined over the set of probability distributions with positive densities and known as the Fisher-Rao metric ( see Definition 4 in Appendix B.1 ; see also Holbrook et al . 2017 ) : Definition 1 ( Fisher information matrix ) . Assume θ 7→ ρθ ( x ) is differentiable for all x on Ω and that∫ ‖∇ρθ ( x ) ‖2 ρθ ( x ) dx < ∞ . Then the Fisher information matrix is defined as the pull-back of the Fisher-Rao metric gF : GF ( θ ) ij=g F ρθ ( ∂iρθ , ∂jρθ ) : = ∫ fi ( x ) fj ( x ) ρθ ( x ) dx , where the functions fi on Ω are given by : fi= ∂iρθρθ . Definition 1 directly introducesGF using the Fisher-Rao metric tensor which captures the infinitesimal behavior of the KL . This approach can be extended to any metric tensor g defined on a suitable space of probability distributions containing PΘ . In particular , when D is the Wasserstein-2 , the Wasserstein information matrix is obtained directly by means of the Wasserstein-2 metric tensor gW ( Otto and Villani , 2000 ; Lafferty and Wasserman , 2008 ) as proposed in Li and Montufar ( 2018a ) ; Chen and Li ( 2018 ) : Definition 2 ( Wasserstein information matrix ) . The Wasserstein information matrix ( WIM ) is defined as the pull-back of the Wasserstein 2 metric gW : GW ( θ ) ij=g W ρθ ( ∂iρθ , ∂jρθ ) : = ∫ φi ( x ) > φj ( x ) dρθ ( x ) , where φi are vector valued functions on Ω that are solutions to the partial differential equations with Neumann boundary condition : ∂iρθ=−div ( ρθφi ) , ∀1≤i≤q . Moreover , φi are required to be in the closure of the set of gradients of smooth and compactly supported functions inL2 ( ρθ ) d. In particular , when ρθ has a density , φi=∇xfi , for some real valued function fi on Ω . The partial derivatives ∂iρθ should be understood in distribution sense , as discussed in more detail in Section 2.3 . This allows to define the Wasserstein natural gradient even when the model ρθ does not admit a density . Moreover , it allows for more generality than the conceptual approach based on ( 2 ) which would require performing a first order expansion of the Wasserstein distance in terms of its linearized version known as the Negative Sobolev distance . We provide more discussion of those two approaches and their differences in Appendix B.3 . From now on , we will focus on the above two cases of the natural gradient∇DL ( θ ) , namely∇FL ( θ ) and∇WL ( θ ) . When the dimension of the parameter space is high , directly using equation ( 4 ) becomes impractical as it requires storing and inverting the matrix G ( θ ) . In Section 2.3 we will see how equation ( 3 ) can be exploited along with Legendre duality to get an expression for the natural gradient that can be efficiently approximated using kernel methods .

Summary: Natural gradient has been proven effective in many statistical learning algorithms. A well-known difficulty in using natural gradient is that it is tedious to compute the Fisher matrix (if one is using Fisher-Rao metric) and the Wasserstein information matrix (if one is using Wasserstein metric). It's important to be able to estimate natural gradient in a practical way, and there have been a few papers looking at this problem but mostly for the case with a Fisher-Rao metric. This paper takes a different and general approach to approximate natural gradient by leveraging the dual formulation for the metric restricted to the Reproducing Kernel Hilbert Space. Some theoretical guarantees of the proposed method is established together to some experimental study.