Over-parameterised Shallow Neural Networks with Asymmetrical Node Scaling: Global Convergence Guarantees and Feature Learning
Authors: Francois Caron
Fadhel Ayed
Paul Jung
Hoil Lee
Juho Lee
Hongseok Yang
Venue: Accepted by TMLR ()

Abstract: We consider gradient-based optimisation of wide, shallow neural networks, where the output of each hidden node is scaled by a positive parameter. The scaling parameters are non-identical, differing from the classical Neural Tangent Kernel (NTK) parameterisation. We prove that for large such neural networks, with high probability, gradient flow and gradient descent converge to a global minimum and can learn features in some sense, unlike in the NTK parameterisation. We perform experiments illustrating our theoretical results and discuss the benefits of such scaling in terms of prunability and transfer learning.

Published in Transactions on Machine Learning Research (02/2025)
Over-parameterised Shallow Neural Networks with Asym- metrical Node Scaling: Global Convergence Guarantees and Feature Learning
François Caron
caron or.
Deportment of Statistics University of Orford, United Kingdom
Fadhel Ayed
fudhel.ayed@gmailcoom com
Huawei Technologies
Puris, France
Paul Jung
peul.jjngg@gmail.comm com
Department of Mathematics Fordham University, USA
Hoil Lee
hoil. lee waa.aaaaamiiiiiite ac.kr
Samsung SDS
South Korea
Juho Lee
juholee@kaist.oom
Kim Jaechal Graduate School of AI KAIST, South Korea
Hongseok Yang
sygggeeokkkymmg@aaassk.. aet.siistss............. oc..r
School of Computing KAIST, South Korea
Reviewed on OpenReview: https: // openreview. net/ forum? id- S5
Abstract
We consider gradient- based optimisation of wide, shallow neural networks, where the output of cach hidden node is scaled by a positive parameter. The scaling paramcters are non- identical, differing from the classical Neural Tangent Kernel (NTK) parameterisation. We prove that for large such neural networks, with high probability, gradient flow and gradient descent converge to a global minimum and can learn features some sense, unlike the NTK parameterisation. We perform experiments illustrating our thcorctical results and discuss the benefits of such scaling in terms of prunability and transfer learning.
1 Introduction
The training of neural networks typically involves the minimisation of non-convex objective function. However, first-order optimisation methods, such as gradient descent (GD) and its variants, often find solutions with low training error. To gain better understanding of this phenomenon, one fruitful direction of research has been to analyse properties of GD training of wver-parameterised, large-width neural networks; that is, neural networks where the number m hidden nodes in a given layer very large. In particular, under a /1/m" scaling of the hidden nodes, Jacot et al. (2018) have shown that, as the number of nodes m tends infinity, the solution obtained by GD achieves zero training error, and coincides with that of kernel regression under so-called limiting Newral Tangent Kernel (NTK). Under the same node scaling, called NTK sealing,

---

Published in Transactions on Machine Learning Research (02/2025)
quantitative theoretical guarantees for the global convergence and generalisation properties have then been obtained for large (but finite) width neural networks (Du et al. 2019b Oymak & Soltanolkotabi 2020 Arora et al. 2019a; Bartlett et al.) 2021). However, it has been noted in a number of articles (Chizat et al. 2019 Yang 2019 Arora et al. 2019a; Yang & Hu 2021) that under NTK scaling, feature learning does not occur and GD training is performed in a lazy-training regime, in contrast with the typical feature-learning
regime exhibited deep neural networks.
Main contributions. We investigate global convergence properties and feature learning in gradient-type training of barg-widlth feedforward neural networks (FFNNs) under more general asymmetrical node scaling. In particular, cach hidden node j = =,....,m! has a fixed node-specific scaling v Vmm. with
1 + (1 Am = 7
(1)
1=1 AR
2 0 are nonnegative fixed scalars with Σ3=1 X3 = 1. Note that = 1 corresponds to the vI/m NTK scaling. If y < 1, the node scaling is necessarily asymmetrical for wirth networks. Two typical examples of the scalars (Xz)121 are (a) X3 6n 2--2 for all j 2 1, and (b) \1 = = Ax = 1//K and X; = 0 for all j > K, for a fixed K.
where 7 [0 and 1 12 \1 12 2
We consider a shallow FFNN with a smooth activation function and without bias, where the first layer weights are trained via gradient flow or descent and empirical risk minimisation under the Eg loss. We show that, under similar assumptions in Du et al. (2019b a) on the data, activation function, and initialisation, when the number of hidden nodes m is sufficiently large: (1) if y >> O, the training error goes to 0 at linear rate with high probability; and (ii) feature learning (in the sense of the definitions given in Section 7.1 occurs if and only if y < 1. We provide numerical experiments which illustrate the theoretical results and demonstrate empirically that such node-scaling is also useful for pruning and for transfer learning.
Organisation of the paper. Section 2 discusscs related work. Section 3 introduces the FFNN model with asymmetrical node scaling, gradient flow or gradient descent updates, and the main assumptions on the data, activation function, and initialisation. Section 4 discusses the properties of the NTK of such model at initialisation, and its infinite-width limit. Sections 5 and derive our main results on the convergence to a global minimum of gradient flow and gradient descent and sketch their proofs. Section 7 gives the main results regarding feature learning. Section S describes our experiments on simulated and real datasets, whose results illustrate our thcoretical results and their potential applications. The Supplementary Matorial contains detailed proofs, as well as results on additional convergence of gradient flow and feature learning when using the ReLU activation function.
Notations. For an integer 2 1, let [n] {1.... n). For multivariate real-valued function f R" - R, 3f
the gradient [v) is the n-dimensional column vector of partial derivatives fvf(vv (ōv, (v), (v)) af (v)) where = (V),.... un) For square matrix B, we denote its minimum and maximum eigenvalues by eigmin (B) and eigmax (B), respectively. For vector v E E2"", we write B diag(v) for the r-by-n diagonal matrix with By U, for E [n].
2 Related work
Large-width FFNNs. The analysis of large-width FFNNs goes back to Neal (9966) who showed the
connection between Gaussian processes and FFNNs in the large-width limit. Recent work has explored this
connection under varying assumptions (Matthews et al. 2018 Lee et al. 2018 Yang 2019
Favaro et al.
2020 Bracale et al. 2021 Lee et al. 2023 Jung et al.] 2023
Large-width FFNNs under NTK scaling, Following the seminal work of Jacot et al. 2018 a number of articles have investigated the benefits of ovar-paramcterisation for gradient descent training, with the "1/vm" mm NTK scaling (Arora et al. 2019b] Du et al.) (2019b a Lee et al., 2019 Zou & Gul 2019 Oymak &
Soltanolkotabi] 2020:
Zou et al.] 2020). Crucially, when the width of the network is large enough with respect.
2

---

Published in Transactions on Machine Learning Research (02/2025)
to the size of the training set, the training loss converges to a global minimum at linear rate under gradient flow or gradient descent. However, under this symmetrical NTK scaling, the hidden-layer features do not. move significantly when the width is large, and the scaling has been coined lazy-training regime for this reason [Chizat et al. 2019 Woodworth et al. 2020
Large-width FFNNs under mean-field scaling. An alternative scaling is the "1/m" mean-field scaling [Rotskoff & Vanden- Lijnden, 2018
Chizat et al. 2019 Sirignano & Spiliopoulos 2020 This scaling is also equivalent, up to the so-called abc-scaling symmetry (Yang & Hu 2021p, p, to the uP parameterisation of Yang & Hu (2021) in the case of
Mei al. 2018 2019
Ghorbani et al., 2020 Chen et al. 2021 Tao al.
2021
shallow networks. Feature learning is known to occur under this mcan-field scaling. Also, Chizat & Bach
(2018) showed that under the same scaling, if the training of the model converges, it converges to global minimum.
Asymmctrical scaling in FFNNs. The idea of using asymmetrical scaling parameters in the context of GD optimisation of deep FFNNs has been previously introduced by Wolinski et al. (2020). The focus of Wolinski et al. (2020) was the (empirical) usefulness in terms of pruning. Indeed, our experiments in Section S are also in line with their findings. The work of Wolinski et al. (2020), however, only considered asymmetrical scaling with y = 0 (no fixed part), and did not investigate global convergence properties under such scaling. The properties of random FFNNs under random asymmetrical node scaling in the large-width limit has also been considered by Lee et al. (2023; but this paper did not investigate the training properties under gradient flow or gradient descent.
3 Problem setup
3.1
Statistical model
We consider shallow FFNN with one hidden layer of m hidden nodes and scalar output. To simplify the analysis, assume that there is no bias term. Let x E R be some input vector, where d 2 1 is the input dimension. The model is defined as
fx W) = Σ (Z W)) with Z,(x: W)
vd WJx x for [m]]
(2)
where fm (m; W) is the scalar output of the FFNN, (x: W) is the pre-activation of the j-th hidden node, σ E - IR is the activation function, W, E Rd is the column vector of weights between node of the hidden layer and the input nodes, and aj R is the weight between the hidden node j and the output node. The Am,'s for j E [m] are nonnegative scaling parameters for the hidden nodes. The parameters to be optimised are contained in W which is an md-dimensional column vector (w (w.......... We assume that a admits a derivative d'.
WT)
For n 2 1, let a: : R" (resp. a' R" - be the vector multivariate function that applies the real-valued function a (resp. o') element- on each of the n input variables. To simplify the analysis, we assume throughout this article that the output weights (ay)je(m) are randomly initialised and fixed afterwards:
id Uniform((-1111) j21. ag
(3)
This simplifying assumption is often made when analysing large shallow networks (soc c.g. (Du et al. 2019b; Bartlett et al., 2021)), and the analysis can also be extended to the case where both layers are trained. The scaling parameters (Am.))g(m) are fixed and parameterised as in Equation (L). By construction, Am.1 > 0 and Σド= Am.j = for all m 2 1. Recall that the case = corresponds to NTK scaling. Also, note that our model covers finite FFNNs: when
have
==0 and X, = 1/K fje
otherwise
3

---

Published in Transactions on Machine Learning Research (02/2025)
for some K < m, the model becomes finite network of width K. In the experiments, we will consider the special case where (Az)521 are the probability masses of Zipl law:
1 2 ((1/a) ju/a j21
(4)
for some a E (0, 1), where ( is the Riemann zeta function. The parameter a tumes how quickly X, decreases with j, smaller values corresponding more rapid decrease and more asymmelry.
3.2 Training
Let D, = {(x, yt) he(n be the training dataset, where n 2 is the number observations. Let X denote the n-by-d matrix whose ith row is x We want to minimise the empirical risk under l loss. Consider the objective function
(m(W) 2
Σ UI - fx W))²
5
which is non-convex in general. For a given dataset Dn, width m 2 1, output weights (ay)je(m) and scaling parameters (emm)) we aim estimate the trainable parameters W by minimising Imm ( (W) using gradient flow or gradient descent. Let Wo be some initialisation. In gradient flow, (Witss is the solution to the following ordinary differential equation (ODE):
dW, WwLm(W)) dt
with lim W = Wo. Let Wy be the value of the parameter w, at time and define Zız(x) Zz(x: W.). Note that VwySm(x: W) o' W)) Jax. Under gradient flow, for j E [m],
dwij_ = dl
E(or (Ui (x;; Wi)) Ww, f' (x;; (x;; W.)
Vimmaa
Σ (v) fm Willo (4)))
Vd
Thus, the derivatives associated with each hidden node j are scaled by VAm.. For input Rd, the output of the FFNN therefore satisfies the ODE
dfm (x; We)
We) (Ww.f((xWW.Ta
dW. dt
11
(.)) W.)
Σ
dt
where Om Rdd Rd + RR the neural tangent kernel for the network fm( W):
Om(x, x'; W)
o (Zz(x; (2)xx
(6)
1=1
The associated neural tangent Gram (NTG) matrix 6. (X; W) is the n-by-n positive semi-definite matrix whose (i,j)-th entry Om (xi, Xj; W). takes the form
Am diag o' Xw
( XX diag o'
Xw
( (X
W
(7)
d
Gradient descent is discretisation of gradient flow. Under gradient descent, the parameters are updated by
- Vm (W.-(W for all s E N,
(8)
where 77 > 0 is learning rate. These updates give rise the family (W (e)s)))( indexed by discrete time steps 0, 1,2,..... rather than continuous times / ' 0.
4

---

Published in Transactions on Machine Learning Research (02/2025)
3.3 Main assumptions
Throughout the paper, assune that the activation function satisfies the following standard condition: for all random variables Z ~ NO0 s²) for some > 0,
[Z)]] < 00 and
Ez (Z)) < 00
This assumption is made all the time, and so do not mention its use explicitly in the paper. The results of this article on global convergence and feature learning use several further assumptions. The first set of these assumptions, which are mild and similar to other assumptions used in the literature, is on the training dataset D, = {(x,, 94) (e)m
Assumption 3.1 (Dataset). (a) All inputs are non-zero and have norms at most 1: 0 < ||x S for all 121 (b) For all * /' and E R, x; # eX.. (c) There C > such that lya S for all
The next assumption concerns the activation function a. Standard activation functions (softplus, tanh, sigmoid, swish) satisfy this assumption, but not the ReLU. However, some of our results, such as global convergence of gradient flow and feature-learning results, also hold in the ReLU case, as shown in Appendix in the Supplementary Material.
Assumption 3.2 (Activation function). The activation function is analytic with [o'(z)) S 1 and [o"(x) s M for some M > 0, but it is not a polynomial.
The last assumption, which standard, is on the initialisation of the weights. ild Assumption 3.3 (Initialisation). For j E N, Woj ~ N/ Ia), where Ia the d-by-d identity matrix.
4 Neural Tangent Kernel at initialisation and its limit
Mean NTG at initialisation and its minimum eigenvalue. Let Wo be random initialisation from Assumption 3.3
Consider the mean NTK at initialisation
xxx E
o
o' (1
E (x, Wo)]
Wo1x
(10)
Vd
The mean NTK, which is also, by the law of large numbers, the limiting NTK under 1/ vm scaling (Jacot ct al. 2018 does not depend on (Am))s21 nor m. Let (x) = O (X, Wo)l be the associated n-by-n mean NTG matrix at initialisation, whose (i,j)-th entry is O*(x;, Let K, cigmin ( (x)X)) be the minimum eigenvalue of the mean NTG matrix initialisation. This minimum eigenvalue plays important role in the analysis of global convergence properties in the lazy-training regime. Many authors (see e.g. (El Karoui 2010: Nguyen et al.) 2021h) have shown that, under some assumptions on the data, activation function, and initialisation, Kn is strictly positive or bounded away from zero. such a result, under the
Assumptions of Section 3.3
(QDu al. 2019b Theorem 3.1) and Duu el al. 2019a Proposition F.1)). When Assump- tions 3.1 3.3 hold, we have Kn
Proposition 4.1
Remark 4.2. Du al. (2019b make the assumption that each x; has unit norm. But their proof holds a).
under the less strict Assumption a.. The above proposition also holds if Assumption 3.2 is replaced by the assumption that is the ReLU function.
Limiting NTG. To give some intuition, we now describe the limiting behaviour of the NTG, for fixed sample size as the width m goes to infinity. The proof, given in Appendix Gin the Supplementary Material, follows from the triangle inequality and the law of large numbers, together with [o'(2) S 1 and Σ Aς 1.
Proposition 4.3. Consider a sequence (Woj)s21 of iid random vectors distributed as in Assumption 3.3 Suppose Assumption 3.2 holds. Then,
θ (x; - 2x (X:
(11)
5

---

Published in Transactions on Machine Learning Research (02/2025)
almost surely as co, where Ôx (X; Wo) - (1 (X; Wo), with 0 (X; Wo) being the
(x;
following random positive semi- remi-definit matriz:
@ (X; Wo)
1
X diag
0' Vd ))
XWoj
xX
diag o'
(
Xwoj
))
(12)
d ==1
vdl
Also, O0 (X; Wo)] = E 2 X; Wo)l = (X)), and
E | X: Wo) ã. (X)) Co( -)))
7)2
(13)
.21
where || || denotes the Frobenius norm, and Co(X) > O is a positive constant equal to
Σ Sn
X5' d
Var 0' (Vammixs))
)
0' 00( (/aoox))))
When 7 =1 (NTK scaling), the NTG converges to a constant matrix, and solutions obtained by gradient flow coincide with that of kernel regression. Whenever 7 < 1, Proposition [1.]shows that the NTG random initialisation, even in the infinite-width limit, contrary to that of NTK scaling, As shown in Equation (13) the departure from the symmetric regime, as measured by the total variance of the limiting random NTG, can be quantified by the nonnegative constant (1 21- Σソ X3 [0, 1]. When this constant is close 0, we approach the kernel regime, and increasing this value leads to departure from the regime. The quantity Σ5²1 X3 E (0, 1] is always strictly positive. More rapid decrease of the Xy as j increases will lead to higher
values of Σソ'1 Xz. For example, when using the Zipf weights in Equation (4), we have ΣØ21 X3 which decreases with a, as shown in Figure [ in the Supplementary Material.
C12/a) 2(1/a)
This section has described the behaviour of the NTG at initialisation in the infinite-width limit, and has provided intuition on the node-scaling paramcters. The next three sections contain results on global convergence and feature learning properties of large, but finite, FFNNs under such asymmetrical scaling.
5 Global convergence analysis for gradient flow
Our global convergence theorem, which is given below, explains what happens during training via gradient flow. Recall that K, is the minimum eigenvalue of the mcan NTG matrix @*(X) initialisation. Our theorem says that with high probability, (i) the loss decays exponentially fast with respect to Kn and training time t, and (ii) the NTG and weights Wij change, respectively, by
n3 Em m.i
n² VV 5= Xmmss
V
and (wwy-wool 0 ( Kin d
n2 1/2
@ X; w,) (X; Woll2=
(
k22 B3g2
Define
[
C2 27
C = sup EswA((,1) cE(0.1)
(14)
Va
Theorem 5.1. (Global convergence, gradient flow)
Let 8 E (0,1). Suppose Assumptions 3.1 to 3.3 hold, and
thot.
22 log
22 C2 2
M
>> 0, and m 2 max
Kidd
where C is the bound on the yi in Assumplion 3.1 properties hold for all t 0:
Then, with probabilily at least 8, the following
(a) cigmin (x: 2 1"; (b) Tim (W.) (1m (mn

---

Published in Transactions on Machine Learning Research (02/2025)
(c) I[a, Woyll - "VAms [2((2)(11 for all j E [m]; K, V 7-8
(d) (X; W.) - mm (X; Wollka
(2'"MP(C44CC)) Na 1200000
Σド=1 XmJ)
(
(2n ki,d9/61/2
Ei X
The above theorem implies that, whenever 7 O, the training error converges to 0 exponentially fast. Additionally, the weight change is bounded by a factor Vmm and the NTG change is bounded by a factor
دー= Am... We have (see Appendix B.2 in the Supplementary Material) that as m + 00,
Am.j + (1 - 7) for all j21 and Σ mj - (1-7)² Σ X3.
If X > 0 (note that we necessarily have xn 0), the upper bound in (c) is therefore vanishing in the infinito-width limit if and only if 7 (lazy-training regime). Similarly, the upper bound in (d) is vanishing if and only 7 1. Although we were not able to obtain matching lower bounds, we show in Section 7 that feature learning arises whenever y << 1.
Remark 5.2. We make two comments on Theorem 5.1 First, although d represents the input dimension, all the occurrences of d in the theorem, such as those in the lower bound of the width m, do not come from the complexity of the input dimension. Instead, it comes from the fact that our model uses the 1/ \/ scaling when computing the pre-activation values of the first layer. If this scalining were removed in our model, the statement of the theorem would not include d (i.e., we would have the theorem with d set to 1). Second, a result similar to Theorem 5.] also holds for the ReLU activation function. See Theorem A. in the Supplementary Matcrial.
Sketch of the proof. We give here a sketch of the proof of Theorem 5.1 (and of Thcorem A.1 its ReLU counterpart, in the Supplementary Material). The detailed prools are given in Appendices and in the Supplementary Material, with secondary lemmas given in Appendices D and there. The structures of the proofs of Theorems 5.1 and A.1 are similar to that of (Du al. 2019b Theorem 3.2), which showed analogous results on the global convergence under NTK scaling. However, there are some key differences which we highlight. below.
Gradient flow converges to global minimum of the objective function if the minimum cigenvalue of the NTG matrix 2m [X; W.) is bounded away from zero, for m sufficiently large, by some positive constant for all / 2 O. In the NTK scaling case (>= 1), sufficiently large, with high probability: (i) the NTG matrix at initialisation is close to its mean, and the minimum cigenvalue is close to that of the mean NTG, (ii) the weights Wyy are nearly constant in time, which implies that (iii) the NTG matrix is nearly constant in time, hence (iv) the minimum eigenvalue of the NTG matrix at time t is close to its value at initialisation, which is bounded away from zero.
)), Du et al. (2019b) showed that the following is satisfied, for m
However, in the case of asymmetrical node scaling (7 < 1), none of the points (i-iv) holds. At initialisation, the random NTG matrix may be significantly different from its mean. Additionally, both the weights and the NTG matrix substantially change over time. This therefore requires a somewhat different approach that we now describe.
Let X 7/m and ( -01--(1111)))) ((1>)\X)/EX) 7e. and note that x 6,0) be the r-by-n symmetric positive semi-definite matrices defined by Equation 17, with Am, replaced by either (m) Note that
X Amy- For i E {1,2), let
(x (X; W.) = O X: W. + + (x:
(15)
with B (X;
-x)
The key idea of the proof to use the above decomposition of the NTG matrix as sum of two terms, and to show that, while the second term may change over time, the first term is close to its mean at initialisation, and does not change much over time. The important points of the proof are follows. For large m, with
7

---

Published in Transactions on Machine Learning Research (02/2025)
high probability: (1) @ (X; Wo) is close to its mean (x*(X) and its minimum eigenvalue therefore lower bounded by (yxm)/2;
The important points of the proof are as follows. For large m, with high probability: (1) @ (X; Wo) is close to its mean (X) and its minimum eigenvalue is therefore lower bounded by (7%m)/2; (ii) while the (X;
weights W. may change significantly over time, O (X; Wt) remains nearly constant over time; (iii) as a result, the minimurn eigenvalue of 6 (X; Wx) can be lower bounded by (7%m)/4; (iv) this implies that the is lower bounded by (7Kw)/4.
minimum eigenvalue of the overall NTG matrix @m
Since showing that the first part of the NTK/NTC in Equation 15 does not change in the limit of m - oo is a key component of the proof, we give here an outline of it in the simplified case of and a being smooth. Let
2 (x, x'; W.)
xx' 7
m Σe of (Wa)x)o'(wii>>
the first part of the NTK. We have, over gradient flow,
d
(1) (x, x'; W:)
m
xo" (Wex))o((w)x'() x' (Wi,x)o" dwei
xx' m7
dt
dt
jml
Y
dwej
< 2M
dt
where the last inequality follows from the triangle inequality and Assumptions 3.1 and 3.2 Furthermore,
dwej dt
v /mm
Σ (v: (x; Wi) (Ze,((()
I v Am.
1Bs - fm Wil
v 2m (We).
where the first inequality follows from the triangle inequality and Assumptions 3.1 and 3.2 and the second inequality follows from Cauchy-Schwarz. The change in the first part of the NTK thus bounded by quantity involving 2 Σラ= Amm.j. Under the scaling in Equation (1), 2 EE= /
Am. = + 0 as + (see Appendix B.2 in the Supplementary Material). Hence, the change in the first part of the NTK/NTG becomes asymptotically small as the width m increases. It is worth noting that for the full NTK, i.e. the sum of both parts in Equation (15), a similar derivation leads to an upper bound of the order of E,=
(2 E Am.s) + +1 1-1 / A X m + co, unless 7 = (symmetric case).
Due to the second term, this quantity does not converges to 0 as
6 Global convergence analysis for gradient descent
Let y E Rn be the vector of outputs (yI....., yn) in the training datasct D, {(x,, (i),::)))e()): For each gradient-descent step 8, let u, T be the outputs step based on the inputs in Dn, that is, u, (f.m (x; W.)....,fll.. (x,, W.)) The following convergence theorem intuitively says that if the learning rate 77, of gradient descent, is sufficiently small and the width of the network is large enough, then with high probability, the training error of the network decays exponentially fast to O. We remark that our node scaling is independent of the learning rate 7). Thus, the only conditions on in this work are in this section where we extend our results from gradient flow to gradient descent.
Theorem 6.1. Consider d E (0, 1). Suppose Assumptions
3.1
3.3 hold, and that y > O, and
2 7Kn
dkndd
?k 28//2
0 < <) < min
8n?
29/2
C1)1/22
8

---

Published in Transactions on Machine Learning Research (02/2025)
where C' and C1 are from Assumption 3.1
and Equation 14
Let B (1 77/k1/2) If
m 2 max (2 am
M2 +()
M ci)
knd
kn -336
11 3
then with probability at least 1 --
6,
ly - =1- u for all (N)
(16)
Note that the condition the learning rate requires n, 0 (()mn/n²)) Thus, the best possible convergence rate from the theorem is (1 (177kn/2)) (1 (Cor² /////>) for some constant Co-
The proof is by induction on the gradient-descent step 8, and is described in detail in Appendix H in the Supplementary Material. It is similar to the proof of (Du et al. 2019a Theorem 5.1), but the two proofs differ significantly because, as in the case of gradient flow, the weights Waj and the Gram matrix θm (X; W.) change during gradient descent, in our case, while they remain nearly constant in the case of Du et al. 2019a)
7 Feature learning analysis
In this section, we present some results about feature learning. We focus on node scalings (Am of the form in Equation [ both asymmetrical (y < 1) and symmetrical (y = 1; that is, Am.j = 1/m),, but we also discuss alternative parameterisstions such as mean-field and jP. For shallow neural works trained by gradient descent, mean-field and pP parameterisations are equivalent (Yang & Hu 202I), and correspond node scalings Am.j 1/m² and lcarning rate 1) = Noo for some f/u > 0. We start with some definitions of feature learning in the context of potentially asymmetric scalings, generalising existing definitions. We then present. feature learning results first under linear activation function, and next under general nonlinear activation function.
7.1 Definitions
Definition 7.1 (Feature learning). Let (Woj)521 be a sequence of random initialisations for nodes j 2 1. We will say that feature learning occurs during training if[
Σ Am (a W.) - -()(x) Wol))² >0
lim inf
(17)
000000
Eis Am, (o(Zx(x; Wol))
almost surely, or in probability, for some E (0, oo] and x. Here t = oo refers to the case that the ratio in the above incquality is the limit as f tends to oc.
The left-hand-side quantity in Equation feature maps between time 0 and time Definition 7.1 matches the definitions of (Yang & Hu 2021 Definitions 3.5, H.2, H.9) or (Frei et al. 2023 node scalings.
corresponds to the relative change in the scaled m-dimensional
Proposition 3.2) in the of symmetrical NTK mean-field
Remark 7.2. As already noted by Yaang & Hu 2021 Remark H.10), Definition 7.1 is a relatively weak notion of fcature learning. It only requires a change in the fcature map for some t (0, oo] and some x, and does not relate the relevance of the learn features for prediction. However, we show empirically in Section S that the feature learning property leads to better performances in terms of prunability and transfer learning.
The previous definition ensures that a change occurs in the feature map. However, may still be the case that the contributions from all the individual nodes remain asymptotically infinitesimally small, such a
"In Yang & Hu 2021), feature learning defined in the largh-width limits for stable and nontrivial parameterisaatons. Noutrivial mcans that the ncural uetwork function fr is not constant in time. Stable means that both the preactivations and activations have (1)) coordinates initialisation and O(1) coordinates throughout training. Both of these properties are satisfied by our model.
9

---

Published in Transactions on Machine Learning Research (02/2025)
way that there are no nodes representing important features in the network. This is problematic if one is interested in pruning the nodes of the network, we show theoretically for linear activation in Section 7.2 and empirically in Section
We introduce below the stronger definition of non-uniform feature learning.
Definition 7.3 (Non-uniform fcature learning). Let (Woj)521 be sequence of random initialisations for nodes 2 1. We will say that non-uniform feature learning occurs during training if
maxjeim] Am.J (o(2)(x; W.) (x) Woll) >0
lim inf
(18)
m0000
Eks Am, (Zxx(( Wo))
almost surely or in probability, for some & E (0, xo] and x.
Note that non-uniform feature learning implics feature learning, but the converse does not hold. For instance, feature learning holds under the mean-field parameterisation, but not non-uniform feature learning (see the next subsection for an illustration with a linear activation).
7.2 Linear activation function
We now describe, with the following theorem, analytic results in the case of a linear activation function. Although for fixed second-layer weights, the NTK does not change this linear-activation case, the evolution of the weights provides useful insights into the differences between the symmetrical and asymmctrical scalings in terms of weight change.
Theorem 7.4. Assume that the activation function (xx z, i.e., the identity map. Let (Am.s) be some mode scalings, not necessarily of the form 1 Let X UDVT be reduced singular value decomposition of X, where U is n x k matrix with orthonormal columns, D is diagonal k x k matriz, V is d x k matrix with orthonormal columns, und k S min(n,), d) is the runk of X. For all j E [m], the difference belween the solution of gradient descenty 33 /flow woot and the initialisation Woz is given by
Vmms a (3 VVT
W Wooj
Woj
Σx m Am,k
(19)
where Bo -E Wo and B Vd UUVD-UTT y is the minémumm-norm solution of
1
argmax ly
X,112
vd
The learrt function is
(m W (xx)
1 Vd
T
m v Am Woos)
x (B +Ta -Vl-VVT),, VVT).))
(20)
Vd
The proof of this theorem is given in Appendix 1.1.1 in the Supplementary Material.
Theorem 7.4 says that along the dimensions spanned by the data, the weight vector of a node j moves by a
quantity proportional to /mm towards the minimum-norm solution. The form of the learnt function Equation 20 implies that the contribution of each hidden node j to the function's output is proportional / Amm.ja, a.a.Wao. soj. The next theorem analyses the asymptotic behaviour of this contribution in the infinite-width
limit and the associated feature lcarning propertics. It shows that, under the scaling of Equation 11). both
feature learning and non-uniform feature learning occur if and only if < 1.
Theorem 7.5 (Feature learning linear activation). Assume the setting of Theorem 7.4 and that Assump-
toon
3.3 holds.
NNote that assume of Assumptions 3.1
3.3 here.
with step-size less than D22 2 E, Am,k) where Daax is the largest entry D.
10

---

Published in Transactions on Machine Learning Research (02/2025)
Under the node scalings (1) both feature learning (Definition 7.1 and non-uniform feature learning (Definition) 7.3) hold if and only if 7 <.
Under the meon-field scaling, feature learning (Definition 7. holds, but non-uniform feature learning (Definition 7.3) does not.
The contribution from the jth node is marginally normally distributed with
Ams _ Boor Amis (Ia Am,k
(Ia-Innma vVT)) for all m 2 and 1 S m.
Am.d
Wa N/
Ex>mm Am.
This implies that, under the mean-field scaling, contributions are asymptotically vanishing. Under the scaling have
Vmmja,Wao Woj ooj - N (a-:((11 (1- (1 7)7 (1 (4-
as m7 + 00.
(21)
As result, if < 1 and N3 0, the contribution from the jth node to the output is non-vanishing in the mnfittte-witth limit.
Before moving on to the case of the nonlincar activation function, we analyse the consoquence of pruning nodes linear network based on the scaling parameters Amj- The result of this analysis is given in the following proposition, which shows the benefit of our asymmetric scaling in pruning.
Proposition 7.6. Assume the setting of Theorem 7.4, and that Assumption 3.3 holds. Assume also the node scaling 11 Let p 1). Consider the following pruned network that is obtained by keeping the [pm] hidden nodes with largest scalings Am. and pruning the other nodes:
[
(x W Wx =
Σ aj
Doj
Σ
a,
Vd
Vd
Then, for all 8 > O, we have the following bound on the pruning error:
fm Woo) >s) s evd ((1B_)+
(
||>||
Pr (fmmo(x; Woo) - (x))
Σ
Ama
+
Σ 1'> Lom
Amis
5' Lame
Since Ej>lom Ami + 7(1 p) as m + xc, the above bound implies the following. If = (no symmetric part), then Pr(1Fm.p(x: V Woo) fm (m W Wx)l > E) + 0 and so the network can be compressed to smaller network vio prowing. Otherwise, the pruming error controlled by the proportion of the symmelric part y in
the infinite-wide width limit:
||x|| lim Pr ([ Wor) fm W Woe) >)) s cvd ((1120) + va) 111-1111111 p)+ Vdx(11 pP)
172 12-900
7.3 Nonlinear activation function
We now analyse feature learning in our model when the activation function is noulinear. Our analysis assumes the following two changes in our setup: Assumption 7.7 (Zeroed initialisation). We assume the model has the following form:
172
fx W) = ( (()3))((((xxW))
Σ
((, (x))))
That is, we subtract from the original model, a duplicate version whose parameters are set W, and are unchanging throughout training. This is commonly used simplification in theoretical analyses of neural networks. It ensures that at initialisation, the model satisfies
Jm (x; Wo) for all x d
11

---

Published in Transactions on Machine Learning Research (02/2025)
Assumption 7.8 (Random outputs). We regard the outputs .1..... Un as random variables, that the probabilities in Equations 17 and [18 refer to the randomness of the outputs as well. We further assume that y1..... are independent and continuous (i.e., the distribution of y/ has density with respect to Lebesgue measure), and that they are also independent from Wo and the ay's. Note that we still treat the imputs #1,..., Tn as deterministic variables. This assumption met if, for example, there exists a true generating function f such that y/ = (x)) + (((z)) where the (xxa) are independent continuous noise variables.
Our analysis on feature learning considers gradient descent with a learning rate that does not depend on m just we did in Section 6 This is in contrast to parameterisations such as mean-field parameterisation where the learning rate has scaling dependent m. However, let us note that key reasoning steps in our proofs also apply to such m-dependent learning rates after minor modifications, allowing to recover existing feature-lcarning results as we will explain shortly after Theorem 7.9 The full proofs of all the theorems in this subsection are given in Appendix [ in the Supplementary Material. Also, the theorems in this subsection have counterparts that hold for the ReLU activation function. Appendix A.2 in the Supplementary Material contains those feature-learning results for the ReLU casc.
We show that if the activation function is continuously differentiable and its derivative is always positive, in the case of sigmoid, then after the first gradient-descent step, (i) both feature learning and non-uniform feature learning in Definitions 7 and 7.3 occur (Theorem 7.9 and (ii) the squared norm of each weight vector Wj changes almost surely by the amount ((A1) in the infinite-width limit (Theorem 7.10p.
Theorem 7.9. Suppose that Assumplions 3.1
7.7 and 7.8 hold. Suppose also that 7 < and that the for all x E R. Let i E [n]. Then, both feature learning and non-uniform feature learning occur after the first gradient-descent step with respect to the input x4 in the almost-sure sense, i.e., the following inequalities hold almost surely:
activation function a continuously differentiable with o' (x)
Ej Ami (x( Wi)) - o( Woll)² 0
lim inf m-000
EE= Am. (a(Z)(x; Wo))
maxje[m] Ama (x; Wi)) (x, wol))² 0.
and
lim inf m-+00
Σラ= Am (o(Z Wol)
Sketch of the proof. The denominator Σ (o, Wo)))² in the condition for non-uniform feature learning converges to positive finite value almost surely as m tends to co. Thus, it is enough to prove that
Since non-uniform feature learning implies feature learning, prove the former only.
lim inf max Am. (x(Z)(x; (1)) - (Z) (o))))) 0 almost surely, m0000
[E(m]
which is implied by
inf (o(Zi(x; W))) - (Z) Wo))) > 0 almost surely.
(22)
Note that the limits from above are not redundant since W₁ depends on m. The sufficient condition in Equation (22) can be simplified further. The assumptions of the theorem allow us to use the inverse function theorem to deduce that the condition in Equation (2) holds whenever
lim inf (x,; Wi) (x 2 > almost surely. 00
(23)
12

---

Published in Transactions on Machine Learning Research (02/2025)
The majority of the detailed proof concerns proving the condition in Equation the following m-independent lower bound: for all m,
23
To that end, compute
x x/xx/ ))
(xxx Wi) (x;; Wo))
7m2m1
Σ
T2
ys' (o'
Woixe X
(24)
d2
Vd
X X; ))²
r/2 7)
Σ Ui' (o'
Woi X
(25)
d²
vd
Then, we show that the right-hand side is almost surely positive, which implies the conclusion of the theorem. The justification of this almost-sure positivity relies on Assumptions 3.1 and 7.8 the positivity of o', and the assumption that y < 1. Concretely, the assumption on y implics that ry 7) > O, while
Assumptions 3.1 and 7.8 and the positivily of o' imply that the squared sum on the right-hand side is almost. surely a positive random variable.
We point out that the proof also reveals that if were to allow the learning rate ni to depend m in particular way, then feature learning could occur even for y (although the condition for non-uniform feature learning might fail). For instance, when v and = 1, Equation (24) still holds, and its
right-hand side is almost surely positive because 1/² // 1/d > 0. This almost-sure positivity and the assumptions on a then imply that
(i(x(Z Wi)) (Z) W wo)) 0,
(26)
which in turn ensures that the condition for feature learning holds, i.e. almost surely,
Σj Am. (n.s(a(Zz)( Wi )) (x; Wo)))² limm
Σド= (o W (Z)(xx Wo)))?
lim inf 10000
ΣØ=1 Am (Z,
lim 000 Tr o Wo))2
B[(a(Zi(x; W - (x; Wo))) > 0.
Elo (oo Wo))
However, note that in this case, by (Yang & Hu 2021 Theorem 3.3), the model becomes unstable and could blow up (cf. the footnote regarding Equation (17)). )- Also, if 1 and were to use 1/m² for Am.j (i.e., mean-field parameterisation), the essentially same argument would apply and lead to feature learning, although such Am.'s i's are not covered by our setup. Just as before, Equation (2 would hold and its right-hand side would be almost surely positive in this case, because 7 /dd = 1/d > 0. This almost-sure positivity implics the inequality in Equation (26), which then gives the condition for feature learning because the lim inf formula. in the definition of feature learning can be simplified to the same ratio of expectations as the one before.
Theorem 7.10. Suppose Assumptions 3.1
3.3 and 7.8 hold. Suppose also that the activation function o is continuously differentiable and satisfies a (x) for all x E R. Then, for all j, the following holds almost surely:
lim Tmminf|| |/ Vwi, ((We) (1 22
ooj X;
w of
Tn 70-000
yile
x
X
Vd
o'
Vd
(27)
In particular, if < 1 and Xy > 0, then the lower bound is positive almost surely so that we have
liminn L(W.)) 0
with probability 1.
Sketch of the proof. The lower bound in Equation 27 is obtained by a relatively straight.orward calculation since the calculation of the gradient is simplified due to Assumption 7.7 The second part of the thcorem follows from the fact that when conditioned on Woj: the absolute value on the right-hand side of Equation (27 is a continuous random variable that depends only on the y ', and it strictly positive with probability one.
13

---

Published in Transactions on Machine Learning Research (02/2025)
L=
0--77
1-02
3.24
0.4, 0.0
10
5-
%
3.20
Figure 1: Results on simulated data. From left to right, 1) training risks, 2) differences in weight norms [[ Wayll with the j's being those neurons which have maximal differences at the end of the training, 3) differences in NTG matrices, and 4) minimum eigenvalues of NTG matrices.
8 Experiments
We use here a (smooth) swish activation function (22) = We obtained quantitatively similar results with the ReLU activation function; Appendix in the Supplementary Material.
8.1 Simulated data.
8.1.1 Illustration of the main results
We first illustrate thoory simulated data R
We generate 100 observations where for E [n], x; is d dimensional and sampled uniformly on the unit sphere, and y/1 _ Σ sin(xx;,5) + e; where
e; ~ 1), 1). We use the FFNN of Section 3.1 with the swish activation function, - 2000 hidden nodes, and Am.j as in Equations T and For each setting, we run GD with a learning rate 1.0 for 50 000 steps, which is repeated five times to get average results. We summarise the results in Figure T which shows the training error and the evolution of clear
We consider the four values (7, a) E (0.5,,,,,) (0.2, (0, 0.4)].
the weights, NTG, and minimum eigenvalue of the NTG as function of the GD iterations. We correspondence between the theory and the empirical results. For 7 > 0, GD achieves near-zero training error. The minimum eigenvalue and training rates increase with y. For y = 1, have the highest minimum cigenvalue and the fastest training rate, but there is no/very little feature learning: the weights and the NTG do not change significantly over the GD iterations. When < 1, there is clear evidence of feature learning: both the weights and the NTG change significantly over time; the smaller the 7 and a, the more feature learning arises.
8.1.2 Improved generalization: Single ReLU experiment.
To illustrate the benefits of asymmetrical scaling, consider here the scenario where the function to learn is single-unit ReLU, setting known to be challenging for the lazy-training regime (Malach et al. 2021 Consider the following data-generating process:
X Na
Y = awo x)
where we (1, d and is the ReLU activation function. In this experiment, sample a training dataset of n 100 samples in dimension d = 10. As previously, train fully connected neural networks composed of single hidden layer, with different node-scaling strategies. The width in all models is set to P = 2000. The generalization error is computed on 5000 samples from the same single ReLU data-generating process. All experiments are repeated 5 times, the training and testing datasets are resampled for cach run. The results are reported in Figure? 2 We notice that at the end of the training, the error near zero for all scalings. Examining the error on the test set, one can see that, as expected, the symmetrical y NTK scaling generalizes poorly in this sctting. The asymmetrical scalings, the other hand, perform significantly
4The code found httpp////ghhiiiccc//a/eeeeeeaaaa
14

---

Published in Transactions on Machine Learning Research (02/2025)
0.025
standard init (torch)
standard Init (torch)
0.5,g 0.2
a: 0.0,1 0.2
0.4,y
044,1 0.0
iid (7 1)
0.3
iid (2-
Training MSE
0.015
Test Test MSE
0.2
0.010
0.005
0.1
0.000
0.0
1
3
Iterations
1e3
Iterations
1c3
Figure 2: Results on simulated data from a single ReLU unit. Evolution of the training error (left) and test error (right) as a function of the training iteration.
better, illustrating the benefits of this strategy in terms of generalization error. For completeness, we also report the results using the standard pytorch initialization. We can see that in this setting of a sparse generating process, the standard initialization also generalizes poorly compared to the asymmetrical model.
8.2
Real data
8.2.1
Regression
We also validate our model on four regression datasots from the UCI repository concrete ((n, d) = (1030, 9)), energy ((n,d) d) (768,8). airfoil ([n,d) (1503, 6)), and plant ((n, d) = (9568, 4)). We split each dataset. into training (40%). test (20%), and validation sets (40%), and use the validation set test transfer learning. We use the same parameters as for the simulated data, but train our FFNNs for 100 000 steps in each run. To further highlight the presence of feature learning in our model, we test the transferability of features learnt from our networks as follows. We first split. the validation set into a held-out training set (50%) and a test set (50%), and extract features of the held-out training set using the FFNNs trained on the original training set. Features are taken to be the outputs of the hidden layers, so each data point in the validation set is represented with a 2000 dimensional vector. Then, we sort feature dimensions with respect to feature importance measured as (Am.) |ey [| [[?)))[m and use the top-k of these to train external model. The chosen external model is a FFNN with a single hidden layer having 64 neurons and ReLU activation, which is trained for 5000 GD steps with a learning rate of 1.0. Our theory suggests that smaller ry and likely lead to better transfer learning. A subset our results appears in Figure 3 see Appendix K for additional results. In line with the simulated data experiments, we observe a stronger presence of feature learning, in terms of weight-norm changes and NTG changes, for smaller values of y and a. Also, we observe that models with smaller values of y have lower risks when small number of features are used for the transfer. The interpretation is that those models are able to learn sufficient number of representative features using relatively fewer neurons.
8.2.2 Classification
We apply our model on two image classification tasks. The first is small-scale using the setting assumed in our theory, while the second is larger-scale using a more realistic setting. In addition to the transferability experiment described before, we test the prunability of the FFNNs. We gradually prune hidden nodes with small feature importance and mcasure risks after pruning. Feature importance measured as above. Our theory suggests that models with smaller y and/or values are likely to be more robust with respect to pruning, as long as 7 1. Wolinski et al. 22000) had similar empirical findings the benefits of asymmetrical scaling for network pruning when 7 = O.
'https://archive.. ics. uci
15

---

Published in Transactions on Machine Learning Research (02/2025)
7.40
v 101
60
9.38
y 10
07,
9.25
9.34
05, 200
2.0
3.22
20
3.05
3.28
3.50
0.3
5.4
3.6
0.0
1.5
0.0
24.
08
1.0
0.2
044 08 Kenstore cre
10
25
60
Aautsma
Paure COTRETAAO
Figure 3: A subset of results for the regression experiments. From left to right, 1) training risks for dataset concrete, : 2) differences in weight norms [[we worll with the j's being the neurons having the maximal difference at the end of the training for dataset energy, 3) differences in NTG matrices for datasel airfoil, and 4) test risks of transferred models for dataset plant.
y:7 0.5
o- 05, y 0.2
10-
03
0.3
B4,,
0.
02
0.6
0.
0.5
0.0
3.1
02
100
540
1000
3500
2030
60
100
Prodins
126
Mamber prined modes
Featum dimndo
Figure 4: A subset of results for MNIST dataset. From left to right, 1) training risks, 2) differences in weight norms, 3) test accuracies of pruned models, and 4) test accuracies of transferred models.
We take subset of size 5000 from the MNIST dataset and train the same models used in the previous experiments. We also test pruning and transfer learning, where we use an additional subset of size 5000 to train an external FFNN having a single hidden layer with 128 nodes. To match our theory, instead of using cross-cntropy loss, we use the MSE loss by treating onc-hot class labels as contimuous-valned targets. The outputs of the models are 10 dimensional, of the outputs. In general, we get similar results in line with previous experiments. The pruning and transfer learning results are displayed in Figure 4 For other results, see Figure 6 in the Supplementary Material.
MNIST.
compute the NTG matrices using only the first dimension
CIFAR. We consider a more challenging image classification task of CIFAR-10 and CIFAR-100. The datasets have 60000 images with 50 for used training and the rest used for testing. There are, respectively, 10 and 100 different classes. We show the benefits of asymmetrical node scaling hold for this more challenging problem. In many applications, fine-tuning or transfer learning to adapt it to the task at hand. We implement this approach on a ResNet-18 model, pre-trained on ImageNet data. With this model, we transform each original image to a vector of dimension 512. We then train shallow FFNNs as described in Section 3.1 with m = 2000 and output dimension 10 (resp. 100). This experiment differs from previous cases as 1) we use stochastic GD with a mini-batch size of 64 instead of full batch GD; 2) we use cross-entropy loss instead of MSE; and 3) both layers are trained. All experiments are five times, with the learning rate 5.0. Figure 5 shows the pruning results for the same four values of pairs (7,a) as above, for CIFAR 100. Similar results are obtained for CIFAR 10; see Appendix K Similar conclusions as before hold here, even though the theory does not apply directly.
uses a large model pre-trained on general task and then performs
9 Discussion, limitations and further work
We have shown that under asymmetrical scaling of the nodes of neural network, it is possible to achieve both zero training error and feature learning, when the width of the neural network is sufficiently large. We considered two definitions of feature learning. The first definition is a minor generalisation of the notion of feature learning from YYagg & Hu 2021). and it defined as a change in the feature map. We proposed a second definition, called non-uniform feature lcarning, which additionally requires that the contributions of some individual nodes remain negliggiie in the asymptotic limit. We showed that under our asymmetric scaling and additional conditions, both definitions hold, whereas for the standard NTK, neither does, and for
16

---

Published in Transactions on Machine Learning Research (02/2025)
y=10
0.6
6-02-
0.6
accuroyy
60
0-03--00 02
0=14 10
0.4
M
Test
0.2
20
0.2
077 05
0- D.
0.0
0
500
100c
1500
2030
550
1030
1500
2030
Ierations
105
Perations
1e5
Number pruned nodes
Number pruned nodes
Figure 5: Results for CIFAR-100. From left to right, 1) test risk through training, 2) differences in weight norms [[ey - Woyll with the j's being the neurons having the maximal difference at the end of training, 3) test risks of pruned models, and 4) test accuracies of pruned models.
the mean field, only the standard dcfinition holds. We demonstrate empirically that having non-uniform feature learning is particularly important when consider transfer learning and pruning. Our definitions of feature learning relate to the change in the feature map. As already noted by YYang & Hu 2021 it is a relatively weak definition of feature learning, as it does not connect the weight change with the generalisation properties. Experimentally, we found that some case (e.g. single ReLU), the asymmetrical, unpruned network provides the best test error, while in others (MNIST and CIFAR), the unpruned symmetrical scaling gave the best test accuracy. An interesting avenue rescarch is investigate thcorctically the generalisation properties of such asymmetrical scaling. We note that the approaches used for the symmetrie NTK Arora et al. 2019a), which rely on the limiting kernel, cannot be applied to setting, due to the evolving kernel.
In this article, we have assumed an iid Gaussian initialisation for the weights (Assumption 3.3 which is a standard assumption in the analysis of large-width neural networks (Du el al. Soltanolkotabi] 2020 Nguyen et al. 2021). Our results rely on the fact that the minimum eigenvalue Ki the mean NTK at initialisation strictly positive; this result was demonstrated by (Du et al. F.1) under the iid Gaussian initialisation. An interesting direction of research would be to investigate whether the results derived in this paper hold under other, possibly non-fid, initialisation schemes. In particular, the
2019b
Oymak &
2019a Proposition
case of orthogonal initialisations would be of particular interest (Hu et al. 2020 Huang et al. 2021).
The asymmetrical parameterisation in Equation (1) is rather general, and only requires the Xy be summable. A natural default choice, taken in this article, is to take power function & j-l/o where < 1. Other parameterisations are also possible, such as X 1/K for j I,...,k K and 0 otherwise. One could also choose other scalings such as (e 1) exp(-j) C/(jlog l/Ojlog((1 (j+2(++1) 1)).
Acknowledgements
We would like lo thank Taeyoung Kim for helpful discussions, and the anonymous reviewers for their useful comments that helped improve the paper. HY was supported by the National Research Foundation of Korea (NRF) grant funded by the Korcan Government (MSIT) (No. RS-2023-00279680) JL acknowledges support from Institute for Information & communications Technology Plamning & Evaluation(IITP) grant funded by the Korea Dnvernmen(MSST)) (RS-2019-T119005,, Artificial Intelligence Graduate School Program(KAAST))
References
S. Arora, S. Du, W. Hu, Z. Li, and R. Wang. Fine-grained analysis of optimization and generalization for overparameterrzze two-layer neural networks. Tn International Conference on Machine Learning, pp- 322 332. PMLR, 2019a.
S. Arora, S. S. Du, W. Hu, Z. Li, R. Salakhutdinov, and R. Wang. On exact computation with an infinitely wide neural net. Advances in Neural Information Processing Systems, 32, 2019b.
P. Bartlett, A. Montanari, and A. Rakhlin. Deep learning: a statistical viewpoint. Acta numerica, 30:87-201, 2021.
17

---

Published in Transactions on Machine Learning Research (02/2025)
D. Bracale, S. Favaro, S. Fortini, and S. Peluchetti. Large-width functional asymptotics for deep Gaussian neural networks. In International Conference on Learning Representations, 2021.
Z. Chen, E. Vanden- Eijnde, and J. Bruna. On feature learning in neural networks with global convergence guarantees. In International Conference Learning Representations 2021.
L. Chizat and F. Bach. On the global convergence of gradient descent for over over-varaminn parameterized models using optimal transport. Advances in newral information processing systems, 31, 2018.
L. Chizat, E. Oyallon, and F. Bach. On lazy training in differentiable programming, Advances in Neural Information Processing Systems, 32, 2019.
S. Du, J. Lee, H. Li, L. Wang, and X. Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pp. 1675-1685. PMLR, 2019a.
S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient descent provably optimizes over wer-parameterized neural networks. In International Conference on Learning Representations, 2019b.
El Karoui. The spectrum kernel random matrices. The Annals of Statistics, 38(1):1 50, 2010.
S. Favaro, S. Fortini, and S. Peluchetti. Stable behaviour of infinitely wide deep neural networks. In International Conference on Artificial Intelligence and Statistics, pp. 1137 1146, 2020.
S. Frei, N. S. Chatterji, and P. L. Bartlett. Random feature amplification: Feature learning and generalization in neural networks. Journal of Machine Learning Research, 24(303):1 49, 2023.
B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. When do neural networks outperform kernel methods? Advances in Newral Informalion Processing Systems, 13:14420-14833 2020.
W. Hu, L. Xiao, and Pennington. Provable benefit of orthogonal initialization in optimizing deep linear networks. In International Conference on Learning Representations, 2020.
W. Huang, W. Du, and R. Xu. On the neural tangent kernel of deep networks with orthogonal initialization. Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence. International Joint. Conferences on Artificial Intelligence Organization, 2021.
A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, pp. 8571 8580, 2018
P. Jung, H. Lee, J. Lee, and H. Yang. a-stable convergence of heavy-tailed infinitely-wide neural networks. Advances in Applied Probability, 55(4), 2023.
H. Lee, F. Ayed, P. Jung, J. Lee, H. Yang, and F. Caron. Deep ncural networks with dependent weights: Gaussian process mixture limit, heavy tails, sparsity and compressibility. Journal of Machine Learning Research, 24:1-78, 2023.
J. Lee, Y. Bahri, R. Novak, S. S. Schoenholz, J. Pennington, and Sohl-Dickstein. Deep neural nerworks as Gaussian processes. In International Conference on Learning Representations, 2018.
J. Lee, L. Xiao, S. Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington. Wide neural networks of any depth evolve as lincar models under gradient descent. Advances in Neural Information Processing Systems, 32, 2019.
E. Malach, P. Kamath, E.I Abbe, and N. Srebro. Quantifying the benefit of using differentiable learning over langent kernels. In International Conference on Machine Learning, pp. 7379 7389. PMLR, 2021.
A. G. de G. Matthews, J. Hron, M. Rowland, R. E. Turner, and Z. Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference Learning Representations, 2018.
S. Mci, A. Montanari, and P.-M. Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, F16633): F7671, 2018.
18

---

Published in Transactions on Machine Learning Research (02/2025)
S. Mei, T. Misiakiewicz, and A. Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pp.
PMTR, 2019.
R. M. Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pp. 29-53. Springer New York, 1996.
Q. Nguyen, M. Mondelli, and G. Montufar. Tight bounds on the smallest eigenvalue of the neural tangent kernel for deep ReLU networks. In International Conference on Machine Learning, pp. 8119 8129. PMLR, 2021.
S. Oymak and M. Soltanolkotabi. Toward moderate overparameterization:: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, (1):84-105, 2020.
G. Rotskoff and E. Vanden-Eijnden. Parameters interacting particles: long time convergence and asymptotic error scaling of neural networks. Advances in Neural Information Processing Systems, 31, 2018.
J. Sirignano and K. Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. SIAM Journal Applied Mathematics, 80(2):72--75, 2020.
L. Tao, J. X. Zhi-Qui Sin, M. Zheng, and Yaoyu. Phase diagram for two-layer ReLU neural networks at witth limit. Journal of Machine Learning Research, 22:1 47, 2021.
P. Wolinski, G. Charpiat, and Y. Ollivier. Asymmetrical scaling layers for stable network pruning. OpenReview Archive, 2020.
B. Woodworth, S. Gunasckar, J. D. Lee, E. Moroshko, P. Savarese, I. Golan, D. Soudry, and N. Srebro. Kernel and rich regimes in overparametrized models. In Conference Learning Theory, pp. 8635-3673 PMLR, 2020.
G. Yang. Wide feedforward or recurrent neural networks of any architecture are Gaussian processes. Advances in Neural Information Processing Systems, pp. 9947-9960, 2019.
G. Yang and E. Hu. Tensor programs iv: Feature learning in infinitewwddth neural networks. In International Conference on Machine Learning, pp- 11727 11737. PMLR, 2021.
D. Zou and Q Gu. An improved analysis of training over-parameterized deep neural networks. Advances in Neural Information Processing Systems, 32, 2019.
D. Zou, Y. Cao, D. Zhou, and Q Gu. Gradient descent optimizes over-parameteriieed deep ReLU networks. Machine learning, (3):3))667-4)) 492,499, 2020.
19


Reviews:
[{"review_id": "FUZFYkfjEp", "reviewer": "Reviewer_3Qfc", "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\n- The idea to consider a per-neuron scaling is original, and leads to some interesting results.\n- The analysis appears to be clean (although I didn't check all the details carefully) \n\n## Weakness\n\n- The thing that I am missing most from the article is the motivation and a discussion of the implications of the results.", "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 authors study two-layer fully connected neural networks. These networks have been studied intensively in two regimes: a mean-field regime, where the network learns features from data, and the NTK [Jacot et al. NeurIPS '18] or \"lazy\" regime [Chizat et al. NeurIPS '19] where an over-parametrised network can fit its training data without learning features, instead performing essentially a type of kernel ridge regression. One way to induce a neural network to behave according to one or the other regime is by scaling the weights of the read-out layer: for a network with $m$ neurons, a scaling of $1/m$ is said to put the network in the mean-field regime, while a scaling of $1 / \\sqrt m$ puts the network in the NTK regime.\n\nHere, the authors consider an alternative scaling where the readout of each hidden neuron has its own scaling \n\\begin{equation}\n    \\lambda_{m, j} = \\gamma \\frac 1 m + (1 - \\gamma) \\frac{\\tilde \\lambda_j}{\\sum_{k=1}^m \\tilde \\lambda_k}\n\\end{equation}\nwhere $\\gamma \\in [0, 1]$ and $1 \\ge \\tilde \\lambda_1  \\ge \\tilde \\lambda_2 \\ge \\ldots  \\ge 0$ are nonnegative fixed scalars that sum to 1 (for example, $\\tilde \\lambda_j \\propto j^{-2}$. \n\nIn this setting, the authors consider the limit of large $m$ and make similar assumptions as some of the first papers to discuss convergence of neural networks in the lazy regime (Du et al., 2019b;a), most notably that second-layer weights are first fixed, and then trained in a second phase, separately from the first-layer weights. Under these assumptions, the authors show that\n- if γ > 0, the training error goes to 0 at a linear rate with high probability; \n- feature learning in a weak sense occurs if and only if γ < 1 (specifically, the authors take feature learning to mean that the representation, i.e. the values of the hidden neurons, has a relative change compared to its representation at initialisation at some point along the training trajectory, for some input x -- their definition does not imply that this change benefits the network in any way.\n\nSome simple experiments suggest that this change in the features is however conductive to generalisation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Please give some hint as to the motivation for considering this regime, as well as for some of the scalings that you give as examples. Is the motivation to have global convergence results with feature learning? If so, what does this type of regime bring to the analysis that is not covered by mean-field methods, which also provide global convergence guarantees? \n2. Please add a Discussion of your results at the end of the article - what are the lessons? What kind of analyses do you anticipate will be possible in this regime?", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ixRGXVc2ER", "reviewer": "Reviewer_nBUn", "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- **Clear Theoretical Contribution**: The paper’s primary theoretical results, e.g. Theorem 5.1 and 6.1, are presented concisely and clearly. Moreover, those results do not require extensive or restrictive assumptions.\n- **Experiments and Empirical Validation**: The experiments demonstrate the relevance of the theoretical findings. For instance, they clearly show how the asymmetrical node scaling achieves 0 training loss wityh gradient descent.\n- **Relevance to Feature Learning and Transfer Learning**: By addressing the limitations of NTK scaling, which typically does not enable feature learning since the NTK is constant, the paper makes an important contribution to understanding how over-parameterized networks might achieve feature learning. \n\n#### Weaknesses\n- **Assumptions and Notations**: Assumption 3.3, which involves a normal distribution, might seem overly restrictive, as other initialization methods (e.g., orthogonal initialization) are common in practice. I beleive this this specific distribution was chosen because it makes the theoretical analysis easier, but how restrictive is this assumptions? Is it possible to generalize it to others?\n- **Clarity and Presentation of Some Concepts**: Even tough most of the theoretical results are well explained, certain results lack sufficient explanation. For instance, in Proposition 7.6, the discution  around cases when $\\gamma=0$ and $\\gamma>0$ is clear but the implications of the inequalities could be made clearer.\n- **Arbitrary Choice of Scaling**: The specific use of the Zipf law in defining node scaling appears arbitrary, as it is unclear why this particular law was selected over other potential distributions.\n- **Gradient Descent Stepsize**: The stepsize for gradient descent is notably small, and there is no discussion on how slow this choice is when compared to the standard stepsize of non-convex gradient descent (1/smoothness constant). A comparison or at least a brief explanation would be appreciated.\n\n### Overall Assessment\n\nThis paper presents an interesting approach with theoretical contributions and empirical validation. However, certain aspects could benefit from further clarification.\n\n**Disclaimer:** I am not an expert in this area, therefore it might be hard for me to judge the impact and also to evaluate this paper w.r.t. existing work in the field.", "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 paper investigates the global convergence and feature learning properties of gradient-based training on wide, shallow neural networks with asymmetrical node scaling, where the output of each hidden node has a non-uniform positive scaling parameter. This scaling method departs from the traditional NTK parameterization, with the intent to achieve feature learning in large-width networks, contrasting with the “lazy training” of NTK. The authors provide theoretical guarantees for convergence under this modified scaling and demonstrate how it allows feature learning. The paper includes numerical experiments that validate the theoretical claims and highlight additional practical benefits, such as enhanced prunability and transfer learning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. **Clarify Definitions and Implications**: The notion of feature learning under asymmetrical scaling could be clarified, especially why it represents a stronger notion than in the standard NTK setting: it is indeed obvious that one definition implies the other, but why this definition ensures that non uniform feeatures learning relate more to the relevance of the learn features for prediction?\n  \n2. **Misc** Add $ ||w - w_0|| $ and $ ||\\theta - \\theta(W_0)|| $) before the big O in section 5 to make the section a bit easier to understand.\n\n3. **Fix Minor Errors**: Address typographical errors, such as “activiation” in Section 3.3.\n\n4. **Explain Assumptions**: Reconsider Assumption 3.3 or provide a justification for the choice of a normal distribution for initialization, especially given that alternative methods exists (like orthogonal initialization, for instance)\n\n5. **Discuss Zipf Law**: Provide reasoning behind the selection of the Zipf law for node scaling, since this choice could appear arbitrary.\n\n6. **Expand on Gradient Descent Stepsize**: Include a discussion about the small stepsize used for gradient descent, particularly how it compares to typical choices in non-convex optimization with smooth functions. This could highlight any advantages or limitations of the approach in practical settings.", "broader_impact_concerns": "This work is purely theoretical, no Broader Impact Concerns section is necessary.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "hRId4WUlnA", "reviewer": "Reviewer_XJbp", "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:\nThe asymmetric scaling formulation provides flexibility in choosing the training regime, while keeping the mathematical analysis tractable based on tools developed by prior works.\n\n# Weaknesses:\n* Theorem 5.1 does not guarantee feature learning, the upper bounds simply suggest that the weights may move significantly from initialization, but the upper bounds may not be sharp. However, feature learning is proved for simpler settings in Section 7.\n* Generally, it is not clear why the notion of feature learning introduced here, i.e. simply the weights moving significantly, is useful on its own. Several works have shown that feature learning is useful because it leads to better generalization guarantees by adapting to the task structure, e.g. [1].\n* Intuitively, I'm not sure I understand why even though the weights move significantly, the first Gram matrix in Equation (15) stays close to initialization.\n\n[1] B. Ghorbani et al. When Do Neural Networks Outperform Kernel Methods? NeurIPS 2020.", "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 considers the gradient flow/descent training dynamics of infinitely wide two layer neural networks under an asymmetric scaling of the neurons. In the considered formulation, the completely symmetric case reduces to the lazy training regime where the neural tangent kernel remains constant and no feature learning regime is achieved. In the asymmetric case, the authors prove that both feature learning and exponential convergence of the training loss can be achieved. Additionally, the authors provide examples where asymmetric scaling can lead to non-uniform feature learning which further enables pruning. Several experiments support the theoretical findings.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Is there an intuitive explanation why the first Gram matrix in Equation (15) does not change significantly even when the weights do change? It might be helpful to highlight this in the main text.\n* In the classification experiments, the best test accuracy is achieved in the lazy regime where $\\gamma = 1$. This is surprising given the general intuition that feature learning achieves better generalization when the task has some hidden low-dimensional structure. Is there an explanation for this observation?\n* The lower bounds on $m$ in Theorem 5.1 and 6.1 are getting better in high dimensions. Does it mean/Is it simply because overfitting is easier in high dimensions when the system becomes underdetermined?\n* Where does the statement $\\sum_{j \\geq \\lfloor \\rho m\\rfloor}\\lambda_{m,j} \\rightarrow \\gamma(1-\\rho)$ in Proposition 7.6 come from? Does it not depend on how $\\lambda_{m,j}$ are chosen?", "broader_impact_concerns": "Not applicable.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agree that the paper is technically sound and provides meaningful theoretical insights into feature learning, which is of interest to the deep learning theory community. Given its solid theoretical contributions and relevance, I conclude that this paper satisfies the acceptance criteria of TMLR.
Decision: Accept (poster)