Generalized Tangent Kernel: A Unified Geometric Foundation for Natural Gradient and Standard Gradient
Authors: Qinxun Bai
Steven Rosenberg
Wei Xu
Venue: Accepted by TMLR ()

Abstract: Natural gradients have been widely studied from both theoretical and empirical perspectives, and it is commonly believed that natural gradients have advantages over standard (Euclidean) gradients in capturing the intrinsic geometric structure of the underlying function space and being invariant under reparameterization. However, for function optimization, a fundamental theoretical issue regarding the existence of natural gradients on the function space remains underexplored. We address this issue by providing a geometric perspective and mathematical framework for studying both natural gradient and standard gradient that is more complete than existing studies. The key tool that unifies natural gradient and standard gradient is a generalized form of the Neural Tangent Kernel (NTK), which we name the Generalized Tangent Kernel (GTK). Using a novel orthonormality property of GTK, we show that for a fixed parameterization, GTK determines a Riemannian metric on the entire function space which makes the standard gradient as “natural" as the natural gradient in capturing the intrinsic structure of the parameterized function space. Many aspects of this approach relate to RKHS theory. For the practical side of this theory paper, we showcase that our framework motivates new solutions to the non-immersion/degenerate case of natural gradient and leads to new families of natural/standard gradient descent methods.

Published in Transactions on Machine Learning Research (02/2025)
Generalized Tangent Kernel: A Unified Geometric Foundation for Natural Gradient and Standard Gradient
Qinxun Bai*
Snnun.bai@hooizon..com
Horizon Robotics, Cupertino, CA, USA
Steven Rosenberg"
sr@maath. bu edu
Department of Mathematics and Statistics, Boston University, Boston, MA, USA
Wei Xu
wei.zu@hoorizon.ott
Horizon Robotics, Cupertino, CA, USA
Reviewed on OpenReview: https: openrevew...........ee. net/ formm id= HOmL5h
Abstract
Natural gradients have been widely studied from both theoretical and empirical perspectives, and it is commonly believed that natural gradients have advantages over standard (Euclidean) gradients in capturing the intrinsic geometric structure of the undorlying function space and being invariant under reparameterization. However, for function optimization, a fundamental theoretical issue regarding the existence of natural gradients on the function space remains underexplored. We address this issue by providing geometric perspective and mathematical framework for studying both natural gradient and standard gradient that is more complete than existing studies. The key tool that unifies natural gradient and standard gradient is gencralized form of the Neural Tangent Kernel (NTK), which we name the Generalized Tangent Kernel (GTK). Using novel orthonormality property of GTK, we show that for a fixed parameterization, GTK determines a Riemannian metric on the entire function space which makes the standard gradient as "natural" as the natural gradient intrinsic structure of the parameterized function space. Many aspects of this approach relate to RKHS theory. For the practical side of this theory paper, we showcase that our framework motivates new solutions to the non-immersion/ tee-immmriioo/deeeeerede /degeneaaee case of natural gradient and leads to new families of natural/standan gradient descent methods.
capturing the
1 Introduction
Since their introduction in Amari 1998
natural gradients have been topic of interest among both thcoretical
researchers (Ollivier 2015 Martens 2014 Li & Montúlar 2018) and practitioners (Kakade 2001 Pascanul
& Bengio, 2013). In recent years, research has examined several important theoretical aspects of natural
gradient for function optimization on neural networks, such as replacing the KL-divergence by the L metric
on the function space (Benjamin et al. 2019). Riemannian metrics for neural networks (Ollivier 2015 and
convergence properties under overparameterizatiin (Zhang et al.] 2019bb.
It is well-known that natural gradient has the favorable property being invariant under change of coordinates. While this invariance property is appealing theoretically, in modern practice of (stochastic) gradient descent, the coordinates are rarely changed during the gradient descent procedure. Therefore, in this work, we focus on the propertics of natural gradient and standard gradient under some fixed function parameterization and coordinates. In this setup, it is widely believed that natural gradients have an advantage over standard (Euclidean) gradients in capturing the intrinsic geometric structure of the parameterized function space.
"Correspondence. These authors contributed equally to this work. Thhooughou paper, we standord gradient to denote the gradient the standard Euclidean metric on parameter space; which default choice for (stochastic) gradient descent in most machine learning practice.

---

Published in Transactions on Machine Learning Research (02/2025)
VF
M
W
Figure 1: Most existing work on natural gradient focuses on the red dashed regime, where øi is mapping from the parameter space W to the parameterized function space o = <(W). Natural gradient decreases the objective F along the gradient flow on & under chosen Riemannian metric, instead of using the standard Euclidean metric on W. However, the functional gradient VF at any point of d need not lie in the tangent space to $.
However, some critical theoretical aspects of natural gradients related to function approximations still largely underexplored. In parametric approaches to machine learning, often maps the parameter domain W C R* to a space = MM,, of differentiable functions from a manifold M to another manifold N. For example, in neural networks, W is the space of network weights, and ø: W + >MR", R") has o(w) equal to the associated network function, where n is the dimension of network inputs and m the dimension of outputs. In the modern formulation of natural gradient, á is an immersion, a Riemannian metric on M induces a pullback or natural metric on W. In supervised training of such neural networks, the empirical loss of a particular type, say mean squared error (MSE) for regression and cross-entropy for classification, is typically applied to the parametric function evaluated on finite set of training data. As shown in Figure 1 most natural gradient focuses the red
existing work (with notable exception of van Oostrun et al. (2023) dashed regime, where natural gradient decreases the objective F along the gradient flow on + = <(W) under the natural metric, instoad of using the standard Euclidean metric on W. However, the functional gradien VF at any point of & need not lie in the tangent space to t. Therefore, a complete study of the natural gradient should analyze both the gradient of F the ambient function space M + and its projection onto b before pulling it back W. This motivates our study of the analytic and differential geometric properties of M that influence the projection and the pullback metric. Morcover, to encompass family of projected functional gradients, including natural gradient and standard gradient, which leads to new understandings of both.
interpret this figure
Based on this geometric viewpoint, we address the relationship between the gradient flow on & and that on M, showing that the inverse of the metric matrix appearing in the natural gradient is in fact a consequence of the orthogonal projection of the functional gradient from M to $ We also show that because of the "pointwise" nature of typical empirical loss, the gradient of the loss function does not in general exist on the infinite-dimension space M for common choices of metrics, such as L and Fisher-Rao. The technical issue e.g., L
is that any such gradient involves delta functions, which are not continuous linear functionals This motivates studying the pullback metric in an RKHS, where by definition delta functions are continuous linear functionals. In particular, for all common empirical loss functions in supervised learning, the gradient is well-defined only in the RKHS context. Using an RKHS framework, we introduce generalized form the Neural Tangent Kernel (NTK) (Jacot et al., 2019), which we name the Generalized Tangent Kernel (GTK), to mathematically express Figure common framework of natural gradient and standard gradient.
In particular, we show that GTK enjoys automatic orthonormality property, and as result, there exists family of projected functional gradients (similar to the projection of F in Figure 1), each of which produces natural gradient that reflects well-defined geometric structure on the immersed function space. Standard gradient is the simplest example of this family in the computational sense that the metric matrix is the identity matrix.
We then study known issue of applying natural gradients modern function approximations, such as neural networks, namely the potential violation of the comerrion//ijectiive assumption on Φ, which leads to
By "functional gradient," mean the gradient of loss function on space functions.
2

---

Published in Transactions on Machine Learning Research (02/2025)
singular metric matrix and therefore an ill-defined natural gradient W. In contrast to existing approaches resorting to the Moore-Penrose generalized inverse (Ollivier 2015; Bernacchia et al. 2018 van Oostrum et al. 2023), ; we use slice methods to reduce the dimension of W and show that our GTK-based framework provides a new solution to this issuc, leading to new variants of natural/standard gradient for exemplar neural nelworks.
To summarize the main theoretical results of the paper, whenever we have a pointwise defined loss function Riemannian Hilbert space/manifole M of functions, the functional gradient exists on M iff each Ty M, M, is an RKHS. For a finite dimensional parameterized function space 4 C M, the existence of natural gradient on & (and thus on the parameter space) is then determined by projection from M to L. The existence of such a projection is unexpectedly universal and flexible, and is encoded by GTK theory, which produces a family of projected RKHS structures on each 2, f L and unifies classic natural gradient, standard gradient, and beyond. We believe this more complete understanding of the geometric foundation of projected functional gradicnts than has appeared in the machine learning literature. Viewing both natural and standard gradients through the lens of GTK also sheds light why standard gradient (and thus SGD) work so well in modern machine learning practice, while the advantage of natural gradient is often minimal
As an outline of the paper, in 2 we introduce the theoretical setup, schematically pictured in Figure 1 which establishes a projection viewpoint for both natural gradient and standard gradient. In $ we provide a careful analysis of the non-existence natural gradients for empirical loss functions on commonly used function spaces. This theoretical issue is addressed in $ by working in an RKHS, which motivates our core concept, the Generalized Tangent Kernel (GTK). The automatic orthonormality property of GTK (Theorem 4.2 produces a family of projections of the functional gradient on M onto the tangent space & via RKHS theory, with natural gradient and standard gradient as special cases. A new variant of natural gradient is also motivated by this RKHS framework; see Appendix G In $5 we study the ill-defined natural gradients for non-immersion function approximations widcly used in machine learning practice, such as neural networks with (piecewise) linear activations. We use a GTK-based framework to study the slice/reducttnn of & for ReLU MLP as an example and derive new variants of standard gradient algorithms. Preliminary experimental results of these new algorithms are also reported. We summarize this work and discuss future directions in Background material, proofs, and preliminary experiments for exemplar new natural gradient variants motivated by the GTK framework are in the Appendices.
Because of the technical nature of the paper, we encourage the reader to refer to Table | in Appendix A for key mathematical notation, and to Appendix B for background materials on Riemannian manifolds, Sobolev spaces, and RKHS theory.
2 The Theoretical Setup: A Geometric Perspective
In supervised learning, we typically want to learn an optimal function f: RR" + R", or more generally f: .M + N, where M, N are n-dimensional, m-dimensional manifolds, respectively. Here an optimal function minimizes some loss function F: M + R, where M is the space of f's.
When function approximation applied, f is parameterized by some parameter vector of finite dimension, and the general situation be described by the following simple diagram:
W
M
(1)
F
R
Here M M) is space of smooth maps from M to N. M is infinite dimensional manifold with reasonable (high Sobolev Fréchet) topology. The parameter space W, also manifold, is usually compact
Forr instance, Zhang al. 2019a) empirically observes that the advantages of natural gradient over standard gradient are minimal except for impracticall ly large training batchsize, even without worrying about the computational cost.
3

---

Published in Transactions on Machine Learning Research (02/2025)
subset of Euclidean space, Φ: W + M is the parametrization of the image <(W) = + M, F is the loss function, and F F @ makes the diagram commute.
In this section, we always assume that Φ is immersion, i.e., at every w E W, the differential dØw Tw W- Te(w) & is injective, where W is the tangent space to W at w, and similarly for To(w) t. Then for fixed Riemannian metric on M, the pullback metric on X, Y E W is given by
( (0') = (()(())
(2)
The case where e is not an immersion is discussed in $5 giving an alternative to using the Moore-Penrose inverse as in van Oostrum et al. (202).. Proofs for this section are in Appendix C
2.1 Smooth Gradients, Orthogonal Projection, and Pullback Metric for Natural Gradient
For metric 9 on M, we want. to study the gradient flow of F on the image $ and this should be equivalent
to studying gradient flow for F on W via the pullback metric (2) The existence of VF on M is carefully
studied in $3 For now, assuming that VF exists on M, we still have to distinguish between the gradient flow lines of F in M and the flow lines on Φ. Indeed, a gradient flow line of F starting in & will not stay in & in general, since VF need not point in directions.
The following lemma establishes the relationship between gradient flow in M and flows on o, and the equivalence belween pullback gradients on the parameter space W and gradients & in the function space.
Lemma 2.1. (1) The gradient flow on & for Fla is given by the flow of P(VF), where P Pa(we) the orthogonal projection from M to Totw) & with respect to Riemannian metric 9 on M.
g (ii) Let ğ = ~ be the pullback metric on W given by (2). Then dél d V F, and (()) is gradient flow line of F with respect to iff o(v(t)) is gradient flow line of Fla.
By this Lemma, to get V on W, we just need to compute P(VF). Starting below, we use the Einstein summation convention of summinng over repeated indices that occur once as a subscript and once as a superscript, so e.g., guo,w, = Engght gooww, (VE, h2) h4 = E. (VE, hi) he-
Proposition 2.1. The orthogonal projection of F onto + is
P(VF) = g4 (VF.
OΦ OΦ Ous
(3)
Ow
where
the pullback metric on W, (g¹) is its inverse matris, {(w*)} are the coordinates on W, and 06/Ow' (66(8/ww))
Proposition 2.1 shows that the appearance of the inverse metric matrix (T") in natural gradicnts is in fact a consequence the orthogonal projection of V from M to $. This projection view becomes clearer when compared with standard gradient where the Euclidean metric is used, which we discuss in the next section.
Remark 2.1. Instead of ( we can first obtain an orthonormal basis (ha)=t,..ddmww) of D say, by applying Gram-Schmidt to the pushforward basis { VF onto & by
de } of T then compute the orthogonal projection of Bust
P(VF)= = Σ (VF,hi)) hy.
(4)
Example 2.2. We give simple example for the resulls in this section. Let Φ: W (0, 2m) x 70 + R3, (0, = (cos 9 sin o. sin sin o, cos be the usual spherical coordinates on S2 Then 0/00- = (- sin 0 sin o, cos 0 sin o, 0), 0/0g- = (cos 0 cos o, sin 0 cos o, - sin Φ), and the induced metric ø²g on W for the standord dot product g on R33 is
sin 6
)
0 1
ブ= ター
0
4

---

Published in Transactions on Machine Learning Research (02/2025)
Then 3
becomes
P(VF) - (VC (- sin sin Φ cos sin 6, 0))
+ VF (- cos f cos o, sin 0 cos Φ, - -in 6)
00
d6
0
0
co) (-Fx sin 0 sin Fg cos 0 sin Φ)
+ (- Fx. cos l cos ø Fyy sin cos o F- sin σ)
00
06
2.2 Surjection and Pushforward Metric for Standard Gradient
In contrast Proposition 2.1 where the gradient on & obtained by projection, standard gradient can be obtained from surjection S from TyM to Ty given by
dΦ 00
S(VF) V.F
(5)
Ow
Ow
While S is similar to the orthogonal projection 33 in that Sl(T,0)1 0, S is not projection: S + S, and S is not the identity map. Thus Lemma 2.1(1) fails, and the flow lines for V(Fl+) are not related to the flow lines of SVF. As will be shown in $4.2 S becomes an orthogonal projection iff we had a Riemannian metric of M in directions tangent to such that:
Qsb
)
dzz.
Ow Ows
In other words, ( } must be an orthonormal basis of T, D, which is in general not true. Nevertheless, under our projection view, standard gradient uses the metric 6 on ④ which, by the following Remark, implicitly uses the pushforward 6.gE of the Euclidean metric from W to 4.
Remark 2.3. For a smooth map o: W + M, the pushforwand of a metric g on W to a metric 0.g on is given by
(X.g)(X, Y) xx (do)) E 0.9 crists iff & is an immersion. Since dΦ( gw- awr delermined by 0 is the pushforward Φ-gF of the Euclidean metric on W.
a = 86 and gE (ow du = dij, the Riemannian metric
3 Natural Gradients of the Empirical Loss
Based on the theoretical setup of $2 in this section, we analyze the non-existence issue of natural gradients for general empirical losses on commonly used function spaces. We consider the most commonly used setting of M in supervised learning, where the larget space N is R", and the loss function is some empirical loss. In this case, M is an infinite dimensional vector space, so we have to specify its topology. For the casiest topology, the induced by the I inner product, we show that the gradient of the loss function does not exist in any strict mathematical sense. More preciscly, follow-your-noss computation of the gradient produces expression containing delta functions, which are not elements computation mathematically rigorous, we have to change the topology product so that delta functions are continuous lincar functionals on the function space. This property of delta functions characterize RKHSs, which motivates the use of RKHS theory in 4
of 1.2. To make this
A typical choice of empirical loss term on M is the mean squared error (MSE)
(x(f) = Σ lf(z,) 1111
for set ((r, C M x Rd of training data. Let f: M: R have components fí: A + R, and let 8, be the "8-vector field', = E R. We formally" and incorrectly consider 8, to be function on M
'Thooughout the paper, 'formal/formalll means that concept not have mathematically sound definition, but be manipulated mechanically.
5

---

Published in Transactions on Machine Learning Research (02/2025)
with the following L inner product,
S/a 6;r= r(x:):
(6c, ())
(8)
for r E C (M). The integral is with respect to some volume form on M.
Take a curve (f:) E M with fo = f. and set fileo E TyM= M. In the following formal computation, we ignore several nontrivial technicalities, which are detailed in Appendix D denoted by d. TyM- + R, satisfics
The differential deg at f.
de (l) = di
Σ [s(()z)) - 3:112= = >> (P(:)) v())((.)))
(9)
Thus we formally get
de. FF (1)) = 0/00//10
M Σ E(P(()) - v²) N2
(10)
By the definition of the gradient, for VLE. lo exist under the I,2 inner product
we must have
(k) = (VE,, M)1 = / E(veR,) -hh,
(11)
M
where "E is the Euclidean dot product. Comparing 11) and 11 the formal gradient of lg at f has jth component given by
(Vln.)) = (Vim.)"= E((xx))
- ys )
(12)
In other words, the gradient of the empirical loss (7). if it exists, must formally be sum of delta functions. Since delta functions are not 1 functions, the L gradienl does not erist on M. Thus we cannot apply Prop. 2.1 to compute the gradient of f on b, even though the gradient the finite dimensional manifold + must exist.
Remark 3.1. We emphasize that the non- non-eristence of the L2 gradient is endemic to supervised learning, because of the discrete nature of empirical loss function. For general differentiable function L L(2, w): RR R RR the corresponding L-empirical loss is ( E, L(J(x,), y:). For V L, the Euclidean gradient in the direction, the formal gradient 12) becomns (VLL.)) = E. L2 v:)) d;, which again is sum of delta functions.
As another common example in machine learning, we consider the Fisher-Rac Rao metric on the space of C probability distributions P {p E (M, p > 0, JM p(z)d (x) = 1} on a measure space (M, jt), where M is compact manifold without boundary. The tangent space at p E P T,P (h E C (M, R) SM hdp 0} Lafferty 1988 $3). Let 6: W + P give parametrized submanifold $ = x(W) C P, with ((w)(z) == p(r, w). The Bether-Rao metric at To(w)) 4b is Pij.o(w) = E[ p(z,(z, w) :wi log p(x, w)), so the Fisher-Rao metric on T,P is
= Sm
hiha du. p²
(p( h2) E[hh
since h(logp) (d/dt)) log(p + th) h/pp
For the MSE loss on P, (e(P)= = E (p(x;;) vil2, as in the 2 case we formally get
gr (VEE,, dl d().
[(((x)) (1))(3)))))3)))
Σ (p(xx) - (1))
M E=1
9p (v²
(p(x.s) - us vi)o,, h).
6

---

Published in Transactions on Machine Learning Research (02/2025)
Under the usual assumption that Sas d,d1 1, the formal gradient is
h Σ (p(xz)) y:)o; - Σ (osaa))((a))) -y),,
VIg = p
where the conslant second lerm ensures JM VL = 0 formally. Again, VlE.p does not exist in TpP, since delta functions are not elements of T,P. Similarly, we use a more general loss function L(z, w): R + R and set &: P + R by {(p) E, L(p(z)), ya), then the formal gradient is p² (00() which again does not exist in T.P..
4 Gradients in RKHS and Generalized Tangent Kernel
The non-existence of the gradient of empirical loss discussed in $ molivates our study of natural gradient in an RKHS, which is detailed in $4.1 In particular, whenever we have a pointwise defined loss function on a Riemannian Hilbert space/ space/maaninie manifold M of functions, the functional gradient exists on M iff cach f E M, is an RKHS. For a finite dimensional parameterized function space & C M, a natural gradient on & (and thus on the parameter space) is then determined by a projected RKHS from TyM onto T, Our study goes further along this line in $4.2 where we introduce the Generalized Tangent Kernel (GTK), the concept of this work, and prove appealing automatic orthonormality property of GTK (Theorem 4.2) The GTK framework provides a unified understanding of projected functional gradients, producing familic les of natural gradients that include both the standard gradient and the classical natural gradient. In fact, a finite dimensional submanifold & C M has many choices of natural gradients, which are in one-to-one correspondence with choices of GTKs on each T, Results in this section have proofs in Appendix and numerical examples in Appendix G
4.1 Natural Gradient in RKHS
In this subsection, we study the gradient of a loss function an RKHS, which by definition a function space where the evaluation/dees functions 6, are continuous.
Let H c M(M.R") be a vector-valued RKHS with kernel Kx. The reproducing property of H (see Appendix B.3
gives
h²",) M(x), vh E H. Then there is "8-vector in the dual space of H that evaluates cach component of h at point i.c.,
(Kx( (1K 1, (Ki(2,, ),,
6 = (x,,), = h² (z),
where Hy is the scalar-valued RKHS corresponding to the j-th component of H. Thus de df) = (Vig. N) defines VlE, as an element of H. In contrast, h + (h)) is not a continuous linear functional on the function spaces discussed in 43 so the functional gradient does not exist.
The next Lemma puts 12 into the RKHS framework for more general loss functions.
Lemma 4.1. For L = L(z,w): R x R + R. differentiable function, the H H-grodient of the L-empirical loss function (x( E. L(Sxa), mh on H C M(M, " is given by
Vli. = Σ KH(a,, V
(13)
where VEL, is the Euclidean gradient of L in the direction.
If = M(M,RT) with a suitable inner product, we know we should project Vll, to the tangent space of Φ- = (o). Proposition 2.1 and Remark 2.1 provide two formulas, based on the pushforward basis { 0g
dum }
and the H--rrthonrrmal basis {h+} of T, respectively. Applying them to (13 gives
do
de Ow
kk Σ PVLL -
(FL,) u:) E Owk
(14)
7

---

Published in Transactions on Machine Learning Research (02/2025)
where IN is the pullback to W of the inner product the H. and
PVLL = - ( E Nol..) hr.
(15)
For this RKHS H. the "6-vector" exists on each tangent space T and is given by the corresponding "projected' evaluation function,
(x) - Py MTyo = (h, K2 DT, for D,
where K K the projected reproducing kernel. As above, Ky can be writlen using either the pushforward basis { 96
&wu } or the orthonormal basis {he} of Ty Combining these results, we have the following kernel form of the projected gradient, as an equivalent alternative to [4] and 15
Proposition 4.1. For H = M"M,R") an RKHS with kernel Ku and L = w)z R x R + R differentäable function, the projection of the H H-gradient of the T--empirical loss function (()) = Σ, ((( to equals
VoLLs:= PVELS = -1
(16)
where K is the projection kernel of KH onto the RKHS T Ky can be represented by the pushforward
basis Dw } of D,
24j dΦ
(z) 2 do
Ks(x,2') 97 Owi
Ous (z)),
(17)
by an orthonormal basis {h,} of T
Ky(x,x²)= = Σ hx(x)) 2 hx(x')
(18)
(Note that T is an RKHS, as it is finite dimensional subspace of the RKHS TyH. ~ H.) In other words, in RKHS theory the projection view of Figure [ and $2.1 produces a family of projected RKHSs parametrized by E W, with the corresponding projected reproducing kernel K,, for f = d(w).
From a practical perspective, although different choices of H. produce different natural gradients, in general, =i) it is expensive to compute either Ix in (17) {44} in (18) For H (RR", RT) Sobolev space (detailed
in Appendix B.2). an explicit computation of (x, and hence of 5, is in Rosenberg 2002).. Because of
this, we introduce a new variant of natural gradient, called the Sobolev natural gradient in Appendix G
where we discuss an RKHS- based approximation of T and provide a practical computational method of 9x
basod on the cronecker-factorce approximation (Martens & Grosse 2015 Grosse & Martens 2016). We also report preliminary experimental results on image classification benchmarks.
Note that the RKHS form of projected gradient 16 in particular the projected kernel forms 17 and [8 is consistent with our observations from Proposition 2.1 and Remark 2.1 namely that the inverse metric matrix of natural gradient appears as a consequence of orthogonal projection from M to 4. This matrix becomes an identity matrix iff an orthonormal basis on each tangent space T is used.
4.2 A Unified Perspective with Generalized Tangent Kernel
In this section, we introduce a general form of the kernel (1). which provides a unified characterization well as a new understanding of projected functional gradients, including both natural gradient and standard gradient. The favorable property of this kernel also enables us to address the con-immersion case of & in $
From $2 the standard gradient formula. (5) for the pushforward basis is similar to the natural gradient:
formula (4h for an orthonormal basis. This formal similarity extends to the corresponding RKHS expressions;
with (16) for the natural gradient, there is kernel form for the standard gradient
VbLY S = Σ ((( v1)
(19)
8

---

Published in Transactions on Machine Learning Research (02/2025)
where (xx,z)= = En De (x) @ d0 (z') happens to be the Neural Tangent Kernel (NTK) (Jacot et al 2019
Oud
Bur
This observation motivates the following generalized form of both Ky and O.
Definition 4.1 (Generalized Tangent Kernel), For Φ: W- M an immersion, where W is a parameter space and M a function space, let { be a basis of the tangent space T at f Φ(w) for some w E W. The Generalized Tongent Kernel of Ty is defined to be
(erk(x, = E1(())
(20)
The Generalized Tangent Kernel (GTK) thus provides a unified characterization of gradients on T with natural gradient and standard gradient as special cases. The natural gradient (16) most closely related to the functional gradient in M (w.r.t. some functional metric), as it is obtained by an orthogonal projection and the corresponding GTK (18) is defined by an orthonormal basis of Ty The standard gradient (19 in contrast, corresponds to NTK, GTK defined by the pushforward basis from W, which in general is not orthonormal for the functional metric. Thus NTK and any GTK besides
18 lose extra information about
the functional gradient in M.
So far, our presentation of the role of the orthogonal projection in natural gradient provides an intuitive justification for the long-belioved statement that, regardless of computational cost, natural gradient has an advantage over standard gradient in capturing the geometric structure of the immersed function space. With the following favorable property of GTK, however, we see the surprising result that the standard gradient is, at least locally, geometrically informative the natural gradient.
Theorem 4.2 (Automatic Orthonormality of GTK). Given a finite dimensional function space F C { + } where X is any set, and any basis {bi} of F, {bi} is an orthonormal basis of the RKHS associated with the reproducing kernel given by the Generalized Tangent Kernel 220
Proof. Given a basis {b,}, let HK denote the RKHS associated with the reproducing kernel K defined by the GTK (20). Then each b, is in the span of {K(x, ), E X] (Manton & Amblard 2015 Lemma 2.1), and can be respresented as b.) Ezex K(I, (I((),, for some function C.: X - R²". By (20)
bi(-) Ec(x) Ed(()) = Σ [(()))
aajbe(), -
be()) be(()= =
(21)
ZEX
L
Since (be) is a basis, the representation where of EBEx ((())
21 must be unique, so of 8%.
Thus by 36
(b by
Σ K(x,)) Σ [((r,)), BEX
ΣΣ FEX x'EX
K(,
FEX
HK
(Σ (gweeeee)
ΣΣ C(t)7) Σ
Σ 8
zEX x'X
T'EX
= Σ 2aaaz= o0)
where 0f equals when i = j and 0 otherwisc.
Theorem 4.2 is the key result in this work, as its direct Corollary expands the scope of natural gradient in an unexpectedly flexible way.
Corollary 4.3 (Natural Gradient for Arbitrary Basis). For f d(w), any choice of the basis {ba} of T produces an orthogonal projection onto T and thus a well-defied natural gradient at w. The natural gradient has the following kernel form,
VOLL := PVLL, = - 1F
(22)
where KETE is defined by 20
9

---

Published in Transactions on Machine Learning Research (02/2025)
&w }, GTK recovers NTK, and (2 recovers the standard gradient formula 19), which is equivalent to the conventional standard gradient on the parameter space W:
As special case of of Corollary 4.3 for {b,} {
V Vuli,o(i)) == L Lo vi) (z)(w)(2.).
(23)
If Φ is a neural network with parameters w, then V is the gradient of the loss function with respect to the network output and V the gradient of the the network output with respect to the network paramcters. (23 can be efficiently computed by backpropagation. Corollary 4.3 guarantees that the standard gradient 23 used everywhere in machine learning is, in fact, a well-defined natural gradient, in the sense that it is the projection (with respect to some specific Riemannian metric) of the functional gradient in the ambient. function space.
To see the flexibility of Corollary 4.3 in defining valid natural gradients, we now give an example of changing a GTK by a change of coordinates on W (which corresponds to a change of Riemannian metric on $). Let (E' (c'.....,5) be new set of coordinates on the parameter space W, given by a full-rank transformation matrix A = (a)), i..., şd = w/w² Then f = ((w) é(A- -g) Ø(E) and 86 ag'
8b dw at, where (a/) = AA- Instead of
choosing the pushforward of { ( } as basis, i.e., {b.} Bwl }, and getting the standard gradient (23) we choose the pushforward of { ) ), i.c., {b,} = I 0d ) and apply Corollary 4.3 again. change of coordinate
computation gives the following altered gradient update rule on W:
O5
0ō
Vell.ōle)
Σ
V La (0(E)(z)), vi) E 0ei (x:)
OEi
ab
ab Ow af
Σ
V (o(E)(z)), v.) E Oio (x:)as
Oo
db dw!
Σ
V (d(w)) E ou (5:)
(24)
(V L. L, (d(w) V A-1 (A-1)"
V La (b (v: (o (A" A)-
We omit in the last two lines as in 23 (4 can also be efficiently computed by standard backpropagation, since the only extra term compared to (23 is the fixed A)- In the next section, we further leverage the flexibility of Corollary 4.3 to propose new approaches to address the ill-defindd natural gradients for non-immersion function approximations.
To the best of our knowledge, this is a new understanding of standard gradient, natural gradient, as well as NTK. It is also worth noting that while GTK is a generalized form of NTK, the theory developed in this scction docs not correspond to known results about NTK.
5 Natural Gradient for Nom-Immersion Function Approximation
If ó is not immersion, the pullback metric matrix (ğt) o" is singular. Therefore, the inverse matrix (G') does not exist, and the formula 3 for natural gradient makes no sense. Neural networks with (piecewise) linear
activations (Dinh et al. 2017) are such examples. Specifically, take a two layer neural network Φe Rn + Rm with activation function v. so o (x) = A2,9 (1b( (1,,, X + b1.0g) ), + b2,0,, where 0 = (01, 02, 0g. 04) and A2, has entries 0,, etc. Then 00, O2 0z. 0.) = Φ 0, 202, A0g. 04) for any X x
Most existing solutions OOlivierr 2015 Bernacchia et al. 2018 use the define natural gradient formula. However, the Moore- Penrose inverse in general loses more information
generalized inverse to
10

---

Published in Transactions on Machine Learning Research (02/2025)
about the function space geometry and gradient than an orthogonal projection from M to Ty Using
Corollary 4.3 we propose an alternative approach in certain explicit cases, consisting of three steps:
For the given @ find a slice (defined in $5.1) of the parameter space W with øls an immersion.
2. Compute basis (ci of each tangent space of S.
3. Choose the basis {b:} on 7, from the pushforward of {ci} and construct the corresponding GTK
from {bi}. By Corollary 4.3 (2) then gives a well-defined natural gradient.
This approach applies to any non-immersion function approximation. While finding the slice S may be difficult depending on the concrete function approximation, once S and its basis {c} are computed, they can be used throughout the training.
This approach is also computationally efficient. Due to the flexibility of Corollary 1.3 we are allowed to choose
the computationally simplest basis on T to construct a GTK and (2
always gives a valid natural gradient.
In the following section, we give two such examples for the ReLU MLP. Our experiments in Section 5.2 show
that the extra computational cost compared to the standard gradient bascling negligible.
Note that naive slice-based approach without leveraging Corollary 4.3 would make little sense in practice.
In particular, the natural gradient (3) for the immersion case is known to be computationally expensive, and
people rely on a series of approximations (Martens & Grosse 2015
Grosse & MMattnn 2016) to make it
acceptable in practice. Restricting the natural gradient to slice S is yet more complicated; in fact, it is unclear if a valid approximation still exists.
5.1 MLP with ReLU Activation
Following Dinh et al. (017)). use ReLU MLP without bias terms concrete example to derive two variants of well-defined natural gradients. Proofs are in Appendix F and results and proofs for a general ReLU MLP with bias terms are in Appendix !
Our setup is a Ê-layer MLP with ReLU activations, with parameters w in the k-th layer, for r = 1,.... fkk where ni is the number of matrix entries in the k-th layer. The space of parameter vectors is W = (w (wl...., , w ) and dim(W) Ee ng. The map o: W Maps(IR", R)" (R") takes a parameter vector to the associated MLP. We abbreviate the subvector for the k-th layer by w = (wf,... wh ), and w = (w .w......u w
We can assume that w 7 0 for all k, since a MLP with some w = 0 is trivial.
Define a multiplicatiiv group G {a = (a.,..., ate) a > 0, I as = 1]. The map G - RR, (ai,..., ar) +
Σ, log a is a bijection to the plane E, x² 0, so G is a manifold of dimension É 1. G acts on W by
(aI...., ae) (w .w.,,.., w²) = (anw", ...., "rnw""...,uu"""" aew").
It is straightforward to check that
(ww', "ul"....") s(xx (w "oo",..."""" w()),
(25)
for all G. Thus is far from injective. As shown below, & is not an immersion.
Proposition 5.1. The dimension of ker(de) w at each w W is at least - 1. Thus & is not an immersion. In fact, the span of
nj Σ
0
2
w
w
:j = 2, 2......
(26)
Ow?
Ow
3r mn
Ōri
w
we
is in ker(dō) Here p is the radial vector in the parameter space of the ith layer.
HHee we treat the standard gradient d were an immersion, so strictly speaking Corollary 4.3 does apply and this gradient valid natural gradient.
11

---

Published in Transactions on Machine Learning Research (02/2025)
Remark 5.1. This issue is fundamentally caused by the homogeneity of neural network activation function, as any (piecewise) linear aclivalion function (ReLU, Leaky ReI.U, Shifted ReLU, Parameterized ReLU, elc.) produces a non-immersive Φ.
We now follow the three steps proposed in the previous section.
Step 1. Find a slice S of W
Note that & is constant along the orbits On = {g w :::g E G}. Thus all the information Φ is contained in slice for G, i.e., a submanifold 8 of W that intersects each orbit once. The following theorem gives an example of S.
Theorem 5.2. Wrile W = R" x... R" A slice S is given by S {Xw*, ..... ,..., w() wi E s5m >>>00 0} where S is the unit sphere in R
E S, it is unlikely that any nonzero in S lies in ker(d)) By dimension count, (T (2')) dim(W). Therefore, we expect ker(dø) Spa in Prop. 5.1 so that øls is immersion.
For w
Step 2. Find a basis {(s of S.
We take a basis (r:= = cg ch k 1,....,.;;1, 1..... ny - 1} of the tangent space TwS for the slice S. In polar coordinates (rx, 1) K2 = 11....,,.. ny - 1, on the 4t layer, r r= Ex 0/Ork, 2 = 0/00 h at x E R The Euclidean metric restricts to metric on S with metric tensor h = ((ccAa)((s,0))) c c.. Details of {cf and h are given in Appendix
k;
Step 3. Define a GTK.
Given the flexibility of Corollary 4.3 we suggest two approaches for constructing the GTK from pushforward of {c
Approach 1: NTK. We push forward the metric h, the Euclidean inner product restricted to S, to TyD Thus we replace in (19) with Σn.,, ((lk)))trs))) @ do(c)).
Approach II: GTK. We can directly use (20) to get KATK = ERR,, dé(c)) & do(c)))
For Approach I, the final expression for the natural gradient is given by the following proposition. Proposition 5.2. The gradient F for F: W + R, is given by
OF a
OF Tw
a awt
S
Wks
wky
w
Duck
Owk,
W1 2
N=
k1=l
1=
k,=1
OF Ows
a wks Ouk
(27)
|w1 /2
Details of the notation and the proof are in Appendix[F] Gradient formulas for Approach are in Appendix L.2
The procedure of training MLP with ReLU activations using Proposition 5.2 is summarized in Algorithm I
(27 may seem complicated, but in fact, it can be easily implemented, and the extra computational cost versus standard gradient computation is negligible. Preliminary experimental results of both Approach T and IT are provided in the next section and PyTorch implementation of Approach I is included in Supplementary Materials.
5.2 Experiments
Following Algorithm [ we implemented a stochastic version of the proposed natural gradient descent for seven- layer MLP with ReLU activations. We test our implementation (slice SGD) versus standard SGD on the CIFAR- 10 image classification benchmark. We use seven-layer MT.P of layer size (2634, 2196, 1758, 1320, 882, 444, 10),
12

---

Published in Transactions on Machine Learning Research (02/2025)
Algorithm Slice SGD for MLP with ReLU activations
Require: stepsize
Init: network weights w (Wi,1..... wtm,)
S,
while stopping criterion not met do
Compute Euclidean gradients
OF Dwn.k
via standard backpropagation
w=
Compute S by 27
Update weights ah+1 w
F
Move wh along its G-orbit back into S by rot+1_ (II | 1/8 + end while
E
100
all layers except the last uses the ReLU activation function. This MLP is trained by standard SGD and our slice SGD methods using a batch size 128 for 200 epochs respectively. To simplify the comparison, we have not used any weight decay or momentum. The same fixed learning rate of 0.01 is used for all methods. Note that these basic settings are commonly used for the SGD baseline, which have been extensively optimized by the ML community. The learning rate is also tuned for the baseline. Our methods simply align everything with the bascline without specific tuning. All experiments are run on a desktop with an Intel 19-7960X 16-core CPU, 64GB memory, and a GeForce RTX 2080Ti GPU.
The testing performance wall-clock training time is shown in Figure ( Though proposed slice SGD formulas seem complicated, the implementation in PyTorrc, as attached in Supplementary Materials, is straightforward. The extra computational cost compared to standard SGD is O(n) scalar multiplications and additions, where n is the number of network parameters. Note that these extra computations are between network weights and their gradients (already) obtained by standard backpropagation, which can be effectively leveraged by modern GPU's. This is fundamentally different from the sequential nature of backpropagation. As shown in the wall-clock axis of Figure 2(a) the difference in time cost versus standard SGD is negligible in practice. We have also tested Approach II described in $5.1 with the corresponding gradient formulas provided in Appendix
I.2 As shown in Figure 2(b) it performs similarly to the Approach I.
60
60
55
Test Accuracy
50
Test l Accuracy
50
45
40
40
35
30
baseline
30
Approach
ours
Approach
20
25
250
500
750
1000
1250
100
150
200
Training time (seconds)
Training Epochs
(a) Ours vs. Baseline
(b) Approach vs. Approach
Figure 2: (a) Testing accuracy vs. wall-clock of the slice SGD (ours) vs standard SGD (baseline) on CIFAR-10. (b) Testing accuracy vs. training epochs of Approach and Approach II of our slice SGD on CIFAR-10. Results averaged five runs of different random seeds, with the shaded area corresponding to the standard deviation.
13

---

Published in Transactions on Machine Learning Research (02/2025)
While our experimental results are still preliminary, feel that they show the potential of our theory, while leaving large room for practical optimization to future work.
6 Discussion
We believe our work is just the beginning of fruitful results in both theory and practice of broader view of natural gradient. From a theoretical viewpoint, the Generalized Tangent Kernel, in particular its automatic orthonormality property, induces family of Riemannian metrics the entire function space M, putting both the natural/pullback metric and the peelidean/pushoowaare metrics reflect (difffrent) geometries on the function space. This motivates the question of which metrics in M lead to better convergence rate for natural gradient descent. Even on a finite dimensional manifold, the relation between a choice of metric and the convergence rate seems to be not well studied. As a first step towards comparing pushforward and pullback metrics, we compare notions of flatness in Appendix H
metric into a unified framework, as both
There are many more topics in Riemannian geometry and machine learning to explore. For example, we have not discussed the relation between the choice of metric and the generalization ability. A related example is Kozachkov et al. (2023), which establishes an interesting connection between Riemannian contraction and an almirithmic stability generalization bound. In another direction, the very intriguing question of which flows on a manifold can be put into gradient flow form for some Riemannian metric (Shoji et al. 2024 is equivalent to determining which flows can be put into GTK form; this will be discussed in future work.
Transformer are known possess the in-context learning capability. It has been shown (Ahn et al. 2023 that transformers actually learn to implement preconditioned gradient descent for the forward pass. Whether there is GTK explanation for this preconditioned gradient descent in forward pass is another interesting question for future study.
From a practical viewpoint, our study and experiments of the new familics of natural/standard gradient descent methods motivated by our theory are still preliminary. We also hope to study slices of parameter spaces for non-immersion function approximations beyond MLPs, such as convolutional neural networks, ResNets, and Transformers. It is therefore appealing from both theoretical and practical perspectives see if general approach to the slice SGD can be designed for different neural network architectures.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:45614 45650, 2023.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 10(2):251-276 1998.
Ari Benjamin, David Rolnick, and Konrad Kording, Measuring and regularizing networks in function space. In International Conference on Learning Representations, 2019.
Alain Berlinet and Christine Thomas-Aguan. Reproducing kernel Hilbert spaces in probability and statistics. Kluwer Academic Publishers, Boston, MA, 2004.
Alberto Bernacchia, Máté Lengyel, and Guillaune Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. Advances in Newral Information Processing Systems, 31, 2018.
Aymeric Dieuleveut, Francis Bach, et al. Nonparametric stochastic approximation with large step-sizes. Annals of Statistics, 44(4):1363- -3999,1199,11 2016.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Joshua Bengio. Sharp minima can generalize for deep nets. Xix preprint Xiv: 1703. 04933, 2017.
Peter B. Gilkey. Invarionce theory, the heat equation, and the Atiyah- Singer inder theorem. Studies in Advanced Mathematics. CRC Press, Boca Raton, FI., second edition, 1995.
14

---

Published in Transactions on Machine Learning Research (02/2025)
Roger Grosse and James Martens. A Kronecker-factordd approximate Fisher matrix for convolution layers. In International Conference on Machine Learning, pp. 573-582. PMLR, 2016.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Xrxiv preprint or
2019.
Sham M. Kakade. natural policy gradient. Advances in Neural Information Processing Systems, 14, 2001.
Leo Kozachkov, Patrick Wensing, and Jean-Jacques Slotine. Generalization dynamical robustness the role of riemannian contraction in supervised learning. Transactions on Machine Learning Research, 2023.
John Lafferty. The density manifold and configuration space quantization. Trans. Amer. Math. Soc., 305(2): 699 741, 1988.
John M. Lee. Introduction to smooth manifolds, volume 218 Graduate Texts in Mathematics. Springer, New York, second edition, 2013.
Wuchen Li and Guido Montúfar. Natural gradient via optimal transport. Information Geometry, 1(2): 181-214, 2018.
Jonathan H Manton and Pierre-Olivier Amblard. A primer on reproducing kernel Hilbert spaces. Foundations and Trendsø in Signal Processing, 8(1-2):1-126, 2015.
James Martens. New insights and perspectives on the natural gradient method. arXiv preprint arXiv:1412.1199 2014.
James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In International Conference on Machine Learning, pp. 2408 2417. PMLR, 2015.
Yann Ollivier. Riemannian metrics for neural networks I: Feedforward networks. Information and Inference: A Journal of the IMA, 4(2):108 153, 2015.
arXiv preprint Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. arXiv:1301.9588, 2013.
Steven Rosenberg. A note on reproducing kernels for Sobolev spaces. arXiv preprint arXiv:2307.0863 2023.
Konrad Schmüdgen. Unbounded self-adjoint operators on Hilbert spaces, volume 265 of Graduate Texts in Mathematics. Springer, Dordrecht, 2012.
Lucas Shoji, Kenta Suzuki, and Leo Kozachkov. Is all learning (natural) gradient descent? arXie preprint 4rXiv:2409.164330 2024.
Jesso van Oostrum, Johannes Müller, and Nihat Ay. Invariance propertics of the natural gradient in overparamel systems. Information Geomelry, 6(1):51-60 67, 2023.
Chaoqi Wang.
com/ httpaa/www...llacccaaaaaaaaaa/aa
2019.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger Grosse. Which algorithmic choices matter which batch sizes? insights from a noisy quadratic model. Advances in Neural Information Processing Systems, 32, 2019a.
Guodong Zhang, James Martens, and Roger Grosse. Fast convergence of natural gradient descent for over-parameterized neural networks. Advances in Neural Information Processing Systems, 32, 2019b.
Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. In International Conference on Learning Representations, 2019c.
15

---

Published in Transactions on Machine Learning Research (02/2025)
A Key Notations
Table 1: Key Notations
M MM, the space of differentiable functions from manifold M to manifold N;
usually M ne k
Φ: W M function approximation, ie., the mapping from parameter space W to function space M
+ ((W) M the parameterized function space, L.e. the image of e M
F: M -> I : the loss function on MM
F Fo d: W R the loss function W from F induced by ø
= [1)) the Riemannian metric on M
0vj o" g the pullback metric on W
the empirical MSE loss on M, given some finite training set
vd the orthogonal projection of the gradient on function space or RKHS to &
W {w (wi...., wh,)): )}: the parameter space for Ł-layer ReLU MLP,,
n& the number of parameters the k-th layer
w (wi, .,.., wh the vector of parameters of the k-th layer; assume that w" of
{a = (an. ae) or 0. I[ a: = 1}: the multiplicative group acting on (MLP) W by w
O [g w gg:ge E G) the orbit of the W under the group action; ø is constant Ow S = (>(w".... ...., w²): : w E 5mi À > 0}: a slice W for the group action,
where s the unit sphere in R*
M(M, vector-valued RKHS with kernel KR
l the more general L-empirical loss on M, where L = L(2, w): R x Rd + R is differentiable A F L. the Euclidean gradient of L(z, w) in the direction
ON (0H,s) the pullback to W of the inner product (:r) on H
K Py Kr the projected reproducing kernel of Kx the tangent space Ty4
O x') Σ, Bø Bhet auat (x') the Neural Tangent Kernel (NTK) (Jacot et al. 2019
(x)
06
Kork(e, = E, (.()) @ (x(2)) the Generalized Tangent Kernel (GTK)
16

---

Published in Transactions on Machine Learning Research (02/2025)
B Background Material
Since our framework involves Riemannian manifolds, RKHS, and Soboley spaces, we briefty review aspects of this material. General references are Lee (2013) for manifolds, Gilkey (9955) for Soboley spaces domains
in R" and on manifolds, and Berlinet & Agnaa 2004 or Manton & Amblard (2015) for RKHSs.
B.1 Riemannian Manifolds
We assume the reader familiar with the concepts of manifolds and local charts. To avoid technicalities, we assume all manifolds are smooth and are embedded in some RV. Let M be an m-dimennional manifold,
denoted M At point p E M", the tangent space T.M R" is the collection of tangent vectors to smooth curves passing through p: TyM- {(d/dt)=o))(()) where y: (-c,c) + M is a smooth map for some > 0 with 7(0) p.
Associated to a smooth map f: 'M + N between manifolds is its differential dfp: T,M + Ty the best lincar approximation to f at p. defined by dfp(X)) = (d/dt)==u)f(((t) where r(t) has 7(0) p. (d) = X. (In many lexts, dfp is denoted by Sapp)) We suggest the reader draw a picture for this definition. The differential is a linear transformation between tangent spaces; for M RT NN R", the differential is the usual Jacobian matrix dfp = (0f%//r"/p)>>x)) pp)mm. The equality d(f g) df o dg for f: M - N, g: W M has a picture proof, and for M, N. W Euclidean spaces is exactly the chain rule. As special case, for f: M + IR, dpp: M - (ppp) R = TR is an element of T M, the dual vector space of linear functionals on T.M.
A manifold M is a Riemannian manifold if it has a positive definite inner product 9p on each TpM, with the condition that Ip depends smoothly on p in a technical sense. Thus 9p: T,M x T.M + R, with (i) g(X, Y) = X). and (ii) x(X,X) 2 0 with g X) iff X = The dot product X 'm Y on M = R" is the basic example of Riemannian metric, and for M C R-, we can set (p(X X 'N Y. However, M has uncountably many Riemannian metrics; for example, we can take (X, = A 'N Y for any positive definite symmetric matrix Ap depending smoothly on p.
A finite dimensional vector space V and its dual space V of linear functionals A: V - IR are isomorphic, just because they have the same dimension, but there is no aanoiicalnnature isomorphism that doesn't involve choosing a basis of V. ("Canonical" mcans "without adding any more data.") In contrast, if V has a positive definite inner product (,,),, then we have a canonical isomorphism V - V* given by v + Ae where (w) = (v. w) for w E V. In particular, on a Ricmannian manifold there is a canonical isomorphism between TpMM and M. Thus for a smooth function f: M + E, the differential dfp E T.M corresponds to a unique tangent vector in M,M, the gradient vector V.p. By the canonical isomorphism above, we have the fundamental formula
(x) = 9p fv.
As expected, for M = RT, the gradient is the usual Euclidean gradient.
In Prop. 5.1 we will use the fact that for f:M M - N and n E N, dfls-1(x) = 0, i.e. the differential vanishes on the level set fn): intuitively, f does not change on the level set. (This assumess that the level set is a submanifold of M.) It follows that any langent vector v E Tm (n) has df(v) = 0. Thus (')) = (..
All this theory is uscless unless chart around p E Mm, then dop; T2U + IRm = Rm is vector space isomorphism. (Apply the chain rule to d(a-11 o a) = dld = Id.) Thus the standard basis {ei,..., emt of T corresponds to a basis of T,M, confusingly donoted (0/ax') [0/ax'Ip..... 0p/...,0/0x" (p): ). or just {0am le..lp....... 0mmlp) In this basis, the Riemannian metric has the local expression as a symmetric matrix:
do explicit calculations in local coordinates. If ap: Up + R is local
Ip = (D:)p, Pi g( lp. Ozs 1p).
Dropping the p index, if X,Y T.M are written as X = X'O Y YY' (using the Einstein summation convention that repeated indices are summed over), then g(X,Y) Y) Dig Xiy.
17

---

since these polynomials in the components of & have the same degree, so the Sobolev s-norm is equivalent to Published in Transactions on Machine Learning Research (02/2025)
The corresponding dual basis of F is also confusingly denoted by {dx' ..... dr ). where dx' (0/0x))) =0; the Kronecker delta function. It follows that the local expression for df, is df,, (0f/0x' ddr'. Thus the coordinate-indeeedder differcntial df, keeps track of all first derivative information of in any local coordinate chart. The positive definite matrix (gt)) is invertible, and its inverse is denoted by (g")). The canonical isomorphism T M - M is given by B = 3,dx² MTM + B,0gs E TpM. As result, the local expression for the gradient of at p is fpp of 1/000/'')pp)p
As final technical point, a smooth function Φ: M - N pulls an inner product g on N back to an inner product denoted o' on M by setting
(6 (x)) - Pe(a) (X), (x)).
(28)
However, Riemannian metric may not pull back to Riemmanian metric: if X + 0 but do(X) = 00, then (Ø" (X, X) = 0. Thus we must require that de is an immersion, which by definition means that dØp is injective for all p E M, in which case the pullback o' g of Riemannian metric is Riemannian metric.
B.2 Sobolev Spaces
One motivation for Sobolev spaces comes from the fact that linear transformations infinite dimensional normed vector spaces may not be continuous; this never happens in finite dimensional lincar algebra. In fact, the simplest differential operator dz on smooth functions on [0, 1] does not exlend to a continuous operator on the Hilbert space - 1]: the functions fn(x) = n n- lim ((ddddz)fmllzz = co.
1/2 sin(mnr) have lim ll = 0, but
There are two approaches to handling differential operators as linear transformations on Banach Danach/Hilberl HHilbert. spaces: (1) develop theory of discontinuous/nm unbounded operators in e.g-, [Schüüdgen (2012); (ii) use Sobolev spaces to make the operators continuous. For (ii), if we define the first Sobolev space If O0, to be the Hilbert space completion of C= with respect to the norm [I := Il + [ldd/dx)fllzz then it is immediate that d/dx: H - T.2 is continuous: 1im [l = 0 implies lim Ilfmllæ== =0
In higher dimensions, let 52 be a bounded open set in RR. For s a positive integer, we define the s-Sobolev space H² to be the Hilbert space completion of C20 with respect to the norm
(12, C) (smooth functions with compact support in ?))
Ilfll = Σ flf... la|5
Here a = (a'....,a") .a)) is a multi-index with |a| = E, ai and 8 /00 ... ...dr"". Thus the norm measures the 1,2 norm all partial derivatives of f up to order s. By standard properties of the Fourier transform,
1111
Σ Ja|Sa
P(E) d.".
R"
where & = (S.,..., (5) and go xitt
conn
There are positive constants C₁ C2 such that
C1 + 1E|)) <
Σ
c2(1+
la|ss
the (renamed) norm
11 / (f()) (1)))) +
(29)
RR
with the associated inner product (f.g). = fx- [f(E) [t df The advantage of 29 is that
we can now define H² for any E R, and we have the basic fact that there is a continuous nondegencrate pairing Ha 8 H- - - C, f f + Ja 7. This implies that the dual space (H*)* to H² is isomorphic to H-
H - (H*)* via JEH (2
(gEH + / 5.g
(30)
18

---

Published in Transactions on Machine Learning Research (02/2025)
By the fundamental Sobolev Embedding Theorem, for (n/2) + +',, H2) is continuously embedded in O' (Q), the space of s' times continuously differentiable functions on S with the sup norm. We always assume satisfies this lower bound for s' = 0. As result, for any I E Ω, the delta function dg is in H2))::: if fi + f in H then fi + f in sup norm, so 6- (fr) = f((z) + f(z) = 6z()). Thus there cxists dg E (M) such that
(2( - ()=2:))) VS l1.
(31)
All this extends to vector-valued functions. For = (f (.... E ("), we can extend the inner product by
(f.9) = E(P.o))
(32)
Then IlfIl = Eド= 1P and we proceed as above to define H, H,, I.)) The delta function generalizes 62, where
= [())) i.e, (5,()) (dz, 1%).
(33)
We also have Sobolev spaces H,MM) on manifolds M, which are defined first in local coordinates on each element in coordinate cover of M, and then patched together using a partition of unity. Finally, we can combine these constructions to form H,MM, RT)) for functions f: M - Rm
B.3 RKHS
The most common Hilbert space is an 2 space, e.g., L 1], the space of real-valued functions defined on [0, which are 2 with respect to e.g., Lebesgue measure. One subtlety of L?2 spaces is that two functions fi, f2 define the same element in [² 1] if fi f2 except on set S of measure zero. As a result, the evaluation maps (often called delta functions) 8, [,, - IR, x E [0, 8,( = [(x), are not well-defined; if f E S, then (fi) # de (f2). Even on well-defined functions, evaluation maps need not be continuous: for example, if (x) = x", then lim llfullza = 0, but [f.)) = 1 A 0. This is a major issuc in $3
The simplest definition of a Reproducing Kernel Hilbert Space (RKHS) a Hilbert space (:, of real-valued functions defined on set X such that the evaluation maps d,; H. E are continuous for all X. Our basic example of an RKHS is (22), n/2, in the previous subscction.
As in 31 by the Riesz representation theorem, for each x E X, there exists dz E H such that
[xx)= = & = (dz, S'H-
(34)
Note that this key equation fails if 8g is not continuous. We usually write dz(y) (x v), so
f(x)
(35)
and call KH the reproducing or Mercer kernel of the RKHS. It is elementary to show that H. is an RKHS iff it has a reproducing kernel, a function Kx: X x X + R with (x( E H for all x E X and with (35
Note that
(x y) (K -), (1(1))))1
(36)
where we consider KH: X + H. More generally, consider a feature map Φ: X - H' from X to a Hilbert space Ht and the associated kernel function K(x, v) = (o(x), dly)rr. Then there is an RKHS Hx with reproducing kernel K. In fact, HK is the Hilbert space completion of Span(K(x, : : E X}. As a consistency check, if we start with an RKHS H with associated KH, and set Φ(x) KH (x,-,, we get K Kn and H'' HK = H. we recover the original RKHS.
As with Sobolev spaces, this theory extends to a Hilbert space H of vector-valued functions f: X - Rm. Let me: I + I be projection to the jth coordinate, and set H; = [m; o f: = (f².....')) E HJ. Assuming ò, in
33 is continuous, the Hilbert space H. has a reproducing kernel K,, and we obtain a reproducing kernel
X X by
(K(x, 11x ((Ki(z,)), )) 1 ........ (Kn (") (f²(s),..,, f(x).
(37)
19

---

Published in Transactions on Machine Learning Research (02/2025)
Notation: a kernel is written finite sum ((z,)) a for a, b, H, then (K(z, f())))
is by definition Σ, (z)) J'H H. If the sum is infinite, convergence issues must be treated.
C Proofs for $2
Proof of Lemma 2. (i) Take X E Tatw) t and curve 7(1) c & with 7(0) (())) = X. Then
(Flo) (r((t) d
(X) dt 1m0
(F) (x(t) FF, = (P(VF),
dt
1-0
Since (XF)+)(X) (V(F))), determines V ((Fl 4). conclude that
P(VF) = V (F)4).
(38)
The result follows.
(ii) Take gradient flow line 7(t) for F, so *(t) V V (x)) Then o is a gradicnt flow line of Fla iff
do(+)(?(t)) (Ø (t) V ((Flo)
Substituting i(t) = VEzc)) and using 38 we need to show that
dø V P(VFo(w))
(39)
Take X W. Because & is an isometry from (W,g) to ( g), we have for each w
= (VF, 1 dF( d(F o (x = dF o (do(X)) = (VF, do()));
= (P(VF), del x)
Since WV) spans T this proves 339 and dø( do(FF) V..
Proof of Proposition general, we can compute the projection in a vector space as the solution to the minimization problem: the projection Pv of a vector v onto a subspace B
Pv argmin wEBBle -oalllyyll y112.
If {by} basis of B, we want argmin respect to a', get
ER v Σ 2/b Taking first derivatives of || Σ b1' with
0 = -2(, bi) + 2 Σ a (bi,
For ğ = (āt)) = ((b,, by)), we get at = gار (v) and thus Po = gB bi)b,,,, where (g¹) (Ø.,) For
v = VF, by = 06/0w' we recover
D Technicalities Ignored in 3
In the formal computation of 3 for the sake of clarity in presentation, ignore several nontrivial technical questions.
Firstly, we ignore the nontrivial question of the norm topologies on M and TyM. Instead, we formally assume an easy topology induced by the 2 inner product, as this is common in the machine learning literature. In this topology, the gradient for the MSE loss function does not exist. We use abstract RKHS theory as a setting where this gradient exists, with a Sobolev topology for large enough Sobolev parameter as a concrete example. This highlights the fundamental issue that norm topologies on infinite dimensional vector spaces can be inequivalent, unlike in finite dimensions.
20

---

Published in Transactions on Machine Learning Research (02/2025)
Secondly, the ceffrrential-raadinn formula
11
dl (h) (Vin, M) S. L Eivens N2
is again formal, since lincar functional 6: H R Hilbert space H satisfies E(h) = (w, h) for some H. iff is contimuus. Since we haven't 't specified the topology on M, can't discuss the contimuity of de As in Remark 3.1 for & a discrete loss function like MSE, & is not continuous on 1,2 but is continuous on high Sobolev spaces.
E Proofs for $4
Proof of Lemma In the notation in BB22 we have
d
aL E 1.j
(0s (a,)).
(VLL,, dh), = dt
L(fe(z)).))==
a
(40)
For L. the Euclidean gradient of L with respect to only, we have
Σ 0h. azi (00
(f(z)).))) ((x,)= L( u))' (K2 ),1)'s
= [Ku( (
h)...
Proof of Proposition 4.1
The projection of KH given by
Ky(æ,x') = (Kx(1,
Σ & 81((())))
which proves
18
For
17 write b, = a.//w/ by = a.0//w0. Using the trace map a @ b (a, (a,b)T,,, we get
du (b,, (,)b))T, ( (0/0w) 0/0w0 ,T, = d.atije.
Thus af = n and (2 x²) E, bz((())) (b(z')))))))))) E.o,0/Ow & a.000ww got 0/O0/ @ 0/0w..
Proof of Corollary 4.3
To produce an orthogonal projection and thus well-defined natural gradient from an arbitrary basis {br} T, there are two steps.
The first step is making {b,} an orthonormal basis of the RKHS T i.e.,
(bi,b;) KerK = 6),
(41)
which is guaranteed by Theorem 4.2
The second step is to define a normal space N2 to Ty inside TyM, so that any E TyM can be decomposed into tangential and normal components, i.c.,
== of V. where of ET NN EN
Since T comes with a Riemannian metric ğ (e.g., Sobolev metric), the casiest choice for the normal space N2 = (u MyM: (v,'"); =0 O, Ve' E Now we have a mixed metric on a tubular neighborhood of + in M, which is given by (11 directions langent to and by ğ in ğ-normal directions. We then exlend this metric to all of M by partition of unity. By Proposition 4.1 the natural/pullback gradient under this mixed metric is given by the following kernel form,
Vllys== := PVL = Σ K FLL.
21

---

Published in Transactions on Machine Learning Research (02/2025)
F Proofs for $5
F.1 Proof of Proposition 5.1
25 is equivalent to
(ww Aw!, ...., x-'w ..) "(w', ,...,,',
w" ...., w ,...."))).
(42)
Note that AA-1=1 1, so this is the two-layer analogue of II [[a= 1.
Differentiate 42 at a fixed w" for = 1 by taking (d/dA)lx=1 of both sides. This yields
aΦ
0,1
Wx 0,j = 0
R
(43)
Ows
1=1
100
w
43
is equivalent to
wos 0,j θ
dowo
Σ
w0.1
dówn
0.
(44)
Owct
Ow
Thus the l - 1 vectors
11 =
w 0.1
w 1. 2
: = 2, 2....I(
(45)
Ow
Ow?
are in ker(d)) w Since wff * 0 for all j, these vectors are nonzero. They are clearly linearly independent, since only the k-th vector involves the k-th layer. Thus the span of the Z is contained in ker(d))
Proof of Theorem 5.2 We first show that every element (w (w"....,ur w²) E W lies the orbit through some slice element.
Let a (a),..., az), with
1 1/8
aj
ww'l
|w 101 j I...,,
Then a G and E S.
We now show that each orbit of G meets S at exactly one point. If az wi a² u22 then in the obvious notation,
af = = = at w abx² for all k.
which implics
(X a
of = (
Thus |wf'l Al 2 |w for all kt af of. Since wf, we the same ray, they must be equal. Therefore, wt W2, a= a
F.2 Details of the basis and the restricted metric on slice S
.......,kkk, l; I...., ny - 1} the tangent space Tw for the slice S. polar coordinates (rk, vo ), k; = 1...., ny - 1, on the tth layer, r r: 0/0rk, ci = 0/0y at E R"s The Fuclidean metric restricts to metric on S with metric tensor h = h(k.k.) di c
We take a basis (r :-: & by : k 1....,
s;
22

---

Published in Transactions on Machine Learning Research (02/2025)
On R" (w ..... )"), introduce spherical coordinates (r, 1b ..... ",.,")...,,, 1) by
w r (1
(46)
w
II sin (v0) cos(bb-1)), k 2, 2....,,
n-1 w
sin(
Thus 1 E [0, n| is the pitch angle from (w* w²) to the w=-axis ihn2 [0, the pitch angle from (w+ ...., wn=:===-----_-m 0) to the wn-1 axis, etc., down to vi2 [0, m]. wta E [0, 2m] is the clockwise angle to the wa-axis
(If the usual polar angle in the plane is 0, then 0 + 2 1/2.) Using
0
Ow 0 a Ow 0
Ous
Orie Ow Ör Dr Owk
gel
0
0 Ör
sin (1ber)
Il (1) e=t
cos(N)-1)00 + cos(sb=-1)8 10
t=2
Il (in(
cos(vit) )ow
II sin( Oww
n-1
Og r
sin(}*)) c
II [si((
Ou
s=t+1
+
cos() cos( (2))
II siin(())
II
sin sin( awt
(47)
Amm2
am8+1
n- 1
r
II ==
sin (55)))
Oulti, e= 2......
2
n-2
0 0---1 rccossin
Il sin(w
Ouv +
cos(2b" [cos(un--)oos(( cos((15-11)
II
(in( Out
k=2
cos cos(1 pp 2) --snnmin-n) 10mn.
From 47 it is straightforward to check that
(0,, 0wy) = (0pr, Oyy)) = O, for i i k. (0,, =
(0pi 0p.) 22 sin (v²).
(48)
Then the Euclidean metric/dot product restricts to metric with metric tensor
0,
k + s.
0,
k # 0, ki # 8j:
h h h(kkk)).....))) c
ck ki
(49)
|w 2 T
svn((
A # 0, k; 8j. 0, =i= 0.
l.
Here the metric is computed at (W1.... we) S with wa (rk, v the spherical coordinates of w in the tt layer. Note that wxll r independent of k. h is diagonal, h is easy to compute.
For computational purposes, we need the right hand side of (47 in Euclidean coordinates. Set
1/2
Σ )2
aq =
23

---

Published in Transactions on Machine Learning Research (02/2025)
in particular On
We have
(ag - (w*)²2 ) 1/2
80
cos(16n1))
sin(t)n1)>
On
w
w++1
cos(1/))
== ,...,,
(50)
II. x-/11 sin (169))
On
8= sin (is"( (wf+1)2 1/2
n-
sin(1/ = (1- cos cos
an IT=24 sin(
Using 50
downward induction from f n to gives
e+1
cos(11))
sin(w =
Og
8 1,..., 1.
(51)
06+1
ag+1
Plugging 51 into
447
gives expressions for 0, Oge in Euclidean coordinates. In particular, in
49 we have
the simplification
,² sin sin?( (v .....siin sin?(v oz. +1
(52)
We can also rewrite
47 as
1
0,
w Deyh,
V (2 +. (w)))
1=1
0a1 w² O' w O
(53)
c+1
0gt )2 +...+ (w Ouo
(w])2 + + (w0))0w0t 2, n
F.3 Proof of Proposition 5.2
For u in the slice S, write = (w:1, ,..., w²) with w* = (wt, ... wh.) Since the gradient is independent of
coordinates, on the gth layer, k= 1...., 1, for any F: W R and any w (w ...,,.., w² E W, we have
mg OF à
A Eucck Frge
Owtt Owk
k==1
0,0 OF θ l44
OF a
OF θ
OF a Oakki au k
+ O.k.
h.0
+ h
k,
Öry Ōr
Ork di
Oupii Ōr
OF 0
aF à
Σ (w/t))
Or Ork
ks
0e k.
The right hand side evaluated Wg. Here h when the context. is clear, and have used 52
(arl o,2)) 1, we abbreviate (((kks))(( by hhikky For in the slice S, write (w .- w²) with
24

---

Published in Transactions on Machine Learning Research (02/2025)
w = (wf, ...,,........................... wha) Using mixture of ordinary summation and summation convention, we have
V 11.00
0 Ōr dr
Σ
OF
i
Ouk
k=1
8
OF
i
1x+1
-1 OF a
(ug)²
+
0rk
dr.
k= k-=1
0F
0
OF 0 Ork dr,
Eue,
F
(54)
Or drk
OF
8 Σ
w w OF a luok awff aw
Ow
kg.k==1
wh w
OF a Dwff Owsy
1
In the last two terms, |w = |w0 ,, so it suffices to compute |
Let (2) denote the component of ck
ki 0
of F Label the first three terms the last line of (54) by
autz
1, 11, III. Then
1. In I, (4)
OF
Bwyy
2. In 11,
Kg
1
OF auk
w
independent of k
independent of R;
In III, by switching the 8 and & indices,
kj
OF Wa Ows
ΣV
Wky
1:
independent of k.kg
G Sobolev Natural Gradient
For H R") a Sobolev space (detailed in Appendix B.2), as shown in (14). in order to compute the Sobolev Natural Gradient on W, we have to efficiently compute the inverse of the Soboley metric tensor In.ij, which is analogous the Fisher information matrix Amari's natural gradient. In SG.1 we discuss RKKHS-based approximation of ØN.ij- In SG.2 provide a practical compulational method of ØH.i based
the Kronecker-factored approximation (Martens & Grosse 2015 Grosse & Martens 2016 In S..4 we
report experimental results on supervised learning benchmarks.
G.1 An approximate computation of TH.i
In general, it is not possible to exactly compute ØH.ij approximation. Following 13) if do gradient descent (Dieuleveut et al. 2016 in the RKHS H with
0b for H H.. We therefore resort to
06
Dw?
25

---

Published in Transactions on Machine Learning Research (02/2025)
sequence of data points (11, 32), (2r, yr), under the initial condition fo(z) 0, we get
T
(r(x) = t =
(55)
Therefore, the gradient descent iterate fo lies in the subspace Kx of H. spanned by {KH( )) for given datasct N' = f From the proof of Proposition 2.1 (in Appendix C for any f E H. the projection is Pkxf = Kij (f, (2 m) -),,,)), where Kij- = (K ( ) (x),, We then approximate gH.
follows,
06
do
do
Oō
TH.is
2
Pxx Owr
Pxx Owo
Owt
H
Oip
OΦ
K K(res) :) K(xd.))
K Kab
K(Fra.))
K(ab. 1), K
Owi
Ous"
OΦ Ows
K (Iras))
K-
K K(ves))
Owi
7
OΦ
Oib
(5a) Kad Ows
(3ie)-
(56)
Ow
Note that by setting Ky to the identity matrix, (56) becomess the Gauss-Newton approximation of the natural gradient metric (Zhang et al. 2019ch. For a practical implementation of (56) used in the next subsection, we only need to invert Kab matrix of size B x B, where B is the mini-batch size of supervised training, which is quite manageable in practice.
Now all we need is to compute Kij = (K(z,, K(rj, +) for all 25 E X. When H is the Sobolev space Hs Rosenberg (203) proves the following Lemma that computes Kij = dz, (rz),
Lemma G.1. For dim(M) + 3,
dz,(x) (x) C, + ||x 2;(l),
(57)
where Cn is some constant only depends on n.
Even with (56), exact computation of the pullback metric for neural networks is in general extremely hard. To efficiently approximate (56) in practice, the next subsection, we adapt the Kronecker-factored approximation
techniques [Martens Grosse 2015
Grosse Martens 2016).
G.2 Cronecker-factoree approximation of 5H
Promecker-factored Approximation Curvature (K-FAC) has been successfully used to approximate the Fisher information matrix for natural gradient/ Moadient/Newton /Newton methods. We now use fully-connected layers example
to show how to adapt K-FAC [Martens & Grosse 2015) to approximate Tn in (56. Approximation techniques
for convolutional layers can be similarly adapted from (Grosse & Martens 2016 We omit standard K-FAC derivations and only focus on critical steps that are adapted for our approximation purposes. For full details
of K-FAC, see [Martens & Grosse 2015 Grosse Martens 2016).
As with the K-FAC approximation to the Fisher matrix, first assume that entries of TH corresponding to different network layers are zero, which makes ÿH block diagonal matrix, with each block corresponding to one layer of the network.
In the notation of Martens & Grosse 2015 the f-th fully-connected layer is defined by
Wiā-1 ā = (b(e),
where We (We by) denotes the matrix of layer bias and weights, à; (a 1)T denotes the activations with appended homogeneous dimension, and the denotes the nonlinear activation function.
26

---

Published in Transactions on Machine Learning Research (02/2025)
(60) For Uz to 56 gives
the block of 9N corresponding to the E-th layer, the argument of (arteess & Grosse 2015 applied
9H E a & Ds
(58)
where Ds
0s Sec
@ denotes the Kronecker product, and Ex defined by
Ex X & YJ X(;) (zy))
Just as K-FAC pushes the expectation of LE az 1a"1 @ Ds;Ds; inwards by assuming the independence of a-- and 8y, we apply the same trick to 558) by assuming the following K --independence between a and Sy:
Ex (ac-1a711 @ Ds Ex [ae-1a711 @ Ex [Ds:Ds]]
(59)
The rest of the computation follows the standard K-FAC for natural gradient. Let A and St be the Kronecker factors
A EK a [a=-1a az- ( (z (xy)
(60)
S. Ex [Ds:Ds]] (s( Krj Ds (zz)-
(61)
Then our natural gradient for the f-th layer can be computed efficiently by solving the lincar system,
vec(V,) (Az-1 1 St) = = (A v SE vec(Ve) = vec(S (vA
FH
where VV) denotes the vector form of the Euclidean gradients of loss with respect to the parameters of the 8-th layer. All Kronecker factors As and St are estimated by moving averages over training batches.
G.3 Summary of Sobolev Natural Gradient Algorithm
Lemma
4.1
provides the following basic formula for our proposed Sobolev Natural Gradient,
PVLL = 9H Σ (FLL(((E)).))
OΦ d$ Ow
E Owk
(62)
where L = L(2,)) R" R R is some loss function and LL denotes the standard Euclidean gradient of I. w.r.t. its first input z. Note that summands on the RHS of (62) are just standard Fuclidean gradient of L w.r.t. model parameters w's evaluated on training pairs (Iz, which can be obtained by standard back-propagation. Then the only extra computation needed is the term zkt which is the inverse of the Sobolev metric tensor IN. Applying this extra inverse of metric tensor is the main computational difference of all natural gradient methods vs. standard Euclidean gradients. For instance, Amari's original natural gradient applies the inverse of the Fisher information matrix.
The approximate computation technique for Jn.s is addressed in G.11 given by the following formula,
Io TH.U Jnn
Oo
(Ia) K K-) Dus (x.),
(63)
where Kac- dx., (Ite) can be computed by Lemma G.1 as following,
Kac C,.
na zell (1 + ||,ra
x
(64)
As discussed in G.1 and detailed in $ G.22 in practice, 64 the only extra computation step of our approach compared to the baseline. In fact, if Kac set to the identity matrix, 663 becomes exactly the Geuso-Newton approximation of Amari's natural gradient metric (Zhang et al. 2019c), which is the baseline of all our experiments. This is another reason that we use the Tororc codebase (Wang. 2019) of (Zhang et al. 2019c to implement our algorithm and follow the experimental setup of (Zhang et al. 2019c in our experiments.
Following the notations GG.2 in Algorithm 2 below, we first copy (with slight tweaking for readability) the Algorithm of (Zhang et al. 2019c), which gives the common algorithmic steps of our approach and the baseline, where we highlight red the key step that our approach differs underneath from the baseline, We
then zoom in on the highlighted step and provide computational details of it in Algorithm g
27

---

Published in Transactions on Machine Learning Research (02/2025)
Algorithm K-FAC with weight decay (Zhang et al. 2019ch
Require: 1) stepsize
Require: 3: weight decay
Require: stats and inverse update intervals Tstats and Tine
Init: (Wo) {Se)en> {Ae-1)E= h =0
while stopping criterion not met do k +- ++1
if mod (mod Tstats) then
Update the factors (St)-n {Ae-1)/ with moving average
end if
if k mod 0 (ıod Tino) then
Calculate the inverses VE- [[A--1]-)) 16=1
end if Vi = Vw.L
We t We (1[As-1)--V((ss--1 We)
end while
Algorithm Update of factors {S+) and {Ac-ikE1
Compute: Kij by (64) for all samples of the current training batch
Compute: the inverse K.j
Update: Ar- by 60
Update: S₃ by 661
G.4
Experimental results
Given that comparisons between natural gradient descent and Euclidean gradient descent have been conducted extensively in the literature (Zhang et al., 2019c a) and that our proposed approach ends up being a new variant of natural gradient, our numerical experiments focus on comparing our Sobolev natural gradient (ours) with the most widely-used natural gradient (baseline) originally proposed by (Amari 1998 in the supervised learning sctting.
Since we borrow the K-FAC approximation techniques in designing an efficient computational method for our Sobolev natural gradient, we follow the settings of (Zhang et al. 2019c), which applies the K-FAC approximation to Amari's natural gradient, to test both gradient methods on the VGG16 ncural network on the CIFAR-10 and the CIFAR-100 image classification benchmarks. Our implementations and testbed are based on the PyTorcc K-FAC codebase (Wang) 2019) provided by one of the co-authors of (Zhang ct al. 2019c). We did a grid search for the optimal combination of the typer-parameters for the bascline: learning rate, weight decay factor, and the damping factor. We end up using learning rate 0.01, weight decay 0.003, and damping 0.03 for the baseline through out our experiments reported in Figure 3 and Figure 4 For learning rate scheduling of the baseline, we follow the suggested default scheduling of (Wang 2019 to decrease the learning rate to 1/10 after every 40% of the total training epochs.
For our method, we use the same learning rate 0.01, weight decay 0.003, and damping 0.03 as the baseline throughout our experiments. We only tune two extra hyper-parameters for our method,
For the learning rate scheduling, decrease our learning rate to 1/5 after the first. 40% of total epochs, and decrease again to 1/5 after another 20% of total epochs.
From 57 the Sobolev kernel Kiz (1z) contains an exponential function with the exponent. -||x; 2y I]. as a result, this kernel reduces to identity matrix if the set of points {2a} are not close to cach other. Given that input data points are by default normalized when loaded from these classification benchmarks, we just introduce for our method a scaling factor further down scale all input data. We fix this scaling factor to be 20 throughout experiments.
28

---

Published in Transactions on Machine Learning Research (02/2025)
We also run the bascline under different combinations of these tuned extra syper-parametents shown in Figure and Figure 6 they cannot bring any benefit to the baseline. All experiments are run on desktop with an Intel 19-7960X 16-core CPU, 64GB memory, and an a GeForce RTX 2080Ti GPU.
When comparing the natural gradient with the SGD, (Zhagg et al., 2019ch trains natural gradients for 100 epochs while training the SGD for 200 epochs, highlighting the training efficiency of natural gradient methods. Following the same philosophy, in comparing the two natural gradient variants, further shorten the training epochs to 50 in all our experiments, and we have found that the final performance is almost par with that trained with 100 epochs. The training and testing behavior of ours vs. the baseline CIFAR-10 and CIFAR-100 is shown in Figure 3 and Figure 4 respectively. While the final testing performance of both natural gradient variants are similar, as expected, our method shows a clear advantage regarding convergence speed, with the margin increased on the more challenging CIFAR-100 benchmark.
100
144
baseine
90
1.2
85
Trakiing A Assurayy
Teis: Accuracy
Traniig Pocuray)
1.0
80
78
0.6
0.4
bascline
0.2
ours
0.0
0
20
20
30
40
50
10
20
10
20
30
40
50
Traiing Epochs
Training
Training Epochs
(a) Testing Accuracy
(b) Training Accuracy
(c) Training Loss
Figure Training and testing behaviors of Sobolev Natural Gradient (ours) Amari's Natural Gradient (baseline) CIFAR-10. Results are averaged over four runs of different. random seeds, with the shaded area corresponding to the standard deviation.
haseline
100
baseline
Durs
ours
Test Aocurayy
traiiing Pocunacyy
80
Training Aocuracy
60
30
10
20
30
50
Training Epochs
Trainin Epoons
Training Cpuchs
(a) Testing Accuracy
(b) Training Accuracy
(c) Training Loss
Figure 4: Training and testing behaviors of Sobolev Natural Gradient (ours) Amari's Natural Gradient (baseline) on CIFAR-100. Results are averaged over four runs of different random seeds, with the shaded area corresponding to the standard deviation.
H Pullback metric vs. pushforward metric: an example involving flatness
Our GTK theory specifies a Riemannian metric on the function space which makes the standard gradient, at least locally, as "natural' as the natural gradient in capturing the intrinsic structure of the parameterized function space. In this section, we compare the pullback metric and the pushforward metric from a purely mathenatical viewpoint, emphasizing that the pullback metric has better compatibility properties. For example, suppose W C W", and Φ: W M extends to Φ': W M. Then the pullback metrics behave
29

---

Published in Transactions on Machine Learning Research (02/2025)
100
baseline
1.2
baseline
Test Aosureyy
Training Racunety
Traking Ansuay
1.0
baselhe_ir _r_sall
bassiin__aan
0.8
0.6
basoling
beseline
baseline_i Ir
08
baseline
0.4
66
beseliee scell
beseline scole
0.2
00
baseine scale
50
0.0
10
20
30
40
10
20
Training Epochs
Training Epochs
Traning Epochs
(a) Testing Accuracy
(b) Training Accuracy
(c) Training Loss
Figure 5: Training and testing behaviors of Amari's Natural Gradient (baseline) on CIFAR-10. baseline_ Ir is the baseline method with the same learning rale scheduling as our method. baseline scale the baseline method with the same imput scaling as our method. baseline_ Ir scale is the baseline method with the same learning rate scheduling and input scaling our method.
100
3.5
beselne
3.0
Accusayy
TraniingPccuaary Acturacy
80
Trainng Ancuacy
2.5
beseline, st_scale
beseline scale
50
2.0
Traning
80
baseline
hasclinc
30
baseiin
baseino_li Ir
1.0
baselin scale
basclino scale
0.5
baselie_J J'_oolle
baseline Ir scale
0
20
30
40
50
0
10
20
80
50
10
20
50
Training Epochts
Treining Epocits
Traning Epochs
(a) Testing Accuracy
(b) Training Accuracy
(c) Training Loss
Figure 6: Training and testing behaviors Amari's Natural Gradient (baseline) on CIFAR-100. baseline Ir is the baseline method with the same learning rate scheduling as our method. baseline_ scale is the baseline method with the same input scaling as our method. baseline Ir scale is the baseline method with the same learning rate scheduling and input scaling as our method.
well: ( = 6"9 for a metric on M. In contrast, there is no relationship in general between (d').gg and 0.GE.
A key point is that pullback metrics are independent of change of coordinates on the parameter space W, while pushforward metrics depend on choice of coordinates. As simple example, if we scale the standard coordinates (x) "......""" x*) on W Ig* to (y² (y'....,y") = (2x'...., ..... 22x"), then in (1) the pushforward metric 2 (80//Oy') (1/4)ø.9E-
0.9E = E, (06/00z)))) ((00/))))0000""))))) of the Euclidean metric on W changes to Σ, More generally, if a²x) is a change of coordinates on W given by diffeomorphism a, the pushforward metrics for the two coordinate systems will be very different.
In contrast, the pullback metric well behaved ("natural' in math terminology) with respect to coordinate changes. If g is a Riemannian metric on M, then (0oa)" g a (0)gg) In particular, for v, w tangent vectors at point in W, we have
(v, () = (da(v), da (w))
(65)
Since v in the x-coordinate chart is identified with da(v) in the coo)-coord chart, (65 says that the inner product of langent vectors to W is independent of chart for pullback metrics.
For the rest of this section, we give a concrete example machine learning where these propertics of the pullback metric have an explicit advantage over the pushforward metric. In particular, we consider the various notions of flatness associated to the basic setup 1 near minimum. In (Dinh et al. 2017 55),
30

---

Published in Transactions on Machine Learning Research (02/2025)
c-flatness of F: W + R near minimum wo is defined by e.g., measuring the Euclidean volume of maximal sets S C W on which F(w) < F(wo) + for all w in S. The authors note that their definitions of flat.ees are not independent of scaling the Euclidean metric by different factors in different directions. Here the pushforward of the Euclidean metric is implicitly used. As above, scaling is the simplest version a change of coordinates of the parameter space. Since this definition of flatness is not coordinate free, has no intrinsic differential geometric meaning.
This issue is discussed in (Dinh et al. 2017 Fig. 5. Under simple scaling of the z-axis, the graph of function f: R E will "look different" Of course, f is independent of scaling, but the graph looks different to an eye with an implicit fixed scale for the z-axis, i.e., from the point of view of a fixed coordinate chart.
We instead obtain the following coordinate free notion of flatness by measuring volumes in the pullback metric ø'g.
Definition H.1. In the notation of 7 the e-flatness of F near local mindmum wo is fg dvole's where S is a maximal connected set such that F(S) C (F(wo), F(wo) + c)-
Since Is d Ss dvolo-g fo(s) dvolg: for any orientation preserving diffeomorphism of W, e-flatness is independent of coordinates on W, and can be measured on S or on M. In contrast, the Euclidean volume of S is not related to the volume of ((5) C M.
MLP with ReLU Activation General Case with Bias Vectors
In Appendix E, we discuss the gradient for an MLP and with no bias vectors. In this Appendix, we give the necessary changes for an MLP with bias vectors.
The setup is MLP with & layers with ReLU activation function, with parameters w in the k-th layer, for r= I,..., n&. where is the number of matrix entries plus the number of components of the bias vector in the k-th layer. The space of parameter vectors is W = {wo (wi.... wh,),, } where each parameter in I. Thus dim(W)- = E. ne.The The map 6: W + Maps([P", M) takes a parameter to the associated MLP.
We abbreviate the subvector for the k-th level by w (w...... wh so w (w'...., w²). We further refine w = (wf,... wh wf,.. wh: := (WAL+ Whe), where the matrix Aş in the w layer has dg entries and the bias vector by has C entries. We may have = 0 for some k. We can assume that w + 0 for all k, since a MLP with some w = is trivial.
Define multiplicative group = {g = (g1...., 96) gi > O. 9: 1). The map G - R (gi,...,
. De) H
Σg loggi gi is bijection to the plane E, x² 0, so G manifold of dimension f- 1. Then G acts on W by
(g1....,gt) 9t) (w !u....,,") w²)
((g) (w], ..., w ),9) (w....., W
(ya (w,..... w2z) 1.922 g192 (w) (w...... w )),
(66)
(gc-1 (wt f ..... w w ), 0102 ge-1 (wt ..... w6-1)
(ge (wf w (wf..... ... [.,)
(e
By Dinh et al. 2017 Appendix B),
o(w) ........""" ww) )
ő(g (w+ ....,.... w²)),
(67)
for all g E G. Thus, in 15.1, Φ is far from injective, and in fact is not immersion.
31

---

Published in Transactions on Machine Learning Research (02/2025)
Proposition 1.1. The dimension of ker(dō) at each w W is at least l - 1. Thus ø is not immersion. In faci, the span of
711 Zy w
Gr
U
(68)
Ju
Owi r-1
Kr
w
Arz
1i--
w
M
is in See(dd)ww.
Proof. For fixed
(2,...,() O:
take 91 0, gr
91 and all other gi 1, and then take (d/dgi)lm== of
both sides of 25 This gives
Ob
dΦ
k1
k2 2
(69)
Ounz
k2
0m,-1
Owe
k==1
kr-1=1
k,=1
where the first sum is over all the neurons in the first layer, and the other sums are over the bias vector neurons. 699 is equivalent to
02
dø
Σ
2
kr Ow,
0.
(70)
Owl
Owz
k=1
kz=1
1r-1=1
R,=1
Thus the l- - vectors in Tw
62
0
i (71)
Z, w
ki
k1
+
2
2
Ow -
W
Rr
k-
A=
are in ker (do) Since Wy * 0 for all j, these vectors are nonzero. They are clearly linearly independent. Thus the span of the Z, is contained in wer(d))w".
As Appendix F] the Z, at form a basis of Tw O by a dimension count. Since e not an immersion for ReLU MLPs, as in $3.2 the natural gradient does not exist, and we must find a slice for the group action. Theorem I.1. Write W = R x " A slice is given by
S = (X(r....... ue): (UI....,.) E S x Sm2 ... S x Sue, A 0} x R
where Sk-1 is the wnil sphere in R& and R&* R$ (0).
The orbit O = (g- w g E G} is contained in (and probably equals) a level set of 6. It is unlikely that any nonzero vector in S lies in ker(dø) Since dim(T...)) + dim(Spant [) = dim(W), we expect that ker(d)) = Sp in Prop. L.I so that Φls is an immersion.
Proof. We prove that every element w = (Wi.... we) E W lies on the orbit through some slice element. Write wk (Ag, by) with A = (Wk,1,..., Wkka)) ) by (wartl...., Wna), where ag the number of entries of the matrix in the hth layer. Set.
of 1(A1, o/2 |(A2, ((2)))(b2)1+
P: = ((Ae-1, 2i22dla b(be-1)) g-= -Aell1 (72)
Then
(oh...., 96) ((A1, (Ae, br))
Al gf b1), A2 22 (2)) ... (g6 Ag-1. G:- ..... (at.a.....ggbbbll).
(73)
E smn S ... x s x S x a
Set
9k =
1/6
(T
Then g (01)...:9t) E G, and g w E S.
32

---

Published in Transactions on Machine Learning Research (02/2025)
For &layer MLP with ReLU activation function, the slice S in Theorem L1 has codimension l - 1 in W, contains all the information in Φ: W , and probably has øls an immersion.
c & & 1...., C k; = -..... ng - } of the tangent space S for the slice in Thm. I.1 We use spherical coordinates (rk, 1/
$5.1 suggests three approaches to treat the non-immersion case. We take a basis
()) ), h; = ,.... n 1, on the 4th layer, for h = =1, I,...: l 1, while for the eth layer we use (re, (re, 4 ,,... 0r, and rectangular coordinates (x)) ,...),) on Rb,,* Then = "E 0/0xx; ch 0/00k k= I...,, - of 0/00// 4 = L.... dgi of 0/0x li as + I...., ne. product restricts to a metric on S with metric tensor
,.. The Euclidean metric/dot
h = ((R,A)))))))) ck c.
(74)
0,
for l # 8, for # 0, ki # 8j.
0,
|wx/ Iallssinn sin (1/ sin' (1x 20 sin (vbket)
for + 0, E; kę 85
|ww/ 2 sin 2 (1b)-1) sin (v/e sin (vb l 1.
for
kū = 8j 5 at:
for L, k; =8 ab
e,
for A: 0, = j= =0
Here the metric is computed / = =().. h is diagonal, h is easy to compute. (74) follows from (48
(wi,..., we) S. Note that | r independent of (except for
We can use this metric in any of Approaches I-III of $5. The implementation of gradient descent in any approach similar to $5, although moving a parameter vector back into the slice is a little more complicated, as we now explain.
To implement gradient descent on W, start with w E and take ul w + a dø- (V Le ((ww)) )), where a is the step size. Move w² along its G-orbit back to w as follows: for w ((A), b1) ...., (Ae, be)), define o by
,..., ń ((A1, 32 1(A2, (/b2) ....., d- ((Ae-1, P22de gibe-1)l O |Ael-1
and set
= ((g'A 9(b1), A2 b2 .., (d sibe-1), (. (S.AA,
(75)
where (I continue with w
1/8
This ensures both that (Boh ..., Bgt1) E G and that be not modified. Now
1.1 Loss function gradient: Approach I
In particular, for Approach we neod:
Proposition I.2. The gradient S F, for F: W + R. is given by
oF wke Tw
i Wk, Dw
6-1
OF
i Ow
S F F.e. F... w
w
wkk
ww
2
Dws.
where
k # l,
1/g
Ok:
k = l.
33

---

Published in Transactions on Machine Learning Research (02/2025)
Proof. For w in the slice S, write (W1, ................................ we) with wk [w] ....,... ) Since the gradient is independent
of coordinates, on the hthh layer, kt = I,..., l 1, for any F: W - R and any w = (WI...., wx) E W, we have
BB OF
a
Buck Far
=1 Owf awkk
OF a
OF a
OF a
h OF a
10.0
+ O
+
Ork Ōrk
Örk dusti
Ouiti Ōrk
Duth d K
OF a
BR-1
+
OF 0
Ork Ork
The right hand side evaluated at wh. As before, 00.0 = (O.m. 0,,)= 1, we abbreviate (k..)) by baik when the context is clear, and we have used 52 For l, we use spherical coordinates (re, 48) on the first ag components and rectangular coordinates on the bias vector components. This gives
or 8 Ow
Euc,& F
owl.
1=
0F
ar-1
14+1
OF a
OF
a
+
Σ
(uo))
+
Σ
Ôr
Ore
Des oui
awt, oul
f==1
4=
To compress the notation, set
{
Tkk- k #l e k E.
T/g
Oke
Using a mixture of ordinary summation and summation convention, we have
SSP = =1.000 a Ör dr
OF Σ
Rel OF
Σ OF à
+
Oy
auk
oy
dul
dwf ow
R=1
(-aet1
OF
h++1
or
nee
or
r-1
(wt)22
Σ
Ork
Òr.
Owf. Ow!
k= k==1
OF ①
OF 0
W
+ 1--
Ork Ork
Or Or:
k=1
k=1s=1
F
wi w OF 8 ky kg
w
ky
w
83
oF
a ow;,
(76)
12 Owh, Owy,
w
Iwo
Out
In the last two terms, |
|wo
so it suffices compute |w*|.
Let (2) denote the component of ko of F Label the first three terms on the last line of III. Then
76 by I, II,
1. I, (K))
oF Bun
2. In II,
kj
n
OF du k.
w
wky
3:
independent of k.k,
independent of
34

---

Published in Transactions on Machine Learning Research (02/2025)
3. In III, by switching the and it indices,
OF Ow
wkk
=
k
w
#=1
independent k:,k
1.2
Loss function gradient: Approach
We show that in Approach II, formulas [4 and (66) are altered only by removing the é-1 factor in the last. term. For (54). we note that the metric in 49 replaced by ((k.k)))(..g))) Thus the proof of Prop. 5.2 changes to
V Bueck Fur OF a
OF
Or Ōrk
Re. Jek
du
and
OF
0
OF a Ousy dukk
Σ Ork
w
+
k=1
k==1
② Ork Ōrk
OF
0F
Wue.WF
Örk Òr, k=
b=
wx. wk, OF 0
By
w 2
OF a Ouf Owy,
Euc.
1 Ow aw
|
ke,ky=1
In the notation of the proof of Prop.
5.2
(22) are unchanged for I and II, while for III
hi
oF
Wky.
w 12
Ow..
The same argument works for
76
35


Reviews:
[{"review_id": "zWaBRQp9LP", "reviewer": "Reviewer_uCC2", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n- The paper is fairly well-written and clear and the idea of GTK seems novel.\n- An effort is made to include all the necessary background in the appendix (which I however, due to time constraints, could not read in its entirety)\n\n# Weaknesses\n\n- More discussion of the literature is would be helpful, for example explaining differences to the ideas of Bai et al.'s \"A Geometric Understanding of Natural Gradient\" (arXiv:2202.06232) or Bastian et al's \"Sobolev gradients and neural networks\".\n- The extra computation cost of $O(n)$ seems not insignificant, and it would be good if the authors could comment a bit more on this issue.\n- I appreciate that numerical examples were conducted, however the advantage of the proposed method as shown in Figure 2(a) seems negligible. The paper would benefit from more numerical experiments or different ones, where differences become more apparent or a discussion on why they are not.\n- I do not understand why in the very beginning of section 2 you equip $\\mathcal{M}$ with a Sobolev or Fréchet topology (which seems a reasonable choice to me), but then in section 3 choose the less natural $L^2$-topology to show non-existence of natural gradients. In other words, if one would instead choose the Sobolev of Fréchet topology or consider these functions as distributions, do natural gradients of empirical losses exist and thus no RKHS approach is necessary? I think at least for the topology induced by viewing the functions as a subset of (tempered) distributions, gradients should exist by the very argument given in the paper.\n\n# Questions\n- In theorem 4.2. should $X$ be instead be $M$? If not, what is $X$, is it a finite set?\n- In Lemma 2.1 (i), what is the \"*flow* of $P(\\nabla F)$\"? Do you mean a gradient flow?\n- In section 2.2. what is meant by \"flow lines for $\\nabla F$ *do not project* to flow lines of $S \\circ \\nabla F$\"? Are you claiming that there is no projection mapping one to the other?\n- In Proposition 4.1, why is $T_f \\Phi$ a RKHS?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The problem at hand is parametric approximation in spaces of functions using gradient flow to decrease some loss functional.\nThe authors introduce the generalized tangent kernel (GTK), which generalizes both the neural tangent kernel (NTK) and natural gradient. The GTK induces a Riemannian structure on the whole space of functions under consideration. It is shown that in many cases, natural gradients of empirical losses do not exist. Analysis of the case that the mapping of parameters to functions from a parametrized set is not injective is also presented, since this case appears frequently in practice.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Apart from the smaller issues I raise below, I think the most important requests are comparisons with similar ideas from the literature, writing the first part of the proof of the central Theorem 4.2 a bit clearer and addressing the last bullet point I listed under weaknesses.\n\n# Small issues and typos\n----\n\nAbstract\n---\n-\"under function approximation\" seems strange to me, choosing a different preposition (e.g. for?) will certainly make this sentence easier to understand. The same goes for the first paragraph of the introduction.\n\nIntroduction\n----\n- In Figure 1, it should be $\\mathscr{W}$ instead of $W$ to match the formatting of the caption.\n- I would instead write \"with the notable exception of van Oostrum\".\n- the term \"functional gradient\" is used without definition and I think (but I am not sure) that it is meant to mean the same as gradient.\n\nSection 2\n----\n- It should be noted somewhere that the Einstein summation convention is used, since not everybody in this community will be familiar with it. However, only in equation (18) it is not used?\n- In the last sentence before subsection 2.1 there is full stop missing at the end of the sentence.\n- In the sentence before equation (2), I do not think that \"Riemannian metric\" is the correct terminology, since the manifold is infinite-dimensional. \n- I do not understand the sentence \"(6) is the pushforward $\\phi_* g_E$ of the Euclidean metric\", how can that formula be $\\phi_* g_E$?\n\nSection 3\n----\n- Even tough no confusion can arise, I would prefer if the summation indices would not be (partially) omitted in formulas (9) - (13), (16).\n- In the new paragraph after equation (8) it should be \" a curve $(f_t)_{t \\ge 0} \\subset \\mathcal{M}$\" instead.\n- In the paragraph after remark 3.1. you first write \"metric on $\\Phi$\" and then \"metric on $T_p \\mathcal{P}$\", which is inconsistent.\n- I would remove the word \"easily\" from the last sentence in this section, because it is redundant.\n- In the example on the bottom of page 5 about the Fisher-Rao metric assumptions on $M$ are missing. If $M$ is not a compact Riemannian manifold in $\\mathbb R^d$ without boundary, then e.g. the identification of the tangent space is not true, see section 3 of Lafferty's 1988 paper \"The Density Manifold and Configuration Space Quantization\".\n\nSection 5\n----\n- In 3. at the beginning of the section, \"Corollary\" is misspelled.\n\nSection 6\n-----\n- \"Euclidean\" is misspelled.\n\n\nBibliography\n----\n- Names of conferences (like NeurIPS) and Journals (like Information Geometry) should be capitalized.\n- The book title by Schmüdgen is missing an 's' at the end.\n\nAppendix B.1\n----\n- When defining functions, I recommend using \\colon for so that the spacing is correct, e.g. $f \\colon M \\to N$ instead of $f : M \\to N$.\n- In my opinion, sentences should not be started with a mathematical symbol (this is not the only place where this occurs, e.g. after equation (1)), which is why I would advise to start the third paragraph by \"A (Hausdorff second countable) space $M$ is a Riemmanian...\".\n- \"a slight perturbation of the standard dot product\" is not very precise, I do not really understand how to understand this.\n- After \"It follows that any tangent vector\" it should by $T_M$, not $T_m$.\n\nAppendix B.2\n-----\n- Two formulas before equation (28), the last differential should be $\\xi^n$, not $x^n$.\n- Before equation (29), the term \"isomorphic\" is misspelled and $\\text{d}\\;\\text{vol}$ is only used there, which is not consistent with the rest of the section and paper (compare e.g. to equation (8)), especially since the Sobolev spaces that are considered two sentences earlier are on $\\mathbb R^n$, not $\\Omega$.\n\nAppendix B.3\n----\n- I recommend getting rid of \"except on a countable set\", since there are sets of Lebesgue measure zero that are uncountably infinite an more generally, one does not have to choose the Lebesgue measure anyway.\n\nAppendix C\n----\n- I think the terminology \"gradient flow\" for the a curve $\\gamma \\colon (0, \\infty) \\to X$ is more common than \"gradient flow line\". \n\nAppendix I\n----\n- I suppose it should be \"function space\" instead of \"functional space\" in the first sentence.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "K01ccWf8wV", "reviewer": "Reviewer_K4VW", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\nThis work answers interesting questions about Natural Gradients. \n\nThe paper is well organized. It introduces problems such as the ill defined gradients, and then provide solutions studying the problem on right spaces. It also provides an analysis of how to perform natural gradient descent for a concrete example (MLP with ReLU activation function), along with an experiment evidence of the fundness of the method.\n\n\n**Weaknesses:**\n\nThe paper is very dense, and not always easy to follow. Also, there are lot of things in appendix, which could maybe help the reading (e.g. some backgrounds...).\n\nA lots of arguments are only formal.\n\nThe experiments seem a bit preliminary and could be improved.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a formal study of natural gradient. First, it describes how to compute natural gradient using projection on the right tangent space. Then, it is shown that for classical settings in Machine Learning, natural gradient are actually not well defined (e.g. for empirical losses using maps on L2). Thus, it is proposed to study natural gradient on RKHS, where they can be well defined. This study motivates the introduction of a general kernel called the generalized Tangent kernel. Finally, the authors study the practical problem of optimizing a MLP with ReLU activation function, and address the problem of the non injectivity of the immersion using a sliced technique.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "While very interesting, some parts of the paper (such as Section 3) feel a bit like a draft with formal arguments for the computations.\n\nIn Section 5.1, it could be nice to add an Algorithm summarizing the procedure for the slice SGD (instead of the Pytorch code in Appendix, which should be in supplementary material).\n\n\n\nTypos:\n- Page 6: $\\langle d_x, h^j\\rangle_{\\mathcal{H}_j}$", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "SgB58Rxeqm", "reviewer": "Reviewer_9Mtm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- Main strength: The paper deals with fundamental mathematical aspects of (natural) gradient flows, and is therefore of broad potential interest. \n\n- Main weakness: The paper is extremely dense, and is not easily accessible to someone  without deep working knowledge of Riemannian manifolds, RKHS, and Sobolev spaces. I found it difficult to follow for the most part.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper demonstrates that standard gradient descent captures the intrinsic geometric structure of the function space in function approximation, rather than the parameter space. The authors introduce a new mathematical object, the Generalized Tangent Kernel (GTK), and show that the GTK defines the Riemannian metric on the function space.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "None of my requested changes are crucial for securing my recommendation for acceptance, but I think they would help strengthen the work. \n\n- As mentioned above, I think this paper could be of potentially broad interest, but it is not presented in a user-friendly way. I understand that the results are highly mathematical in nature, so this might be difficult. But I think interspersing simple, down-to-earth examples throughout the paper and clearly highlighting main results would help the reader along. This is not crucial for securing my recommendation for acceptance, but it would help strengthen the work. \n\n- It could be interesting to connect your results with the very recent findings of [Shoji et al (2024)](https://arxiv.org/abs/2409.16422), who show that when an optimizer decreases a loss function, this implies that the iterates of the optimizer can be put into natural gradient form. Is it possible that a similar converse result holds for GTK, in the function space? This would be very interesting.\n\n- It would be nice to connect this work to recent work on in-context learning (e.g., [Anh et al, 2023](https://arxiv.org/abs/2306.00297))--do your result apply to this setting as well? Or only supervised learning? \n\n- \"It is even more interesting to investigate the relation between choices of metrics and the generalization ability, which we have not touched in current work.\" This question has been studied in [Kozachkov et al (2023)](https://openreview.net/pdf?id=Sb6p5mcefw), with connections to NTK in section 4.3.\n\nIt would be great to see the above three papers (and related works) mentioned/discussed!", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper has been well-received by all the reviewers, and the revised version has successfully addressed their concerns. However, a minor issue remains regarding the lack of practical applications or examples in the current version. Please include all the promised changes.
Decision: Accept (poster)