Title: KERNELIZED WASSERSTEIN NATURAL GRADIENT

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.
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 ν

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 ρ →F(ρ) over probability distributions ρ from a parametric model P Θ . The set P Θ contains probability distributions defined on an open sample space Ω ⊂ R d and parametrized by some vector θ ∈ Θ, where Θ is an open subset of R q . 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 model M t 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:

EQUATION

)

The linear model M t is an approximation of the cost function L in the neighborhood of θ t and is simply obtained by a first order expansion: M t (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 θ →ρ θ 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 divergence D between probability distributions and use it as a new penalization term:

-argmin

EQUATION

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 G D (θ)u where G D (θ) contains second order information about the model as measured by D. This leads to the following expression for the generalized natural gradient ∇ D L(θ t ) where the dependence in D is made explicit:

EQUATION

From (3), it is possible to express the generalized natural gradient by means of the Euclidean gradient:

∇ D L(θ t )=G D (θ t ) -1 ∇L(θ t ).

The parameter updates are then obtained by the new update rule:

EQUATION

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 distance W 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 distance W .

2.2. INFORMATION MATRIX VIA DIFFERENTIAL GEOMETRY:

When D is the Kullback-Leibler divergence or relative entropy F , then (3) defines the Fisher-Rao natural gradient ∇ F L(θ) (Amari, 1985) and G F (θ) is called the Fisher information matrix. G F (θ) is well defined when the probability distributions in P Θ 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 g F 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 θ → ρ θ (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 g F :

G F (θ) ij =g F ρ θ (∂ i ρ θ ,∂ j ρ θ ):= f i (x)f j (x)ρ θ (x)dx,

where the functions f i on Ω are given by:

f i = ∂iρ θ ρ θ .

Definition 1 directly introduces G F 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 g W (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 g W :

G W (θ) 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 in L 2 (ρ θ ) d . In particular, when ρ θ has a density, φ i =∇ x f i , for some real valued function f i 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 ∇ D L(θ), namely ∇ F L(θ) and ∇ W L(θ). 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.

2.3. LEGENDRE DUALITY FOR METRICS:

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 (Ω)→R q such that:

f(x)dρ θ+ u (x)-f(x)dρ θ (x)= ∇ρ θ (f) u+ δ( ,f,u) ∀f ∈C ∞ c (Ω), ∀u∈R q

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 →ρ θ (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:

EQUATION

Moreover, defining U θ (f):=∇L(θ)+∇ρ θ (f), the Fisher-Rao natural gradient satisfies:

∇ F L(θ)=-argmin u∈R q sup f∈C ∞ c (Ω) f(x)dρ θ (x)=0

defined over a latent space Z and (θ,z) → 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 θ :

EQUATION

A different expression for ∇ρ θ is obtained in the case of implicit models when θ →h θ (z) is differentiable for ν-almost all z and ∇h θ is square integrable under ν:

EQUATION

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:

EQUATION

and the Wasserstein natural gradient satisfies:

EQUATION

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.

3. KERNELIZED WASSERSTEIN NATURAL GRADIENT:

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 ∇ W L(θ), 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.

3.1. GENERAL FORMULATION AND MINIMAX THEOREM:

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 : Ω × Ω → R such that k(x, .) ∈ H for all x ∈ Ω and satisfying the Reproducing property for all functions f in H:

EQUATION

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:

EQUATION

The kernelized Wasserstein natural gradient is obtained by solving (11) and is denoted by ‹ ∇ W L(θ). Here, is a positive real numbers, λ is non-negative while D(θ) is a diagonal matrix in R q 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:

EQUATION

where f * is the unique solution to the quadratic optimization problem:

EQUATION

)

When λ = 0, f * might not be well defined, still, we have:

‹ ∇ W L(θ) = lim j→∞ 1 D(θ) -1 U θ (f j )

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 that H enjoys some denseness properties.

3.2. NYSTRÖM METHODS FOR THE KERENALIZED NATURAL GRADIENT:

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 consider N samples (Z n ) 1≤n≤N from the latent distribution ν which are used to produce N samples (X n ) 1≤n≤N from ρ θ using the map h θ , i.e., X n =h θ (Z n ). 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),

EQUATION

where U θ (f) is given by: U 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 →∂ i k(X n ,x) with 1≤n≤N and 1≤i≤d and ∂ i k(y,x) denotes the partial derivative of k w.r.t. y i . This requires to solve a system of size Nd×Nd which can be prohibitive when both N 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 subspace H M of H 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

θ (f)= ÷ ∇L(θ)+ 1 N N n=1 ∇h θ (Z n )) ∇ x f(X n ). (

EQUATION

where (Y m ) 1≤m≤M are basis points drawn uniformly from (X n ) 1≤n≤N with M ≤ N. This further reduces the computational cost when M 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 (∂ i k(Y m ,.)) 1≤i≤d for each basis point Y m . Hence, we consider M indices (i m ) 1≤m≤M uniformly drawn form {1,...,d} and define the Nyström subspace H M to be:

H M :=span{x →∂ im k(Y m ,x)|1≤m≤M}.

An estimator for the kernelized Wasserstein natural gradient (KWNG) is then given by:

EQUATION

)

By definition of the Nyström subspace H M , the optimal solution f * is necessarily of the form:

f * (x)= M m=1 α m ∂ im k(Y m ,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:

EQUATION

where C and K are matrices in R M×Nd and R M×M given by

EQUATION

)

while T is a matrix in R M×q obtained as the Jacobian of θ →τ(θ)∈R M , i.e., T :=∇τ(θ), with

(τ(θ)) m = 1 N N n=1 ∂ im k(Y m ,h θ (Z n )).

In ( 18), we used the notation ∂ i+d k(y,x) for the partial derivative of k w.r.t. x i . 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.

3.3. PRACTICAL CONSIDERATIONS:

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 matrix T admits a second expression of the form T =CB where B is the Jacobian matrix of (h θ (Z n )) 1≤n≤N . Although this expression is impractical to compute in general, it suggests that C can be 'simplified'. This simplification can be achieved in practice by computing the SVD of CC =USU T and pre-multiplying T by S † U T . 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 of CC T of the form CC =USU T , then ( 17) is equal to:

EQUATION

where P :=S † S and T :=S † U T T .

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 when N 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 CC T in the case of ( 19). This gives an overall cost of O(dNM 2 +qM 2 +M 3 ).

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 (X n ) 1≤n≤N and the basis points (Y m ) 1≤m≤M and σ 0 is fixed a priori. Another choice is the median heuristic Garreau et al. (2018) .

3.4. THEORY:

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 in H 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:

EQUATION

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 c means 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:= 1 2+c . Under Assumption 1 and Assumptions (A) to (G) listed in Appendix A.2, for N 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:

ÿ ∇ W L(θ)-∇ W L(θ) 2 =O N -2 4+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, with M required to be of order dNlog(N) in the limit when c→∞ in which case the Nyström approximation loses its computational advantage.

4. EXPERIMENTS:

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

4.1. CONVERGENCE ON SYNTHETIC MODELS:

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. . An optimal step-size γ t is used: γ t = 0.1 for both KWNG and WNG while γ t =0.0001 for EG.

Figure 1 shows the evolution of the the relative error w.r.t. the sample-size N, 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.

4.2. APPROXIMATE INVARIANCE TO PARAMETRIZATION:

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 10 7 . 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 ).

ACKNOWLEDGEMENT:

GM has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement n o 757983).

A PRELIMINARIES:

A.1 NOTATION

We recall that Ω is an open subset of R d while Θ is an open subset of parameters in R q . Let Z ⊂ R p be a latent space endowed with a probability distribution ν over Z. Additionally, (θ,z) →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 write B to denote the jacobian of h θ w.r.t. θ viewed as a linear map from R q to L 2 (ν) d without explicit reference to θ: Bu(z)=∇h θ (z).u; ∀u∈R q . 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(θ) using N samples from ρ θ .

We also consider a Reproducing Kernel Hilbert Space H 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:

∂ i f(x)= f,∂ i k(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 L 2 (ν) the space of real valued and square integrable functions under ν and by L 2 (ν) d the space of square integrable vector valued functions under ν and with values in R d .

A.2 ASSUMPTIONS:

We make the following set of assumptions: (F) The following quantity is finite:

κ 2 =sup x∈Ω 1≤i≤q ∂ i ∂ i+q k(x,x).

(G) For all 0≤δ ≤1, it holds with probability at least 1-δ that ÷ ∇L(θ)-∇L(θ) N -1 2 .

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 θ (Z i ))

where Z i 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(θ) 2 dν(z)= 1 N ∇ θ l(h θ (z)) 2 dν(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.

A.3 OPERATORS DEFINITION:

Differential operators. We introduce the linear L operator and its adjoint L :

L:H→L 2 (ν) d L :L 2 (ν) d →H f →(∂ i f •h θ ) 1≤i≤d v → d i=1 ∂ i k(h θ (z),.)v i (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 for A in terms of outer product ⊗ and its regularized version for a given λ>0,

A= d i=1 ∂ i k(h θ (z),.)⊗∂ i k(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) that A is also a compact operator under Assumption (E).

Assume now we have access to N samples (Z n ) 1≤n≤N as in the main text. We define the following objects:

Â:= 1 N N n=1 d i=1 ∂ i k(h θ (Z n ),.)⊗∂ i k(h θ (Z n ),.), Âλ := Â+λI.

Furthermore, if v is a continuous function in L 2 (ν) d , then we can also consider an empirical estimator for L v:

' L v := 1 N N n=1 d i=1 ∂ i k(h θ (Z n ),.)v i (Z n ).

Subsampling operators. We consider the operator Q M defined from H to R M by:

EQUATION

where (e m ) 1≤m≤M is an orthonormal basis of R M . Q M admits a singular value decomposition of the form Q M =UΣV , with V V :=P M being the orthogonal projection operator on the Nyström subspace H M . Similarly to Rudi et al. (2015) ; Sutherland et al. (2017) , we define the projected inverse function G M (C) as:

G M (C)=V (V CV ) -1 V .

We recall here some properties of G M 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. G M (A)P M =G M (A) 2. P M G M (A)=G M (A) 3. G M (A λ )A λ P M =P M 4. G M (A λ )=(P M AP M +λI) -1 P M 5. A 1 2 λ G M (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:

ÿ ∇ W L(θ)=( D(θ)+G M,N ) -1 ÷ ∇L(θ)

where G M,N is given by:

G M,N :=( ' L B) G M ( Âλ ) ' 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 G M,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 of G M,N defined as :

G M :=(L B) G M (A λ )L B (22) B BACKGROUND IN INFORMATION GEOMETRY B.1 FISHER-RAO STATISTICAL MANIFOLD

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 g F is defined for all ρ∈P as an inner product over T ρ of the form:

g F ρ (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.

B.2 WASSERSTEIN STATISTICAL MANIFOLD:

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 location x that is mapped to an element of mass of ρ at location y using the coupling π:

x-y 2 dπ(x,y)

The squared Wasserstein 2 distance between ρ and ρ is defined as the smallest transport cost over all possible couplings:

W 2 2 (ρ,ρ )= inf π∈Π(ρ,ρ )

x-y 2 dπ(x,y).

A dynamical formulation of W 2 was provided by the celebrated Benamou-Brenier formula in Benamou and Brenier (2000) :

W 2 2 (ρ,ρ )=inf φt 1 0 φ l (x) 2 dρ l (x)dl

where the infimum is taken over the set of vector fields φ : [0,1] × Ω → R d . Each vector field of the potential determines a corresponding probability distribution ρ l as the solution of the continuity equation:

EQUATION

)

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 that W 2 (ρ,ρ ) 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) 2 dρ l (x) Hence, W 2 2 (ρ,ρ ) can be written as:

W 2 2 (ρρ )=inf ρ l 1 0 |∂ l ρ l | 2 dl.

In fact, ∂ l ρ l can be seen as an element in the tangent space T ρ l P 2 to P 2 at point ρ l . To ensure that ( 23) is well defined, T ρ P 2 can be defined as the set of distributions σ satisfying σ(1)=0.

EQUATION

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 ρ l P 2 . This inner product defines in turn a metric tensor g W on P 2 called the Wasserstein metric tensor (see Otto and Villani (2000) ; Ambrosio et al. ( 2004)): Definition 5. The Wasserstein metric g W is defined for all ρ∈P 2 as the inner product over T ρ P 2 of the form:

g W ρ (σ,σ ):= φ σ (x) φ σ (x)dρ(x), ∀σ,σ ∈T ρ P 2

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. L 2 (ρ) d .

Definition 5 allows to endow P 2 with a formal Riemannian structure with W 2 being its geodesic distance:

W 2 2 (ρ,ρ )=inf ρ l 1 0 g ρ l (∂ l ρ l ,∂ l ρ l )dl.

B.3 NEGATIVE SOBOLEV DISTANCE AND LINEARIZATION OF THE WASSERSTEIN DISTANCE:

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) :

W 2 2 (ρ θ ,ρ θ+u )= ρ θ+u -ρ θ 2 H -1 (ρ θ ) +o( u 2 ).

Further performing a Taylor expansion of ρ θ+u -ρ θ H -1 (ρ θ ) in u leads to a quadratic term u G W (θ t )u where we call G W (θ 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 G W (θ 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 W 2 (ρ,ρ+ (ρ -ρ))= ρ-ρ 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:

W 2 (δ θ ,δ θ + (δ θ -δ θ ))= √ θ-θ

Hence, 1 W 2 (δ θ ,δ θ + (δ θ -δ θ )) 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 W 2 (δ θ ,δ θ+ 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 ∇ x f(x) 2 dρ(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 ∇ x f(θ) 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 G W (θ)u which can be computed in closed form:

EQUATION

In this case, choosing f(θ)=0 and ∇f(θ)=u achieves the supremum which is simply given by 1 2 u 2 . Equation ( 26) can be seen as a limit case of (25) when →0:

1 2 u G W (θ)u:= sup f∈C ∞ c (Ω) f(θ)=0 lim →0 1 (f(θ)-f(θ+ u))- 1 2 ∇ x f(θ) 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.

C PROOFS C.1 PRELIMINARY RESULTS:

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 ḠW (ψ) and ḠF (ψ) their corresponding Wasserstein and Fisher information matrices whenever they are well defined.

Proof of Proposition 1. Here we only consider the case when ∇ D L(θ) is either given by the Fisher natural gradient ∇ F L(θ) or the Wasserstein Natural gradient ∇ W L(θ). 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 ) G D (θ s ) -1 ∇ θ L(θ s )

By the chain rule, we have that

∇ θ L(θ s )=∇ θ Ψ(θ s )∇ ψ L( ψ s ), hence: ˙ ψ s =-∇ θ Ψ(θ s ) G D (θ 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 =-G D ( ψ 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 ḠW (ψ) and ḠF (ψ) are related to G W (θ) and G F (θ) by the relation:

ḠW (ψ)=∇ ψ Ψ -1 (ψ) G W (θ)∇ ψ Ψ -1 (ψ) ḠF (ψ)=∇ ψ Ψ -1 (ψ) G F (θ)∇ ψ Ψ -1 (ψ)

Proof. Let v ∈ R q and write u = ∇ θ Ψ -1 (ψ)v, then by the dual formulations of G W (θ) and G F (θ) in Proposition 2 we have that:

1 2 v ∇ ψ Ψ -1 (ψ) G F (θ)∇ ψ Ψ -1 (ψ)v = sup f∈C ∞ c (Ω) f(x)dρ θ (x)=0 ∇ρ θ (f) ∇ θ Ψ -1 (ψ)v- 1 2 f(x) 2 dρ θ (x)dx, Now recalling that ∇ ψ ρψ =∇ θ ρ θ ∇ ψ Ψ -1 (ψ)

by Lemma 11, it follows that:

1 2 v ∇ ψ Ψ -1 (ψ) G F (θ)∇ ψ Ψ -1 (ψ)v = sup f∈C ∞ c (Ω) f(x)dρ θ (x)=0 ∇ ψ ρψ (f) v- 1 2 f(x) 2 dρ θ (x)dx,

Using again Proposition 2 for the reparametrized model ρψ , we directly have that:

1 2 v G F (ψ)v = sup f∈C ∞ c (Ω) f(x)dρ θ (x)=0 ∇ ψ ρψ (f) v- 1 2 f(x) 2 dρ θ (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∈R q 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 →∇ρ θ (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∈R q , 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) → 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:

EQUATION

Recalling that the set of smooth and compactly supported functions

C ∞ c (∞) is dense in L 2 (ρ θ )

and that the objective function in ( 27) is continuous and coercive in f, it follows that (27) admits a unique solution f * in L 2 (ρ θ ) which satisfies the optimality condition:

f(x)(∇ρ θ (x)) udx= f(x)f * (x)ρ θ (x)dx ∀f ∈L 2 (ρ θ )

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 G F (θ)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 (e s ) 1≤s≤q be an orthonormal basis of R q . Under Assumptions (B) and (C), there exists an optimal solution Φ = (φ s ) 1≤s≤q with φ s in L 2 (ρ θ ) d satisfying the PDE:

∂ s ρ θ =-div(ρ θ φ s )

The elliptic equations also imply that L ∇h θ =L (Φ•h θ ). Moreover, the Wasserstein information matrix G W (θ) on P Θ at point θ can be written as G W (θ) = Φ Φ where the inner-product is in L 2 (ρ θ ) d and satisfies:

1 2 u G W (θ)u= sup f∈C ∞ c (Ω) f(x)dρ θ (x)=0 ∇ρ θ (f) u- 1 2 ∇ x f(h θ (z)) 2 dν(z).

for all u∈R q .

Proof. Let (e s ) 1≤s≤q be an orthonormal basis of R q . For all 1 ≤ s ≤ q, we will establish the existence of an optimal solution φ s in L 2 (ρ θ ) d satisfying the PDE:

EQUATION

Consider the variational problem:

EQUATION

where S is a Hilbert space obtained as the closure in L 2 (ρ θ ) d of functions of the form φ

= ∇ x f with f ∈C ∞ c (Ω): S :={∇ x f | f ∈C ∞ c (Ω)} L2(ρ d θ )

. We have by Assumption (C) that:

φ(h θ (z)) ∂ θs h θ (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)) ∂ θs h θ (z))dν(z)

In particular, for any f ∈C ∞ c (Ω), it holds that:

∇ x f(h θ (z)) φ * s (h θ (z))dν(z)= ∇ x f(h θ (z)) ∂ θs h θ (z))dν(z)

which is equivalent to (28) and implies directly that L T ∇h θ = L T Φ • h θ where Φ := (φ * s ) 1≤s≤q . The variational expression for 1 2 u G W u follows by noting that (29) admits the same optimal value as

sup f∈C ∞ c (Ω) f(x)dρ θ (x)=0 ∇ρ θ (f) u- 1 2 ∇ x f(h θ (z)) 2 dν(z).

That is because S is by definition the closure in L 2 (ρ θ ) d of the set of gradients of smooth and compactly supported functions on Ω.

Proof of Proposition 3. This is a consequence of Proposition 13.

C.2 EXPRESSION OF THE ESTIMATOR:

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∈H M 1 N N n=1 ∇f(X n ) 2 +λ f 2 H + 1 R(f) D -1 R(f)+ 2 R(f) D -1 ÷ ∇L(θ) with R(f) given by R(f)= 1 N N n=1 ∇f(X n ) B(Z n ). Now, recalling that any f ∈H M can be written as f = M m=1 α m ∂ im k(Y m ,.

), and using the reproducing property

∂ i f(x)= f,∂ i k(x,.

) H (Steinwart and Christmann, 2008, Lemma 4.34) , it is easy to see that:

1 N N n=1 ∇f(X n ) 2 = 1 N 1≤n≤N 1≤i≤d ( M m=1 α m ∂ im ∂ i+d k(Y m ,X n )) 2 . f 2 H = 1≤m,m ≤M α m α m ∂ im ∂ i m +d k(Y m ,Y m ) R(f)= 1 N 1≤n≤N 1≤i≤d 1≤m≤M α m ∂ im ∂ i+d k(Y m ,X n )B i (Z n )

The above can be expressed in matrix form using the matrices defined in Proposition 5:

1 N N n=1 ∇f(X n ) 2 =α CC α; f 2 H =α Kα; R(f)=α CB.

Hence the optimal solution f * is of the form f * = M m=1 α * m ∂ im k(Y m ,.), with α * obtained as a solution to the finite dimensional problem in R M :

min α∈R M α ( CC + λK +CBD -1 B C )α+2α CBD -1 ÷ ∇L(θ)

It is easy to see that α * are given by:

α * =-( CC T + λK +CBD -1 B T C T ) † CBD -1 ÷ ∇L(θ).

Now recall that the estimator in Proposition 5 is given by: ÿ

∇ W L(θ) = 1 D -1 U θ ( f * ). Hence, 1 D -1 ( ÷ ∇L(θ)-B T C T α * )

The desired expression is obtained by noting that CB =T using the chain rule.

C.3 CONSISTENCY RESULTS:

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 in H. More precisely: Assumption 2. For all 1 ≤ i ≤ q, there exits functions f i ∈ H such that φ i = ∇f i . Additionally, f i are of the form f i =A α v i for some fixed α≥0, with v i ∈H and A being the differential covariance operator defined on H by A:f

→ d i=1 ∂ i k(h θ (z),.)∂ i f(h θ (z))dν(z).

The parameter α characterizes the smoothness of f i 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, α +1 2 ). Under Assumption 2 and Assumptions (A) to (G) listed in Appendix A.2, for N 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:

ÿ ∇ W L(θ)-∇ W L(θ) 2 =O N -2b 2b+1 .

In the worst case where α = 0, the proposed estimator needs at most M ∼ (d √ N log(N)) to achieve a convergence rate of N -1 2 . The smoothest case requires only M ∼(dN

1 3 log(N))

to achieve a rate of N -2 3 . 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 advantage 1 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-δ:

' ∇ W L-∇ W L =O(N -b 2b+1 )

provided that M ∼ dN Proof. Here for simplicity we assume that D(θ) = I without loss of generality and we omit the dependence in θ and write ∇ W L and ∇L instead of ∇ W L(θ) and ∇L(θ) and ∇ W L(θ). We also define Ĝ = I + G M,N and G = I + G W . By Proposition 9, we know that ' ∇ W L = Ĝ-1 ∇L. We use the following decomposition:

' ∇ W L-∇ W L ≤ Ĝ-1 ( ∇L-∇L) + Ĝ-1 (G M,N -G W )G -1 W ∇L + Ĝ-1 G -1 W ∇L To control the norm of Ĝ-1 we write Ĝ-1 = G -1 2 (H + I) -1 G -1 2

, where H is given by

H := G -1 2 (G M,N -G W )G -1

2 . Hence, provided that µ := λ max (H), the highest eigenvalue of H, is smaller than 1, it holds that:

(H +I) -1 ≤(1-µ) -1 .

Moreover, since G W is positive definite, its smallest eigenvalue η is strictly positive. Hence, G -1 ≤(η+ ) -1 . Therefore, we have Ĝ-1 ≤(η+ )(1-µ)) -1 , which implies:

' ∇ W L-∇ W L ≤(η+ ) -1 Ç 1-µ +η -1 ∇L G M,N -G W + η -1 ∇L å .

Let 0 ≤ δ ≤ 1. We have by Assumption (G) that ∇L -∇L = O(N -1 2 ) with probability at least 1-δ. Similarly, by Proposition 16 and for N large enough, we have with probability at least

1-δ that G M,N -G W =O(N -b 2b+1

) where b is defined in Proposition 16. Moreover, for N large enough, one can ensure that µ≤ 1 2 so that the following bound holds with probability at least 1-δ:

' ∇ W L-∇ W L (η+ ) -1 Ä 2N -1 2 +η -1 ∇L (N -b 2b+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 for N large enough that:

G M,N -G W =O(N -b 2b+1 ).

provided that M ∼ dN 

G M,N -G W ≤ G M -G W + G M -G M,N

were G M is defined in ( 22) and is obtained by taking the number of samples N to infinity while keeping the number of basis points M fixed.

The estimation error G M -G M,N is controlled using Proposition 17 where, for any 0≤δ ≤1, we have with probability at least 1-δ and as long as N ≥M(1,λ,δ):

G M,N -G M ≤ B √ Nλ (a N,δ + 2γ 1 κ+2γ 1 λ+κ √ Nλ )+ 1 Nλ a 2 N,δ .

In the limit where N → ∞ and λ → 0, only the dominant terms in the above equation remain which leads to an error G M,N -G M = O((Nλ) -1 2 ). Moreover, the condition on N can be expressed as λ -1 logλ -1 N.

To control the error approximation error G M -G W 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 if M ≥M(d,λ,δ), it holds with probability at least 1-δ: 

G M -G W λ min(1,α+

G M -G W λ 1 2+c

Let's set b := min(1,α + 1 2 ) 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λ -1 logλ -1 and λ -1 logλ -1 N it holds with probability as least 1-δ that

G M,N -G W =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:

G M,N -G W =O(N -b 2b+1 ).

Proposition 17. For any 0≤δ ≤1, provided that N ≥M(1,λ,δ), we have with probability as least 1-δ:

G M,N -G M ≤ B √ Nλ (2a N,δ + 2γ 1 κ+2γ 1 λ+κ √ Nλ )+ 1 Nλ a 2 N,δ .

with:

a N,δ := … 2σ 2 1 log 2 δ + 2alog 2 δ √ N Proof.

For simplicity, we define E = ' L B -L B. By definition of G M,N and G M we have the following decomposition:

G M,N -G M =E G M ( Âλ )E E0 +E G M ( Âλ )L B E1 +B LG M ( Âλ )E E2 -B LG M (A λ )P M ( Â-A)P M G M ( Âλ )L B E3

The first three terms can be upper-bounded in the following way:

E 0 = E Â-1 2 λ Â1 2 λ G M ( Âλ ) Â1 2 λ Â-1 2 λ E ≤ E 2 Â-1 λ ≤1/λ Â1 2 λ G M ( Âλ ) Â1 2 λ ≤1 E 1 = E 2 = E A -1 2 λ A 1 2 λ G M (A λ )A 1 2 λ A -1 2 λ L B ≤ B E Â-1 2 λ ≤1/ √ λ Â1 2 λ G M ( Âλ ) Â1 2 λ ≤1 A -1 2 λ L ≤1 Â-1 2 λ A 1 2 λ

For the last term E 3 , we first recall that by definition of G M (A λ ) we have:

G M (A λ )P M ( Â-A)P M G M (A λ )=G M (A λ )( Â-A)G M (A λ ).

Therefore, one can write:

E 3 = B LA -1 2 λ A 1 2 λ G M (A λ )A 1 2 λ A -1 2 λ ( Â-A)A -1 2 λ A 1 2 λ Â-1 2 λ Â1 2 λ G M ( Âλ ) Â1 2 λ Â-1 2 λ A 1 2 λ A -1 2 λ L B ≤ B 2 LA -1 2 λ 2 ≤1 A 1 2 λ G M (A λ )A 1 2 λ ≤1 Â1 2 λ G M ( Âλ ) Â1 2 λ ≤1 A 1 2 λ Â1 2 λ 2 A -1 2 λ ( Â-A)A -1 2 λ ≤ B 2 A 1 2 λ Â1 2 λ 2 A -1 2 λ ( Â-A)A -1 2 λ

We recall now (Rudi et al., 2015, Proposition 7.) which allows to upper-bound A

1 2 λ Â1 2 λ by (1-η) -1 2 where η =λ max (A 1 2 λ (A-Â)A 1 2

λ ) provided that η <1. Moreover, (Rudi et al., 2015, Proposition 8.) allows to control both η and A

-1 2 λ ( Â-A)A -1 2 λ

under Assumption (F). Indeed, for any 0≤δ ≤1 and provided that 0<λ≤ A it holds with probability 1-δ that:

A -1 2 λ ( Â-A)A -1 2 λ ≤2γ 1 1+κ/λ 3N + … 2γ 1 κ Nλ ; η ≤ 2γ 2 3N + … 2γ 2 κ

Nλ where γ 1 and γ 2 are given by:

γ 1 =log( 8T r(A) λδ ); γ 2 =log( 4T r(A) λδ ).

Hence, for N ≥M(1,λ,δ) we have that (1-η) -1 2 ≤2 and one can therefore write:

E 3 ≤4 B 2 (2γ 1 1+κ/λ 3N + … 2γ 1 κ Nλ ) E 1 = E 1 ≤ 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σ 2 1 log 2 δ + 2alog 2 δ √ N ):= 1 √ N a N,δ .

Finally, we have shown that provided that N ≥M(1,λ,δ) then with probability greater than 1-δ one has:

G M,N -G M ≤ B √ Nλ (2a N,δ + 2γ 1 κ+2γ 1 λ+κ √ Nλ )+ 1 Nλ a 2 N,δ .

Proposition 18. Let 0≤λ≤ A and define M(d,λ,δ):= 128 9 log 4T r(A) λδ (dκλ -1 +1). Under Assumption 1 and Assumption (F), for any δ ≥0 such that M ≥M(d,λ,δ) the following holds with probability 1-δ:

G M -G W λ 1 2+c

Proof. We consider the error G M -G W . Recall that G W is given by G W = Φ Φ 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 G M -G W . Let's call E =Φ•h θ -LF κ

We consider the following decomposition:

G M -G W =(L Φ•h θ ) G M (A λ )L Φ•h θ -Φ Φ =E LG M (A λ )L E E1 -E E E2 +F κ L LG M (A λ )-I L Φ•h θ E3 +E L G M (A λ )L L-I F κ E4

First we consider the term E 1 one simply has:

E 1 ≤κ 2 LA -1 2 λ ≤1 A 1 2 λ G M (A λ )A 1 2 λ ≤1 A -1 2 λ L ≤1 ≤κ 2

The second term also satisfies E 1 ≤ κ 2 by definition of F κ . For the last two terms E 3 and E 4 we use Lemma 20 which allows to control the operator norm of L(G M (A λ )L L-I). Hence, for any δ ≥0 and M such that M ≥M(d,λ,δ) and for κ≤1 it holds with probability 1-δ that:

E 3 √ λκ -c ; E 4 √ λκ -c

We have shown so far that G M -G W (κ 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 G M -G W λ 1 2+c

which shows the desired result.

Proposition 19. Let 0≤λ≤ A and define M(d,λ,δ):= 128 9 log 4T r(A) λδ (dκλ -1 +1). Under Assumption 2 and Assumption (F), for any δ ≥0 such that M ≥M(d,λ,δ) the following holds with probability 1-δ:

G M -G W λ min(1,α+ 1 2 )

Proof. Recall that G W is given by G W =Φ Φ with Φ defined in Proposition 13. By Assumption 2, we have that Φ=∇(A α V ) with V :=(v s ) 1≤s≤q ∈H q . Hence, one can write

G M -G W =(L Φ•h θ ) G M (A λ )L Φ•h θ -Φ Φ =V (A α (AG M (A λ )A-A)A α V

we can therefore directly apply Lemma 20 and get G M -G W λ min(1,α+ 1 2 ) with probability 1-δ for any δ ≥0 such that M ≥M(d,λ,δ).

Lemma 20. Let 0 ≤ λ ≤ A , α ≥ 0 and define M(d,λ,δ) := 128 9 log 4T r(A) λδ (dκλ -1 + 1). Under Assumption (F), for any δ ≥0 such that M ≥M(d,λ,δ) the following holds with probability 1-δ:

L(G M (A λ )L L-I)A α λ min(1,α+ 1 2 )

Proof. We have the following identities:

L(G M (A λ )L L-I)A α =L(G M (A λ )A λ -I -λG M (A λ ))A α =LA -1 2 λ A 1 2 λ (G M (A λ )A λ P M -I)A α E1 -λLA -1 2 λ A 1 2 λ G M (A λ )A 1 2 λ A -1 2 λ A α E3 +LA -1 2 λ A 1 2 λ G M (A λ )A 1 2 λ A 1 2 λ (I -P M )A α E2 .

For the first E 1 we use (Sutherland et al., 2017, Lemma 1 

(iii)) which implies that G M (A λ )A λ P M =P M . Thus E 1 = LA -1 2 λ A 1 2

λ (P M -I)A α . Moreover, by Lemma 21 we have that A 1 2 λ (I -P M ) ≤ 2

√:

λ with probability 1 -δ for M > M(d,λ,δ). Therefore, recalling that (I -P M ) 2 = I -P M since P M is a projection, one can further write:

E 1 ≤ LA -1 2 λ ≤1 A 1 2 λ (P M -I) 2 ≤λ A -1 2 λ A α E 2 ≤ LA -1 2 λ ≤1 A 1 2 λ G M (A λ )A 1 2 λ ≤1 A 1 2 λ (P M -I) 2 ≤4λ A -1 2 λ A α E 3 ≤λ LA -1 2 λ ≤1 A 1 2 λ G M (A λ )A 1 2 λ ≤1 A -1 2 λ A α It remains to note that A -1 2 λ A α ≤ λ α-1 2 when 0 ≤ α ≤ 1 2 and that A -1 2 λ A α ≤ A α-1 2 for α > 1 2

which allows to conclude.

C.4 AUXILIARY RESULTS:

Lemma 21. Let 0 ≤ λ ≤ A . Under Assumption (F), for any δ ≥ 0 such that M ≥M(d,λ,δ):= 128 9 log 4T r(A) λδ (κλ -1 +1) the following holds with probability 1-δ:

A 1 2 λ (I -P M ) ≤2 √ λ

Proof. The proof is an adaptation of the results in Rudi et al. (2015) ; Sutherland et al. (2017) . Here we recall Q M defined in (21). Its transpose Q M sends vectors in R M to elements in the span of the Nyström basis points, hence P M and Q M have the same range, i.e.: range(P M )= range(Q M ). We are in position to apply (Rudi et al., 2015, Proposition 3.) which allows to find an upper-bound on A 1 2

λ (P M -I) in terms of Q M :

A 1 2 λ (P M -I) ≤ √ λ A 1 2 λ (Q M Q M +λI) -1 2 .

For simplicity we write ÂM :=Q M Q M and E 2 :=A λ . We also denote by β =λ max (E 2 ) the highest eigenvalue of E 2 . We can therefore control A 1 2 λ ( ÂM +λI) -1 2 in terms of β using (Rudi et al., 2015, Proposition 7) provided that β <1: ∂ i k(y,.)⊗∂ i k(y,.)dp θ (y)=A.

We also have by Assumption (F) that v m ,A -1 λ v m ≤ 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 γ =log 4T r(A) λδ provided that λ≤ A . Thus by choosing M ≥ 128γ 9 (dκλ -1 +1) we have that β ≤ 3 4 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 ≤ 2alog 2 δ N + 2σ 2 1 log 2 δ N

Proof. denote by v n = d i=1 ∂ i k(X n ,.)B i (Z n ), we have that E[v n ] = 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[ z n -L B p H ]≤ 1 2 p!σ 2 1 a p-2 . To simplify notations, we write x and x instead of h θ (z) and h θ (z ). We have that: 

E[ z n -L B p H ]=

E2:

We used the convexity of the norm and the triangular inequality to get the last line. We introduce the notation γ i (x):=∂ i k(x,.)-∂ i k(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 E 1 can be upper-bounded as follows:

E 1 = T r(B(z)B(z) Γ(x)) p 2 ≤ B(z) 2 T r(Γ(x) 2 ) 1 2 p 2 .

Published as a conference paper at ICLR 2020 Moreover, we have that T r(Γ(x) 2 ) 1 2 = ( 1≤i,j≤d γ i (x), γ j (x) 2 H ) 1 2 ≤ d i=1 γ i (x) 2 . We further have that γ i (x) ≤ ∂ i ∂ i+d k(x,x) 1 2 + ∂ i ∂ i+d k(h θ (z),h θ (z)) 1 2 dν(z) and by Assumption (F) it follows that γ i (x) ≤ 2 √ κ. Hence, one can directly write that: E 1 ≤ (2 √ κd) p B(z) p dν(z). Recalling Assumptions (B) and (C) we get:

E 1 ≤2 p-1 (2 √ κd) p C(θ) p (1+ 1 2 p!ζ p-2 σ 2 )

Similarly, we will find an upper-bound on E 2 . To this end, we introduce the matrix Q(x ,x") whose components are given by Q(x ,x") i,j =∂ i ∂ i+d k(x ,x"). One, therefore has: Once again, we have that T r(Q(x ,x") 2 )

E 2 = T

1 2 ≤ ( d i=1 ∂ i ∂ i+d k(x ,x )) 1 2 ( d i=1 ∂ i ∂ i+d k(x",x"))

1 2 ≤ dκ thanks to Assumption (F). Therefore, it follows that: 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[ z n -L B p H ]≤ 1 2 p!σ 2 1 a p-2 . Hence one concludes using Bernstein's inequality for a sum of random vectors (see for instance Rudi et al., 2015, Proposition 11) .

D EXPERIMENTAL DETAILS D.1 NATURAL WASSERSTEIN GRADIENT FOR THE MULTIVARIATE NORMAL MODEL:

Multivariate Gaussian. Consider a multivariate gaussian with mean µ ∈ R d and covariance matrix Σ ∈ R d ×R d parametrized using its lower triangular components s = T (Σ). We denote by Σ = T -1 (s) the inverse operation that maps any vector s∈R 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(θ),∇ s L(θ), 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 (∇ s L(θ)). 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 =Σ 1 2 Z +µ,Z ∼N (0,I d )

:

The estimator proposed bySutherland et al. (2017) also requires M to grow linearly with the dimension d although such dependence doesn't appear explicitly in the statement ofSutherland et al.

2017, Theorem 2.

Figures and Tables:

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 as M increases and N =5000. A gaussian kernel is used with a fixed bandwidth σ =1.

Figure 2: Left (a): Training error per iteration for KWNG, WNG, and EG. Right (b): projection of the sequence of updates obtained using KWNG, WNG and EG along the first two PCA directions of the WNG trajectory. The dimension of the sample space is fixed to d = 10. Exact valued for the gradient are used for EG and WNG. For KWNG, N = 128 samples and M = 100 basis points are used. The regularization parameters are set to: λ = 0 and = 10 -10. An optimal step-size γ t is used: γ t = 0.1 for both KWNG and WNG while γ t =0.0001 for EG.

Figure 3: Test accuracy and Training accuracy for classification on Cifar10 (top) and Cifar100 (bottom) in both the ill-conditioned case (left side) and well-conditioned case (right side) for different optimization methods. on Cifar10 Results are averaged over 5 independent runs except for KFAC and eKFAC.

A) Ω is a non-empty open subset of R d . (B) There exists positive constants ζ and σ such that z p dν(z)≤ 1 2 p!ζ 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+d k(x,x)dp θ (x)<∞ for all 1≤i≤d.

2b+1 log N, λ ∼ N 1 2b+1 and N -b 2b+1 where b := min(1, α + 1 2 ) when Assumption 2 holds and b= 1 2+c when Assumption 1 holds instead.

logN where b := min(1,α+ 1 2 ) when Assumption 2 holds and b = 1 2+c when Assumption 1 holds instead. Proof. To control the error G M,N -G W we decompose it into an estimation error G M,N -G M and approximation error G M -G W :

Now we need to make sure that β <1 for M large enough. To this end, we will apply(Rudi et al., 2015,  Proposition 8.)  to ÂM . Denote by v m = √ d∂ im k(Y m ,.). Hence, by definition of ÂM it follows that ÂM = 1 M M m=1 v m ⊗v m . Moreover, (v m ) 1≤m≤M are independent and identically distributed and satisfy: E[v m ⊗v m ]= q i=1

i k(x,.)B i (z)-d i=1 ∂ i k(x ,.)B i (z )dν(z ) k(x,.)-∂ i k(x ,.))B i (z)dν(z )

r((B(z)-B(z ))(B(z)-B(z")) Q(x ,x")dν(z )dν(z")

2 ≤( √ dκ) p | B(z)-B(z ) dν(z)| p dν(z) ≤3 p-1 ( √ dκ) p C(θ) p (2 p + z p dν(z

symmetric matrix in R d ×R d . 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 R d and S and S are symmetric matrices in R d ×R d the metric evaluated at u and v is given by:u G(θ)v =m m +T r(AΣA )where A and A are symmetric matrices that are solutions to the Lyapunov equation: S =AΣ+ΣA, S =A Σ+ΣA .

Miss-specified case. Here we work under Assumption 1 which allows to use Proposition 18. Hence, for any 0≤δ ≤1 and if M ≥M(d,λ,δ), it holds with probability at least 1-δ:

Table content: <table/>

Peer review (iclr2020):
Thank you for your revision and for the rebuttal. This a strong submission with insightful angle on natural gradients and with provable guarantees. The authors improved a lot the manuscript and incorporated reviewers feedback. I am increasing my score to 8. ### Summary of the paper: The paper provides a way to estimate the natural Wasserstein gradient using Kernel estimators. The idea is neat and novel. Natural Wasserstein Gradient similar to the so called natural fisher gradients preconditions the gradient using a matrix that uses the local curvature of the manifold of the parametric distribution. Authors give variational forms of the Fisher information matrix of an explicit model , using the variational form of the chi squared or the Fisher Rao divergence. Similarly authors give a variational form of the wasserstein natural gradient . Let theta be the parameter of the parametric implicit model, theta in R^q. For a descent direction , the variational form is obtained via finding an objective S, , where is a form of *Wasserstein information matrix*. Authors then propose to learn the function f in an RKHS and propose to find the descent direction by solving where r(u) is a quadratic regularizer on u. The sup problem has a closed form solution and can be approximated using Nystrom approximation and randomization on dimensions. The problem in u has also a closed form solution , and one used u as the proxy to the natural descent. Authors under some assumption show that the estimated natural W gradient in RKHS is concentrated around the true one. Experiments on synthetic data and in classification on CIFAR 10 and CIFAR 100 shows that the preconditioning of the gradients that the method offers allows faster convergence in both well conditioned and ill conditioned initialization of the weights of the neural network. Review : The paper is not easy to follow and the high level intuition how the method works is not well explained. It would be easier for the reader, to motivate the natural wasserstein descent from how one defines natural Fisher descent , where one seeks a first order approximation of as we perturb in the parameter space and this well known that this epsilon . Hence natural Gradient descent is : Now for the wasserstein distance one has also similarly: and it is known that as epsilon goes to zero we have: Now replacing with the implicit model as epsilon goes to zero we get the expression given in the paper using a simple taylor expansion: in a sense the paper is proposing to linearize around the perturbation in the parameter space of the implicit model and this can be done using , as pointed and used in many recent works. then the paper proposes to approximate in RKHS which was already proposed in Mroueh et al in Sobolev Descent. AISTATS 2019. We encourage the authors to layout in the beginning the derivations form this point of view which will make the paper easier to digest, the expression in Equation 7 seems mysterious and pulled out of a hat, but it is easier to understand by going to perturbation analysis usually done on KL for Fisher Natural gradient and to do it also here starting from the linearization of with , and how to approximate it in RKHS as it was already proposed in the literature in Mroueh et al Sobolev Descent. I read carefully the proofs of Proposition 1, 2, 3. I did not ready full the proofs of the concentration of the estimator , but they seem sensible as they follow usual bounding strategies in this context. Questions: - There is nothing special about the wasserstein natural gradient flow variational form and implicit model, once can apply the same to the variational form of Fisher, that would be probably more efficient? It would be great to baseline this one ? -the constraint is not imposed in the kernelized version? - the method comes disappointing since it seems that the preconditioning that the Wasserstein gradient gives is not enough and is need where D is diagonal depends on T. Have you tried with ? it might be that the scaling of the gradients is coming only from that ? - Can you give timings for computing each gradient update and how it compares to regular SGD or diagonal approximation of Fisher natural gradient? - Does one need preconditioned gradient if the network was self normalized (like batch norm or spectral norm etc)? Overall assessment: That is a good theoretical work with provable guarantees. The computational complexity of each gradient estimate is large which makes the method not quite appealing in practice.