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: The authors propose an approximate of the natural gradient under Wasserstein metric when optimizing some cost function over a parametric family of probability distributions. The authors leverage the dual formulation and restrict the feasible space to a RKHS. The authors show a trade-off between accuracy and computational cost with theoretical guarantees for the proposed method, and empirically verify it for classification tasks.