KERNELIZED WASSERSTEIN NATURAL GRADIENT
Year: 2020

Abstract:
Many machine learning problems can be expressed as the optimization of some cost functional over a parametric family of probability distributions. It is often beneficial to solve such optimization problems using natural gradient methods. These methods are invariant to the parametrization of the family, and thus can yield more effective optimization. Unfortunately, computing the natural gradient is challenging as it requires inverting a high dimensional matrix at each iteration. We propose a general framework to approximate the natural gradient for the Wasserstein metric, by leveraging a dual formulation of the metric restricted to a Reproducing Kernel Hilbert Space. Our approach leads to an estimator for gradient direction that can trade-off accuracy and computational cost, with theoretical guarantees. We verify its accuracy on simple examples, and show the advantage of using such an estimator in classification tasks on Cifar10 and Cifar100 empirically.

heading: 1 INTRODUCTION
text: 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.

heading: 2 NATURAL GRADIENT DESCENT
text: 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.

heading: 2.1 GENERAL FORMULATION
text: 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 .

heading: 2.2 INFORMATION MATRIX VIA DIFFERENTIAL GEOMETRY
text: 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.

heading: 2.3 LEGENDRE DUALITY FOR METRICS
text: In this section we provide an expression for the natural gradient defined in (3) as the solution of a saddle-point optimization problem. It exploits Legendre duality for metrics to express the quadratic term u>G(θ)u as a solution to a functional optimization problem over C∞c (Ω), the set of smooth and compactly supported functions on Ω. The starting point is to extend the notion of gradient∇ρθ which appears in Definitions 1 and 2 to the distributional sense of Definition 3 below.
Definition 3. Given a parametric family PΘ of probability distributions, we say that ρθ admits a distributional gradient at point θ if there exists a linear continuous map∇ρθ :C∞c (Ω)→Rq such that:∫
f(x)dρθ+ u(x)− ∫ f(x)dρθ(x)= ∇ρθ(f)>u+ δ( ,f,u) ∀f∈C∞c (Ω), ∀u∈Rq
where δ( ,f,u) depends on f and u and converges to 0 as approaches 0. ∇ρθ is called the distributional gradient of ρθ at point θ.
When the distributions in PΘ have a density, written x 7→ρθ(x) by abuse of notation, that is differentiable w.r.t. θ and with a jointly continuous gradient in θ and x then∇ρθ(f) is simply given by ∫ f(x)∇θρθ(x)dx as shown in Proposition 12 of Appendix C.1. In this case, the Fisher-Rao natural gradient admits a formulation as a saddle point solution involving∇ρθ and provided in Proposition 2 with a proof in Appendix C.1. Proposition 2. Under the same assumptions as in Definition 1, the Fisher information matrix admits the dual formulation:
1 2 u>GF (θ)u := sup
f∈C∞c (Ω)∫ f(x)dρθ(x)=0
∇ρθ(f)>u− 1
2
∫ f(x)2dρθ(x)dx. (5)
Moreover, defining Uθ(f):=∇L(θ)+∇ρθ(f), the Fisher-Rao natural gradient satisfies:
∇FL(θ)=−argmin u∈Rq sup f∈C∞c (Ω)∫ f(x)dρθ(x)=0
Uθ(f)>u− 1
2
∫ f(x)2dρθ(x)dx,
Another important case is when PΘ is defined as an implicit model. In this case, any sample x from a distribution ρθ in PΘ is obtained as x=hθ(z), where z is a sample from a fixed latent distribution ν
defined over a latent space Z and (θ,z) 7→hθ(z) is a deterministic function with values in Ω. This can be written in a more compact way as the push-forward of ν by the function hθ:
PΘ :={ρθ :=(hθ)#ν |θ∈Ω}. (6)
A different expression for∇ρθ is obtained in the case of implicit models when θ 7→hθ(z) is differentiable for ν-almost all z and∇hθ is square integrable under ν:
∇ρθ(f)= ∫ ∇hθ(z)>∇xf(hθ(z))dν(z). (7)
Equation (7) is also known as the re-parametrization trick (Kingma et al., 2015) and allows to derive a dual formulation of the Wasserstein natural gradient in the case of implicit models. Proposition 3 below provides such formulation under mild assumptions stated in Appendix A.2 along with a proof in Appendix C.1. Proposition 3. Assume PΘ is defined by (6) such that ∇ρθ is given by (7). Under Assumptions (B) and (C), the Wasserstein information matrix satisfies:
1 2 u>GW (θ)u= sup f∈C∞c (Ω) ∇ρθ(f)>u− 1 2
∫ ‖∇xf(x)‖2dρθ(x) (8)
and the Wasserstein natural gradient satisfies:
∇WL(θ)=−argmin u∈Rq sup f∈C∞c (Ω)
Uθ(f)>u− 1
2
∫ ‖∇xf(x)‖2dρθ(x). (9)
The similarity between the variational formulations provided in Propositions 2 and 3 is worth noting. A first difference however, is that Proposition 3 doesn’t require the test functions f to have 0 mean under ρθ. This is due to the form of the objective in (8) which only depends on the gradient of f . More importantly, while (8) is well defined, the expression in (5) can be infinite when∇ρθ is given by (7). Indeed, if the ρθ doesn’t admit a density, it is always possible to find an admissible function f∈C∞c (Ω) with bounded second moment under ρθ but for which ∇ρθ(f) is arbitrarily large. This is avoided in (8) since the quadratic term directly penalizes the gradient of functions instead. For similar reasons, the dual formulation of the Sobolev distance considered in Mroueh et al. (2019) can also be infinite in the case of implicit models as discussed in Appendix B.3 although formally similar to (8). Nevertheless, a similar estimator as in Mroueh et al. (2019) can be considered using kernel methods which is the object of Section 3.

heading: 3 KERNELIZED WASSERSTEIN NATURAL GRADIENT
text: In this section we propose an estimator for the Wasserstein natural gradient using kernel methods and exploiting the formulation in (9). We restrict to the case of the Wasserstein natural gradient (WNG), denoted by∇WL(θ), as it is well defined for implicit models, but a similar approach can be used for the Fisher-Rao natural gradient in the case of models with densities. We first start by presenting the kernelized Wasserstein natural gradient (KWNG) in Section 3.1, then we introduce an efficient estimator for KWNG in Section 3.2. In Section 3.4 we provide statistical guarantees and discuss practical considerations in Section 3.3.

heading: 3.1 GENERAL FORMULATION AND MINIMAX THEOREM
text: Consider a Reproducing Kernel Hilbert Space (RKHS) H which is a Hilbert space endowed with an inner product 〈.,.〉H along with its norm ‖.‖H. H has the additional property that there exists a symmetric positive semi-definite kernel k : Ω×Ω 7→ R such that k(x,.) ∈ H for all x ∈ Ω and satisfying the Reproducing property for all functions f inH:
f(x)=〈f,k(x,.)〉H, ∀x∈Ω. (10)
The above property is central in all kernel methods as it allows to obtain closed form expressions for some class of functional optimization problems. In order to take advantage of such property for estimating the natural gradient, we consider a new saddle problem obtained by restricting (9) to functions in the RKHS H and adding some regularization terms:‹∇WL(θ):=−min
u∈Rq sup f∈H Uθ(f)>u−
1
2
∫ ‖∇xf(x)‖2dρθ(x)+ 1
2 ( u>D(θ)u−λ‖f‖2H). (11)
The kernelized Wasserstein natural gradient is obtained by solving (11) and is denoted by ‹∇WL(θ). Here, is a positive real numbers, λ is non-negative whileD(θ) is a diagonal matrix in Rq with positive diagonal elements whose choice will be discussed in Section 3.3. The first regularization term makes the problem strongly convex in u, while the second term makes the problem strongly concave in f when λ>0. When λ=0, the problem is still concave in f . This allows to use a version of the minimax theorem (Ekeland and Témam, 1999, Proposition 2.3, Chapter VI) to exchange the order of the supremum and minimum which also holds true when λ=0. A new expression for the kernelized natural gradient is therefore obtained: Proposition 4. Assume that >0 and λ>0, then the kernelized natural gradient is given by:‹∇WL(θ)= 1 D(θ)−1Uθ(f∗), (12)
where f∗ is the unique solution to the quadratic optimization problem:
inf f∈H J (f):=
∫ ‖∇xf(x)‖2dρθ(x)+ 1 Uθ(f)>D(θ)−1Uθ(f)+λ‖f‖2H. (13)
When λ=0, f∗ might not be well defined, still, we have: ‹∇WL(θ)= limj→∞ 1 D(θ)−1Uθ(fj) for any limiting sequence of (13).
Proposition 4 allows to compute the kernelized natural gradient directly, provided that the functional optimization (13) can be solved. This circumvents the direct computation and inversion of the metric as suggested by (11). In Section 3.2, we propose a method to efficiently compute an approximate solution to (13) using Nyström projections. We also show in Section 3.4 that restricting the space of functions to H can still lead to a good approximation of the WNG provided thatH enjoys some denseness properties.

heading: 3.2 NYSTRÖM METHODS FOR THE KERENALIZED NATURAL GRADIENT
text: We are interested now in finding an approximate solution to (13) which will allow to compute an estimator for the WNG using Proposition 4. Here we considerN samples (Zn)1≤n≤N from the latent distribution ν which are used to produceN samples (Xn)1≤n≤N from ρθ using the map hθ, i.e.,Xn=hθ(Zn). We also assume we have access to an estimate of the Euclidean gradient∇L(θ) which is denoted by ÷∇L(θ). This allows to compute an empirical version of the cost function in (13),
Ĵ (f):= 1 N N∑ n=1 ‖∇xf(Xn)‖2+ 1 Ûθ(f)>D(θ)−1Ûθ(f)+λ‖f‖2H, (14)
where Ûθ(f) is given by: Ûθ(f)=÷∇L(θ)+ 1N∑Nn=1∇hθ(Zn))>∇xf(Xn). (14) has a similar structure as the empirical version of the kernel Sobolev distance introduced in Mroueh et al. (2019), it is also similar to another functional arising in the context of score estimation for infinite dimensional exponential families (Sriperumbudur et al., 2017; Sutherland et al., 2017; Arbel and Gretton, 2017). It can be shown using the generalized Representer Theorem (Schölkopf et al., 2001) that the optimal function minimizing (14) is a linear combination of functions of the form x 7→∂ik(Xn,x) with 1≤n≤N and 1≤i≤d and ∂ik(y,x) denotes the partial derivative of k w.r.t. yi. This requires to solve a system of sizeNd×Nd which can be prohibitive when bothN and d are large. Nyström methods provide a way to improve such computational cost by further restricting the optimal solution to belong to a finite dimensional subspaceHM ofH called the Nyström subspace. In the context of score estimation, Sutherland et al. (2017) proposed to use a subspace formed by linear combinations of the basis functions x 7→∂ik(Ym,x):
span{x 7→∂ik(Ym,x) |1≤m≤M ; 1≤i≤d}, (15) where (Ym)1≤m≤M are basis points drawn uniformly from (Xn)1≤n≤N with M ≤ N . This further reduces the computational cost whenM N but still has a cubic dependence in the dimension d since all partial derivatives of the kernel are considered to construct (15). Here, we propose to randomly sample one component of (∂ik(Ym,.))1≤i≤d for each basis point Ym. Hence, we consider M indices (im)1≤m≤M uniformly drawn form {1,...,d} and define the Nyström subspaceHM to be:
HM :=span{x 7→∂imk(Ym,x)|1≤m≤M}. An estimator for the kernelized Wasserstein natural gradient (KWNG) is then given by:ÿ ∇WL(θ)= 1 D(θ)−1Ûθ(f̂∗), f̂∗ :=argmin
f∈HM Ĵ (f). (16)
By definition of the Nyström subspaceHM , the optimal solution f̂∗ is necessarily of the form: f̂∗(x)=∑M m=1αm∂imk(Ym,x), where the coefficients (αm)1≤m≤M are obtained by solving a finite dimensional quadratic optimization problem. This allows to provide a closed form expression for (17) in Proposition 5. Proposition 5. The estimator in (16) is given by:ÿ ∇WL(θ)= 1 Å D(θ)−1−D(θ)−1T> ( TD(θ)−1T>+λ K+
N CC>
)† TD(θ)−1 ã÷∇L(θ), (17) where C andK are matrices in RM×Nd and RM×M given by
Cm,(n,i) =∂im∂i+dk(Ym,Xn), Km,m′=∂im∂im′+dk(Ym,Ym′), (18)
while T is a matrix in RM×q obtained as the Jacobian of θ 7→τ(θ)∈RM , i.e., T :=∇τ(θ), with
(τ(θ))m= 1
N N∑ n=1 ∂imk(Ym,hθ(Zn)).
In (18), we used the notation ∂i+dk(y,x) for the partial derivative of k w.r.t. xi. A proof of Proposition 5 is provided in Appendix C.2 and relies on the reproducing property (10) and its generalization for partial derivatives of functions. The estimator in Proposition 5 is in fact a low rank approximation of the natural gradient obtained from the dual representation of the metric (9). While low-rank approximations for the Fisher-Rao natural gradient were considered in the context of variational inference and for a Gaussian variational posterior (Mishkin et al., 2018), (17) can be applied as a plug-in estimator for any family PΘ obtained as an implicit model. We next discuss a numerically stable expression of (17), its computational cost and the choice of the damping term in Section 3.3. We then provide asymptotic rates of convergence for (17) in Section 3.4.

heading: 3.3 PRACTICAL CONSIDERATIONS
text: Numerically stable expression. When λ=0, the estimator in (17) has an additional structure which can be exploited to get more accurate solutions. By the chain rule, the matrixT admits a second expression of the form T=CB whereB is the Jacobian matrix of (hθ(Zn))1≤n≤N . Although this expression is impractical to compute in general, it suggests thatC can be ’simplified’. This simplification can be achieved in practice by computing the SVD of CC>=USUT and pre-multiplying T by S†UT . The resulting expression is given in Proposition 6 and falls into the category of Ridgless estimators (Liang and Rakhlin (2019)). Proposition 6. Consider an SVD decomposition ofCCT of the formCC>=USUT , then (17) is equal to:ÿ ∇WL(θ)= 1 Å D(θ)−1−D(θ)−1T̃> ( T̃D(θ)−1T̃>+ N P )† T̃D(θ)−1
ã÷∇L(θ), (19) where P :=S†S and T̃ :=S†UTT .
Choice of damping term. So far, we only required D(θ) to be a diagonal matrix with positive coefficients. While a natural choice would be the identity matrix, this doesn’t necessarily represent the best choice. As discussed by Martens and Sutskever (2012, Section 8.2), using the identity breaks the self-rescaling properties enjoyed by the natural gradient. Instead, we consider a scale-sensitive choice by setting (D(θ))i=‖T̃.,i‖ where T̃ is defined in Proposition 6. When the sample-size is limited, as it is often the case whenN is the size of a mini-batch, larger values for might be required. That is to prevent the KWNG from over-estimating the step-size in low curvature directions. Indeed, these directions are rescaled by the inverse of the smallest eigenvalues of the information matrix which are harder to estimate accurately. To adjust dynamically during training, we use a variant of the Levenberg-Marquardt heuristic as in Martens and Sutskever (2012) which seems to perform well in practice; see Section 4.
Computational cost. The number of basis points M controls the computational cost of both (17) and (19) which is dominated by the cost of computing T and C, solving an M×M linear system and performing an SVD of CCT in the case of (19). This gives an overall cost of O(dNM2+qM2+M3). In practice, M can be chosen to be small (M ≤20) while N corresponds to the number of samples in a mini-batch. Hence, in a typical deep learning model, most of the computational cost is due to computing T as the typical number of parameters q is of the order of millions. In fact, T can be computed using automatic differentiation and would require performing M backward passes on the model to compute the gradient for each component of τ . Overall, the proposed estimator can be efficiently implemented and used for typical deep learning problems as shown in Section 4.
Choice of the kernel. We found that using either a gaussian kernel or a rational quadratic kernel to work well in practice. We also propose a simple heuristic to adapt the bandwidth of those kernels to the data by setting it to σ=σ0σN,M , where σN,M is equal to the average square distance between samples (Xn)1≤n≤N and the basis points (Ym)1≤m≤M and σ0 is fixed a priori. Another choice is the median heuristic Garreau et al. (2018).

heading: 3.4 THEORY
text: In this section we are interested in the behavior of the estimator in the limit of large N and M and when λ> 0; we leave the case when λ= 0 for future work. We work under Assumptions (A) to (G) in Appendix A.2 which state that Ω is a non-empty subset, k is continuously twice differentiable with bounded second derivatives, ∇hθ(z) has at most a linear growth in z and ν satisfies some standard moments conditions. Finally, we assume that the estimator of the euclidean gradient ÷∇L(θ) satisfies Chebychev’s concentration inequality which is often the case in Machine learning problem as discussed in Remark 1 of Appendix A.2. We distinguish two cases: the well-specified case and the miss-specified case. In the well-specified case, the vector valued functions (φi)1≤i≤q involved in Definition 2 are assumed to be gradients of functions inH and their smoothness is controlled by some parameter α≥0 with worst case being α=0. Under this assumption, we obtain smoothness dependent convergence rates as shown in Theorem 14 of Appendix C.3 using techniques from Rudi et al. (2015); Sutherland et al. (2017). Here, we will only focus on the miss-specified which relies on a weaker assumption: Assumption 1. There exists two constants C>0 and c≥0 such that for all κ>0 and all 1≤i≤q, there is a function fκi satisfying:
‖φi−∇fκi ‖L2(ρθ)≤Cκ, ‖f κ i ‖H≤Cκ−c. (20)
The left inequality in (20) represents the accuracy of the approximation of φi by gradients of functions in H while the right inequality represents the complexity of such approximation. Thus, the parameter c characterizes the difficulty of the problem: a higher value of cmeans that a more accurate approximation of φi comes at a higher cost in terms of its complexity. Theorem 7 provides convergences rates for the estimator in Proposition 5 under Assumption 1: Theorem 7. Let δ be such that 0≤δ≤1 and b := 12+c . Under Assumption 1 and Assumptions (A) to (G) listed in Appendix A.2, forN large enough,M∼(dN 1 2b+1 log(N)), λ∼N 1 2b+1 and .N− b 2b+1 , it holds with probability at least 1−δ that:
‖ÿ ∇WL(θ)−∇WL(θ)‖2 =O(N− 24+c). A proof of Theorem 7 is provided in Appendix C.3. In the best case where c=0, we recover a convergence rate of 1√
N as in the well specified case for the worst smoothness parameter value α=0. Hence, Theorem 7
is a consistent extension of the well-specified case. For harder problems where c>0 more basis points are needed, withM required to be of order dN log(N) in the limit when c→∞ in which case the Nyström approximation loses its computational advantage.

heading: 4 EXPERIMENTS
text: This section presents an empirical evaluation of (KWNG) based on (19). Code for the experiments is available at https://github.com/MichaelArbel/KWNG.

heading: 4.1 CONVERGENCE ON SYNTHETIC MODELS
text: To empirically assess the accuracy of KWNG, we consider three choices for the parametric model PΘ: the multivariate normal model, the multivariate log-normal model and uniform distributions on hyper-spheres. All have the advantage that the WNG can be computed in closed form (Chen and Li, 2018; Malagò et al., 2018). While the first models admit a density, the third one doesn’t, hence the Fisher natural gradient is not defined in this case. While this choice of models is essential to obtain closed form expressions for WNG, the proposed estimator is agnostic to such choice of family. We also assume we have access to the exact Euclidean Gradient (EG) which is used to compute both of WNG and KWNG.
2 4 6 8 10 d
0.0% 0.5% 1.0% 1.5% 2.0% 2.5% 3.0% 3.5% 4.0% Re la tiv e er ro r
(a)
0 1000 2000 3000 4000 5000 N
0.0% 0.5% 1.0% 1.5% 2.0% 2.5% 3.0% 3.5% 4.0% (b)
101 102 103 M
0.0% 0.5% 1.0% 1.5% 2.0% 2.5% 3.0% 3.5% 4.0% (c)
d = 10 d = 8 d = 6 d = 4 d = 2
Figure 1: Relative error of KWNG averaged over 100 runs for varying dimension form d= 1 (yellow) to d=10 (dark red) for the hyper-sphere model. (a): box-plot of the relative error as d increases while
N = 5000 and M =
ö
d
√
N
ù
. (b) Relative error as the sample size N increases and M =
ö
d
√
N
ù
. (c):
Relative error asM increases andN=5000. A gaussian kernel is used with a fixed bandwidth σ=1.
Figure 1 shows the evolution of the the relative error w.r.t. the sample-sizeN , the number of basis points M and the dimension d in the case of the hyper-sphere model. As expected from the consistency results provided in Section 3.4, the relative error decreases as the samples size N increases. The behavior in the number of basis points M shows a clear threshold beyond which the estimator becomes consistent and where increasing M doesn’t decrease the relative error anymore. This threshold increases with the dimension d as discussed in Section 3.4. In practice, using the rule M = ö d √ N ù
seems to be a good heuristic as shown in Figure 1 (a). All these observations persist in the case of the normal and log-normal model as shown in Figure 4 of Appendix D.1. In addition we report in Figure 5 the sensitivity to the choice of the bandwidth σ which shows a robustness of the estimator to a wide choice of σ.
We also compare the optimization trajectory obtained using KWNG with the trajectories of both the exact WNG and EG in a simple setting: PΘ is the multivariate normal family and the loss function L(θ) is the squared Wasserstein 2 distance between ρθ and a fixed target distribution ρθ∗ . Figure 2 (a), shows the evolution of the loss function at every iteration. There is a clear advantage of using the WNG over EG as larger step-sizes are allowed leading to faster convergence. Moreover, KWNG maintains this properties while being agnostic to the choice of the model. Figure 2 (b) shows the projected dynamics of the three methods along the two PCA directions of the WNG trajectory with highest variance. The dynamics of WNG seems to be well approximated by the one obtained using KWNG.

heading: 4.2 APPROXIMATE INVARIANCE TO PARAMETRIZATION
text: We illustrate now the approximate invariance to parametrization of the KWNG and show its benefits for training deep neural networks when the model is ill-conditioned. We consider a classification task on two datasets Cifar10 and Cifar100 with a Residual Network He et al. (2015). To use the KWNG
estimator, we view the input RGB image as a latent variable z with probability distribution ν and the output logits of the network x :=hθ(z) as a sample from the model distribution ρθ∈PΘ where θ denotes the weights of the network. The loss function L is given by:
L(θ):= ∫ y(z)>log(SM(Uhθ(z)))dν(z),
where SM is the Softmax function, y(z) denotes the one-hot vector representing the class of the image z and U is a fixed invertible diagonal matrix which controls how well the model is conditioned. We consider two cases, the Well-conditioned case (WC) in which U is the identity and the Ill-conditioned case (IC) where U is chosen to have a condition number equal to 107. We compare the performance of the proposed method with several variants of SGD: plain SGD, SGD + Momentum, and SGD + Momentum + Weight decay. We also compare with Adam Kingma and Ba (2014), KFAC optimizer (Martens and Grosse, 2015; Grosse and Martens, 2016) and eKFAC (George et al., 2018) which implements a fast approximation of the empirical Fisher Natural Gradient. We emphasize that gradient clipping by norm was used for all experiments and was crucial for a stable optimization using KWNG. Details of the experiments are provided in Appendix D.2. Figure 3 shows the training and test accuracy at each epoch on Cifar10 in both (WC) and (IC) cases. While all methods achieve a similar test accuracy in the (WC) case on both datasets, methods based on the Euclidean gradient seem to suffer a drastic drop in performance in the (IC) case. This doesn’t happen for KWNG (red line) which achieves a similar test accuracy as in (WC) case. Moreover, a speed-up in convergence in number of iterations can be obtained by increasing the number of basis points M (brown line). The time cost is also in favor of KWNG (Figure 6). On Cifar100, KWNG is also less affected by the ill-conditioning, albeit to a lower extent. Indeed, the larger number of classes in Cifar100 makes the estimation of KWNG harder as discussed in Section 4.1. In this case, increasing the batch-size can substantially improve the training accuracy (pink line). Moreover, methods that are used to improve optimization using the Euclidean gradient can also be used for KWNG. For instance, using Momentum leads to an improved performance in the (WC) case (grey line). Interestingly, KFAC seems to also suffer a drop in performance in the (IC) case. This might result from the use of an isotropic damping term D(θ) = I which would be harmful in this case. We also observe a drop in performance when a different choice of damping is used for KWNG. More importantly, using only a diagonal pre-conditioning of the gradient doesn’t match the performance of KWNG (Figure 7).

heading: ACKNOWLEDGEMENT
text: GM has received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (grant agreement no 757983).

heading: A PRELIMINARIES

heading: A.1 NOTATION
text: We recall that Ω is an open subset of Rd while Θ is an open subset of parameters in Rq. Let Z ⊂Rp be a latent space endowed with a probability distribution ν over Z. Additionally, (θ,z) 7→hθ(z)∈Ω is a function defined over Θ×Z. We consider a parametric set of probability distributions PΘ over Ω defined as the implicit model:
PΘ :={ρθ :=(hθ)#ν ; θ∈Θ}, where by definition, ρθ =(hθ)#ν means that any sample x from ρθ can be written as x=hθ(z) where z is a sample from ν. We will writeB to denote the jacobian of hθ w.r.t. θ viewed as a linear map from Rq to L2(ν)d without explicit reference to θ: Bu(z)=∇hθ(z).u; ∀u∈Rq. As in the main text, L :Θ→R is a loss functions which is assumed to be of the form L=F(ρθ), with F being a real valued functional over the set of probability distributions. ∇L(θ) denotes the euclidean gradient of L w.r.t θ while ÷∇L(θ) is an estimator of∇L(θ) usingN samples from ρθ. We also consider a Reproducing Kernel Hilbert SpaceH of functions defined over Ω with inner product 〈.,.〉H and norm ‖.‖H and with a kernel k : Ω×Ω→R. The reproducing property for the derivatives (Steinwart and Christmann, 2008, Lemma 4.34) will be important: ∂if(x)=〈f,∂ik(x,.)〉H for all x∈Ω. It holds as long as k is differentiable.
C∞b (Ω) denotes the space of smooth bounded real valued functions on Ω, and C ∞ c (Ω)⊂C∞b (Ω) denotes the subset of compactly supported functions. For any measured space Z with probability distribution ν, we denote by L2(ν) the space of real valued and square integrable functions under ν and by L2(ν)d the space of square integrable vector valued functions under ν and with values in Rd.

heading: A.2 ASSUMPTIONS
text: We make the following set of assumptions:
(A) Ω is a non-empty open subset of Rd. (B) There exists positive constants ζ and σ such that ∫ ‖z‖pdν(z)≤ 12p!ζ
p−2σ2 for any p≥2. (C) For all θ∈Θ there exists C(θ) such that ‖∇θhθ(z)‖≤C(θ)(1+‖z‖) for all z∈Z. (D) k is twice continuously differentiable on Ω×Ω. (E) For all θ∈Θ it holds that ∫ ∂i∂i+dk(x,x)dpθ(x)<∞ for all 1≤i≤d.
(F) The following quantity is finite: κ2 =sup x∈Ω 1≤i≤q ∂i∂i+qk(x,x). (G) For all 0≤δ≤1, it holds with probability at least 1−δ that ‖÷∇L(θ)−∇L(θ)‖.N−12 . Remark 1. Assumption (G) holds if for instance ÷∇L(θ) can be written as an empirical mean of i.i.d. terms with finite variance: ÷∇L(θ)= 1
N N∑ i=1 ∇θl(hθ(Zi))
where Zi are i.i.d. samples from the latent distribution ν where ∫ ∇θl(hθ(z))dν(z)=L(θ). This is often
the case in the problems considered in machine-learning. In, this case, the sum of variances of the vector÷∇L(θ) along its coordinates satisfies:∫ ‖÷∇L(θ)−∇L(θ)‖2dν(z)= 1
N
∫ ‖∇θl(hθ(z))‖2dν(z):= 1
N σ2
One can then conclude using Cauchy-Schwarz inequality followed by Chebychev’s inequality that with probability 1−δ: ‖÷∇L(θ)−∇L(θ)‖≤ σ√ δN
Moreover, Assumption (C) is often satisfied when the implicit model is chosen to be a deep networks with ReLU non-linearity.

heading: A.3 OPERATORS DEFINITION
text: Differential operators. We introduce the linear L operator and its adjoint L>:
L :H→L2(ν)d L> :L2(ν)d→H
f 7→(∂if◦hθ)1≤i≤d v 7→ ∫ d∑
i=1
∂ik(hθ(z),.)vi(z)dν(z)
This allows to obtain the linear operator A defined in Assumption 2 in the main text by composition A :=L>L. We recall here another expression forA in terms of outer product⊗ and its regularized version for a given λ>0,
A= ∫ d∑ i=1 ∂ik(hθ(z),.)⊗∂ik(hθ(z),.)dν(z) Aλ :=A+λI.
It is easy to see that A is a symmetric positive operator. Moreover, it was established in Sriperumbudur et al. (2017) thatA is also a compact operator under Assumption (E).
Assume now we have access toN samples (Zn)1≤n≤N as in the main text. We define the following objects:
Â := 1
N N∑ n=1 d∑ i=1 ∂ik(hθ(Zn),.)⊗∂ik(hθ(Zn),.), Âλ :=Â+λI.
Furthermore, if v is a continuous function in L2(ν)d, then we can also consider an empirical estimator for L>v: ‘L>v := 1
N N∑ n=1 d∑ i=1 ∂ik(hθ(Zn),.)vi(Zn).
Subsampling operators. We consider the operatorQM defined fromH to RM by:
QM :=
√ q
√ M M∑ m=1 em⊗∂imk(Ym,.) (21)
where (em)1≤m≤M is an orthonormal basis ofRM . QM admits a singular value decomposition of the form QM =UΣV
>, with V V > :=PM being the orthogonal projection operator on the Nyström subspaceHM . Similarly to Rudi et al. (2015); Sutherland et al. (2017), we define the projected inverse function GM(C) as:
GM(C)=V (V >CV )−1V >. We recall here some properties of GM from (Sutherland et al., 2017, Lemma 1): Lemma 8. Let A :H→H be a positive operator, and define Aλ=A+λI for any λ>0. The following holds:
1. GM(A)PM =GM(A)
2. PMGM(A)=GM(A)
3. GM(Aλ)AλPM =PM
4. GM(Aλ)=(PMAPM+λI)−1PM
5. ‖A 1 2 λGM(Aλ)A 1 2 λ‖
Estimators of the Wasserstein information matrix. Here we would like to express the estimator in Proposition 5 in terms of the operators introduced previously. We have the following proposition: Proposition 9. The estimator defined in Proposition 5 admits the following representation:ÿ ∇WL(θ)=( D(θ)+GM,N)−1÷∇L(θ) whereGM,N is given by:
GM,N :=(’L>B)>GM(Âλ)’L>B.
Proof. This a direct consequence of the minimax theorem (Ekeland and Témam, 1999, Proposition 2.3, Chapter VI) and applying (Sutherland et al., 2017, Lemma 3).
The matrix GM,N is in fact an estimator of the Wasserstein information matrix defined in Definition 2. We will also need to consider the following population version ofGM,N defined as :
GM :=(L >B)>GM(Aλ)L>B (22)

heading: B BACKGROUND IN INFORMATION GEOMETRY

heading: B.1 FISHER-RAO STATISTICAL MANIFOLD
text: In this section we briefly introduce the non-parametric Fisher-Rao metric defined over the set P of probability distributions with positive density. More details can be found in Holbrook et al. (2017). By abuse of notation, an element ρ∈P will be identified with its density which will also be denoted by ρ. Consider Tρ, the set of real valued functions f defined over Ω and satisfying∫
f(x)2
ρ(x) dx<∞;
∫ f(x)ρ(x)dx=0.
We have the following definition for the Fisher-Rao metric: Definition 4 (Fisher-Rao metric). The Fisher-Rao metric gF is defined for all ρ∈P as an inner product over Tρ of the form:
gFρ (f,g):=
∫ 1
ρ(x) f(x)g(x)dx, ∀f,g∈Tρ
Note that the choice of the set Tρ is different from the one considered in Holbrook et al. (2017) which replaces the integrability condition by a smoothness one. In fact, it can be shown that these choices result in the same metric by a density argument.

heading: B.2 WASSERSTEIN STATISTICAL MANIFOLD
text: In this section we review the theory of Wasserstein statistical manifold introduced in Li and Montufar (2018a); Chen and Li (2018). By analogy to the Fisher-Rao metric which allows to endow the parametric model PΘ with the structure of a Riemannian manifold, it is also possible to use a different metric that is derived from the Wasserstein 2 distance. We first start by briefly introducing the Wasserstein 2 distance. Given two probability distributions ρ and ρ′, we consider the set of all joint probability distributions Π(ρ,ρ′) between ρ and ρ′ usually called the set of couplings between ρ and ρ′. Any couplingπ defines a way of transporting mass from ρ to ρ′. The cost of such transport can be measured as the expected distance between an element of mass of ρ at locationx that is mapped to an element of mass of ρ′ at location y using the couplingπ:∫
‖x−y‖2dπ(x,y)
The squared Wasserstein 2 distance between ρ and ρ′ is defined as the smallest transport cost over all possible couplings:
W22 (ρ,ρ ′)= inf
π∈Π(ρ,ρ′)
∫ ‖x−y‖2dπ(x,y).
A dynamical formulation ofW2 was provided by the celebrated Benamou-Brenier formula in Benamou and Brenier (2000):
W22 (ρ,ρ ′)=inf
φt ∫ 1 0 ∫ ‖φl(x)‖2dρl(x)dl
where the infimum is taken over the set of vector fields φ : [0,1]×Ω→ Rd. Each vector field of the potential determines a corresponding probability distribution ρl as the solution of the continuity equation:
∂lρl+div(ρlφl)=0, ρ0 =ρ,ρ1 =ρ ′. (23)
When Ω is a compact set, a Neumann condition is added on the boundary of Ω to ensure that the total mass is conserved. Such formulation suggests thatW2(ρ,ρ′) corresponds in fact to the shortest path from ρ to ρ′. Indeed, given a path ρl from ρ to ρ′, the infinitesimal displacement direction is given by the distribution ∂lρl. The length |∂lρl| of this direction is measured by: |∂lρl|2 := ∫ ‖φl(x)‖2dρl(x) Hence, W22 (ρ,ρ′) can be written as:
W22 (ρρ ′)=inf
ρl ∫ 1 0 |∂lρl|2dl.
In fact, ∂lρl can be seen as an element in the tangent space TρlP2 to P2 at point ρl. To ensure that (23) is well defined, TρP2 can be defined as the set of distributions σ satisfying σ(1)=0.
|σ(f)|≤C‖∇f‖L2(ρ), ∀f∈C ∞ c (Ω) (24)
for some positive constant C. Indeed, the condition in (24) guarantees the existence of a vector field φσ that is a solution to the PDE: σ=−div(ρφσ). Moreover, |∂lρl|2 can be seen as an inner product of ∂lρl with itself in TρlP2. This inner product defines in turn a metric tensor gW on P2 called the Wasserstein metric tensor (see Otto and Villani (2000); Ambrosio et al. (2004)): Definition 5. The Wasserstein metric gW is defined for all ρ∈P2 as the inner product over TρP2 of the form:
gWρ (σ,σ ′):= ∫ φσ(x) >φσ′(x)dρ(x), ∀σ,σ′∈TρP2
where φσ and φσ′ are solutions to the partial differential equations:
σ=−div(ρφσ), σ′=−div(ρφσ′).
Moreover, φσ and φσ′ are required to be in the closure of gradient of smooth and compactly supported functions w.r.t. L2(ρ)d.
Definition 5 allows to endow P2 with a formal Riemannian structure withW2 being its geodesic distance:
W22 (ρ,ρ ′)=inf
ρl ∫ 1 0 gρl(∂lρl,∂lρl)dl.

heading: B.3 NEGATIVE SOBOLEV DISTANCE AND LINEARIZATION OF THE WASSERSTEIN DISTANCE
text: To device the Wasserstein natural gradient, one can exploit a Taylor expansion of W which is given in terms of the Negative Sobolev distance ‖ρθ+u−ρθ‖H−1(ρθ) as done in Mroueh et al. (2019):
W22 (ρθ,ρθ+u)=‖ρθ+u−ρθ‖2H−1(ρθ)+o(‖u‖ 2).
Further performing a Taylor expansion of ‖ρθ+u−ρθ‖H−1(ρθ) in u leads to a quadratic term u>GW (θt)u where we call GW (θt) the Wasserstein information matrix. This two steps approach is convenient conceptually and allows to use the dual formulation of the Negative Sobolev distance to get an estimate of the quadratic term u>GW (θt)u using kernel methods as proposed in Mroueh et al. (2019) for learning non-parametric models. However, with such approach, ‖ρθ+u−ρθ‖H−1(ρθ) needs to be well defined for u small enough. This requirement does not exploit the parametric nature of the problem and can be restrictive if ρθ+u and ρθ do not share the same support as we discuss now.
As shown in (Villani, 2003, Theorem 7.26) and discussed in (Arbel et al., 2019, Proposition 17 and 18), the Wasserstein distance between two probability distributions ρ and ρ′ admits a first order expansion in terms of the Negative Sobolev Distance:
lim →0
1
W2(ρ,ρ+ (ρ ′−ρ))=‖ρ−ρ′‖H−1(ρ)
when ρ′ admits a bounded density w.r.t. ρ. When such assumption fails to hold, there are cases when this first order expansion is no longer available. For instance, in the simple case when the parametric family consists of dirac distributions δθ located at a value θ, the Wasserstein distance admits a closed form expression of the form:
W2(δθ,δθ+ (δθ′−δθ))= √ ‖θ−θ′‖
Hence, 1 W2(δθ,δθ+ (δθ′−δθ)) diverges to infinity. One can consider a different perturbation of the model δθ+ u for some vector u which the one we are interested in here. In this case, the Wasserstein distance admits a well-defined asymptotic behavior:
lim →
1
W2(δθ,δθ+ u)=‖u‖.
On the other hand the Negative Sobolev Distance is infinite for any value of . To see this, we consider its dual formulation as in Mroueh et al. (2019):
1 2 ‖ρ−ρ′‖H−1(ρ) = sup
f∈C∞c (Ω)∫ f(x)dρ(x)=0
∫ f(x)dρ(x)− ∫ f(x)dρ′(x)− 1
2
∫ ‖∇xf(x)‖2dρ(x)
Evaluating this expression for δθ and δθ+ u for any value >0 and for any u that is non zero, on has: 1
2 ‖δθ−δθ+ u‖H−1(δθ) = sup
f∈C∞c (Ω) f(θ)=0
1 (f(θ)−f(θ+ u))− 1
2 ‖∇xf(θ)‖2 (25)
One can always find a function f such that∇f(θ)=0, f(θ)=0 and−f(θ+ u) can be arbitrarily large, thus the Negative Sobolev distance is infinite. This is not the case of the metric u>GW (θ)u which can be computed in closed form:
1 2 u>GW (θ)u= sup
f∈C∞c (Ω) f(θ)=0
∇f(θ)>u− 1 2 ‖∇xf(θ)‖2 (26)
In this case, choosing f(θ)=0 and∇f(θ)=u achieves the supremum which is simply given by 12‖u‖ 2. Equation (26) can be seen as a limit case of (25) when →0: 1
2 u>GW (θ)u := sup
f∈C∞c (Ω) f(θ)=0
lim →0
1 (f(θ)−f(θ+ u))− 1
2 ‖∇xf(θ)‖2
However, the order between the supremum and the limit cannot be exchanged in this case, which makes the two objects behave very differently in the case of singular probability distributions.

heading: C PROOFS

heading: C.1 PRELIMINARY RESULTS
text: Here we provide a proof of the invariance properties of the Fisher and Wasserstein natural gradient descent in the continuous-time limit as stated in Proposition 1. Consider an invertible and smoothly differentiable re-parametrization Ψ, satisfying ψ = Ψ(θ). Denote by ρ̄ψ = ρΨ−1(ψ) the re-parametrized model and ḠW (ψ) and ḠF (ψ) their corresponding Wasserstein and Fisher information matrices whenever they are well defined.
Proof of Proposition 1. Here we only consider the case when∇DL(θ) is either given by the Fisher natural gradient ∇FL(θ) or the Wasserstein Natural gradient ∇WL(θ). We will first define ψ̃s := Ψ(θs) and show that in fact ψ̃s=ψs at all times s>0. First, let’s differentiate ψ̃s in time:
˙̃ ψs=−∇θΨ(θs)>GD(θs)−1∇θL(θs)
By the chain rule, we have that∇θL(θs)=∇θΨ(θs)∇ψL̄(ψ̃s), hence: ˙̃ ψs=−∇θΨ(θs)>GD(θs)−1∇θΨ(θs)∇ψL̄(ψ̃s).
It is easy to see that ∇θΨ−1(ψ̃s) = (∇ψΨ−1(ψs))−1 by definition of Ψ and ψ̃s, hence by Lemma 10 one can conclude that:
˙̃ ψs=−GD(ψ̃s)−1∇ψL̄(ψ̃s).
Hence, ψ̃s satisfies the same differential equation as ψs. Now keeping in mind that ψ0 = ψ̃0 = Ψ(θ0), it follows that ψ0 =ψ̃0 =Ψ(θ0) by uniqueness of differential equations.
Lemma 10. Under conditions of Propositions 2 and 3, the informations matrices ḠW (ψ) and ḠF (ψ) are related toGW (θ) andGF (θ) by the relation:
ḠW (ψ)=∇ψΨ−1(ψ)>GW (θ)∇ψΨ−1(ψ) ḠF (ψ)=∇ψΨ−1(ψ)>GF (θ)∇ψΨ−1(ψ)
Proof. Let v∈Rq and write u=∇θΨ−1(ψ)v, then by the dual formulations of GW (θ) and GF (θ) in Proposition 2 we have that:
1 2 v>∇ψΨ−1(ψ)>GF (θ)∇ψΨ−1(ψ)v= sup
f∈C∞c (Ω)∫ f(x)dρθ(x)=0
∇ρθ(f)>∇θΨ−1(ψ)v− 1
2
∫ f(x)2dρθ(x)dx,
Now recalling that∇ψρ̄ψ=∇θρθ∇ψΨ−1(ψ) by Lemma 11, it follows that:
1 2 v>∇ψΨ−1(ψ)>GF (θ)∇ψΨ−1(ψ)v= sup
f∈C∞c (Ω)∫ f(x)dρθ(x)=0
∇ψρ̄ψ(f)>v− 1
2
∫ f(x)2dρθ(x)dx,
Using again Proposition 2 for the reparametrized model ρ̄ψ, we directly have that:
1 2 v>GF (ψ)v= sup
f∈C∞c (Ω)∫ f(x)dρθ(x)=0
∇ψρ̄ψ(f)>v− 1
2
∫ f(x)2dρθ(x)dx,
The result follows by equating both expression. The same procedure can be applied for the case the Wasserstein information matrix using Proposition 3.
Lemma 11. The distributional gradients∇ψρ̄ψ and∇θρθ are related by the expression:
∇ψρ̄ψ=∇θρθ∇ψΨ−1(ψ)
Proof. The proof follows by considering a fixed direction u∈Rq and a test function f∈C∞c (Ω) and the definition of distributional gradient in Definition 3:
∇ρ̄ψ(f)>u= lim →0
1 ∫ f(x)dρ̄ψ+ u(x)− ∫ f(x)dρ̄ψ(x)
= lim →0
1 ∫ f(x)dρΨ−1(ψ+ u)(x)− ∫ f(x)dρΨ−1(ψ)(x)
Now by differentiability of Ψ−1, we have the following first order expansion:
Ψ−1(ψ+ u)=Ψ−1(ψ)+ ∇Ψ−1(ψ)>u+ υ( )
where υ( ) converges to 0 when →0. Now using again the definition Definition 3 for ρΨψ one has:
1 Å∫ f(x)dρΨ−1(ψ+ u)(x)− ∫ f(x)dρΨ−1(ψ)(x) ã =∇ρΨ−1(ψ)(f)>∇Ψ−1(ψ)>u
+ υ( )+δ( ,f,(∇Ψ−1(ψ)u+υ( )))
The last two terms converge to 0 as →0, hence leading to the desired expression.
Proposition 12. When ρθ admits a density that is continuously differentiable w.r.t θ and such that x 7→∇ρθ(x) is continuous, then the distributional gradient is of the form:
∇ρθ(f)= ∫ f(x)∇ρθ(x)dx, ∀f∈C∞c (Ω)
where∇ρθ(x) denotes the gradient of the density of ρθ(x) at x.
Proof. Let >0 and u∈Rq, we define the function ν( ,u,f) as follows:
ν( ,u,f)= ∫ f(x)
Å 1
(ρθ+ u)−ρθ−∇ρ>θ u
ã dx
we just need to show that ν( ,u,f)→0 as →0. This follows form the differentiation lemma (Klenke, 2008, Theorem 6.28) applied to the function (θ,x) 7→ f(x)ρθ(x). Indeed, this function is integrable in x for any θ′ in a neighborhood U of θ that is small enough, it is also differentiable on that neighborhood U and satisfies the domination inequality:
|f(x)∇ρθ(x)>u|≤|f(x)| sup x∈Supp(f),θ∈U ∇ρθ(x)>u|.
The inequality follows from continuity of (θ,x)∇ρθ(x) and recalling that f is compactly supported. This concludes the proof.
We fist provide a proof of the dual formulation for the Fisher information matrix.
Proof of Proposition 2 . Consider the optimization problem:
sup f∈C∞c (Ω)∫ f(x)dρθ(x)=0
Å∫ f(x)∇ρθ(x)dx ã> u− 1
2
∫ f(x)2ρθ(x)dx (27)
Recalling that the set of smooth and compactly supported functions C∞c (∞) is dense in L2(ρθ) and that the objective function in (27) is continuous and coercive in f , it follows that (27) admits a unique solution f∗ in L2(ρθ) which satisfies the optimality condition:∫
f(x)(∇ρθ(x))>udx= ∫ f(x)f∗(x)ρθ(x)dx ∀f∈L2(ρθ)
Hence, it is easy to see that f∗=(∇ρθ)>u/ρθ and that the optimal value of (27) is given by:
1
2
∫ ((∇ρθ(x))>u)2
ρθ(x) dx.
This is equal to u>GF (θ)u by Definition 1.
The next proposition ensures that the Wasserstein information matrix defined in Definition 2 is well-defined and has a dual formulation.
Proposition 13. Consider the model defined in (6) and let (es)1≤s≤q be an orthonormal basis of Rq. Under Assumptions (B) and (C), there exists an optimal solution Φ = (φs)1≤s≤q with φs in L2(ρθ)d satisfying the PDE:
∂sρθ=−div(ρθφs)
The elliptic equations also imply that L>∇hθ=L>(Φ◦hθ). Moreover, the Wasserstein information matrix

heading: GW (θ) on PΘ at point θ can be written as GW (θ) = Φ>Φ where the inner-product is in L2(ρθ)d and
text: satisfies:
1 2 u>GW (θ)u= sup
f∈C∞c (Ω)∫ f(x)dρθ(x)=0
∇ρθ(f)>u− 1
2
∫ ‖∇xf(hθ(z))‖2dν(z).
for all u∈Rq.
Proof. Let (es)1≤s≤q be an orthonormal basis of Rq. For all 1≤ s≤ q, we will establish the existence of an optimal solution φs in L2(ρθ)d satisfying the PDE:
∂sρθ=−div(ρθφs) (28)
Consider the variational problem:
sup φ∈S
∫ φ(hθ(z)) >∂θshθ(z)− 1
2 ‖φ‖2L2(ρθ) (29)
where S is a Hilbert space obtained as the closure in L2(ρθ)d of functions of the form φ=∇xf with f∈C∞c (Ω):
S :={∇xf | f∈C∞c (Ω)}L2(ρdθ).
We have by Assumption (C) that:∫ φ(hθ(z)) >∂θshθ(z)dν(z)≤≤C(θ)  ∫ (1+‖z‖2)dν(z) ∫ ‖φ‖L2(ρθ).
Moreover, by Assumption (B), we know that »∫
(1+‖z‖2)dν(z)<∞. This implies that the objective in (29) is continuous in φwhile also being convex.s It follows that (29) admits a unique solution φ∗s∈S which satisfies for all φ∈S:∫
φ(hθ(z)) >φ∗s(hθ(z))dν(z)= ∫ φ(hθ(z)) >∂θshθ(z))dν(z)
In particular, for any f∈C∞c (Ω), it holds that:∫ ∇xf(hθ(z))>φ∗s(hθ(z))dν(z)= ∫ ∇xf(hθ(z))>∂θshθ(z))dν(z)
which is equivalent to (28) and implies directly that LT∇hθ =LTΦ◦hθ where Φ := (φ∗s)1≤s≤q. The variational expression for 12u >GWu follows by noting that (29) admits the same optimal value as
sup f∈C∞c (Ω)∫ f(x)dρθ(x)=0
∇ρθ(f)>u− 1
2
∫ ‖∇xf(hθ(z))‖2dν(z).
That is because S is by definition the closure in L2(ρθ)d of the set of gradients of smooth and compactly supported functions on Ω.
Proof of Proposition 3. This is a consequence of Proposition 13.

heading: C.2 EXPRESSION OF THE ESTIMATOR
text: We provide here a proof of Proposition 5
Proof of Proposition 5. Here, to simplify notations, we simply write D instead of D(θ). First consider the following optimization problem:
inf f∈HM
1
N N∑ n=1 ‖∇f(Xn)‖2+λ‖f‖2H+ 1 R(f)>D−1R(f)+ 2
R(f)>D−1÷∇L(θ)
withR(f) given byR(f)= 1N ∑N n=1∇f(Xn)>B(Zn). Now, recalling that any f∈HM can be written
as f= ∑M m=1αm∂imk(Ym,.), and using the reproducing property ∂if(x)=〈f,∂ik(x,.)〉H (Steinwart and Christmann, 2008, Lemma 4.34) , it is easy to see that:
1
N N∑ n=1 ‖∇f(Xn)‖2 = 1 N ∑ 1≤n≤N 1≤i≤d ( M∑ m=1 αm∂im∂i+dk(Ym,Xn)) 2.
‖f‖2H= ∑
1≤m,m′≤M αmαm′∂im∂im′+dk(Ym,Ym′)
R(f)= 1 N ∑ 1≤n≤N 1≤i≤d
1≤m≤M
αm∂im∂i+dk(Ym,Xn)Bi(Zn)
The above can be expressed in matrix form using the matrices defined in Proposition 5:
1
N N∑ n=1 ‖∇f(Xn)‖2 =α>CC>α; ‖f‖2H=α>Kα; R(f)=α>CB.
Hence the optimal solution f̂∗ is of the form f̂∗= ∑M m=1α ∗ m∂imk(Ym,.), with α
∗ obtained as a solution to the finite dimensional problem in RM :
min α∈RM α>( CC>+ λK+CBD−1B>C>)α+2α>CBD−1÷∇L(θ) It is easy to see that α∗ are given by:
α∗=−( CCT+ λK+CBD−1BTCT )†CBD−1÷∇L(θ). Now recall that the estimator in Proposition 5 is given by: ÿ ∇WL(θ) = 1 D−1Uθ(f̂∗). Hence, 1 D −1(÷∇L(θ)−BTCTα∗) The desired expression is obtained by noting thatCB=T using the chain rule.

heading: C.3 CONSISTENCY RESULTS
text: Well-specified case. Here, we assume that the vector valued functions (φi)1≤i≤q involved in Definition 2 can be expressed as gradients of functions inH. More precisely: Assumption 2. For all 1≤ i≤ q, there exits functions fi ∈H such that φi =∇fi. Additionally, fi are of the form fi=Aαvi for some fixed α≥0, with vi∈H andA being the differential covariance operator defined onH byA :f 7→ ∫∑d i=1∂ik(hθ(z),.)∂if(hθ(z))dν(z).
The parameter α characterizes the smoothness of fi and therefore controls the statistical complexity of the estimation problem. Using a similar analysis as Sutherland et al. (2017) we obtain a convergence rate for the estimator in Proposition 5
the following convergence rates for the estimator in Proposition 5: Theorem 14. Let δ be such that 0 ≤ δ ≤ 1 and b := min(1, α + 12). Under Assumption 2 and Assumptions (A) to (G) listed in Appendix A.2, forN large enough,M∼(dN 1 2b+1 log(N)), λ∼N− 1 2b+1 and .N− b 2b+1 , it holds with probability at least 1−δ that:
‖ÿ ∇WL(θ)−∇WL(θ)‖2 =O(N− 2b2b+1). In the worst case where α= 0, the proposed estimator needs at most M ∼ (d √ N log(N)) to achieve a convergence rate ofN− 1 2 . The smoothest case requires onlyM∼(dN 13 log(N)) to achieve a rate ofN−23 . Thus, the proposed estimator enjoys the same statistical properties as the ones proposed by Sriperumbudur et al. (2017); Sutherland et al. (2017) while maintaining a computational advantage1 tNow we provide a proof for Theorem 14 which relies on the same techniques used by Rudi et al. (2015); Sutherland et al. (2017).
Proof of Theorem 14 . The proof is a direct consequence of Proposition 15 under Assumption 2.
Proof of Theorem 7. The proof is a direct consequence of Proposition 15 under Assumption 1.
Proposition 15. Under Assumptions (A) to (G) and for 0≤ δ ≤ 1 and N large enough, it holds with probability at least 1−δ: ‖’∇WL−∇WL‖=O(N− b2b+1 ) provided that M ∼ dN 1 2b+1 logN , λ ∼ N 1 2b+1 and . N−
b 2b+1 where b := min(1,α + 12) when
Assumption 2 holds and b= 12+c when Assumption 1 holds instead. 1 The estimator proposed by Sutherland et al. (2017) also requires M to grow linearly with the dimension d although such dependence doesn’t appear explicitly in the statement of Sutherland et al. 2017, Theorem 2.
Proof. Here for simplicity we assume that D(θ) = I without loss of generality and we omit the dependence in θ and write∇WL and∇L instead of∇WL(θ) and∇L(θ) and∇WL(θ). We also define Ĝ = I+GM,N and G = I+GW . By Proposition 9, we know that ’∇WL= Ĝ−1 ∇̂L. We use the following decomposition:
‖’∇WL−∇WL‖≤‖Ĝ−1 (∇̂L−∇L)‖+‖Ĝ−1 (GM,N−GW )G−1W ∇L‖+ ‖Ĝ−1 G−1W ∇L‖ To control the norm of Ĝ−1 we write Ĝ −1 = G −12 (H + I)−1G −12 , where H is given by H := G −12 (GM,N − GW )G −12 . Hence, provided that µ := λmax(H), the highest eigenvalue of H, is smaller than 1, it holds that: ‖(H+I)−1‖≤(1−µ)−1.
Moreover, since GW is positive definite, its smallest eigenvalue η is strictly positive. Hence, ‖G−1 ‖≤(η+ )−1. Therefore, we have ‖Ĝ−1 ‖≤(η+ )(1−µ))−1, which implies: ‖’∇WL−∇WL‖≤(η+ )−1Ç‖∇̂L−∇L‖ 1−µ +η−1‖∇L‖‖GM,N−GW‖+ η−1‖∇L‖ å .
Let 0 ≤ δ ≤ 1. We have by Assumption (G) that ‖∇̂L−∇L‖ = O(N−12 ) with probability at least 1−δ. Similarly, by Proposition 16 and for N large enough, we have with probability at least 1−δ that ‖GM,N−GW‖=O(N− b 2b+1 ) where b is defined in Proposition 16. Moreover, forN large enough, one can ensure that µ≤ 12 so that the following bound holds with probability at least 1−δ:
‖’∇WL−∇WL‖.(η+ )−1Ä2N−12 +η−1‖∇L‖(N− b2b+1 + )ä. Thus by setting .N− b 2b+1 we get the desired convergence rate.
Proposition 16. For any 0≤δ≤1, we have with probability as least 1−δ and forN large enough that:
‖GM,N−GW‖=O(N− b 2b+1 ).
provided that M∼dN 1
2b+1 logN where b :=min(1,α+ 12) when Assumption 2 holds and b= 1 2+c when Assumption 1 holds instead.
Proof. To control the error ‖GM,N−GW‖ we decompose it into an estimation error ‖GM,N−GM‖ and approximation error ‖GM−GW‖:
‖GM,N−GW‖≤‖GM−GW‖+‖GM−GM,N‖
wereGM is defined in (22) and is obtained by taking the number of samplesN to infinity while keeping the number of basis pointsM fixed.
The estimation error ‖GM−GM,N‖ is controlled using Proposition 17 where, for any 0≤δ≤1, we have with probability at least 1−δ and as long asN≥M(1,λ,δ):
‖GM,N−GM‖≤ ‖B‖√ Nλ
(aN,δ+ √
2γ1κ+2γ1 λ+κ√ Nλ )+ 1 Nλ a2N,δ.
In the limit where N →∞ and λ→ 0, only the dominant terms in the above equation remain which leads to an error ‖GM,N −GM‖=O((Nλ)− 1 2 ). Moreover, the condition on N can be expressed as λ−1logλ−1.N .
To control the error approximation error ‖GM −GW‖ we consider two cases: the well-specified case and the miss-specified case.
• Well-specified case. Here we work under Assumption 2 which allows to use Proposition 19. Hence, for any 0≤δ≤1 and ifM≥M(d,λ,δ), it holds with probability at least 1−δ:
‖GM−GW‖.λmin(1,α+ 1 2 )
• Miss-specified case. Here we work under Assumption 1 which allows to use Proposition 18. Hence, for any 0≤δ≤1 and ifM≥M(d,λ,δ), it holds with probability at least 1−δ:
‖GM−GW‖.λ 1 2+c
Let’s set b := min(1,α+ 12) for the well-specified case and b = 1 2+c for the miss-specified case. In the limit where M →∞ and λ→ 0 the condition on M becomes: M ∼ dλ−1 logλ−1. Hence, when M∼dλ−1logλ−1 and λ−1logλ−1.N it holds with probability as least 1−δ that
‖GM,N−GW‖=O(λb+(λN)− 1 2 ).
One can further choose λ of the form λ = N−θ. This implies a condition on M of the form dNθ log(N).M and Nθ log(N).N . After optimizing over θ to get the tightest bound, the optimal value is obtained when θ = 1/(2b+1) and the requirement on N is always satisfied once N is large enough. Moreover, one can choose M∼dN 1 2b+1 logN so that the requirement on M is satisfied for N large enough. In this case we get the following convergence rate:
‖GM,N−GW‖=O(N− b 2b+1 ).
Proposition 17. For any 0≤δ≤1, provided thatN≥M(1,λ,δ), we have with probability as least 1−δ:
‖GM,N−GM‖≤ ‖B‖√ Nλ
(2aN,δ+ √
2γ1κ+2γ1 λ+κ√ Nλ )+ 1 Nλ a2N,δ.
with:
aN,δ := … 2σ21log 2
δ + 2alog2δ√ N
Proof. For simplicity, we define E = ’L>B −L>B. By definition of GM,N and GM we have the following decomposition:
GM,N−GM =E>GM(Âλ)E︸ ︷︷ ︸ E0 +E>GM(Âλ)L>B︸ ︷︷ ︸ E1 +B>LGM(Âλ)E︸ ︷︷ ︸ E2
−B>LGM(Aλ)PM(Â−A)PMGM(Âλ)L>B︸ ︷︷ ︸ E3
The first three terms can be upper-bounded in the following way:
‖E0‖=‖E>Â −12 λ Â 1 2 λGM(Âλ)Â 1 2 λ Â −12 λ E‖
≤‖E‖2‖Â−1λ ‖︸ ︷︷ ︸ ≤1/λ ‖Â 1 2 λGM(Âλ)Â 1 2 λ‖︸ ︷︷ ︸ ≤1
‖E1‖=‖E2‖=‖E>A −12 λ A 1 2 λGM(Aλ)A 1 2 λA −12 λ L >B‖
≤‖B‖‖E‖‖Â− 1 2 λ ‖︸ ︷︷ ︸ ≤1/ √ λ ‖Â 1 2 λGM(Âλ)Â 1 2 λ‖︸ ︷︷ ︸ ≤1 ‖A− 1 2 λ L >‖︸ ︷︷ ︸ ≤1 ‖Â− 1 2 λ A 1 2 λ‖
For the last term E3 , we first recall that by definition of GM(Aλ) we have:
GM(Aλ)PM(Â−A)PMGM(Aλ)=GM(Aλ)(Â−A)GM(Aλ). Therefore, one can write:
‖E3‖=‖B>LA −12 λ A 1 2 λGM(Aλ)A 1 2 λA −12 λ (Â−A)A −12 λ A 1 2 λ Â −12 λ Â 1 2 λGM(Âλ)Â 1 2 λ Â −12 λ A 1 2 λA −12 λ L >B‖
≤‖B‖2‖LA− 1 2 λ ‖ 2︸ ︷︷ ︸
≤1
‖A 1 2 λGM(Aλ)A 1 2 λ‖︸ ︷︷ ︸ ≤1 ‖Â 1 2 λGM(Âλ)Â 1 2 λ‖︸ ︷︷ ︸ ≤1 ‖A 1 2 λ Â 1 2 λ‖ 2‖A− 1 2 λ (Â−A)A −12 λ ‖
≤‖B‖2‖A 1 2 λ Â 1 2 λ‖ 2‖A− 1 2 λ (Â−A)A −12 λ ‖
We recall now (Rudi et al., 2015, Proposition 7.) which allows to upper-bound ‖A 1 2 λ Â 1 2 λ‖ by (1−η)− 1 2 where η=λmax(A 1 2 λ (A−Â)A 1 2 λ ) provided that η<1. Moreover, (Rudi et al., 2015, Proposition 8.) allows to control both η and ‖A− 1 2
λ (Â−A)A −12 λ ‖ under Assumption (F). Indeed, for any 0≤δ≤1 and provided
that 0<λ≤‖A‖ it holds with probability 1−δ that:
‖A− 1 2 λ (Â−A)A −12 λ ‖≤2γ1 1+κ/λ
3N +
… 2γ1κ
Nλ ; η≤ 2γ2 3N +
… 2γ2κ
Nλ where γ1 and γ2 are given by:
γ1 =log( 8Tr(A)
λδ ); γ2 =log(
4Tr(A)
λδ ).
Hence, forN≥M(1,λ,δ) we have that (1−η)−12 ≤2 and one can therefore write:
‖E3‖≤4‖B‖2(2γ1 1+κ/λ
3N +
… 2γ1κ
Nλ )
‖E1‖=‖E1‖≤ 2‖B‖√ λ ‖E‖
The error ‖E‖ is controlled by Proposition 22 where it holds with probability greater or equal to 1−δ that:
‖E‖≤ 1√ N (
… 2σ21log 2
δ + 2alog2δ√ N ):= 1√ N aN,δ.
Finally, we have shown that provided thatN≥M(1,λ,δ) then with probability greater than 1−δ one has:
‖GM,N−GM‖≤ ‖B‖√ Nλ
(2aN,δ+ √
2γ1κ+2γ1 λ+κ√ Nλ )+ 1 Nλ a2N,δ.
Proposition 18. Let 0≤λ≤‖A‖ and defineM(d,λ,δ):= 1289 log 4Tr(A) λδ (dκλ −1+1). Under Assumption 1 and Assumption (F), for any δ≥0 such thatM≥M(d,λ,δ) the following holds with probability 1−δ:
‖GM−GW‖.λ 1 2+c
Proof. We consider the error ‖GM −GW‖. Recall that GW is given by GW = Φ>Φ with Φ defined in Proposition 13. Let κ be a positive real number, we know by Assumption 1 that there exists Fκ := (fκs )1≤s≤q with f κ s ∈H such that ‖Φ−Fκ‖L2(ρθ) ≤ Cκ and ‖fκs ‖H ≤ Cκ−c for some fixed positive constant C. Therefore, we use Fκ to control the error ‖GM−GW‖. Let’s callE=Φ◦hθ−LFκ We consider the following decomposition:
GM−GW =(L>Φ◦hθ)>GM(Aλ)L>Φ◦hθ−Φ>Φ =E>LGM(Aλ)L>E︸ ︷︷ ︸
E1
−E>E︸ ︷︷ ︸ E2
+F>κ ( L>LGM(Aλ)−I ) L>Φ◦hθ︸ ︷︷ ︸
E3
+E>L ( GM(Aλ)L>L−I ) Fκ︸ ︷︷ ︸
E4
First we consider the term E1 one simply has:
‖E1‖≤κ2‖LA −12 λ ‖︸ ︷︷ ︸ ≤1 ‖A 1 2 λGM(Aλ)A 1 2 λ‖︸ ︷︷ ︸ ≤1 ‖A− 1 2 λ L >‖︸ ︷︷ ︸ ≤1 ≤κ2
The second term also satisfies ‖E1‖≤κ2 by definition of Fκ. For the last two terms E3 and E4 we use Lemma 20 which allows to control the operator norm of L(GM(Aλ)L>L−I). Hence, for any δ≥0 and M such thatM≥M(d,λ,δ) and for κ≤1 it holds with probability 1−δ that:
‖E3‖. √ λκ−c; ‖E4‖. √ λκ−c
We have shown so far that ‖GM−GW‖.(κ2+2κ−c √ λ). One can further optimize over κ on the interval [0,1] to get a tighter bound. The optimal value in this case is κ∗=min(1,(cλ 1 2 ) 1 2+c ). By considering λ>0 such that (cλ 1 2 ) 1 2+c )≤1, it follows directly that ‖GM−GW‖.λ 1 2+c which shows the desired result.
Proposition 19. Let 0≤λ≤‖A‖ and defineM(d,λ,δ):= 1289 log 4Tr(A) λδ (dκλ −1+1). Under Assumption 2 and Assumption (F), for any δ≥0 such thatM≥M(d,λ,δ) the following holds with probability 1−δ:
‖GM−GW‖.λmin(1,α+ 1 2 )
Proof. Recall thatGW is given byGW =Φ>Φ with Φ defined in Proposition 13. By Assumption 2, we have that Φ=∇(AαV ) with V :=(vs)1≤s≤q∈Hq. Hence, one can write
GM−GW =(L>Φ◦hθ)>GM(Aλ)L>Φ◦hθ−Φ>Φ =V >(Aα(AGM(Aλ)A−A)AαV
we can therefore directly apply Lemma 20 and get ‖GM−GW‖.λmin(1,α+ 1 2 ) with probability 1−δ for any δ≥0 such thatM≥M(d,λ,δ).
Lemma 20. Let 0 ≤ λ ≤ ‖A‖, α ≥ 0 and define M(d,λ,δ) := 1289 log 4Tr(A) λδ (dκλ −1 + 1). Under Assumption (F), for any δ≥0 such thatM≥M(d,λ,δ) the following holds with probability 1−δ:
‖L(GM(Aλ)L>L−I)Aα‖.λmin(1,α+ 1 2 )
Proof. We have the following identities:
L(GM(Aλ)L>L−I)Aα=L(GM(Aλ)Aλ−I−λGM(Aλ))Aα
=LA −12 λ A 1 2 λ (GM(Aλ)AλPM−I)A α︸ ︷︷ ︸
E1
−λLA− 1 2 λ A 1 2 λGM(Aλ)A 1 2 λA −12 λ A α︸ ︷︷ ︸ E3
+LA −12 λ A 1 2 λGM(Aλ)A 1 2 λA 1 2 λ (I−PM)A α︸ ︷︷ ︸
E2
.
For the first E1 we use (Sutherland et al., 2017, Lemma 1 (iii)) which implies that GM(Aλ)AλPM =PM . Thus E1 =LA −12 λ A 1 2 λ (PM−I)Aα. Moreover, by Lemma 21 we have that ‖A 1 2 λ (I−PM)‖≤2 √ λ with probability 1− δ for M >M(d,λ,δ). Therefore, recalling that (I −PM)2 = I −PM since PM is a projection, one can further write:
‖E1‖≤‖LA −12 λ ‖︸ ︷︷ ︸ ≤1 ‖A 1 2 λ (PM−I)‖ 2︸ ︷︷ ︸ ≤λ ‖A− 1 2 λ A α‖
‖E2‖≤‖LA −12 λ ‖︸ ︷︷ ︸ ≤1 ‖A 1 2 λGM(Aλ)A 1 2 λ‖︸ ︷︷ ︸ ≤1 ‖A 1 2 λ (PM−I)‖ 2︸ ︷︷ ︸ ≤4λ ‖A− 1 2 λ A α‖
‖E3‖≤λ‖LA −12 λ ‖︸ ︷︷ ︸ ≤1 ‖A 1 2 λGM(Aλ)A 1 2 λ‖︸ ︷︷ ︸ ≤1 ‖A− 1 2 λ A α‖
It remains to note that ‖A− 1 2
λ A α‖≤ λα−12 when 0≤α≤ 12 and that ‖A −12 λ A α‖≤ ‖A‖α−12 for α> 12 which allows to conclude.

heading: C.4 AUXILIARY RESULTS
text: Lemma 21. Let 0 ≤ λ ≤ ‖A‖. Under Assumption (F), for any δ ≥ 0 such that M≥M(d,λ,δ):= 1289 log 4Tr(A) λδ (κλ −1+1) the following holds with probability 1−δ:
‖A 1 2 λ (I−PM)‖≤2 √ λ
Proof. The proof is an adaptation of the results in Rudi et al. (2015); Sutherland et al. (2017). Here we recall QM defined in (21). Its transposeQ>M sends vectors in RM to elements in the span of the Nyström basis
points, hence PM andQ>M have the same range, i.e.: range(PM)= ¯range(Q>M). We are in position to apply (Rudi et al., 2015, Proposition 3.) which allows to find an upper-bound onA 1 2
λ (PM−I) in terms ofQM :
‖A 1 2 λ (PM−I)‖≤ √ λ‖A 1 2 λ (Q > MQM+λI) −12‖.
For simplicity we write ÂM :=Q>MQM andE2 :=A −12 λ (A−ÂM)A −12 λ . We also denote by β=λmax(E2) the highest eigenvalue of E2. We can therefore control ‖A 1 2 λ (ÂM +λI) −12‖ in terms of β using (Rudi et al., 2015, Proposition 7) provided that β<1:
‖A 1 2 λ (PM−I)‖≤ √ λ
1√ 1−β .
Now we need to make sure that β<1 forM large enough. To this end, we will apply (Rudi et al., 2015, Proposition 8.) to ÂM . Denote by vm = √ d∂imk(Ym,.). Hence, by definition of ÂM it follows that ÂM = 1 M ∑M m=1vm⊗vm. Moreover, (vm)1≤m≤M are independent and identically distributed and satisfy:
E[vm⊗vm]= ∫ q∑
i=1
∂ik(y,.)⊗∂ik(y,.)dpθ(y)=A.
We also have by Assumption (F) that 〈vm,A−1λ vm〉≤ dκ λ almost surely and for all λ>0. We can therefore apply (Rudi et al., 2015, Proposition 8.) which implies that for any 1≥δ≥0 and with probability 1−δ it holds that:
β≤ 2γ 3M +
… 2γdκ
Mλ
with γ=log4Tr(A)λδ provided that λ≤‖A‖. Thus by choosingM≥ 128γ 9 (dκλ −1+1) we have that β≤ 34 with probability 1−δ which allows to conclude.
Proposition 22. There exist a>0 and σ1>0 such that for any 0≤δ≤1, it holds with probability greater of equal than 1−δ that:
‖’L>B−L>B‖≤ 2alog2δ N +
  2σ21log 2 δ
N
Proof. denote by vn= ∑d i=1∂ik(Xn,.)Bi(Zn), we have that E[vn]=L>B. We will apply Bernstein’s inequality for sum of random vectors. For this we first need to find a>0 and σ1>0 such that E[‖zn− L>B‖pH]≤ 1 2p!σ 2 1a p−2. To simplify notations, we write x and x′ instead of hθ(z) and hθ(z′). We have that:
E[‖zn−L>B‖pH]= ∫ ∥∥∥∥∥ d∑
i=1
∂ik(x,.)Bi(z)− ∫ d∑
i=1
∂ik(x ′,.)Bi(z ′)dν(z′) ∥∥∥∥∥ p dν(z)
≤2p−1 ∫ ∥∥∥∥∥ d∑
i=1
∫ (∂ik(x,.)−∂ik(x′,.))Bi(z)dν(z′) ∥∥∥∥∥ p
dν(z)︸ ︷︷ ︸ E1
+2p−1 ∫ ∥∥∥∥∥ ∫ d∑
i=1
∂ik(x,.)(Bi(z)−Bi(z′))dν(z′) ∥∥∥∥∥ p
dν(z)︸ ︷︷ ︸ E2
We used the convexity of the norm and the triangular inequality to get the last line. We introduce the notation γi(x):=∂ik(x,.)− ∫ ∂ik(hθ(z
′),.)dν(z′) and by Γ(x) we denote the matrix whose components are given by Γ(x)ij :=〈γi(x),γj(x)〉H. The first term E1 can be upper-bounded as follows:
E1 = ∫ ∣∣∣Tr(B(z)B(z)>Γ(x))∣∣∣p2 ≤ ∫ ∣∣∣‖B(z)‖2Tr(Γ(x)2)12 ∣∣∣p2 .
Moreover, we have that Tr(Γ(x)2) 1 2 = ( ∑ 1≤i,j≤d〈γi(x),γj(x)〉2H) 1 2 ≤ ∑d i=1 ‖γi(x)‖2. We further
have that ‖γi(x)‖≤∂i∂i+dk(x,x) 1 2 + ∫ ∂i∂i+dk(hθ(z),hθ(z)) 1 2 dν(z) and by Assumption (F) it follows that ‖γi(x)‖ ≤ 2 √ κ. Hence, one can directly write that: E1 ≤ (2 √ κd)p ∫ ‖B(z)‖pdν(z). Recalling Assumptions (B) and (C) we get:
E1≤2p−1(2 √ κd)pC(θ)p(1+ 1
2 p!ζp−2σ2)
Similarly, we will find an upper-bound on E2. To this end, we introduce the matrix Q(x′,x”) whose components are given byQ(x′,x”)i,j=∂i∂i+dk(x′,x”). One, therefore has:
E2 = ∫ ∣∣∣∣∫ ∫ Tr((B(z)−B(z′))(B(z)−B(z”))>Q(x′,x”)dν(z′)dν(z”)∣∣∣∣p2 dν(z) ≤ ∫ ∣∣∣∣∫ ∫ ‖B(z)−B(z′)‖‖B(z)−B(z”)‖Tr(Q(x′,x”)2)12 dν(z′)dν(z”)∣∣∣∣p2 dν(z)
Once again, we have that Tr(Q(x′,x”)2) 1 2 ≤ ( ∑d i=1∂i∂i+dk(x ′,x′)) 1 2 ( ∑d i=1∂i∂i+dk(x”,x”)) 1 2 ≤ dκ thanks to Assumption (F). Therefore, it follows that:
E2≤( √ dκ)p ∫ | ∫ ‖B(z)−B(z′)‖dν(z)|pdν(z)
≤3p−1( √ dκ)pC(θ)p(2p+ ∫ ‖z‖pdν(z)+ Å∫ ‖z‖dν(z) ãp )
≤3p−1( √ dκ)pC(θ)p(2p+ 1
2 p!ζp−2σ2+
Å∫ ‖z‖dν(z) ãp ).
The second line is a consequence of Assumption (C) while the last line is due to Assumption (B). These calculations, show that it is possible to find constants a and σ1 such that E[‖zn−L>B‖pH]≤ 1 2p!σ 2 1a p−2. Hence one concludes using Bernstein’s inequality for a sum of random vectors (see for instance Rudi et al., 2015, Proposition 11).

heading: D EXPERIMENTAL DETAILS

heading: D.1 NATURAL WASSERSTEIN GRADIENT FOR THE MULTIVARIATE NORMAL MODEL
text: Multivariate Gaussian. Consider a multivariate gaussian with mean µ ∈ Rd and covariance matrix Σ∈Rd×Rd parametrized using its lower triangular components s=T(Σ). We denote by Σ=T−1(s) the inverse operation that maps any vector s∈R d(d+1) 2 to its corresponding symmetric matrix in Rd×Rd. The concatenation of the mean µ and s will be denoted as θ : θ= (µ,s). Given two parameter vectors u= (m,T(S)) and v = (m′,T(S′)) where m and m′ are vectors in Rd and S and S′ are symmetric matrices in Rd×Rd the metric evaluated at u and v is given by:
u>G(θ)v=m>m′+Tr(AΣA′)
whereA andA′ are symmetric matrices that are solutions to the Lyapunov equation:
S=AΣ+ΣA, S′=A′Σ+ΣA′.
A and A′ can be computed in closed form using standard routines making the evaluation of the metric easy to perform. Given a loss function L(θ) and gradient direction ∇θL(θ) = ∇µL(θ),∇sL(θ), the corresponding natural gradient∇Wθ L(θ) can also be computed in closed form:
∇Wθ L(θ)=(∇µL(θ),T(Σ(A+diag(A))+(A+diag(A))Σ)),
where A = T−1(∇sL(θ)). To use the estimator proposed in Proposition 5 we take advantage of the parametrization of the Gaussian distribution as a push-forward of a standard normal vector:
X∼N (µ,Σ)⇐⇒X=Σ12Z+µ,Z∼N (0,Id)
error as d increases with N =5000 and the number of basis points is set to M =
ö
d
√
N
ù
. (b) Relative
error as the sample size N increases and the number of basis points is set to M =
ö
d √ N
ù
. Right (c):
Relative error asM increases andN fixed to 5000.
is set to M=
ö
d
√
N
ù
. Left: uniform distributions on a hyper-sphere, middle: multivariate normal, and
right: multivariate log-normal.
D.2 CLASSIFICATION ON CIFAR10 AND CIFAR100
Architecture. We use a residual network with one convolutional layer followed by 8 residual blocks and a final fully connected layer. Each residual block consists of two 3×3 convolutional layers each and ReLU nonlinearity. We use batch normalization for all methods. Details of the intermediate output shapes and kernel size are provided in Table 1.
Hyper-parameters. For all methods, we used a batch-size of 128. The optimal step-size γ was selected in {10,1,10−1,10−2,10−3,10−4} for each method. In the case of SGD with momentum, we used a momentum parameter of 0.9 and a weight decay of either 0 or 5×10−4. For KFAC and EKFAC, we used a damping coefficient of 10−3 and a frequency of reparametrization of 100 updates. For KWGN we set M=5 and λ=0 while the initial value for is set to =10−5 and is adjusted using an adaptive scheme based on the Levenberg-Marquardt dynamics as in (Martens and Grosse, 2015, Section 6.5). More
precisely, we use the following update equation for after every 5 iterations of the optimizer:
←ω , if r> 3 4 ←ω−1 , if r< 1 4 .
Here, r is the reduction ratio:
r= max tk−1≤t≤tk Å 2 L(θt))−L(θt+1) ∇WL(θ)>∇L(θ)> ã where (tk)k are the times when the updates occur. and ω is the decay constant chosen to ω=0.85.

heading: D.3 ADDITIONAL EXPERIMENTS
text: Figure 6 reports the time cost in seconds for each methods, while Figure 7 compares KWNG to diagonal conditioning.

Review summaries:
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.

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.