Kernelized Wasserstein Natural Gradient
Authors: M Arbel
A Gretton
W Li
G Montufar
Venue:  (2020)
Keywords: kernel methods
natural gradient
information geometry
Wasserstein metric

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 \texttt{Cifar10} and \texttt{Cifar100} empirically. 

Published as a conference paper at ICLR 2020
KERNELIZED WASSERSTEIN NATURAL GRADIENT
Michael Arbel, Arthur Gretton
Gatshy Computational Neuroscience Unit University College London imichael. n. arbel, arthur. gretton @ggaall com
Wuchen Li
University of California, Los Angeles
wcli@math. ucla.edu
Guido Montúfar
University of California, Los Angeles, and Max Planck Institute for Mathematics in the Sciences montu faremir wwg.ppp......o..
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 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 requires inverting a high dimensional matrix at each iteration. We propose general framework to approximate the natural gradient for the Wasserstein metric, by leveraging a dual formulation of the metric restricted to 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 simple examples, and show the advantage of using such an estimator classification tasks Cifarl0 and Cifar100 empirically.
INTRODUCTION
The success machine learning algorithms relies on the quality of an underlying optimization method. Many of the current state-of-the-ar methods rely on variants of Stochastic Gradient Descent (SGD) such as AdaGrad (Duchi et al., 201 1), 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 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 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 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 al., 2012; Ioffe and Szegedy, 2015; Salimans and Kingma, 2016) is often critical for the success of such optimization methods.
aligned with the coordinates of the vector parameters.
The natural gradient method (Amari, 1998) takes a related but different perspective. Rather than re- parametrizing the model, the natural gradient method tries to make the optimizer itself invariant to re- parameterizations by directly operating on the manifold probability distributions. This requires endowing the parameter space with suitable notion of proximity formalized by metric. An important metric in this context is the Fisher information metric (Fisher and Russell, 1922; Rao, 1992), which induces the Fisher-Ree Rao natural gradient (Amari, 1985). Another important metric in probability space is the Wasserstein metric (Vil- lani, 2009; Otto, 2001), which induces the Wasserstein natural gradient (Li and Montufar, 201 8a;;; 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 challenging in practice. Indeed, each parameter update requires inverting the metric tensor. This becomes infcasible for current deep learning models, which typically have millions of parameters. This has motivated research into finding efficient algorithms estimate the natural gradient (Martens and Grosse, 2015; Grosse and Martens, 2016; George al., 2018;

---

Published as a conference paper at ICLR 2020
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 ct al. (2019) proposed to estimate the metric based on 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 be solved 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 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 has the advantage of being well defined even when the model doesn't 't admit 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 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, illustrate the benefits of our proposed estimator for solving classification problems when the model has an ill-conditioned parametrization.
The paper is organized follows. In Section 2. after 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 present experiments 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
is often possible to formulate learning problems as the minimization of some cost functional F(p) over probability distributions p from parametric model Pe. The set Pe contains probability distributions defined on an open sample space Ω C IR and parametrizee by some vector 0 E θ where is an open subset of R7, The learning problem can thus be formalized as finding optimal value O that locally minimizes a loss function L(0): = =(po)= defined over the parameter space θ One convenient way to solve this problem approximately is by gradient descent, which uses the Faclidean gradient of L w.r.t. the parameter vector 0 to produce a sequence of updates Og according to the following rule:
0++1=0-- 0 (..))
Here the step-size Yt is a positive real number. The Euclidean gradient can be viewed the direction in parameter space that leads to the highest decrease of some linear model M, 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:
1 glul.
C((,)) = argmin (1)
uER
The linear model M, is an approximation of the cost function L in the neighborhood of 9, and is simply obtained by a first order expansion: M*(u)= = C().) + V (O.)) Tu. u. The quadratic term |//l| penalizes
the change in the parameter and ensures that the solution remains in the neighhorhood where the linear model is still 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-propagationn. The performance of this algorithm is highly dependent on the parametrization of the model Pe. however (Martens, 2010; Sutskever et al., 2013). To obtain algorithm that is robust to parametrization, one can take advantage of the structure of the cost function L(0) which is obtained as the composition of the functional F and the model 0~+ po and define generalized natural gradient (Amari and Cichocki, 2010). We first provide
2

---

Published as a conference paper at ICLR 2020
conceptual description of the general approach to obtain such gradient. The starting point is to choose a divergence D between probability distributions and use as a new penalization term:
-argmin MM(())+ D ((po..Po.ta). weRv
(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 quadratic term u Go where GD( contains second order information about the model as measured by D. This leads to the following expression for the generalized natural gradient V ((e( where the dependence in D made explicit:
VD c(0,) :- argmin M,(u)+ + GD((.)u.. (un(()))
(3)
NER
From (3), it possible to express the generalized natural gradient by means of the Euclidean gradient: The parameter updates are then obtained by the new update rule:
(4)
Equation (4) leads a descent algorithm which invariant to parametrization in the continuous-tim limit:
Proposition 1. Les VI be invertible and smoothly differentiable re- re-parametrization +b vr(() and denote by C(t) : C(V) (16)). Consider the continuous-time natural gradient flows:
0, = - L(O.),
V/ -==-V -V==VC
4/n=(((0)
Then vh and 0, are related by the equation bb (0.)) at all times $20.
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: ((²) = Vn (()) which holds in general for any covariant gradient. We provide a proof Proposition in Appendix C.I 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 1) is either the Kullback-Leible divergence F, or the squared Wasserstein-2 distance W.
2.2 INFORMATION MATRIX VIA DIFFERENTIAL GEOMETRY
When D is the Kullback-Leinn divergence or relative entropy F, then (3) defines the Fisher-Rao natural gradient V (()) (Amari, 1985) and Gr (6) is called the Fisher information matrix. ( is well defined when the probability distributions in Pe all have positive densities, and when some additional differentiabil- ity and integrability assumptions pe are satisfied. In fact, it has 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-Raa metric (see Definition 4 in Appendix B.1; also Holbrook et al. 2017):
Definition 1 (Fisher information matrix). Assume 0 + px(x) is differentiable for all x on S2 and that Polooo) dx < 20. Then the Fisher information matrix is defined as the pull-back of the Fisher-Rao pola) metric gf:
where the functions fr on Q are given by: fo
0.pe
Pe
Definition 1 directly introduces GF 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 Pe. In particular, when D is the W the Wasserstein information matrix is obtained directly by means of the Wasserstein-2 metric tensor g (Otto and Villani, 2000; Lafferty and Wasserman, 2008) as proposed in Li and Montufar (2018a); Chen and Li (2018):
3

---

Published as a conference paper at ICLR 2020
Definition 2 (Wasserstein information matrix), The Wasserstein information matrix (WIM) is defined as the pull-back of the Wasserstein 2 metric g
Gw (0), (0:p0.000)=0.0)====010)))))))))) (o(0)) (,(2)d)((z))))
where &i are vector valued functions Q8 that are solutions to the partial differential equations with Neumann boundary condition:
0,ps= dio--pi)(
visisq.
Moreover, d; are required to in the closure of the set of gradients of smooth and compactly supported functions in L2(pa)". In particular, when pa has a density, o, b-=Vo v some real valued function on St
The partial derivatives 0,ps should be understood in distribution sense, as discussed in more detail in Section 2.3. This allows define the Wasserstein natural gradient even when the model pe 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 Appendix B.3. From now on, we will focus on the above two cases the natural gradient V ((0) namely V [()) and VW [(0). When the dimension of the parameter space is high, directly using equation (4) becomes impractical as it requires storing and inverting the matrix G(O). In Section 2.3 will see how equation (3) can be exploited along with Legendre duality get expression for the natural gradient that can be efficiently approximated using kemel methods.
2.3 LEGENDRE DUALITY FOR METRICS
In this section we provide an expression for the natural gradient defined (3) as the solution saddle-point optimization problem. It exploits Legendre duality for metrics to express the quadratic term u G(f)u as a solution to a functional optimization problem over C (S2), the set of smooth and compactly supported functions on Ω. The starting point is to extend the notion of gradient Upp which appears in Definitions and to the distributional sense of Definition 3 below.
Definition 3. Given a parametric family Pe of probability distributions, we say that pa admits a distributional gradient at point if there exists linear continuous map V CC2 (S2) +R9 such that:
J [(c)ddp+oa (x)) (aaaiaaaaed--odddldddeeeeee-ee-- [(a)doo(t)==VVaddl uu ee
VSEC? VUER"
where f(c.f.u) depends on f and and converges to 0 as approaches O. Vp is called the distributional gradient of pa at point
When the distributions in Pe have a density, written x Po (I)) by abuse of notation, that differentiable w.r.t. and with jointly continuous gradient in @ and then ppp(f is simply given by f fz (epo (2) as shown in Proposition 12 of Appendix C.I. In this case, the Fisher-Rao natural gradient admits a formu- lation as saddle point solution involving Vp and provided in Proposition 2 with proof in Appendix C.I.
Proposition 2. Under the same assumptions as in Definition 1, the Fisher information matrix admits the dual formulation:
1
Vpo(f) u / / f(=))
u (G
sup SEC? (0) S ((())do(a)==
(5)
2
Moreover defining Us
the Fisher-Rao natural gradient satisfies:
sup uo(s) a- 1 / [xx) SEC, (0)) [ ((x)ddo())=
VFC(0)=
argmin wERv
Another important case is when Pe is defined as an implicit model. In this case, any sample x from distribution Po in Pe is obtained ho(z), where is sample from fixed latent distribution
4

---

Published as a conference paper at ICLR 2020
defined over latent space 2 and (0,2) + ho(2) is deterministic function with values in Q. This can be written in a more compact way as the
of v by the function hy:
Pe {po == (he #2 [E.)
(6)
different expression for Up is obtained in the case of implicit models when θ ha(2)))) is differentiable for y-almost all z and Wh is square integrable under
p / Vho(=)"Vaaa V
Equation (7) is also known the re-parametrization trick (Kingma al., 2015) and allows to derive a dual formulation the Wasserstein natural gradient the case of implicit models. Proposition 3 below provides such formulation under mild assumptions stated in Appendix A.2 along with proof in Appendix C.I.
Proposition 3. Assume Pe is defined by (6) such that V is given by (7). Under Assumptions (B) and (C), the Wasserstein information matrix satisfies:
2" Gw sup Ppolf) T 1 (ECx(())
u- 2/ [Vaf(())
(8)
and the Wasserstein natural gradient satisfies:
sup uo(t) / (V NE (Q)
w
(9)
The similarity between the variational formulations provided in Propositions and 3 is worth noting, A first difference however, that Proposition 3 doesn't 't require the test functioons f have mean under pe- This is due to the form of the objective in (8) which only depends on the gradient of f. More importantly (5) can be infinite when Vp is given by (7). Indeed, if the Po doesn't It admit density, it is always possible to find an admissible function f / (Q) with bounded second moment under Po but for which Vps(f) is arbitrarily large. This is avoided in (8) since the quadratic tern directly penalizes the gradient of functions instead. For similar reasons, the dual formulation of the Sobolev distance considered Mrouch et al. (2019) can also be infinite in the case implicit models discussed in Appendix B.3 although formally similar to (8). Nevertheless, similar estimator as in Mroueh et al. (2019) can be considered using kernel methods which the object of Section 3.
while (8) well defined, the expression
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 V C(0), as it is well defined for implicit models, but a similar approach can be used for the Fisher-Ran 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 Reproducing Kernel Hilbert Space (RKHS) H. which is Hilbert space endowed with an inner product (..)) along with its norm 11-1 H has the additional property that there exists a symmetric positive semi-definite kernel & Ω x Ω + E such that k(x,,) E H for all E l and satisfying the Reproducing property for all functions f in H:
(10) The above property is central all kernel methods as 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 new saddle problem obtained by restricting (9) to functions the RKHS H and adding some regularization terms:
f() = (f.((f.f))
1 To [(() = min sup UT( u
2. 1111///))))))))p((())))))))))))) 1 D(0)=~>1//13 (11)
wER* fEM
5

---

Published as a conference paper at ICLR 2020
The kernelized Wasserstein natural gradient is obtained by solving (11) and is denoted by ((O). Here, é is a positive real numbers, À is non- non-negative while D is a diagonal matrix in R2 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 X= the problem is still concave in f- This allows use 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 - A new expression for the kernelized natural gradient is therefore obtained:
Proposition 4. Assume that >> 0 and X O, then the kemelized natural gradient is given by:
A (( = - --00=----)~=0
(12)
e
where f is the unique solution to the quadratic optimization problem:
inf /V=))
(13)
FER FER
When = I might not be well defined, still, we have: A - -D0) for any limiting sequence (13).
Proposition 4 allows to compute the kernelized natural gradient directly, provided that the functional optimization (13) can be solved. This cireumvents the direct computation and inversion of the metric as suggested by (1I). 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 (Zn)iensN from the latent distribution E which are used to produce N samples (Xn)igngN from pa using the map hgg i... Xx, = ho (Z,)). We
also assume we have access to an estimate of the Euclidean gradient V VZ((0) which is denoted by V.(()). This allows to compute an empirical version of the cost function in (13),
N
J(): = Σ V (xx,)
+
- D(0) (00))))) (f)+1)
(14)
(Z,) Vxx((x). (14) has 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 al., 2001) that the optimal function minimizing (14) is a linear combination of functions of the form x++ (,,k(X, with Sn S N and 1 < S d and (,k(y, denotes the partial derivative of i w.r.t y/- This requires solve 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 finite dimensional subspace Has of H called the Nyström subspace. In the context of score estimation, Sutherland et al. (2017) proposed to use subspace formed by linear combinations of the basis functions x x 0.I.(mm,.))
where (f) is given by: [[())
(15) where (Ynn. ( are basis points drawn uniformly from (X,)ssneN with M N. This further reduces the computational cost when M & N but still has 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 (a,k(Ym,.)) (YYn))))ssss sisd for each basis point Yn. Hence, we consider M indices (im Sm uniformly drawn form {1 (I....d)) d) and define the Nyström subspace H.s to be:
span(r+00kk
ISmSM; 1sisd),
:= span( ( k(Ym 2 SmSM).
An estimator for the kernelized Wasserstein natural gradient (KWNG) is then given by:
W - (77)
=== argminn fEHu
(16)

---

Published as a conference paper at ICLR 2020
By definition of the Nyström subspace Hy, the optimal solution f is necessarily of the form: Em=1 Om k(Y,,,z), where the coefficients (0m))smSSM obtained by solving finite dimensional quadratic optimization problem. This allows to provide a closed form expression for (17) in Proposition 5.
-M
Proposition 5. The estimator in (16) is given by:
N V [(0) ( (D()) D(0)) TT (DD(O) + e CCTT) (TD(0)-))Z(()) (17)
where C and K are matrices in RxMxN and TR. MRM given by Cm.(n.)) = =ô... dctak(Yymmt Xn),
Kmm = -.., 0im, 4ake (Y:Ymmmm))
(18)
while T is matrix in TR gxxx xa obtained as the Jacobian of 0 (()) ERM Le., : :VV)()), with
N
(T(0)) N kYm
Σ 1=
In (18), we used the notation (..2)( for the partial derivative of ls w.r.t Ty. A proof of Proposition 5 is provided in Appendix C.2 and relies on the reproducing property (10) and its generalization for partial deriva- tives 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 poste- rior (Mishkin al., 2018), (17) can be applied as a plug-in estimator for any family Pe obtained as an im- plicit 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 X= the estimator in (17) has an additional structure which can be exploited to get more accurate solutions. By the chain rule, the matrix T addmis a second expression of the form T'= CB where B is the Jacobian matrix of (hæ(Z,))) (1,)),,))) 1 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 CCT =USUT and pre- seellllpppyyig T by STUT. The resulting expression is given in Proposition and falls into the category of Ridgless estimators (Liang and Rakhlin (2019)).
Proposition 6. Consider an SVD decomposition of CCT the form CCT =USUT, then (17) is equal to:
V - - (TD(O)- + x )
(19)
where StS and T:=SUFT.
Choice of damping term. So far, we only required D(0) to be diagonal matrix with positive coefficients. While a natural choice would be the identity matrix, this doesn't '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(0); = -|.T..ill where 7 is defined in Proposition 6. When the sample-size is limited, as it is often the case when N is the size of mini-batch, larger values for might be required. That is prevent the KWNG from ever-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-Mann Marquardt heuristic as in Martens and Sutskever (2012) which seems to perform well in practice; see Section
Computational cost. The number of basis points M controls the computational cost of both (17) and (19) which dominated by the cost of computing 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 M2N In practice, MM can be chosen to be small (M < 20) while N corresponds to the number of samples in mini-batch. Hence, in typical deep learning model, most of the computational cost is due computing as the typical number of parameters is of the order of millions. In fact, T can be computed using automatic differentiation and would require performing M backward passes the model to compute the gradient for each component of r. Overall, the proposed estimator can be efficiently implemented and used for typical deep learning problems as shown in Section 4.
7

---

Published as a conference paper at ICLR 2020
Choice of the kernel. We found that using either gaussian kernel or rational quadratic kemell to work well in practice. We also propose a simple heuristic to adapt the bandwidth of those kemels to the data by setting it to == N,M. where N,N, is equal to the average square distance between samples (Xn) and the basis points (Ym Igmgg and do is fixed a priori. Another choice 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 À >> O; we leave the case when À for future work. We work under Assumptions (A) to (G) in Appendix A.2 which state that 52 is a non-empty subset, k is continuously twice differentiable with bounded second derivatives, (z) (z) has most a linear growth in 2 and v/ satisfies some standard
moments conditions. Finally, we assume that the estimator of the euclidean gradient VL(()) satisfies Chebychev's concentration inequality which is often the case in Machine learning problem as discussed in Remark 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 (di) iiيc, involved in Definition 2 are assumed to be gradients functions in H. and their smoothness is controlled by some parameter a 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, will only focus the miss-specified which relies weaker assumption:
A ssumption 1. There exists two constants C'> O and e² O such that for all x>> O and all 1 ISiSg there a function f satisfying:
S CK.
(20)
The left inequality in (20) represents the accuracy of the approximation of de by gradients of functions in H while the right inequality represents the complexity of such approximation. Thus, the parameter e characterizes the difficulty of the problem: a higher value e means that a more accurate approximation of o; comes higher cost in terms of its complexity. Theorem 7 provides convergences rates for the estimator Proposition 5 under Assumption I:
Theorem 7. Let be such that 05651 and b:= 1. Under Assumption I and Assumptions (A) to (G) 2+c
listed in Appendix A.2, for N large enough, M ~ (dN log())), A~ NaTx and e e zin, it holds with probubility at least 1 8 that:
LI0) [(0))2=0(~~2)
A proof of Theorem 7 is provided in Appendix C.3. In the best case where c= 0, we recover a convergence rate of V/ as in the well specified case for the worst smoothness parameter value == Hence, Theorem a consistent extension of the well-specifted case. For harder problems where e> 0 more basis points are needed, with M required be of order d. (N) in the limit when oo 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: //gittub...////// Somm/McchaelAAnn 1/KNNG.
4.1 CONVERGENCE ON SYNTHETIC MODELS
To empirically assess the accuracy of KWNG, we consider three choices for the parametric model Pe: 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 '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 Fuclidean Gradient (EG) which is used to compute both of WNG and KWNG.
8

---

Published as a conference paper at ICLR 2020
4.0%
4.0%
(6)
4.0%
10
3.5%
3.5%
3.5%
3,0%
2.5%
Reloive
2.0%
2.0%
2.0%
1.5%
155
1.5%
1.0%
1.0%
1.0%
0.0%
10
0.0%
0
1000
2000 3000
4000
5000
0.0%
10
10a
10
N
Figure 1: Relative error of KWNG averaged over 100 runs for varying dimension form d = (yellow)
d = 10 (dark red) for the hyper-sphere model. (a): box-plot of the relative error d increases while
N = 5000 and M = dvN (b) Relative error as the sample size N increases and M = dv/ (c):
Relative error as M increases and N = 5000. A gaussian kemeel is used with fixed bandwidth o =
per iteration
(b) Projection of trajectories
10
200
Cost function
10-4
10"31
100
10-22
200
10-32
KIWNG
-300
GD
Inital value
10
WNG
4000
Optimal value
101
102
10a
10
10a
0
200
400
600
800
1000
Iteradions
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 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: x= 0 and = 10-10, An optimal step-size "/ is used: y = 0.1 for both KWNG and WNG while 7/2 = 0.0001 for EG.
Figure 2:
Figure 1 shows the evolution of the the relative enror w.r.t the sample-size N, the number of basis points M and the dimension d in the case the hyper-sphere model. As expected from the consistency results provided in Section 3.4, the relative error decreases 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 as discussed in Section 3.4. In practice, using the rule = dvNN seems to be good (a). All these observations persist in the case of the normal and log- lor-normal model as shown in Figure 4 of Appendix D.1. In addition we report in Figure the sensitivity the choice of the bandwidth a which shows a robustness of the estimator to a wide choice of a.
heuristic as shown in Figure
We also compare the optimization trajectory obtained using KWNG with the trajectories of both the exact WNG and EG in simple setting: Pe is the multivariate normal family and the loss function C()) is the squared Wasserstein distance between pe and fixed target distribution po-. 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 00 with Residual Network He et al. (2015). To use the KWNG

---

Published as a conference paper at ICLR 2020
accuracy: (IC)
[al Training acturacy. (IC)
acturscy. (WC)
[cl Training accuracy: IWC)
100
100
1 accurscy
accunnc
Reccurccy
Adam
60
60
90
20
500
SPAC
80
200
Epochs
Epochs
Epochs
Epochs
[b) Test accuracyy (IC)
(a) Training accuracy, (CC)
[d] Test couracy: (WC)
[c] Training accuracy: (WC)
100
55 3S.0cuccyy
50
accurrrycy
%: accurrcyy
60
70
Stacuncy
20
KWNCCMti S1+M
KWNGGMU 51
40
SCD
40
30
70
5
RPAC
#FFAC
KVNGGN-518
30
100
200 Epocts
200
100
200
300
100
200
300
100
200
200
Epochs
Epachs
Epochs
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 Cifarlo Results are averaged over 5 independent runs except for KFAC and eKFAC.
estimator, we view the input RGB image as a latent variable with probability distribution v and the output logits of the network (2) as a sample from the model distribution Oa E Pe where denotes the weights of the network. The loss function L is given by:
u(>) log (o0:=JuttTTogSSMMWWeee (og(ToosUUUUU(()))dde
where SM is the Softmax function, y(2) denotes the one-hot vector representing the class of the image 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 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 each epoch on Ci farlll in both (WC) and (IC) cases. While all methods achieve similar test accuracy in the (WC) case on both datasets, methods based on the Euclidean gradient seem to suffer drastic drop in performance in the (IC) case. This doesn't 't happen for KWNG (red line) which achieves a similar test accuracy as in (WC) case. Moreover, a speed-up 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 Ci Fa 00, KWNG is also less affected by the ill ill-conditioning. albeit a lower extent. Indeed, the larger number of classes in Cifar100 makes the estimation of KWNG harder 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 also suffer drop in performance in the (IC) case. This might result from the use of an isotropic damping term D(0) = I which would be harmful in this case. We also observe a drop in performance when different choice of damping is used for KWNG. More importantly, using only diagonal pre-conditioning of the gradient doesn't 't match the performance of KWNG (Figure 7).
LCCNOWLEDGEMENT
GM has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement nº 757983).
10

---

Published as a conference paper at ICLR 2020
REFERENCES
Amari, S. and Cichocki, A. (2010). Information geometry of divergence functions. Bulletin of the Polish Academy of Sciences: Technical Sciences, 58(No 11:183-195.
Amari, S.-i. (1985), Differentil-GGoomtriiee Methods in Statistics. Lecture Notes Statistics. Springer- Verlag, New York.
Amari, S.-i. (1998). Natural Gradient Works Efficiently in Leamnin. Neural Computation, 10(2):551-276
Ambrosio, L.,, Gigli, N., and Savaré, G. (2004). Gradient flows with metric and differentiable structures, and applications to the Wasserstein space. Atti della Accademia Nazionale dei Lincei. Classe di Scienze Fisiche, Matematiche e Naturali. Rendiconti Lincei. Matematica e Applicazioni, 15(3-4)33273333
Arbel, M. and Gretton, A. (2017). Kernel Conditional Exponential Family. arXiv: [stat]. arXiv: 1711.05363.
Arbel, M., Korba, A., Salim, A., and Gretton, A. (2019). Maximum Mean Discrepancy Gradient Flow. 19X:::1906004370 [cs, stat]. arXiv: 1906.04370.
Benamou, J.-D. and Brenier, Y. (2000). A computational fluid mechanics solution to the Monge- Kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393.
Bernacchia, A., Lengyel, M., and Hennequin, G. (2018). Exact natural gradient in deep linear networks and its application to the nonlinear case. In Bengio. S., Wallach, H., Larochelle, H., Grauman, K.. Cesa-Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems 31, pages 5941-5950. Curran Associates, Inc.
Chen, Y. and Li, W. (2018). Natural gradient in Wasserstein statistical manifold. arXiv:1885.083300 [cs, mathj. arXiv: 1805.08380.
Duchi, J. Hazan, E., and Singer, Y. (2011). Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. Journal of Machine Learning Research, 12(ul))2121-2155)
Ekeland, L. and Témam, R. (1999). Convex Analysis and Variational Problems. Classics in Applied Mathematics. Society for Industrial and Applied Mathematics.
Fisher, R. A. and Russell, E. J. (1922). On the mathematical foundations of theoretical statistics. Philosophical Transactions of the Royal Society of London. Series A, Containing Papers of Mathematical or Physical Character, 604)534-0
Garreau, D., Jitkrittum, W., and Kanagawa, M. (2018). Large sample analysis of the median heuristic. 17X7::1707.07666 [math, stat]. arXiv: 772.99269.
George, T. Laurent, C., Bouthillier, X., Ballas, N., and Vincent, P. (2018). Fast Approximate Natural facttoeer-factored Eigenbasis. a7Xiv1186.03888 [cs, stat]. arXiv: 1806.03884.
Gradient Descent
Grosse, R. and Martens, (2016). A Cronecker-factorre Approximate Fisher Matrix for Convolution Layers. In Proceedings of the 33rd International Conference on International Conference Machine Learning Volume 48, ICML' 16, pages 573-582, JMLR.org. event-place: New York, NY, USA.
He, K., Zhang, X., Ren, S., and Sun, J. (2015). Deep Residual Learning for Image Recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770-778.
Heskes, T. (2000), On "Natural" Learning and Pruning Multilayered Perceptrons. Neural Computation, 12(4):811991.
Hinton, G., Srivastava, N., and Swersky, K. (2012). Lecture 6a overview of mini-batch gradient descent.
Holbrook, A., Lan, S., Streets, J., and Shahbaba, B. (2017). The nonparametric Fisher geometry and the chi-square process density prior. arXiv: 170.003177 [stat]. arXiv:
Loffe, S. and Szegedy, C. (2015). Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In Proceedings of the 32Nd International Conference on International Confer- Machine Learning- Volume 37, ICML 15, pages 456. JMLR.org. event-place: Lille, France.

---

Published as a conference paper at ICLR 2020
Kingma, D. P. and Ba, J. (2014). Adam: A Method for Stochastic Optimization. arXhv:1412.6980 [cs]. arXiv: 1412.6980.
Kingma, D. P., Salimans, T., and Welling, M. (2015). Variational Dropout and the Local Reparamettrrization Trick. ArXiv, abs/1506.02557,, arXiv: 1506.02557.
Klenke, A. (2008). Probability Theory: Comprehensive Course. World Publishing Corporation.
Lafferty, J. and Wasserman, L. (2008). Rodeo: sparse, greedy nonparametrie regression. The Annals of Statistics, (6(1):28-4)) 63.
LeCun, Y. A., Bottou, L.., Orr, G. B., and Müller, K.-R. (2012). Efficient BackProp. In Montavon, G., Orr, G. B., and Müller, K.-R., editors, Neural Networks: Tricks of the Trade: Second Edition, Lecture Notes in Computer Science, pages 9-48. Springer Berlin Heidelberg, Berlin, Heidelberg.
Li, W. (2018). Geometry of probability simplex via optimal transport. arXiv:1803.06360 [math]. arXiv: 1803.06360.
Li, W., Lin, A. T., and Montufar, G. (2019). Affine natural proximal learning.
Li, W. and Montufar, G. (2018a). Natural gradient via optimal transport. arXiv: [cs, math]. arXiv: 1803.07033.
Li, W. and Montufar, G. (2018b). Ricci curvature for parametric statistics via optimal transport. arXiv [cs, math, stat]. arXiv: 1807.07095.
Liang, T. and Rakhlin, (2019). Just Interpolate: Kernel "Ridgeless" Regression Can Generalize. 1 [cs, math, stat). arXiv: 1808.00387.
Malagò, L., Montrucchio, L., and Pistone, G. (2018). Wasserstein Riemannian Geometry of Positive Definite Matrices. arXiv:1801.0920 09960999990 [math, stat]. arXiv: 1801.09269.
Martens, J. (2010). Deep Learning via Hessian-free Optimization. Proceedings of the 27th International Conference International Conference on Machine Learning, ICML pages 735-742, USA. Omnipress. event-place: Haifa, Israel.
Martens, J. and Grosse, (2015). Optimizing Neural Networks with Kronecker-factored Approximate [cs, stat]. arXiv: 1503.05671.
Curvature.
Martens, J. and Sutskever, (2012). Training Deep and Recurrent Networks with Hessian-Free Optimization. In Montavon, G., Orr, G. B., and Müller, K.-R., editors, Neural Networks: Tricks of the Trude: Second Edition, Lecture Notes Computer Science, pages 479-535. Springer Berlin Heidelberg. Berlin, Heidelberg.
Mishkin, A., Kunstner, F., Nielsen, D., Schmidt, M., and Khan, M. E. (2018). SLANG: Fast Structured Covariance Approximations for Bayesian Deep Leariing with Natural Gradient. arXiv:1811.04504 [cs, stat]. arXiv: 1811 1814004
Modin, K. (2017). Geometry of Matrix Decompositions Seen Through Optimal Transport and Information Geometry. Journal of Geometric Mechanics, 93):3335-390 390. arXiv: 1601.01875.
Mroueh, Y., Sercu, T., and Raj, (2019). Sobolev Descent. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2976-201 2985.
Ollivier, Y., Amold, L., Auger, A., Hansen, (2011). Infommation-eommetiic optimization algorithms: A unifying picture via invariance principles. J. Mach. Learn. Res,, 18:18:1-18:65.
Otto, F. (2001). The Geometry of Dissipative Evolution Equations: The Porous Medium Equation. Communications in Partial Differential Equations, 26(1-2):111-174
Otto, F. and Villani, C. (2000). Generalization of an Inequality by Talagrand and Links with the Logarithmic Soboley Inequality. Journal of Functional Analysis, 173(2):361-00)
12

---

Published as a conference paper at ICLR 2020
Rao, C. R. (1992). Information and the Accuracy Attainable in the Estimation of Statistical Parameters. In Kotz, S. and Johnson, N. L., editors, Breakthroughs in Statistics: Foundations and Basic Theory, Springer Series in Statistics, pages 235-244-- 245. Springer New York, New York, NY.
Rudi, A., Camoriano, R., and Rosasco, L. (2015). Less more: Nyström computational regularization.
Salimans, T. and Kingma, (2016). Weight Normalization; Simple Reparameterization to Accelerate Training of Deep Neural Networks. arXiv: [cs]. arXiv: 1602.07868.
Schölkopf, B., Herbrich, R., and Smola, A. (2001). A generalized representer theorem. In Helmbold, D. and Williamson, B., editors, Computational Learning Theory, pages 416-426, Berlin, Heidelberg. Springer Berlin Heidelberg.
Simonyan, K. and Zisserman, A. (2014). Very Deep Convolutional Networks for Large-Scale Image Recognition. arXiv:140911556 [cs]. arXiv: 1409.1556.
Sriperumbudur, B., Fukumizu, K., Kumar, R., Gretton, A., and Hyvärinen, A. (2017). Density estimation in infinite dimensional exponential families. Journal of Muchine Learning Research.
Steinwart, I. and Christmann, A. (2008). Support Vector Machines. Springer Publishing Company, Incorporated, 1st edition.
Sutherland, D. J., Strathmann, H., Arbel, M., and Gretton, A. (2017). Efficient and principled score estimation.
Sutskever, I., Martens, J., Dahl, G., and Hinton, G. (2013). On the importance of initialization and momentum deep learning. In International Conference on Machine Learning, pages 1139-1147.
Villani, C. (2003). Topics Optimal Transportation. American Mathematical Soc. Google-Books-ID. R nWgiq890EC.
Villani, C. (2009). Optimal transport: Old and new. Technical report.
13

---

Published as a conference paper at ICLR 2020
A PRELIMINARIES
A.1 NOTATION
We recall that Ω is open subset of Rd while θ is open subset of parameters in R7 Let 2 R be a latent space endowed with a probability distribution // over Z. Additionally, (0,z) (2( E S2 is function defined over e 2. We consider parametric set of probability distributions Pe over 52 defined as the implicit model:
Pe = {po :- (
E
where by definition, Pe = (ha)# means that any sample x from pe can be written - ha(=) where 2 is sample from v. We will write B to denote the jacobian of ha w.r.t A viewed as a linear map from I2 to [2( without explicit reference 0:
VuER".
As in the main text, L: R is a loss functions which is assumed to be of the form L= with F being real valued functional over the set of probability distributions. VL(() denotes the euclidean gradient w.r.& 9 while (()) an estimator of V (() using N samples from pe.
We also consider a Reproducing Kernel Hilbert Space H of functions defined over 52 with inner product (..)) and norm I1-1|z and with kemel k Ω Q2 + The reproducing property for the derivatives (Steinwart and Christmann, 2008, Lemma 4.34) will be important: fz = ( (f.O:k(z..)) for all x E holds as long as l differentiable.
(92) (I2) denotes the space of smooth bounded real valued functions on Ω and C (92) C Cg denotes the subset of compactly supported functions. For any measured space 2 with probability distribution v, we denote by (22(v) the space of real valued and square integrable functions under v and by [2))) the space of square integrable vector valued functions under v and with values in Rd
A.2 ASSUMPTIONS
We make the following set of assumptions:
(A) Ω is non-empty open subset of IRd.
(B) There exists positive constants and a such that f|l=///d><=) 2plce- (plcc for any p22. 2.
(C) For all A there exists C(0) such that [oohe())ll for all 2 Z.
(D) k is twice continuously differentiable on 2
( For all @ it holds that
< for all 1sisd.
The following quantity is finite: x² =sup wen 0,0++gh(z..0 15154
(G) For all 058 it holds with probability at least 1 à that V
Remark 1. Assumption (G) holds if for instance VL(0) can be written as an empirical mean of i..d.
terms with finite variance:
1 N Σ (ol(h(Z)))
CC(I)
where Z; are i..dd samples from the latent distribution v where (z)() This is often the case the problems considered in machine- leariig In, this case, the sum of variances of the vector
V ((0) along its coordinates satisfies:
f [IV /
One can then conclude using Cauchy-Schware inequality followed by Chebychev's inequality that with probubilit 1-ô:
VC(O))
VEN
Moreover, Assumption (C) is often satisfied when the implicit model is chosen to be a deep networks with R non-linearity
14

---

Published as a conference paper at ICLR 2020
A.3
OPERATORS DEFINITION
Differential operators. We introduce the linear L operator and its adjoint L
L:H La(v
L :L2(u) H
wー> J Eak(ho(s)))uut)tee
(=(00foha)isssA
This allows to obtain the linear operator A defined in Assumption 2 in the main text by composition A: L. We recall here another expression for A in terms of outer product @ and its regularized version for given A>
A 1
A+==
is easy to see that A is symmetric positive operator. Moreover, was established Sriperumbudur et al. (2017) that A is also compact operator under Assumption (E).
Assume now we have access to samples (Zn) as in the main text. We define the following objects:
N
À
A+==
1==
Furthermore, is continuous function in [2(1)d, then we can also consider an empirical estimator for L L.v v:
N d ack(ho (z))o(Z)).)))((z)))
LT
Subsampling operators.
We consider the operator Qu defined from H to R&M by:
Qm: VM
V
Σ cm B. m.....)
(21)
where ( ar is an orthonormal basis of RM Qa admits singular value decomposition of the form QM=UEVT, with VVT := Par being the orthogonal projection operator on the Nyström subspace HM- Similarly Rudi et al. (2015); Sutherland al. (2017), we define the projected inverse function GM (C) as:
((CC)=((VTVV)-VV
We recall here some properties of GM from (Sutherland al., 2017, Lemma 1): Lemma 8. Let A:Hl be positive operator, and define Ax = A+ for any X> 0. The following holds:
1. (A) -MPP)
2. (M)
3. (Ar(A))A = PM
4. APM
5. (A)GM(ANAA2I
Estimators of the Wasserstein information matrix. Here we would like to express the estimator in Proposition 5 in terms the operators introduced previously. We have the following proposition:
Proposition 9. The estimator defined in Proposition 5 admits the following representation:
W ((0) = (e))
where GM,N is given by:
(T.T::((TT)) Gnr GT (nx B.
15

---

Published as a conference paper at ICLR 2020
Proof. This 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 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 GM,N defined as:
( (M)
(22)
B BACKGROUND IN INFORMATION GEOMETRY
B. FISHER-RAA STATISTICAL MANIFOLD
In this section we briefly introduce the non con-parametric Fisher-Rao metric defined over the set P of probability distributions with positive density, More details can 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 Tp, the set of real valued functions f defined over and satisfying
f f(x)² p(x)
dx <o; f f())(())d==
We have the following definition for the Fisher-Rao metric:
Definition 4 (Fisher-Rao metric). The Fisher-Ran metric of is defined for all PEP as an inner product over To of the form:
1 0h (1=)): plz) - VS.ET,
Note that the choice of the set Tp is different from the one considered in Holbrook et al. (2017) which replaces the integrabilityy condition by smoothness one. In fact, it can be shown that these choices result in the same metric by density argument.
B.2 WASSERSTEIN STATISTICAL MANIFOLD
In this section we review the theory of Wasserstein statistical manifold introduced Li and Montufar (2018a); Chen and Li (2018). By analogy to the Fisher-Raa Rao metric which allows to endow the parametric model Pe with the structure of a Riemannian manifold, it is also possible to use a different metric that derived from the Wasserstein 2 distance. We first start by briefly introducing the Wasserstein 2 distance. Given two probability distributions and p', we consider the set all joint probability distributions [I(p,p)) between p and o' usually called the set of couplings between p and p'. Any coupling defines way of trans- porting mass from to p'. The cost of such transport can be measured as the expected distance between an el- ement of mass of p location that is mapped element of mass of pl at location y using the coupling
f || w/ (xr
The squared Wasserstein 2 distance between p and p' is defined as the smallest transport cost over all possible couplings:
W3(a))) inf / ||=- dd TETT(
dynamical formulation of W2 was provided by the celebrated Benamou-Brenier formula in Benamou and Brenier (2000):
W3 W f &x
where the infimum taken over the set of vector fields Φ: [0011] x - Rd. Each vector field of the potential determines corresponding probability distribution Pl as the solution of the continuity equation:
Oipt diocroi( =0
Po P
(23)
16

---

Published as a conference paper at ICLR 2020
When is compact set, Neumann condition is added the boundary of S2 ensure that the total mass conserved. Such formulation suggests that W2(p,') corresponds fact to the shortest path from to p'. Indeed, given a path Pt from p to p'. the infinitesimal displacement direction is given by the distribution Otpr. The length [ótpr] of this direction measured by: [ōrpi == (d) Hence, W2(p,t)) can be written as:
W inf
In fact, 0it can be seen as an element in the tangent space p, P2 to P2 at point Pt. To ensure that (23) is well defined, TpPr can be defined as the set of distributions o satisfying d(1)=0.
VSECt
(24)
for some positive constant C. Indeed, the condition in (24) guarantees the existence of vector field d, that is a solution to the PDE: a= div(p(a).
Moreover, | l can be as inner product of aip: with itself in Tpp P2- This inner product defines in turn metric tensor g on P2 called the Wasserstein metric tensor (sce Otto and Villani (2000); Ambrosio ct al. (2004)):
Definition 5. The Wasserstein metric g defined for all pE P2 the inner product over P2 of the form:
30 (a,,o):= (x) (1))p(((), Va,o TTpp
where do and oe are solutions to the partial differential equations:
a= a=-div(pbo); o' = -iiv((p
Moreover, do and the are required to be in the closure of gradient of smooth and compactly supported functions w.rl L2(p)
Definition 5 allows endow P2 with formal Riemannian structure with W2 being its geodesic distance:
Pi Wi(ppp')=int inf ( (o. (Orp)
B.3 NEGATIVE SOBOLEV DISTANCE AND LINEARIZATION OF THE WASSERSTEIN DISTANCE
To device the Wasserstein natural gradient, one can exploit Taylor expansion of W which is given in terms of the Negutive Sobolev distance [|po+.: - Po (po) as done in Mroueh et al. (2019);
W (po.po+w)= 11p00= - +o(lll)))))(( )
Further performing a Taylor expansion of poolu-Pooll---- where we call (0,) 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-(0,) using kernel methods as proposed in Mroueh et al. (2019) for learning non-parametric models. However, with such approach, l|po+u PollH-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 Pa-. and pa do not share the same support as we discuss now.
po ( in leads a quadratic term u Gw (0.
As shown (Villani, 2003, Theorem 7.26) and discussed (Arbel et al., 2019, Proposition 17 and 18), the Wasserstein distance between two probability distributions p and p' admits first order expansion in terms of the Negative Sobolev Distance:
lim Walp-ptdd-pp)- - -100-10 c00
when pl admits bounded density w.r.t p. 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 dg located at value 6, the Wasserstein distance admits closed form expression the form:
W2000,00 c6lor - = 0'1l
17

---

Published as a conference paper at ICLR 2020
Hence, Wog + c(óg. - -66) diverges to infinity. One consider different perturbation of the model deteu for some vector u which the one we are interested in here. In this case, the Wasserstein distance admits well-defined asymptotic behavior.
lim (o ||u]].
On the other hand the Negative Sobolev Distance is infinite for any value of c. To see this, we consider its dual formulation as Mroueh et (2019):
f [(o))p())) ((o)))))))-)))0)))00-))))))))))) (//////())))()))))
sup sect (2)
J [())dp())=
Evaluating this expression for dg and doten for any value e>> 0 and for any u that is non zero, on has:
1 (f(0) - f(0 f(++())) - [
2e 1600
00tealll--(66) = sup
(25)
sect (12)
=(0)=0
One always find function f such that (0()==0, ((0) 0 and (00+ can be arbitrarily large, thus the Negative Sobolev distance is infinite. This is not the case the metric u Go(00u which can be computed in closed form:
1 " Gw(f)u= sup Vf(0) T u 1101"1=00000
|V
(26)
2
FEC
f(9)=0
In this case, choosing - and VS(0)= =u achieves the supremum which is simply given by 2|lw|1 Equation (26) can seen as a limit case of (25) when c++0:
1 T (w(:)u:=
1 (0+t)) |V==
2 u
sup lim (n) /(0)=0
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.I PRELIMINARY RESULTS
Here provide 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 e-parametrization V, satisfying 1o- (t().. Denote by Pv: (v-)(w) the re-parametrized model and Gw(c) and (+()) their corresponding Wasserstein and Fisher information matrices whenever they are well defined.
Proof of Proposition 1. Here we only consider the case when VDC(0) either given by the Fisher natural gradient V [(0) or the Wasserstein Natural gradient V C(0). We will first define e.. == (0,) and show that in fact vh- u/, at all times > O. First, let's differentiate 3/ in time:
v V.V. G.
By the chain rule, we have that
hence:
v. = -3=
It is easy to see that V. one can conclude that:
- (1.) (V, v - by definition of v and C. hence by Lemma 10
v. - GD D(ba).
Hence, N. satisfies the same differential equation as th.- Now keeping in mind that v/o w (B(o).. it follows that eho = =/ (0()) by uniqueness of differential equations.
18

---

Published as a conference paper at ICLR 2020
Lemma 10. Under conditions of Propositions 2 and 3, the informations matrices Gw(e) and Gr(b) are related to Gw (0) and (0) by the relation:
Gw(b) V V (16) Gw (V) Vr( V (6) Gr
Proof. Let E R and write VoN Vo then by the dual formulations of Gw(O) and Gr(0) in Proposition 2 we have that:
Vool //2 (c))dpo(z))d
GOV
sup fEcr (0) ((
Now recalling that VyPp Vapo (+) by Lemma 11, it follows that:
VOPa(() ( _ (())) ((=))
1 T w. (+) (w)c=
sup JECT (9) (()))po())==
Using again Proposition for the reparametriieed model Pv. we directly have that:
1 u G=(D)u=
Vepo(f) リー / f(z) dox)
2
sup fECt (9) [ (=))
The result follows by equating both expression. The same procedure can be applied for the case the Wasserstein information matrix using Proposition
Lemma 11. The distributional gradients V uPb and Vapa are related by the expression:
Pp V sp8
v (xb)
Proof. The proof follows by considering fixed direction E and test function f (S)) and the definition of distributional gradient in Definition 3:
(P u= 1iimimmm f(=)dpo+=c((t) (000--(())ppppp f(x)dpu(() <-
1
=lim ((x))po- -00
(())dpe-
Now by differentiability of V - we have the following first order expansion:
T -ー ( = (VV) ( (w cu(e)
ut
where u(e) converges to when e = Now using again the definition Definition 3 for Pae one has:
(f (())dp)-(0))))(0)> f())dpu--()()))) = [f
The last two terms converge to as e e, hence leading to the desired expression.
Proposition 12. When admits density that is continuously differentiable w.r. and such that x+ Vp)(z) continuous, then the distributional gradient is of the form:
ppo(f)= F(x)Vool
Vfec
where (x)()) denotes the gradient of the density of po(x)
19

---

Published as a conference paper at ICLR 2020
Proof. Let <>> and EER, define the function v(c,u.f) follows:
( (100aa)) Up )
we just need to show that f(E,u,,) -0 ~0 This follows form the differentiation lemma (Klenke, 2008, Theorem 6.28) applied to the function (0,z) f(z) Indeed, this function is integrable in for any f' in neighborhood U of & that is small enough, it is also differentiable that neighborhood U and satisfies the domination inequality:
po((r) ul]
sup
ofe SESup((f))
The inequality follows from continuity of ( and recalling that f is compactly supported. This concludes the proof.
We fist provide proof of the dual formulation for the Fisher information matrix.
Proof of Proposition 2. Consider the optimization problem:
((Lo)Vo(o)))) u J (x))
"n-1 u-2/ 2 f(x)) (ox)))(()d)
sup fECE (0)
(27)
Recalling that the of smooth and compactly supported functions C2 (co) is dense in L2(ps) and that the objective function (27) is continuous and coercive in f. follows that (27) admits a unique solution f" in Le(po) which satisfies the optimality condition:
J [(=)(Voo(())) udr udr= 1 (4)
Hence, is easy see that (Vpp) T u/ p9 and that the optimal value of (27) is given by:
((Vpo(x))Tu)) dr.
2 po(x)
This is equal to u Gr()) 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>)igusg be an orthonormal basis of TR. Under Assumptions (B) and (C), there exists an optimal solution = (ds) with o, in L2(po) satisfying the PDE:
0.pe = -ioo=
The elliptic equaations also imply that Wh = L T (4ohh)). Moreover, the Wasserstein information matrix Gw (0) on Pe at point 0 can be written as Gw = 1 Th where the inner-product is in L2(pg) and satisfies:
1
vpolS) u u- 11-1/11100000 -//Iv//(h((+))Padee))
2
Gw(f)u=
sup SEC (0) f ((=))dp(0))=
for all u uER?.
Proof. Let (ex) be an orthonormal basis of 12 For all 1 15 1 g. we will establish the existence of optimal solution d; in (2( satisfying the PDE:
O.po dio(po
(28)
20

---

Published as a conference paper at ICLR 2020
Consider the variational problem:
sup dES. / o(ha(z)) (o,ho(()) 2
( (00
(29)
where S is Hilbert space obtained as the closure in L2(e)) of functions of the form Φ= = Vn with CC))
S==(Vaf (1))
We have by Assumption (C) that:
f (((hole)) To, (0t 1 (1+1))))dd(()) |lolllatoo
Moreover, by Assumption (B), know that [[L++l22/)dddl co. This implies that the objective in (29) is continuous in ø while also being convex.s It follows that (29) admits unique solution o ES which satisfies for all Φ
f ((hh(z))) of (ho()))))) ((h(()))
( Top) 0n, (o.ho()))dl2)
In particular, for any f C (9), holds that:
f V f(h(())))) T of (ho(2))dr(zz=
(0..h(z))ddv2)
which is equivalent to (28) and implies directly that Vh = L he where == (6)) The variational expression for u Gwu follows by noting that (29) admits the same optimal value as
[pols) 2 [Vaa(ho(z))]Paate)
sup
fEC2 (O) J
(=))
That is because S is by definition the closure in L2(p))² of the set of gradients of smooth and compactly supported functions on
Proof of Proposition 3. This is consequence of Proposition 13.
C.2 EXPRESSION OF THE ESTIMATOR
We provide here a proof of Proposition
Proof of Proposition 5. Here, to simplify notations, we simply write instead of D(0). First consider the following optimization problem:
Σ ((X.)) "R() 2 (R() D DL'() SEH
inf
with R() given by R()- I
B(Zn). Now, recalling that any S can be written as f f= Mf am k(Ym..), and using the reproducing property (f (Steinwart and Christmann, 2008, Lemma 4.34), it is easy to see that:
N
1
1
Σ
am O ((tck((n,
N
ISnEN IsiSd
m=1
l
Σ
Om P.m Yom
15m.m SM
1
R(f)= N
Σ
ISASN Isidd
15mCM
21

---

Published as a conference paper at ICLR 2020
The above can be expressed in matrix form using the matrices defined in Proposition 5:
N [= CCT oç
1 Ko;
R())=aCC
Hence the optimal solution 7 of the form f Σ Em=1a in Drm k(mm..), with a" obtained as a solution to the finite dimensional problem in R:
min a T (eCCT +CK Ja (CBD---
MERM
It is easy to see that a" are given by:
--==(CC" ("
Now recall that the estimator in Proposition 5 given by: (W)(()) = uo("). Hence,
(VZ(0) a²) 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 (0:)ı involved in Definition 2 can be expressed as gradients of functions H. More precisely:
Assumption 2. For all 1 S sisq q. there exits functions fi E H such that d; Vv Additionally, fs are of the form f. = e, for some fixed a 20, with u, EH and A being the differential covariance operator defined on H by A:
(h)((d))))
The parameter a characterizes the smoothness fi and therefore controls the statistical complexity of the estimation problem. Using a similar analysis Sutherland et al. (2017) we obtain convergence rate for the estimator Proposition 5
the following convergence rates for the estimator in Proposition 5:
Theorem 14. Let be such that S 8 S and b = min(1, + 2). Under Assumption and Assumptions (A) (G) listed in Appendix A.2. for N large enough, M (dN Nit log(N)), \~N xh and e पी,, it holds with probability at least 1-6 5 that:
[W)) -000000 [[00) -0
24 25-1
In the worst case where = 0, the proposed estimator needs at most M ~ (dv (dVNNlog(N) achieve convergence rate N- The smoothest case requires only . ~ (dN log(NN) achieve rate N- 3. Thus, the proposed estimator enjoys the same statistical properties as the ones proposed by Sriperumbudur al. (2017); Sutherland et al. (2017) while maintaining computational advantage! tNow we provide 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 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 < 8 and N large enough, holds with probability at least 1 8-8:
FEC-v"cl=01 JPFT
provided that M dN DTTT log N and S N wi where b: min(1, a + 4) when
Assumption 2 holds and b= 1 when Assumption holds instead.
24c
The estimator proposed by Sutherland et al. (2017) also requires M to grow linearly with the dimension d although such dependence doesn't appear explicilly in the statement of Sutherland et al. 2017, Theorem
22

---

Published as a conference paper at ICLR 2020
Proof. Here for simplicity assume that D without loss of generality and omit the dependence in and write V C and VC instead of V [(0) and (2( and V [(0). We also define
G, =e GMrNN and G. = + +ww By Proposition 9, we know that J G, VV... use following decomposition:
W I| V2I Gw GWANN-GGWWWWWLD VCl GWW GW/
To control the norm of G-1 we write G-² Go + I) where H is given by H == GT (GM. Gw) Hence, provided that μ == Amax (H), the highest eigenvalue of H, is smaller than it holds that:
Moreover, since Gw is positive definite, its smallest eigenvalue is strictly positive. Hence, (()+t) Therefore, we have (>++) (
which implies:
VC V +1 V II 1-p
(1)+e)
Let 0 We have by Assumption (G) that IIC - VLI O(N-b) with probability at least 1-6. Similarly, by Proposition 16 and for large enough, we have with probability least 1-8 that |[G.,N-Gw -w - (() where 6 is defined in Proposition 16. Moreover, for N large enough, one can ensure that م so that the following bound holds with probability at least 1 ó-6:
(1
2 +e)).
Thus by setting e N I get the desired convergence rate.
Proposition 16. For any 05851, we have with probability as least 1-8 and for large enough that:
IGM.N-GGW|I=0 2-)).
provided that M ~ dN 2 log.N where b:=min(1a++ when Assumption 2 holds and b= 24c when Assumption 1 holds instead.
Proof. To control the error |GM.N Gw we decompose into an estimation error [GM,, Gas and FCM-Gwy Gw
approximation error ||GM
IGGM,NGGwll
Gw
were GM is defined in (22) and is obtained taking the number of samples N infinity while keeping the number of basis points M fixed.
The estimation error ||GM GM,M is controlled using Proposition 17 where, for any 058-1, we have with probability at least 5 and as long as N2 M(1,),6):
1 [B] V
X+
(av.st 2
NX oi
In the limit where N + 0o and X + 0, only the dominant terms in the above equation remain which leads to an error ||GM,N Gm ((NA)-((N)))))) =). Moreover, the condition on N can be expressed as X -og
To control the emror approximation error || Gw we consider two cases: the well-specified case and the miss- miss-specified case.
Well-specified case. Here we work under Assumption 2 which allows to use Proposition 19, Hence, for any 0<85 1 and M > M (d,2,6), holds with probability least 1-8:
||Gas - Gw mmin(1,a+))
23

---

Published as a conference paper at ICLR 2020
Miss siss-speeified case. Here work under Assumption 1 which allows to use Proposition 18. Hence, for any 05ō5 1 and M (d,,66)) holds with probability at least 1-6:
Let's set b:= min(1, + ()) for the well-specified case and = 24c for the miss-specified case. In the limit where M - do and À - 0 the condition on M becomes: M ~ dA "looA- Hence, when M dd "logA 11 and A 'logÀ S N holds with probability as least 1-6 that
GG Gw [00]==0[>>=(AN))
One further choose À of the form X N This implies condition on M of the form d log( loV))o((N) S M and log S N. After optimizing over to get the tightest bound, the optimal value is obtained when 0 (2b (2b+ and the requirement N is always satisfied once N is large enough. Moreover, one can choose .M iddN=n log so
requirement on M is satisfied for N large enough. In this case get the following convergence rate:
BCM.N-GWl1-2
2²1).
Proposition 17. For any 0.05 1, provided that 22 M(1,, have with probability as least 1 5:
|B||
d+ )+
1
(2a / 2 2n1
2x
NX
VNX
NA
with:
2 2alog
ON
2001og
VN
Proof. For simplicity, we define E -LTB-LLB. By definition of GMMNN and Gm we have the following decomposition:
GM,N MAW-GM-TT GM= E GM(AWE+E (N EA B+B BIB LOM
Eo
E
- -B (AAAAAAAAAAAAMMMAAAA
The first three terms can be upper-bounded in the following way:
| A (
< <| A (AN)
21/A
21
e || || AT A (M I Bl|
| ||I Gas (Ax) I1| ||| 1
<|B
<1/VA
For the last term C3 we first recall that by definition of (AA) we have:
Therefore, one can write:
| - LA
A
< <|IBIi
|
(
1
21
< <113111442 1 + (A-1)A
24

---

Published as a conference paper at ICLR 2020
We recall now (Rudi ct al., 2015, Proposition 7.) which allows to upper-bound || by (1-n)*
where 7 Amax (A ( 4)) provided that v v/< 1. Moreover, (Rudi et al., 2015, Proposition 8.) allows to control both 17 and ||A (A- A under Assumption (F). Indeed, for any 0 0 and provided that 0< A |||||||]]]]]] holds with probability 1 that:
1-x/1
27k NA
212 17.5
272
7 2
+
3N
3N
NA
where 71 and 722 are given by:
87r(AA)
(A)(A)
71 log(
);
72
log
Ad
Hence, for N2 22M((1,66) we have that (1-1)) 2 and one can therefore write:
1+x/>
2
||E/|| S <4)
3N
NX
| B||
|
V\
The error | is controlled by Proposition 22 where holds with probability greater or equal to 5 that:
1
log 2
2alog 1
2
|
UN.5
VN
VN
Finally, we have shown that provided that N M (1.X.8) then with probability greater than 1 one has:
B1 (2aN + 2
A+
|
NX N
VNX
Proposition 18. Leto [ and define (ddA,,)) == 128 log 26
4Tv(A) (dRA-1+1 Under Assumption
and Assumption (F), any 8 620 such that M 2 PM (d.2,5) the following holds with probability 6-ó:
|GGM-GWIISAF
Proof. We consider the error 1|G--Gwll Recall that Gw is given by Gw = ① 4 with & defined in Proposition 13. Let be positive real number, know by Assumption 1 that there exists F = hsf with f E H such that 1|6 |1 F= S Cs and || S Cks for some fixed positive constant C. Therefore, we use F* to control the error || - Gw ||. Let's call E E= LF We consider the following decomposition:
G Gw =LL (oho) Gar (Ax)L poho T
=E (AN) E-E
E FT (L LUM 1)L woho+ E L(9M(AA) L L-1)
C,
E
First we consider the term é, one simply has:
||,, x ILLA M IIAA I| & L ITIS Sk2
The second term also satisfies ||E_|| S x2 by definition of Fg. For the last two terms E and E, we use Lemma 20 which allows to control the operator norm of L(g L T L-1). Hence, for any 2 and M such that M > (d.X.8) and for < it holds with probability 1-6 that:
|| I
|
We have shown so far that Gww [0.1] to get a tighter bound. The optimal value in this case is x² = =min(1(cAA) * By considering À A 0 such that (cx²) xेz 11 it follows directly that I| Gw II उमं which shows the desired result. n
Gw 1-
+2 VX). One can further optimize over x on the interval
25

---

Published as a conference paper at ICLR 2020
Proposition 19. Ler0 - ||||| and define (d.A,,6) log 4Tv( and Assumption (F), for any 6 220 such that M M2M Md,,,, the following holds with probability 1 -::
41), Under Assumption
|G - (GM-Gwlle X
Proof. Recall that Gw is given by Gw + d with & defined in Proposition 13. By Assumption 2, we have that d V) with V (vs) H Hence, one can write
Gm - Gw = (I. boho) T (Ax)L toho TT =V (AGM A)AV
can therefore directly apply Lemma 20 and get |G Gwll & (min(1,a+t with probability 1- for any 820 such that M 2 (d,A,6).
Lemma 20. Let 0 < A ||| a 2 0 and define M(d,X,6) == 128 log 4Tr(A) (dmA-1+1) +). Under Assumption (F), for any d 2 such that M 2 M (d.2,8) the following holds with probability ó-
A6
|L (Ax) L- A) mm
Proof.
We have the following identities:
=LA + AR - - ALA ARGM(AN)
LA 1 (A -
For the first E, we use (Sutherland et al., 2017, Lemma (ii)) which implies that PM
Thus e, = LA A (PM - T).a. Moreover, by Lemma 21 we have that ||A (T PM)) - 2 with probability 1 1 8 for M > 6(d, Therefore, recalling that (I Pas)2 Par since Par is projection, further write:
||e, || S ILA
(PM- A-
1
||e < |LAx ||IA
9M (AN)
(PM- A
24A
}
9N A 1 A
| S S
51
21
||A remains to note that which allows conclude.
1 A < - when 0 05 a $ and that A -114|=1 for a
C.4 AUXILIARY RESULTS
Lemma 21. Let 0 < X < ||A||
Under Assumption (F), for any 8 2 0 such that +1) the following holds with probability 6:
128 log Ad M
(Tr(A) (mA-1+
Proof. The proof is an adaptation of the results in Rudi et al. (2015); Sutherland al. (2017). Here we recall Qu defined in (21). Its transpose Q sends vectors in R.M elements in the span of the Nyström basis
26

---

Published as a conference paper at ICLR 2020
points, hence Pas and Qh have the same range, i.e.: range(Pus) Frangee We are in position to ap- ply (Rudi 2015, Proposition 3.) which allows to find upper-bound on A (Pas - in terms of Q.:
For simplicity we write Aas = Qi QM and F2:= A A We also denote by 3 Xmaz (F2) the highest eigenvalue of Eg. We therefore control || +1) in terms of 3 using (Rudi ct al., 2015, Proposition 7) provided that <<1:
(PM
1 V1-3
Now need make sure that < 1 for M large enough. To this end. we will apply (Rudi al., 2015, Proposition 8.) to Am- Denote by Vmn vdo... k(Ym:.). Hence, by definition of Aas it follows that Am m Um Um Moreover, (Pm)iSmSM are independent and identically distributed and satisfy:
M
4
Elum 3mml = JaM..)aaAA...)d =A
We also have Assumption (F) that (Um--Ay ' pm) S de almost surely and for all X> O. We can therefore apply (Rudi et al., 2015, Proposition 8.) which implies that for any 12 5 2 0 and with probability 1 8 it holds that:
27
2gydic
3M
MA
with y= log (Tr(A)
Ad provided that > ||. Thus by choosing M2 1287
(dxA (dxA-1+1) we have that BS 9
with probability 8 which allows to conclude.
Proposition 22. There exist a> 0 and σ > 0 such that for any 05Ō<1, it holds with probability greater of equal than 1 that:
2alog N
2oflog 2
| B- Bl|
Proof. denote by Un - E (, we have that E[u,]] = L LTB. B. We will apply Bernstcin' inequality for of random vectors. For this first need to find > and a >> such that B[[ BlP S Eppofap-2 To simplify notations, we write x. and x' instead of ho(z) and ho(²"). We have that:
Bil - L BBBR / Σ Bc(())))((=)) 1 (.M
du(z)
52--1
(AM(),,)-O((aaa.. -Odd) du(z)
+2
(BB))))
du(z)
We used the convexity of the norm and the triangular inequality to get the last line. We introduce the notation 3(a):==0kkc.)-1004(( (dx()) and by I'(x) we denote the matrix whose components are given by (())))== The first term E, can be upper-bounded as follows:
e = / (()
27

---

Published as a conference paper at ICLR 2020
Moreover, have that Tr([(x))) Eis I Ei= l/((1)11 We further
have that (x)[l = _ (h that (x:)] < 2/K. Hence, can directly write that: e₁ (2Vkd)) [[]B(=)) Assumptions (B) and (C) we get:
and by Assumption (F) follows
Recalling
(2 2l
Similarly, we will find an upper-bound on Ez. To this end, we introduce the matrix Q(2', whose components are given by
x²"",, One, therefore has:
E - ~ (=(B6))-- T((2)))))(())))) Q(1' x )dul " dv(z)
de(z)
Once again, we have that Tr(Q(e', <>ا 0,0 00 (x"",, sdn dx thanks to Assumption (F). Therefore, follows that:
6 (VVdx)) B(
( [I=](db(())
<3 <3 (Van) [=]]dd(()++
<3 (V) - p 202 ( [ellddu()) )
The second line a consequence Assumption (C) while the last line due to Assumption (B). These calculations, show that it is possible to find constants a and 0 such that E|]|=, - L B s f Hence one concludes using Bernstein's 's inequality for a sum of random vectors (see for instance Rudi 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 μa E IRd and covariance matrix Ee Rd x Ed parametrized using its lower triangular components s= s= T()). We denote by Σ -
the inverse operation that maps any vector E R 2 to its corresponding symmetric matrix in I2 x Rd The concatenation the mean j and s will denoted as 0 0 (ft,8). Given two parameter vectors (m, T(S)) and = (m',T(S") where m and m' are vectors in Re and S and S' are symmetric matrices in Fd x xRd the metric evaluated at u and is given by:
TT G(0)=== T +Tr(AAA))
where A and are symmetric matrices that are solutions to the Lyapunov equation:
S= -'AA'+EAA'
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 VgL(() = ((,) V.C(O), the corresponding natural gradient V/ C()) can also be computed in closed form:
W 8
(A- ())
where (V, To use the estimator proposed in Proposition 5 take advantage of the parametrization the Gaussian distribution as push-forward a standard normal vector:
28

---

Published as a conference paper at ICLR 2020
874
0%
074
076
4%
034
094
0%
1000
20000
3000
4000
5000
199
M
10%
(a)
100%
(2)
10%
10
8%
894
8%
856
2000 4200
1.0
Figure 4: Evolution of the relative error of KWNG averaged over 100 runs for varying dimension form d=1 1 (yellow) d= 10 (dark red). For each run, a random value for the parameter and for the Euclidean gradient VL( is sampled from centered Gaussian with variance 0.1. In all cases, A= 0 and 10-5. Top row: multivariate normal model, bottom row: multivariate tog-normal. Left (a): box-plot of the relative
error as d increases with N = 5000 and the number of basis points is set to M = N (b) Relative
error the sample size N increases and the number of basis points is set to M dv N Right (c):
Relative error as M increases and N fixed to 5000.
10%
sphere
gaussian
10%
log, nog_normal
Reiatve error
Relative error
Relative error
8%
6%
6%
4%
4%
495
2%
2
-2
-2
2
Figure 5: Relative eтrоr of the KWNG for varying bandwidth of the kemel. Results are averaged over 100 runs for varying dimension form d=1 (yellow) to d= 10 (dark red). For each run, random value for the parameter and for the Euclidean gradient V (()) is sampled from a centered Gaussian with variance 0.1. In all cases, X >= 10-10 The sample size is fixed to N = 5000 and the number of basis points set to M- dvN Left: uniform distributions on a hyper-sphere, middle: multivariate normal, and right: multivariate log-normal.
D.2 CLASSIFICATION ON CIFARI0 AND CIFAR10C
Architecture. We use residual network with one convolutional layer followed by 8 residual blocks and 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 Table
Hyper-parameters, For all methods, we used a batch-size of 128. The optimal step-size y was selected in 110,1,10-1,10 1,022 1,033, 10-4) 4] for each method. In the case of SGD with momentum, we used momentum parameter of 0.9 and a weight decay of either 0 or 5 x 10 For KFAC and EKFAC, we used 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 e is set to c= 10 and is adjusted using an adaptive scheme based the Levenherg-Marquaat dynamics in (Martens and Grosse, 2015, Section 6.5). More
29

---

Published as a conference paper at ICLR 2020
Kerel size
Output shape
Z
x2x32x3
Conv
3x
Residual block
3 x2
64
Residual block
3x3 x2
128
Residual block
3x3 x2
256
Residual block
3x33 x2
512
Linear
Number of classes
Table Network architecture.
precisely, we use the following update equation for e after every iterations the optimizer:
6 e we.
if >>>
e 6t w
if r <
Here, r is the reduction ratio:
2 [(0.)) -(0)))-C(0++) ((0++1) Gw L(0) VL(0)
max S1Sh
where (t+) are the times when the updates occur, and the decay constant chosen =0055.
D.3
ADDITIONAL EXPERIMENTS
Figure 6 reports the time cost in seconds for each methods, while Figure 7 compares KWNG to diagonal conditioning.
30

---

Published as a conference paper at ICLR 2020
(a) Training accuracy: (IC)
(b) Test accuracy: (IC)
100
55
90
90
KWNG(M=20)
SGD
80
accuracy
SGD+Momenturi
% accuracy
SGD+Momentum+WD
70
70
KFAC
60
60
eKFAC
50
40
40
30
30
10a
10a
10a
10
10a
10a
104
10a
Time (s)
Time (s)
(c) Training accuracy: (WC)
(d) Test accuracy: (WC)
100
90
90
80
accuracy
80
accuracy
70
KWNG(M=5)
KWNG(M=20)
60
60
SGD
50
50
SGD+Moone Momentum+WD
KFAC
40
eKFAC
30
30
102
10a
104
10a
102
10a
10a
10a
Time (s)
Time (s)
Figure 6: Training accuracy (left) and test accuracy (right) as a function of time for classification on Cifar10 in both the ill-conditioned case (top) and well-conditioned case (bottom) for different optimization methods.
(a) Training accuracy: (IC)
(b) Test accuracy: (IC)
100
Diagonal conditioning:
80
Diogonal conditioning: D [T]
KWNG
% aacuracy
KWNG: D !71
60
60
KWNG: D 1T1
KWNG: D |7| (rq-kemee)
%
40
20
20
100
200
300
100
200
300
Epochs
Epochs
Figure 7: KWNG vs Diagonal conditioning in the ill-conditioned case Cifarl0. In red and blue, the euclidean gradient is preconditioned using diagonal matrix D either given by D or D==1[T.<II where T and T are defined in Propositions and The rest of the traces are obtained using the stable version of KWNG in Proposition with different choices for the damping term D D= >= and [[T., .l. All use gaussian kemee except the yellow traces which uses rational quadratic kernel.
31


Reviews:
[{"review_id": "r1xIbmXaKr", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n \n###\nSummary of the paper: \n\nThe 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. \n\nAuthors 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 $u$, the variational form is obtained via finding an objective S,  $\\sup_{f\\in C^c_{\\infty}} S(f, u ) = u^{\\top}G_{W}u$, where $G_{W}$ is a form of  \"Wasserstein information matrix\".\n\nAuthors then propose to learn the function f in an RKHS and propose to find the descent direction by solving\n $\\min_{u} <u, \\nabla_{\\theta} Loss (p_{\\theta})>  + \\sup_{f\\in RKHS}  S(f, u ) + r(u)- \\lambda ||f||^2_{rkhs}$\nwhere r(u) is a quadratic regularizer on u. \n\nThe 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. \n\nAuthors under some assumption  show that the  estimated natural W gradient in RKHS  is concentrated around the true one. \n\nExperiments 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.\n \nReview : \n\nThe paper is not easy to follow and the high level intuition how the method works is not well explained. \n\nIt 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 $KL(p_{\\theta},p_{\\theta+ \\epsilon u})$ as we perturb in the parameter space and this well known that this epsilon $u^{\\top}F u$. \nHence natural Gradient descent is :\n$\\min _{u} <u, \\nabla_{\\theta}Loss(p_{\\theta})> + KL(p_{\\theta},p_{\\theta+  u}) \\approx min _{u} <u, \\nabla_{\\theta}Loss(p_{\\theta})> + u^{\\top}F u$  \n\nNow for the wasserstein distance one has also similarly:\n$\\min _{u} <u, \\nabla_{\\theta}Loss(p_{\\theta})> + W^2_2(p_{\\theta},p_{\\theta+  u})$\n\nand it is known that as epsilon goes to zero we have: \n\n $W^2_2(p_{\\theta},p_{\\theta+ \\epsilon u})/\\epsilon = ||p_{\\theta} - p_{\\theta+ \\epsilon u}||^2_{H^{-1}(p_{\\theta})}+o(\\epsilon) = \\sup_{f} \\int f (p_{\\theta} - p_{\\theta+\\epsilon u}) - \\frac{1}{2} \\mathbb{E}_{p_{\\theta}}||\\nabla_x f(x)||^2+o(\\epsilon) $\n\nNow replacing with the implicit model as epsilon goes to zero we get the expression given in the paper using a simple taylor expansion:\n$=   \\sup_{f} \\int <\\nabla_{\\theta} h_{\\theta}^{\\top}\\nabla_x f(h_{\\theta}),u > d\\nu - \\frac{1}{2}\\mathbb{E} _{p_{\\theta}}||\\nabla_x f(x)||^2$\n\nin a sense the paper is proposing to linearize $W^2_2$ around the perturbation in the parameter space of the implicit model and this can be done using  $|| .||_{H^{-1}(q)}$ , as pointed and used in many recent works.  then the paper proposes to approximate $|| .||_{H^{-1}(q)}$ in RKHS which was already proposed in Mroueh et al  in Sobolev Descent.  AISTATS 2019. \n\nWe 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 $W_2$ with $||.||_{H^{-1}(q)}$  , and how to approximate it in RKHS as it was already proposed in the literature in Mroueh et al Sobolev Descent. \n\n\nI 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. \n\nQuestions:\n\n- 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 ? \n-the constraint $\\int f(x)p_{\\theta}(x)=0$ is not imposed in the kernelized version?\n- the method comes disappointing since it seems that the preconditioning that the Wasserstein gradient gives is not enough and $r(u)=u^{\\top}D u$ is need where D is diagonal depends on T. Have you tried with $D=Identity$? it might be that the scaling of the gradients is coming only from that $D^{-1}$?\n\n- Can you give timings for computing each gradient update and how it compares to regular SGD or diagonal approximation of Fisher natural gradient?\n\n- Does one need preconditioned gradient if the network was self normalized (like batch norm or spectral norm etc)?   \n\n\nOverall assessment: \n\nThat 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. ", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "S1xohUaAFS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nI find this work interesting with some important merit, as it tackles an important problem in statistical learning. My main concern, however, is the problems related to RKHS from a practical point of view. For example, solving optimization problem (11) is difficult and the paper makes a range of further approximations to be able to arrive at an approximate solution. Also, selecting the kernel and its bandwidth is crucial in practice. From a practical point of view, I suspect that more evidence is needed to justify if the proposed method can really offer a method of choice.  \n\nHaving said that, I believe this paper provides an important first (and alternative) step towards an important problem. The paper is also well written and well structured. I have a few further comments below\n1) In the abstract and introduction, the invariant property of natural gradient is mentioned several times without a detailed explanation why/what it is. Adding a brief explanation of this property is appreciated.\n2) The sentence on line 8 in Introduction reads \".. It can be not alleviated by using adaptive step size...\". This is when the authors are talking about the adaptive learning methods. Is this a too strong comment about the adaptive learning methods? Can the authors know for sure that these methods cannot be used here?\n3) Equations (1) and (2): Are they correct? should the minus sign just be in front of the first term \\mathcal{M}_t(u) only?\n4) Page 4, first line after Def 3: \"one covers the usual gradient \\nabla\\rho_\\theta(x)\". It is not very clear (to me) how to get this. Can the authors please elaborate more on this?\n\n \n     \n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "S1gMw5s2YH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe motivation of the natural gradient is well-motivated. Although the choice of Wasserstein metric is sound, especially for models that do not admit a density, it seems that there are no supporting experiments for this choice over Fisher Information metric. In general, the writing is fine. The flow idea is clear. However, the content is quite dense. Some assumptions just pump out without careful judgment. The idea to restrict into RKHS and use low-rank approach is interesting to approximate for the natural gradient under Wasserstein metric. Overall, I lean to the acceptance side.\n\nBelow are some of my concerns:\n\n1) It seems that the natural gradient under Wasserstein metric is well-motivated for models which do not admit a density (to compare with the natural gradient under Fisher information metric). However, it seems that there is no supporting experiments about it yet. For models in the experiments, it is better to show a comparison between natural gradient under Wasserstein metric and Fisher information metric w.r.t. time consumption and accuracy.\n\n2) In proposition 3 and theorem 5, they require some assumptions. The authors should place those assumptions into the main text instead of only putting it in the appendix, and should give more discussions about those assumptions. Especially, for assumption (D), why one can have this assumption? It seems that this assumption (D) has a strong influence to the complexity in Theorem 5? More detail discussion is required.\n\n3) For the relaxation in Equation (9), it seems that the authors do not simply add some regularization terms. How does it relate to the original Equation (7)? What is the meaning of the 3rd term in Equation (9)? and how's about the 2nd term?\n\n4) For the experiments, the authors evaluate the multivariate normal model and the multivariate log-normal model which are very special cases under Wasseserstein information matrix where one can compute in closed-form. The authors should show some general models, especially models which do not admit a density. For the experiments in Section 4.2, the authors should add the natural baseline: natural gradient under Fisher information metric. It is unclear to me why one needs natural gradient under Wasserstein metric over Fisher information metric for this setup? What is the benefit to use natural gradient under Wasserstein metric?\n\n ", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This is a very interesting paper which extends natural gradient to output space metrics other than the Fisher-Rao metric (which is motivated by approximating KL divergence). It includes substantial mathematical and algorithmic insight. The method is shown to outperform various other optimizers on a neural net optimization problem that's artificially made ill-conditioned; while it's not clear how practically meaningful this setting is, it seems like a good way to study optimization. I think this paper will be of interest to a lot of researchers and could open up new research directions, so I recommend acceptance as an Oral.

Decision: Accept (Spotlight)