Global Convergence in Shallow Neural Networks
Global Convergence in Shallow Neural Networks
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.
1 Introduction
The training of neural networks typically involves the minimisation of a 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 a better understanding of this phenomenon, one fruitful direction of research
has been to analyse properties of GD training of over-parameterised, large-width neural networks; that is,
neural networks where the number m of hidden nodes in a given layer is 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 to
p
infinity, the solution obtained by GD achieves zero training error, and coincides with that of kernel regression
under a so-called limiting Neural Tangent Kernel (NTK). Under the same node scaling, called NTK scaling,
1
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;a; 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 in deep neural networks.
Main contributions. We investigate global convergence properties and feature learning in gradient-type
training of large-width feedforward neural networks (FFNNs) under a more general p asymmetrical node scaling.
In particular, each hidden node j = 1, . . . , m has a fixed node-specific scaling λm,j with
1 λ
ej
λm,j = γ · + (1 − γ) · Pm (1)
m k=1 λk
e
large-width networks. Two typical examples of the scalars (λ ej )j≥1 are (a) λ
ej = 6π −2 j −2 for all j ≥ 1, and
(b) λ
e1 = . . . = λ
eK = 1/K and λ
ej = 0 for all j > K, for a fixed K.
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 ℓ2 loss. We show that,
under similar assumptions as in Du et al. (2019b;a) on the data, activation function, and initialisation, when
the number of hidden nodes m is sufficiently large: (i) if γ > 0, the training error goes to 0 at a 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 γ < 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 discusses 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 a model
at initialisation, and its infinite-width limit. Sections 5 and 6 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 8 describes our experiments on simulated and real datasets,
whose results illustrate our theoretical results and their potential applications. The Supplementary Material
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 n ≥ 1, let [n] = {1, . . . , n}. For a multivariate real-valued function f : Rn → R,
the gradient ∇v f (v) is the n-dimensional column vector of partial derivatives ∇v f (v) = ( ∂v
∂f
1
(v), . . . , ∂v
∂f
n
(v))⊤
where v = (v1 , . . . , vn )⊤ . For a square matrix B, we denote its minimum and maximum eigenvalues by
eigmin (B) and eigmax (B), respectively. For a vector v ∈ Rn , we write B = diag(v) for the n-by-n diagonal
matrix with Bii = vi for i ∈ [n].
2 Related work
Large-width FFNNs. The analysis of large-width FFNNs goes back to Neal (1996) 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 over-parameterisation for gradient descent training, with the
“1/ m” NTK scaling (Arora et al., 2019b; Du et al., 2019b;a; Lee et al., 2019; Zou & Gu, 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 a 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-Eijnden, 2018; Mei et al., 2018; 2019; Chizat et al., 2019; Sirignano & Spiliopoulos, 2020;
Ghorbani et al., 2020; Chen et al., 2021; Tao et al., 2021). This scaling is also equivalent, up to the so-called
abc-scaling symmetry (Yang & Hu, 2021), to the µP parameterisation of Yang & Hu (2021) in the case of
shallow networks. Feature learning is known to occur under this mean-field scaling. Also, Chizat & Bach
(2018) showed that under the same scaling, if the training of the model converges, it converges to a global
minimum.
Asymmetrical 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 on the (empirical) usefulness in terms of pruning. Indeed, our experiments in
Section 8 are also in line with their findings. The work of Wolinski et al. (2020), however, only considered
asymmetrical scaling with γ = 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
We consider a shallow FFNN with one hidden layer of m ≥ 1 hidden nodes and a scalar output. To simplify
the analysis, we assume that there is no bias term. Let x ∈ Rd be some input vector, where d ≥ 1 is the
input dimension. The model is defined as
m
X 1
fm (x; W) = λm,j aj σ(Zj (x; W)) with Zj (x; W) = √ wj⊤ x for j ∈ [m] (2)
p
j=1
d
where fm (x; W) is the scalar output of the FFNN, Zj (x; W) is the pre-activation of the j-th hidden node,
σ : R → R is the activation function, wj ∈ Rd is the column vector of weights between node j of the hidden
layer and the input nodes, and aj ∈ R is the weight between the hidden node j and the output node. The
λm,j ’s for j ∈ [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 (w1⊤ , . . . , wm ) . We assume that σ admits a
⊤ ⊤
derivative σ .′
For n ≥ 1, let σ : Rn → Rn (resp. σ ′ : Rn → Rn ) be the vector-valued multivariate function that applies the
real-valued function σ (resp. σ ′ ) element-wise on each of the n input variables. To simplify the analysis, we
assume throughout this article that the output weights (aj )j∈[m] are randomly initialised and fixed afterwards:
iid
aj ∼ Uniform({−1, 1}), j ≥ 1. (3)
This simplifying assumption is often made when analysing large shallow networks (see e.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
Pm (λm,j )j∈[m] are fixed and parameterised as in Equation (1). By construction, we have
λm,1 > 0 and j=1 λm,j = 1 for all m ≥ 1. Recall that the case γ = 1 corresponds to NTK scaling. Also,
note that our model covers finite FFNNs: when
(
1/K if j ∈ [K]
γ = 0 and λj =e
0 otherwise
3
Published in Transactions on Machine Learning Research (02/2025)
for some K ≤ m, the model becomes a finite network of width K. In the experiments, we will consider the
special case where (λ
ej )j≥1 are the probability masses of a Zipf law:
ej = 1 1
λ 1/α
, j≥1 (4)
ζ(1/α) j
for some α ∈ (0, 1), where ζ is the Riemann zeta function. The parameter α tunes how quickly λ
ej decreases
with j, smaller values corresponding to more rapid decrease and more asymmetry.
3.2 Training
Let Dn = {(xi , yi )}i∈[n] be the training dataset, where n ≥ 1 is the number of observations. Let X denote
the n-by-d matrix whose ith row is xi⊤ . We want to minimise the empirical risk under ℓ2 loss. Consider the
objective function
n
1X
Lm (W) = (yi − fm (xi ; W))2 (5)
2 i=1
which is non-convex in general. For a given dataset Dn , width m ≥ 1, output weights (aj )j∈[m] , and scaling
parameters (λm,j )j∈[m] , we aim to estimate the trainable parameters W by minimising Lm (W) using gradient
flow or gradient descent. Let W0 be some initialisation. In gradient flow, (Wt )t>0 is the solution to the
following ordinary differential equation (ODE):
dWt
= −∇W Lm (Wt )
dt
with limt→0 Wt = W0 . Let p
wtj be the value of the parameter wj at time t, and define Ztj (x) = Zj (x; Wt ).
Note that ∇wj fm (x; W) = λm,j aj σ ′ (Zj (x; W)) · √1d x. Under gradient flow, for j ∈ [m],
n
! p n
!
dwtj X λm,j aj X
= (yi − fm (xi ; Wt ))∇wj fm (xi ; Wt ) = √ (yi − fm (xi ; Wt ))σ ′ (Ztj (xi ))xi .
dt i=1
d i=1
Thus, the derivatives associated with each hidden node j are scaled by λm,j . For an input x ∈ Rd , the
p
where Θm : Rd × Rd → R is the neural tangent kernel for the network fm (x; W):
m
x⊤ x′ X
Θm (x, x ; W) =
′
λm,j σ ′ (Zj (x; W))σ ′ (Zj (x′ ; W)). (6)
d j=1
The associated neural tangent Gram (NTG) matrix Θ b m (X; W) is the n-by-n positive semi-definite matrix
whose (i, j)-th entry is Θm (xi , xj ; W). It takes the form
m
b m (X; W) = 1
X Xwj Xwj
Θ λm,j diag σ ′ √ XX⊤ diag σ ′ √ . (7)
d j=1 d d
Gradient descent is a discretisation of gradient flow. Under gradient descent, the parameters are updated by
where η > 0 is a learning rate. These updates give rise to the family (Ws )s∈N∪{0} indexed by discrete time
steps s = 0, 1, 2, . . ., rather than continuous times t ≥ 0.
4
Published in Transactions on Machine Learning Research (02/2025)
Throughout the paper, we assume that the activation function σ satisfies the following standard condition:
for all random variables Z ∼ N (0, s2 ) for some s > 0,
This assumption is made all the time, and so we 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 Dn = {(xi , yi )}i∈[n] .
Assumption 3.1 (Dataset). (a) All inputs are non-zero and have norms at most 1: 0 < ∥xi ∥ ≤ 1 for all
i ≥ 1. (b) For all i ̸= i′ and c ∈ R, xi ̸= cxi′ . (c) There is C > 0 such that |yi | ≤ C for all i ≥ 1.
The next assumption concerns the activation function σ. 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 A
in the Supplementary Material.
Assumption 3.2 (Activation function). The activation function is analytic, with |σ ′ (x)| ≤ 1 and |σ ′′ (x)| ≤ M
for some M > 0, but it is not a polynomial.
Mean NTG at initialisation and its minimum eigenvalue. Let W0 be a random initialisation from
Assumption 3.3. Consider the mean NTK at initialisation
x⊤ x′ 1 ⊤ 1 ⊤ ′
Θ∗ (x, x′ ) = E [Θm (x, x′ ; W0 )] = E σ ′ √ w01 x σ ′ √ w01 x . (10)
d d d
√
The mean NTK, which is also, by the law of large numbers, the limiting NTK under 1/ m scaling (Jacot
et al., 2018), does not depend on (λm,j )j≥1 nor m. Let Θ b ∗ (X) = E[Θ
b m (X; W0 )] be the associated n-by-n
mean NTG matrix at initialisation, whose (i, j)-th entry is Θ∗ (xi , xj ). Let κn = eigmin (Θ b ∗ (X)) be the
minimum eigenvalue of the mean NTG matrix at initialisation. This minimum eigenvalue plays an 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., 2021)) have shown that, under some assumptions on the data, activation
function, and initialisation, κn is strictly positive or bounded away from zero. such a result, under the
Assumptions of Section 3.3.
Proposition 4.1 ((Du et al., 2019b, Theorem 3.1) and (Du et al., 2019a, Proposition F.1)). When Assump-
tions 3.1 to 3.3 hold, we have κn > 0.
Remark 4.2. Du et al. (2019b;a) make the assumption that each xi has unit norm. But their proof holds
under the less strict Assumption 3.1(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 a fixed
sample size n, as the width m goes to infinity. The proof, given in Appendix C in the Supplementary
P Material,
follows from the triangle inequality and the law of large numbers, together with |σ ′ (z)| ≤ 1 and j≥1 λ
ej = 1.
Proposition 4.3. Consider a sequence (w0j )j≥1 of iid random vectors distributed as in Assumption 3.3.
Suppose Assumption 3.2 holds. Then,
Θ
b m (X; W0 ) → Θ
b ∞ (X; W0 ) (11)
5
Published in Transactions on Machine Learning Research (02/2025)
(2)
Also, E[Θ
b ∞ (X; W0 )] = E[Θ
b ∞ (X; W0 )] = Θ
b ∗ (X), and
h i X
E ∥Θb ∞ (X; W0 ) − Θb ∗ (X)∥2 = C0 (X)(1 − γ)2
F
e2
λ j (13)
j≥1
where ∥ · ∥F denotes the Frobenius norm, and C0 (X) > 0 is a positive constant equal to
X x⊤ xi′ 2
1 ⊤
1 ⊤
i
Var σ √ w01 xi σ √ w01 xi′ .
′ ′
′
d d d
1≤i,i ≤n
When γ = 1 (NTK scaling), the NTG converges to a constant matrix, and solutions obtained by gradient
flow coincide with that of kernel regression. Whenever γ < 1, Proposition 4.3 shows that the NTG is random
at 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 measuredPby the total variance of the limiting random NTG,
can be quantified by the nonnegative constant (1 − γ)2 j≥1 λ e2 ∈ [0, 1]. When this constant is close to 0, we
j
approach the kernel regime, and increasing this value leads to a departure from the regime. The quantity
P e2 ∈ (0, 1] is always strictly positive. More rapid decrease of the λ
λ ej as j increases will lead to higher
j≥1 j
values of
P e . For example, when using the Zipf weights in Equation (4), we have P
λ 2 e2 = ζ(2/α)2 ,
λ
j≥1 j j≥1 j ζ(1/α)
which decreases with α, as shown in Figure S.1 in the Supplementary Material.
This section has described the behaviour of the NTG at initialisation in the infinite-width limit, and
has provided intuition on the node-scaling parameters. The next three sections contain results on global
convergence and feature learning properties of large, but finite, FFNNs under such asymmetrical scaling.
Our global convergence theorem, which is given below, explains what happens during training via gradient
flow. Recall that κn is the minimum eigenvalue of the mean NTG matrix Θ b ∗ (X) at initialisation. Our
theorem says that with high probability, (i) the loss decays exponentially fast with respect to κn and training
time t, and (ii) the NTG and weights wtj change, respectively, by
P qP
m
m n2 2
j=1 λm,j
1/2
!
n3 j=1 λ2m,j nλm,j
∥Θm (X; Wt )−Θm (X; W0 )∥2 = O
b b + and ∥wtj −w0j ∥ = O .
κ2n d3 γ 2 κn d2 γ κn d1/2 γ
Define " 2 #
cz
C1 = sup Ez∼N (0,1) σ √ . (14)
c∈(0,1] d
Theorem 5.1. (Global convergence, gradient flow) Let δ ∈ (0, 1). Suppose Assumptions 3.1 to 3.3 hold, and
that
2 n log 2n 210 n3 M 2 (C 2 + C1 ) 215 n4 M 2 (C 2 + C1 )
3
γ > 0, and m ≥ max δ
, , ,
κn d κ3n d3 γ 2 δ κ4n d4 γ 2 δ
where C is the bound on the yi ’s in Assumption 3.1. Then, with probability at least 1 − δ, the following
properties hold for all t ≥ 0:
(a) eigmin (Θ
b m (X; Wt )) ≥ γκn ;
4
(b) Lm (Wt ) ≤ e−(γκn t)/2 Lm (W0 );
6
Published in Transactions on Machine Learning Research (02/2025)
√ q
n λm,j 27 (C 2 +C1 )
(c) ∥wtj − w0j ∥ ≤ κn d1/2 γ2δ for all j ∈ [m];
7 3 2 2 5 2 2
+C1 )1/2
qP
b m (X; W0 )∥2 ≤ 2 n M2 (C +C1 ) Pm m
(d) ∥Θ
b m (X; Wt ) − Θ κn d3 γ 2 δ · j=1 λ2m,j + 2 n κMn (C
2
d γδ 1/2 · λ 2
j=1 m,j .
The above theorem implies that, whenever γ > 0, the p training error converges to 0 exponentially fast.
Additionally, the weight change is bounded by a factor λm,j and the NTG change is bounded by a factor
qP
m
j=1 λm,j . We have (see Appendix B.2 in the Supplementary Material) that as m → ∞,
2
m
X ∞
X
λm,j → (1 − γ)λ
ej for all j ≥ 1 and λ2m,j → (1 − γ)2 e2 .
λ j
j=1 j=1
If λ
ej > 0 (note that we necessarily have λ e1 > 0), the upper bound in (c) is therefore vanishing in the
infinite-width limit if and only if γ = 1 (lazy-training regime). Similarly, the upper bound in (d) is vanishing
if and only if γ = 1. Although we were not able to obtain matching lower bounds, we show in Section 7 that
feature learning arises whenever γ < 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/ d
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.1 also holds for the ReLU activation function. See Theorem A.1 in the
Supplementary Material.
Sketch of the proof. We give here a sketch of the proof of Theorem 5.1 (and of Theorem A.1, its ReLU
counterpart, in the Supplementary Material). The detailed proofs are given in Appendices F and G in the
Supplementary Material, with secondary lemmas given in Appendices D and E there. The structures of
the proofs of Theorems 5.1 and A.1 are similar to that of (Du et 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 a global minimum of the objective function if the minimum eigenvalue of the
NTG matrix Θ b m (X; Wt ) is bounded away from zero, for m sufficiently large, by some positive constant for
all t ≥ 0. In the NTK scaling case (γ = 1), Du et al. (2019b) showed that the following is satisfied, for m
sufficiently large, with high probability: (i) the NTG matrix at initialisation is close to its mean, and the
minimum eigenvalue is close to that of the mean NTG, (ii) the weights wtj 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.
However, in the case of asymmetrical node scaling (γ < 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.
(1) (2) ej )/ Pm λ (1) (2)
Let λm,j = γ/m and λm,j = ((1 − γ)λ k=1 k , and note that λm,j + λm,j = λm,j . For k ∈ {1, 2}, let
e
b (k)
Θ m be the n-by-n symmetric positive semi-definite matrices defined by Equation (7), with λm,j replaced by
(k)
either λm,j . Note that
Θ
b m (X; Wt ) = Θ
b (1) (X; Wt ) + Θ
m
b (2) (X; Wt )
m (15)
b (1)
with E[Θ m (X; W0 )] = γ Θ (X).
b∗
The key idea of the proof is to use the above decomposition of the NTG matrix as a 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 as follows. For large m, with
7
Published in Transactions on Machine Learning Research (02/2025)
b (1)
high probability: (i) Θ m (X; W0 ) is close to its mean γ Θ (X) and its minimum eigenvalue is therefore lower
b∗
bounded by (γκn )/2;
The important points of the proof are as follows. For large m, with high probability: (i) Θ b (1)
m (X; W0 ) is
close to its mean γ Θ
b ∗ (X) and its minimum eigenvalue is therefore lower bounded by (γκn )/2; (ii) while the
weights Wt may change significantly over time, Θ b (1)
m (X; Wt ) remains nearly constant over time; (iii) as a
(1)
result, the minimum eigenvalue of Θm (X; Wt ) can be lower bounded by (γκn )/4; (iv) this implies that the
b
minimum eigenvalue of the overall NTG matrix Θ b m (X; Wt ) is lower bounded by (γκn )/4.
Since showing that the first part of the NTK/NTG in Equation (15) does not change in the limit of m → ∞
is a key component of the proof, we give here an outline of it in the simplified case of d = 1 and σ being
smooth. Let
m
′ γ
X
Θ(1)
m (x, x ′
; W t ) = xx σ ′ (wtj x)σ ′ (wtj x′ )
m j=1
be the first part of the NTK. We have, over gradient flow,
(1) m
dΘm (x, x′ ; Wt ) γ X ′′ dw
tj
= xx′ xσ (wtj x)σ ′ (wtj x′ ) + x′ σ ′ (wtj x)σ ′′ (wtj x′ )
dt m j=1 dt
m
γ X dwtj
≤ 2M ,
m j=1 dt
where the last inequality follows from the triangle inequality and Assumptions 3.1 and 3.2. Furthermore,
n
dwtj X
= λm,j (yi − fm (xi ; Wt ))σ ′ (Ztj (xi ))xi
p
dt i=1
n
X
|yi − fm (xi ; Wt )|
p
≤ λm,j
i=1
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 p is thus bounded by a
γ Pm γ Pm
quantity involving m λm,j . Under the scaling in Equation (1), m λm,j → 0 as m → ∞
p
j=1 j=1
(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 P Equation√ (15), a similar derivation leads to an upper bound of the order of
m
γ Pm λj λm,j
( m j=1 λm,j ) + (1 − γ) Pm . Due to the second term, this quantity does not converges to 0 as
p j=1
e
λk
e
k=1
m → ∞, unless γ = 1 (symmetric case).
8
Published in Transactions on Machine Learning Research (02/2025)
where C and C1 are from Assumption 3.1 and Equation (14). Let β = (1 − ηγκn /2)1/2 . If
2 n log 2n 25 η 2 n3 M 2 (C 2 + C1 ) 211 η 2 n4 M 2 (C 2 + C1 )
3
m ≥ max δ
, , ,
κn d κn d3 (1 − β)2 δ κ2n d4 (1 − β)2 δ
Note that the condition on the learning rate requires η = O(γκn /n2 ). Thus, the best possible convergence
rate from the theorem is (1 − (ηγκn /2)) = (1 − (C0 γ 2 κ2n /n2 )) for some constant C0 .
The proof is by induction on the gradient-descent step s, 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 wsj and the Gram matrix Θ
b m (X; Ws )
change during gradient descent in our case, while they remain nearly constant in the case of (Du et al., 2019a).
In this section, we present some results about feature learning. We focus on node scalings (λm,j )j∈[m] of the
form in Equation (1), both asymmetrical (γ < 1) and symmetrical (γ = 1; that is, λm,j = 1/m), but we also
discuss alternative parameterisations such as mean-field and µP . For shallow neural networks trained by
gradient descent, mean-field and µP parameterisations are equivalent (Yang & Hu, 2021), and correspond to
node scalings λm,j = 1/m2 and learning rate η = η0 m for some η0 > 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 a linear activation function, and next under a general nonlinear
activation function.
7.1 Definitions
Definition 7.1 (Feature learning). Let (w0j )j≥1 be a sequence of random initialisations for nodes j ≥ 1.
We will say that feature learning occurs during training if1
Pm 2
j=1 λ m,j σ(Z j (x; Wt )) − σ(Z j (x; W 0 ))
lim inf 2 >0 (17)
m→∞ Pm
k=1 λm,k σ(Zk (x; W0 ))
almost surely, or in probability, for some t ∈ (0, ∞] and x. Here t = ∞ refers to the case that the ratio in the
above inequality is the limit as t tends to ∞.
The left-hand-side quantity in Equation (17) corresponds to the relative change in the scaled m-dimensional
feature maps between time 0 and time t. Definition 7.1 matches the definitions of (Yang & Hu, 2021,
Definitions 3.5, H.2, H.9) or (Frei et al., 2023, Proposition 3.2) in the case of symmetrical NTK or mean-field
node scalings.
Remark 7.2. As already noted by (Yang & Hu, 2021, Remark H.10), Definition 7.1 is a relatively weak notion
of feature learning. It only requires a change in the feature map for some t ∈ (0, ∞] and some x, and does
not relate to the relevance of the learn features for prediction. However, we show empirically in Section 8
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, it may still be the case
that the contributions from all the individual nodes remain asymptotically infinitesimally small, in such a
1 In (Yang & Hu, 2021), feature learning was defined in the large-width limits for stable and nontrivial parameterisations.
Nontrivial means that the neural network function fm is not constant in time. Stable means that both the preactivations and
activations have Θ(1) coordinates at 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, as we show theoretically for a linear activation in Section 7.2,
and empirically in Section 8. We introduce below the stronger definition of non-uniform feature learning.
Definition 7.3 (Non-uniform feature learning). Let (w0j )j≥1 be a sequence of random initialisations for
nodes j ≥ 1. We will say that non-uniform feature learning occurs during training if
2
maxj∈[m] λm,j σ(Zj (x; Wt )) − σ(Zj (x; W0 ))
lim inf 2 >0 (18)
m→∞ Pm
k=1 λm,k σ(Zk (x; W0 ))
Note that non-uniform feature learning implies 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).
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 in this linear-activation case, the evolution
of the weights provides useful insights into the differences between the symmetrical and asymmetrical scalings
in terms of weight change.
Theorem 7.4. Assume that the activation function is σ(x) = x, i.e., the identity map. Let (λm,j ) be some
node scalings, not necessarily of the form (1)2 . Let X = UDV⊤ be a reduced singular value decomposition of
X, where U is an n × k matrix with orthonormal columns, D is a diagonal k × k matrix, V is a d × k matrix
with orthonormal columns, and k ≤ min(n, d) is the rank of X. For all j ∈ [m], the difference between the
solution of gradient descent3 /flow w∞j and the initialisation w0j is given by
p
λm,j
w∞j − w0j = Pm aj (β ∞ − VV⊤ β 0 ) (19)
λ
k=1 m,k
Pm p √
where β 0 = j=1 λm,j aj w0j , and β ∞ = d VD−1 U⊤ y is the minimum-norm solution of
1 1
argmax ∥y − √ Xβ∥2 .
β 2 d
The proof of this theorem is given in Appendix I.1.1 in the Supplementary Material.
Theorem 7.4 says that alongp the dimensions spanned by the data, the weight vector of a node j moves by a
quantity proportional to λm,j towards the minimum-norm solution. The form of the learnt function in
Equation (20) implies that the contribution of each hidden node j to the function’s output is proportional to
λm,j aj w∞j . The next theorem analyses the asymptotic behaviour of this contribution in the infinite-width
p
limit and the associated feature learning properties. It shows that, under the scaling of Equation (1), 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-
tion 3.3 holds.
2 Note that we assume none of Assumptions 3.1 to 3.3 here.
P
3 with 2
a step-size less than d(Dmax λ
k m,k
)−1 , where Dmax is the largest entry of 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 γ < 1.
• Under the mean-field scaling, feature learning (Definition 7.1) holds, but non-uniform feature learning
(Definition 7.3) does not.
The contribution from the jth node is marginally normally distributed with
λm,j λm,j
VV⊤ for all m ≥ 1 and 1 ≤ j ≤ m.
p
λm,j aj w∞j ∼ N P β ∞ , λm,j Id − P
k λm,k k λm,k
This implies that, under the mean-field scaling, contributions are asymptotically vanishing. Under the scaling
(1), we have
d
λm,j aj w∞j → N (1 − γ)λ ej β , (1 − γ)λej Id − (1 − γ)λej VV⊤ (21)
p
∞ as m → ∞.
Before moving on to the case of the nonlinear activation function, we analyse the consequence of pruning
nodes of a linear network based on the scaling parameters λm,j . 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 (1). Let ρ ∈ (0, 1). Consider the following pruned network that is obtained by keeping the ⌊ρm⌋ hidden
nodes with largest scalings λm,j and pruning the other nodes:
⌊ρm⌋ ⊤ ⌊ρm⌋
Xp w∞j x 1 X
fem,ρ (x; W∞ ) = λm,j aj √ = √ x⊤
p
λm,j aj w∞j .
j=1
d d j=1
Then, for all ε > 0, we have the following bound on the pruning error:
∥x∥ √ X s X
Pr fem,ρ (x; W∞ ) − fm (x; W∞ ) > ε ≤ √ ∥β ∞ ∥ + d λm,j + d λm,j .
ε d j>⌊ρm⌋ j>⌊ρm⌋
λm,j → γ(1 − ρ) as m → ∞, the above bound implies the following. If γ = 0 (no symmetric
P
Since j>⌊ρm⌋
part), then Pr(|fem,ρ (x; W∞ ) − fm (x; W∞ )| > ε) → 0 and so the network can be compressed to a smaller
network via pruning. Otherwise, the pruning error is controlled by the proportion of the symmetric part γ in
the infinite-width limit:
∥x∥ √
lim Pr fem,ρ (x; W∞ ) − fm (x; W∞ ) > ε ≤ √ ∥β ∞ ∥ + d γ(1 − ρ) + dγ(1 − ρ) .
p
m→∞ ε d
We now analyse feature learning in our model when the activation function is nonlinear. Our analysis assumes
the following two changes in our setup:
Assumption 7.7 (Zeroed initialisation). We assume the model has the following form:
Xm Xm
fm (x; W) = λm,j aj σ(Zj (x; W)) − λm,j aj σ(Zj (x; W0 )) .
p p
j=1 j=1
That is, we subtract from the original model, a duplicate version whose parameters are set to W0 and are
unchanging throughout training. This is a commonly used simplification in theoretical analyses of neural
networks. It ensures that at initialisation, the model satisfies
fm (x; W0 ) = 0 for all x ∈ Rd .
11
Published in Transactions on Machine Learning Research (02/2025)
Assumption 7.8 (Random outputs). We regard the outputs y1 , . . . , yn as random variables, so that the
probabilities in Equations (17) and (18) refer to the randomness of the outputs as well. We further assume
that y1 , . . . , yn are independent and continuous (i.e., the distribution of yi has a density with respect to
Lebesgue measure), and that they are also independent from W0 and the aj ’s. Note that we still treat
the inputs x1 , . . . , xn as deterministic variables. This assumption is met if, for example, there exists a true
generating function f ∗ such that yi = f ∗ (xi ) + ϵi (xi ) where the ϵi (xi ) are independent continuous noise
variables.
Our analysis on feature learning considers gradient descent with a learning rate η that does not depend on m
just as we did in Section 6. This is in contrast to parameterisations such as mean-field parameterisation where
the learning rate has a scaling dependent on m. However, let us note that key reasoning steps in our proofs
also apply to such m-dependent learning rates after minor modifications, allowing us to recover existing
feature-learning results as we will explain shortly after Theorem 7.9. The full proofs of all the theorems in
this subsection are given in Appendix I 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 case.
We show that if the activation function is continuously differentiable and its derivative is always positive, as
in the case of sigmoid, then after the first gradient-descent step, (i) both feature learning and non-uniform
feature learning in Definitions 7.1 and 7.3 occur (Theorem 7.9) and (ii) the squared norm of each weight
vector wj changes almost surely by the amount Ω(λ e1 ) in the infinite-width limit (Theorem 7.10).
Theorem 7.9. Suppose that Assumptions 3.1, 3.3, 7.7 and 7.8 hold. Suppose also that γ < 1 and that the
activation function σ is continuously differentiable with σ ′ (x) > 0 for all x ∈ R. Let i ∈ [n]. Then, both
feature learning and non-uniform feature learning occur after the first gradient-descent step with respect to
the input xi in the almost-sure sense, i.e., the following inequalities hold almost surely:
Pm 2
j=1 λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))
lim inf 2 >0
m→∞ Pm
λ
j=1 m,j σ(Z j (xi ; W 0 ))
2
maxj∈[m] λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))
and lim inf 2 > 0.
m→∞ Pm
j=1 λ m,j σ(Z j (xi ; W 0 ))
Pm Since non-uniform feature learning implies feature learning, we prove the former only.
Sketch of the proof.
The denominator j=1 λm,j (σ(Zj (xi ; W0 )))2 in the condition for non-uniform feature learning converges to
a positive finite value almost surely as m tends to ∞. Thus, it is enough to prove that
2
lim inf max λm,j (σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))) >0 almost surely,
m→∞ j∈[m]
which is implied by
2
lim inf (σ(Z1 (xi ; W1 )) − σ(Z1 (xi ; W0 ))) > 0 almost surely. (22)
m→∞
Note that the limits from above are not redundant since W1 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 (22) holds whenever
2
lim inf (Z1 (xi ; W1 ) − Z1 (xi ; W0 )) > 0 almost surely. (23)
m→∞
12
Published in Transactions on Machine Learning Research (02/2025)
The majority of the detailed proof concerns proving the condition in Equation (23). To that end, we compute
the following m-independent lower bound: for all m,
n ⊤ !2
η 2 λm,1 X
2
w xi ′
Z1 (xi ; W1 ) − Z1 (xi ; W0 ) = yi′ σ ′ √01
xi⊤′ xi (24)
d2 ′ d
i =1
n !2
η 2 (1 − γ)λ
⊤
e1 X w01 xi′
≥ yi′ σ ′
√ ⊤
xi′ xi . (25)
d2 d
i′ =1
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 σ ′ , and
the assumption that γ < 1. Concretely, the assumption on γ implies that η 2 (1 − γ)λ e1 /d2 > 0, while
Assumptions 3.1 and 7.8 and the positivity of σ imply that the squared sum on the right-hand side is almost
′
with probability 1.
Sketch of the proof. The lower bound in Equation (27) is obtained by a relatively straightforward
calculation since the calculation of the gradient is simplified due to Assumption 7.7. The second part of the
theorem follows from the fact that when conditioned on w0j , the absolute value on the right-hand side of
Equation (27) is a continuous random variable that depends only on the yi ’s, and so it is strictly positive
with probability one.
13
Published in Transactions on Machine Learning Research (02/2025)
0.08 25
γ = 1.0
12
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
2
00
0.0 0.0 0.0 0.0 0.1 0.1
20
0.0
0.06 α = 0.7, γ = 0.5
00
α = 0.5, γ = 0.2
kw tj − w 0j k
0.0
Training risk
15
08
0.04 α = 0.4, γ = 0.0
00
0.0
10
6 0
00
0.0
4
0.02 5
04
00
0.0
2
0
02
0.00
00
0
0.0
0 1 2 3 4 5 0 1 2 3 4 5 0 1 2 3 4 5 0 1 2 3 4 5
0.0
Iterations 1e4 Iterations 1e4 Iterations 1e4 Iterations 1e4
Figure 1: Results on simulated data. From left to right, 1) training risks, 2) differences in weight norms
∥wtj − w0j ∥ 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 σ(z) = z/(1 + e−z ). We obtained quantitatively similar
results with the ReLU activation function; see Appendix L in the Supplementary Material.
We first illustrate our theory on simulated data.4 We generate n = 100 observations where for i ∈ [n], xi
Pd
is d = 50 dimensional and sampled uniformly on the unit sphere, and yi = d5 j=1 sin(πxi,j ) + εi where
iid
εi ∼ N (0, 1). We use the FFNN of Section 3.1, with the swish activation function, m = 2000 hidden nodes,
and λm,j as in Equations (1) and (4). We consider the four values (γ, α) ∈ {(1, −), (0.5, 0.7), (0.2, 0.5), (0, 0.4)}.
For each setting, we run GD with a learning rate of 1.0 for 50 000 steps, which is repeated five times to get
average results. We summarise the results in Figure 1, which shows the training error and the evolution of
the weights, NTG, and minimum eigenvalue of the NTG as a function of the GD iterations. We see a clear
correspondence between the theory and the empirical results. For γ > 0, GD achieves near-zero training
error. The minimum eigenvalue and training rates increase with γ. For γ = 1, we have the highest minimum
eigenvalue 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 γ and α, the more feature
learning arises.
To illustrate the benefits of asymmetrical scaling, we consider here the scenario where the function to learn
is a single-unit ReLU, a setting known to be challenging for the lazy-training regime (Malach et al., 2021).
Consider the following data-generating process:
X ∼ N (0, Id )
Y = σ(w0T X)
where w0 = (1, .., 1) ∈ Rd and σ is the ReLU activation function. In this experiment, we sample a training
dataset of n = 100 samples in dimension d = 10. As previously, we train fully connected neural networks
composed of a 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 each run.
The results are reported in Figure 2. We notice that at the end of the training, the error is near zero for all
scalings. Examining the error on the test set, one can see that, as expected, the symmetrical γ = 1 NTK
scaling generalizes poorly in this setting. The asymmetrical scalings, on the other hand, perform significantly
4 The code can found at [Link]
14
Published in Transactions on Machine Learning Research (02/2025)
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.1 Regression
We also validate our model on four regression datasets from the UCI repository5 : concrete ((n, d) = (1030, 9)),
energy ((n, 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 to 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 m = 2000 dimensional vector. Then, we sort feature dimensions with respect to
feature importance measured as (λm,j ∥wtj ∥2 )j∈[m] and use the top-k of these to train an 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 γ and α likely
lead to better transfer learning. A subset of 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 γ and α. Also, we observe that
models with smaller values of γ have lower risks when a small number of features are used for the transfer.
The interpretation is that those models are able to learn a 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 measure risks after pruning. Feature importance is measured as above. Our
theory suggests that models with smaller γ and/or α values are likely to be more robust with respect to
pruning, as long as γ < 1. Wolinski et al. (2020) had similar empirical findings on the benefits of asymmetrical
scaling for network pruning when γ = 0.
5 [Link]
15
Published in Transactions on Machine Learning Research (02/2025)
0.20 60 0.18
γ = 1.0 3.0 γ = 1.0
kw tj − w 0j k
Training risk
Test risk
30 0.12
0.10 1.5
20 0.10
1.0
0.05 10 0.08
0.5
0 0.06
0.00 0.0
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 20 40 60 80 100
Iterations 1e6 Iterations 1e6 Iterations 1e6 Feature dimension
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 ∥wtj − w0j ∥ 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 dataset airfoil,
and 4) test risks of transferred models for dataset plant.
1.0
0.9 0.9
40
0.8 0.8
0.8
0.7
Training accuracy
30 0.7
Test accuracy
Test accuracy
kw tj − w 0j k
0.6 0.6
0.6
γ = 1.0 20 0.5
0.4 0.5
0.4 α = 0.7, γ = 0.5
10 0.3 0.4
α = 0.5, γ = 0.2 0.3
0.2 0.2
α = 0.4, γ = 0.0 0 0.2
0.1
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 500 1000 1500 2000 20 40 60 80 100
Iterations 1e6 Iterations 1e6 Number of pruned nodes Feature dimension
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.
MNIST. We take a 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-entropy loss, we use the MSE loss by treating one-hot class labels as continuous-valued
targets. The outputs of the models are 10 dimensional, so we compute the NTG matrices using only the
first dimension of the outputs. In general, we get similar results in line with our previous experiments. The
pruning and transfer learning results are displayed in Figure 4. For other results, see Figure S.6 in the
Supplementary Material.
CIFAR. We consider a more challenging image classification task of CIFAR–10 and CIFAR–100. The
datasets have 60 000 images with 50 000 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, one uses a large model pre-trained on a general task and then performs
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 run five times, with the learning rate 5.0. Figure 5 shows the pruning results
for the same four values of pairs (γ, α) 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.
We have shown that under an asymmetrical scaling of the nodes of a 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 (Yang & Hu, 2021), and it is defined as a change in the feature map. We proposed a
second definition, called non-uniform feature learning, which additionally requires that the contributions of
some individual nodes remain non-negligible 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)
Figure 5: Results for CIFAR–100. From left to right, 1) test risk through training, 2) differences in weight
norms ∥wtj − w0j ∥ 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 definition holds. We demonstrate empirically that having non-uniform
feature learning is particularly important when we consider transfer learning and pruning. Our definitions of
feature learning relate to the change in the feature map. As already noted by (Yang & 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 in 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 of research is to investigate theoretically the generalisation
properties of such asymmetrical scaling. We note that the approaches used for the symmetric NTK (Arora
et al., 2019a), which rely on the limiting kernel, cannot be applied to our 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 et al., 2019b;a; Oymak &
Soltanolkotabi, 2020; Nguyen et al., 2021). Our results rely on the fact that the minimum eigenvalue κn of the
mean NTK at initialisation is strictly positive; this result was demonstrated by (Du et al., 2019a, Proposition
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-iid, initialisation schemes. In particular, the
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 λ
ej to be summable.
A natural default choice, taken in this article, is to take a power function ∝ j −1/α
where 0 < α < 1. Other
parameterisations are also possible, such as λ ej = 1/K for j = 1, . . . , K and 0 otherwise. One could also
choose other scalings such as (e − 1) exp(−j) or C/(j log2 (j + 1)).
Acknowledgements
We would like to 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 Korean Government (MSIT) (No. RS-2023-00279680). JL acknowledges support
from Institute for Information & communications Technology Planning & Evaluation(IITP) grant funded by
the Korea government(MSIT) (RS-2019-II190075, Artificial Intelligence Graduate School Program(KAIST)).
References
S. Arora, S. Du, W. Hu, Z. Li, and R. Wang. Fine-grained analysis of optimization and generalization
for overparameterized two-layer neural networks. In 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-Eijnden, and J. Bruna. On feature learning in neural networks with global convergence
guarantees. In International Conference on Learning Representations, 2021.
L. Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using
optimal transport. Advances in neural 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-parameterized neural
networks. In International Conference on Learning Representations, 2019b.
N. El Karoui. The spectrum of 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 Neural Information Processing Systems, 33:14820–14830, 2020.
W. Hu, L. Xiao, and J. 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.
In 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. α-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 neural 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 J. Sohl-Dickstein. Deep neural networks 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 linear models under gradient descent. Advances in Neural Information
Processing Systems, 32, 2019.
E. Malach, P. Kamath, E.l Abbe, and N. Srebro. Quantifying the benefit of using differentiable learning over
tangent 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 on Learning Representations, 2018.
S. Mei, 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, 115(33):E7665–E7671, 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. 2388–2464. PMLR, 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. F. 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(1):84–105,
2020.
G. Rotskoff and E. Vanden-Eijnden. Parameters as 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 on Applied Mathematics, 80(2):725–752, 2020.
L. Tao, J. X. Zhi-Qin, M. Zheng, and Z. Yaoyu. Phase diagram for two-layer ReLU neural networks at
infinite-width limit. Journal of Machine Learning Research, 22:1–47, 2021.
Joel A Tropp. User-friendly tail bounds for sums of random matrices. Foundations of computational
mathematics, 12(4):389–434, 2012.
P. Wolinski, G. Charpiat, and Y. Ollivier. Asymmetrical scaling layers for stable network pruning. OpenReview
Archive, 2020.
19
Published in Transactions on Machine Learning Research (02/2025)
20
Published in Transactions on Machine Learning Research (02/2025)
Although we assume a smooth activation function in the main text of the paper (Assumption 3.2), some
of the results remain true when we drop this assumption and use the ReLU activation function instead.
In this section, we explain these results for ReLU. Throughout the section, we assume a weak derivative
σ ′ (x) = 1{x>0} of the ReLU activation function σ.
Our global convergence theorem under gradient flow in the main text (Theorem 5.1) has a counterpart for
the ReLU case, which is given below. This counterpart says that when we train the network with the ReLU
activation, with high probability, the loss decays exponentially fast with respect to κn and the training time
t, and the weights wtj and the NTG matrix respectively change by
P qP
1/2 m 3/2 m 3/2
n3/2
!
nλm,j n2 j=1 λm,j λ
j=1 m,j
O and O + .
κn d1/2 κn d3/2 1/2 5/4
κn d
Theorem A.1 (Global convergence, gradient flow, ReLU). Consider δ ∈ (0, 1). Let D0 = 2C 2 + (2/d).
p
Assume Assumptions 3.1 and 3.3, and the use of the ReLU activation function. Also, assume γ > 0 and
Then, with probability at least 1 − δ, the following properties hold for all t ≥ 0:
(a) eigmin (Θ
b m (X; Wt )) ≥ γκn ;
4
(b) Lm (Wt ) ≤ e−(γκn t)/2 Lm (W0 );
3
nD0
(c) ∥wtj − w0j ∥ ≤ κn d21/2 λm,j for all j ∈ [m];
p
γδ 1/2
1/2
qP
26 n3/2 D0
9 2 Pm 3/2
m 3/2
(d) ∥Θb m (X; Wt ) − Θ
b m (X; W0 )∥2 ≤ 2 n D0
κn d3/2 γδ 5/2 · j=1 λm,j + 1/2 5/4 1/2 5/4 · j=1 λm,j .
κn d γ δ
The proof of the theorem is given in Appendix F, and uses Lemmas E.1 to E.3 and Proposition D.1.
The theorem guarantees that whenever γ > 0, the training
p error converges to 0 exponentially fast. Also, it
implies that the weight change is bounded by a factor λm,j , and the NTG change is bounded by a factor
qP
m 3/2
j=1 λm,j . As we show in Appendix B.2, as m tends to ∞,
m ∞
3/2 e3/2 .
X X
λm,j → (1 − γ)λ
ej for every j ≥ 1, and λm,j → (1 − γ)3/2 λ j
j=1 j=1
Thus, when λ ej > 0 (note that we necessarily have λ e1 > 0), the upper bound in (c) is vanishing in the
infinite-width limit if and only if γ = 1 (NTK regime); similarly, the upper bound in (d) is vanishing if and
only if γ = 1. In fact, both feature learning and non-uniform feature learning in high-probability versions of
Definitions 7.1 and 7.3 occur whenever γ < 1, as we will show in the next subsection.
We present feature-learning results for the ReLU activation. The proofs of the theorems in this subsection
appear in Appendix J
We start with a result that corresponds to Theorem 7.9 in the smooth-activation case. The result says that if
γ < 1 and the activation function is ReLU, then after the first step of gradient descent, both feature learning
and non-uniform feature learning occur in a slightly weaker sense than that of Definitions 7.1 and 7.3 where
we have substituted the almost-sure conditions with corresponding high-probability conditions.
21
Published in Transactions on Machine Learning Research (02/2025)
Theorem A.2. Suppose that Assumptions 3.1, 3.3, 7.7 and 7.8 hold. Suppose also that γ < 1 and that the
ek > 0, then with probability at least 1 − (1/2)k , the following inequalities
activation function σ is ReLU. If λ
hold for all i ∈ [n]:
Pm 2
j=1 λ m,j σ(Z j (xi ; W 1 )) − σ(Z j (x i ; W 0 ))
lim inf 2 >0 (S.1)
m→∞ Pm
j=1 λ m,j σ(Z j (xi ; W 0 ))
2
maxj∈[m] λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))
and lim inf 2 >0 (S.2)
m→∞ Pm
λ
j=1 m,j σ(Z j (x i ; W 0 ))
As we mentioned already, the proof of Theorem A.2 appears in Appendix J. Here we explain the key steps
of the proof. Note that the condition for non-uniform feature learning in Equation (S.2) implies that for
feature learning in Equation (S.1). Thus, we focus on proving the former condition. The crux of proving the
condition in Equation (S.2) lies in the derivation of the following lower bound:
2
lim inf max λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))
m→∞ j∈[m]
( )!
e2 (1 − γ)λ
η 2 c2 (1 − γ)2 λj
ej (w⊤ xi )2
0j
≥ max 1{w⊤ xi ≥0} · min ,
j∈[k] 0j d2 d
where c is a positive real-valued continuous random variable that depends only on the outputs y1 , . . . , yn . In
particular, c does not depend on W0 nor m, and moreover c2 > 0 almost surely. The assumptions of the
theorem and the properties of c imply that the above lower bound is strictly positive if w0j
⊤
xi > 0 for some j,
and this latter condition happens with probability at least 1 − (1/2) , which gives the claim of the theorem.
k
Pm 2
j=1 λ m,j σ(Z j (xi ; W1 )) − σ(Zj (xi ; W 0 ))
lim inf 2 >0
m→∞ Pm
j=1 λm,j σ(Zj (xi ; W0 ))
2
maxj∈[m] λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))
and lim inf 2 > 0.
m→∞ Pm
λ
j=1 m,j σ(Z j (x i ; W 0 ))
Our next result about the ReLU activation function is a counterpart of Theorem 7.10 in the smooth-activation
case. It says that for all j, if γ < 1 and λ
ej > 0, then with probability at least 1/2, the first step of gradient
descent induces a non-zero change in the squared norm of the weight vector wj in the infinite-width limit.
The result also suggests that the change in the squared norm is proportional to λ ej .
22
Published in Transactions on Machine Learning Research (02/2025)
Theorem A.4. Suppose Assumptions 3.1, 3.3, 7.7 and 7.8 hold. Suppose also that the activation function σ
is ReLU. Then, for all j, the following holds almost surely:
n X n
2 (1 − γ)λej X
lim inf ∇wtj L(Wt ) t=0 ≥ yi yi′ xi⊤ xi′ 1{w⊤ x ≥0} 1{w ⊤ x ≥0}
′
. (S.3)
m→∞ d i=1 ′
0j i 0j i
i =1
A.3 Discussion
Theorem A.1 is the counterpart of Theorem 5.1 for the global convergence of gradient flow with the ReLU
activation function. Despite empirical evidence from Appendix L suggesting that similar convergence
results could potentially be applicable to GD in the ReLU context, we have yet to substantiate this with
a comprehensive proof. The proof of the global convergence of GD with smooth activation provided in
Appendix H relies on a Taylor approximation. This necessitates the activation function σ to be twice
differentiable. It is worth noting that, in the symmetric NTK case, the global convergence of GD with the
ReLU activation has been shown by Du et al. (2019b, Section 4). Their proof, however, critically relies on
the fact that the weights remain stationary throughout the iterations of GD, which is not the scenario we are
dealing with here when γ > 0. As such, it remains a compelling open question to determine whether the
global convergence of GD can be proven within our specific framework when employing the ReLU activation
function.
The following proposition summarises a number of useful properties on the scaling parameters defined by (1).
Proposition B.2. For all m ≥ 1,
m
X
λm,j = 1, (S.4)
j=1
m
√ X √
(S.5)
p
γm ≤ λm,j ≤ m.
j=1
23
Published in Transactions on Machine Learning Research (02/2025)
Proof. Equation (S.4) follows from the definition of λm,j as shown below:
m m
! m
X X γ λ
ej X λ
ej
λm,j = + (1 − γ) Pm = γ + (1 − γ) m e
= γ + (1 − γ) = 1.
m
P
k=1 λk k=1 λk
e
j=1 j=1 j=1
In Equation (S.5), the upper bound follows from Cauchy-Schwarz and Equation (S.4), and the lower bound
from the definition of λm,j :
v v
m m
um um √
r uX uX
√ X γ X
γm = 1 = 1 · m.
p
≤ λm,j ≤ t λm,j t
j=1
m j=1 j=1 j=1
For Equation (S.6), we note the following bounds on the sum of the λrm,j for all r > 1:
1/r 1/r r
m r m m m
X (1) r X (2) r
(2)
X X r
λm,j ≤ (λm,j ) ≤ λm,j + λm,j
j=1 j=1 j=1 j=1
Pm (1) r
where the second inequality uses the Minkowski inequality. But as m → ∞, the term j=1 (λm,j ) =
γ r m−(r−1) → 0. Furthermore, as m → ∞,
m m
X (2)
r (1 − γ)r X r X
λm,j = Pm r λ̃j → (1 − γ)r λ̃rj
j=1 k=1 λ̃k j=1 j≥1
γ(m−⌊ρm⌋) Pm e
By sandwiching, m → γ(1 − ρ). Additionally, the series j=1 λj converges to 1. Thus, its tail
Pm
converges to 0 and j=⌊ρm⌋+1 λej → 0.
ζ(2/α)
Figure S.1 shows the value of e2 = as a function of α, when using Zipf weights Equation (4).
P
j≥1 λ j ζ(1/α)2
This proposition holds also under the ReLU activation case. In what follows, we will give a proof that works
for both the smooth activation function and ReLU.
24
Published in Transactions on Machine Learning Research (02/2025)
P∞ e2
Figure S.1: Value of j=1 λ j as a function of α, where (λj )j≥1 are defined as in Equation (4), As α → 1, it
e
converges to 0, which corresponds to the kernel regime.
It is sufficient to look at the convergence of individual entries of the NTG matrix; that is, to show that, for
each pair 1 ≤ i, i′ ≤ n,
m
xi⊤ xi′
γ X ′
Θm (xi , xi′ ; W0 ) = × σ (Zj (xi ; W0 ))σ ′ (Zj (xi′ ; W0 ))
d m j=1
m
(S.8)
(1 − γ) X e ′
+ Pm λj σ (Zj (xi ; W0 ))σ ′ (Zj (xi′ ; W0 ))
k=1 λk j=1
e
tends to
∞
(1 − γ) ⊤ X e ′
γΘ∗ (xi , xi′ ) + xi xi′ λj σ (Zj (xi ; W0 ))σ ′ (Zj (xi′ ; W0 )) (S.9)
d j=1
almost surely as m → ∞. Using the fact that |σ ′ (z)| ≤ 1 and the triangle inequality, the modulus of the
difference between the RHS of Equation (S.8) and Equation (S.9) is upper bounded by
m
xi⊤ xi′ γ 1
X
σ ′ (Zj (xi ; W0 ))σ ′ (Zj (xi′ ; W0 )) − E[σ ′ (Z1 (xi ; W0 ))σ ′ (Z1 (xi′ ; W0 ))]
d m j=1
∞
! m
1 X X
+ (1 − γ) Pm −1 ej +
λ λ
ej
λ
ej
j=1 j=1 j=m+1
m
x⊤ xi′ 1
X
= i γ σ (Zj (xi ; W0 ))σ (Zj (x ; W0 )) − E[σ ′ (Z1 (xi ; W0 ))σ ′ (Z1 (xi′ ; W0 ))]
′ ′
i′
d m j=1
!
m
X
+ 2(1 − γ) 1 − λ
ej
j=1
which
P∞ etends to 0 almost surely as m tends to infinity using the law of large numbers and the fact that
j=1 λj = 1.
The following proposition is a corollary of Lemma 4 in (Oymak & Soltanolkotabi, 2020). It holds under both
the ReLU and smooth activation cases. A proof is included for completeness.
25
Published in Transactions on Machine Learning Research (02/2025)
23 n log n
Proposition D.1. Let δ ∈ (0, 1). Assume Assumptions 3.1 and 3.3, γ > 0, and m ≥ κn d δ . Also, assume
that the activation function satisfies Assumption 3.2 or it is ReLU. Then, with probability at least 1 − δ,
eigmin (Θ
b m (X; W0 )) ≥ eig (Θ b (1) (X; W0 )) > γκn > 0.
min m
2
Proof. We follow here the proof of Lemma 4 in (Oymak & Soltanolkotabi, 2020).
m
b m (X; W) = 1
X
Θ λm,j Aj
d j=1
m m
1 X (1) 1 X (2)
= λm,j Aj + λ Aj
d j=1 d j=1 m,j
where √ √
Aj = diag(σ ′ (Xwj / d))XX⊤ diag(σ ′ (Xwj / d)).
b (1)
Let Θ m (X; W) = d
Pm (1) γ Pm b (1)
j=1 λm,j Aj = md j=1 Aj . Note that eigmin (Θm (X; W)) ≥ eigmin (Θm (X; W)) a.s.,
1 b
and
E[Θb (1) (X; W0 )] = γ Θ
m
b ∗ (X)
where Θ
b ∗ (X) is defined in Equation (10). We have, for all j ≥ 1,
√
∥Aj ∥2 = eigmax (Aj ) ≤ eigmax (diag(σ ′ (Xwj / d))2 ) eigmax (XX⊤ ) ≤ eigmax (XX⊤ )
(S.10)
≤ trace(XX⊤ ) ≤ n.
At initialisation, A1 , A2 , . . . , Am are independent random matrices. Using matrix Chernoff inequalities (see
Proposition B.1), we obtain, for all ϵ ∈ [0, 1),
Pr eigmin (Θ b m (X; W0 )) ≤ (1 − ϵ)γκn ≤ ne−ϵ2 mκn d/(2n) .
b m (X; W0 )) ≤ γκn ≤ δ.
Pr eigmin (Θ
2
The proof technique used to prove Theorems 5.1 and A.1 is similar to that of (Du et al., 2019b) (NTK
scaling). In particular, we provide in this section Lemmas similar to Lemmas 3.2, 3.3 and 3.4 in (Du et al.,
2019b), but adapted to our setting. Lemma E.1 is an adaptation of Lemma 3.3. Lemmas E.2 and E.4 are
adaptations of Lemma 3.2, respectively for the ReLU and smooth activation cases. Lemmas E.3 and E.5 are
adaptations of Lemma 3.4, respectively for the ReLU and smooth activation cases.
E.1 Lemma on exponential decay of the empirical risk and scaling of the weight changes
The following lemma is an adaptation of Lemma 3.3 of (Du et al., 2019b), and applies to both the ReLU and
smooth activation cases. It shows that, if the minimum eigenvalue of the NTG matrix is bounded away from
0, gradient flow converges to a global minimum exponentially fast. Recall that y = (y1 , . . . , yn )⊤ ∈ Rn .
Lemma E.1. Let t > 0 and ζ > 0. Assume Assumption 3.1 and eigmin (Θ b m (X; Ws )) ≥ ζ for all 0 ≤ s ≤ t.
2
Also, assume that the activation function satisfies Assumption 3.2 or it is ReLU. Then,
26
Published in Transactions on Machine Learning Research (02/2025)
d
us = Θ
b m (X; Ws )(y − us ).
ds
It follows that
dLm (Ws ) b m (X; Ws )(y − us ) ≤ − ζ (y − us )⊤ (y − us ) = −ζLm (Ws ).
= −(y − us )⊤ Θ
ds 2
Using Grönwall’s inequality, we obtain
2 n 2
dwsj aj X ′
= σ (Zsj (xi ))xi · (yi − fm (xi ; Ws ))
p
λm,j √
ds d i=1
d n
!2
λm,j X X ′
= σ (Zsj (xi ))xik · (yi − fm (xi ; Ws ))
d
k=1 i=1
d n
! n !
λm,j X X 2 X
≤ xik σ (Zsj (xi )) (yi − fm (xi ; Ws ))
′ 2 2
d
k=1 i=1 i=1
n
! d n !
λm,j X ′ XX
= σ (Zsj (xi )) (yi − fm (xi ; Ws ))
2 2 2
xik
d i=1 k=1 i=1
n
! n d !
λm,j X XX
≤ (yi − fm (xi ; Ws )) 2 2
xik
d i=1 i=1 k=1
nλm,j
≤ ∥y − us ∥2
d
nλm,j
≤ ∥y − u0 ∥2 e−ζs .
d
Integrating and using Minkowski’s integral inequality, we obtain
Z t Z t
d d
∥wtj − w0j ∥ = wsj ds ≤ wsj ds
0 ds 0 ds
r Z t
nλm,j
≤ ∥y − u0 ∥ e−ζs/2 ds
d 0
r
nλm,j 2
≤ ∥y − u0 ∥ .
d ζ
From now on, the proofs for the ReLU and smooth-activation cases slightly differ.
27
Published in Transactions on Machine Learning Research (02/2025)
E.2 Lemma bounding the NTK change and minimum eigenvalue - ReLU case
The next lemma and its proof are similar to Lemma 3.2 in (Du et al., 2019b) and its proof. Recall that
0 < ∥xi ∥ ≤ 1 for every i ∈ [n], and the w0j are iid N (0, Id ).
Lemma E.2. Let δ ∈ (0, 1), and cm,j > 0 for every j ∈ [m]. Assume that Assumptions 3.1 and 3.3 holds
and the activation function is ReLU. Then, with probability at least 1 − δ, the following holds. For every
W = (w1⊤ , . . . , wm ) , if it satisfies
⊤ ⊤
δ 2 cm,j
∥w0j − wj ∥ ≤ for all j ∈ [m],
4
we have
v
m m
2n u
uX
n X (k) (k)
Θ
b (s) (X; W)
m − Θ
b (s) (X; W0 )
m ≤ λm,j cm,j + t λm,j cm,j for all k ∈ [2]
2 d j=1 d j=1
and
v
m m
2nγ t
uX
b (1) (X; W0 )) − nγ
X
eigmin (Θ
b m (X; W)) ≥ eig (Θ cm,j + (S.12)
u
min m cm,j .
dm j=1 dm1/2 j=1
(k)
Define ∇W fm (X; W) to be the n-by-(md) matrix whose i-th row is the md-dimensional row vector
(k)
(∇W fm (xi ; W))⊤ .
Note that for all k ∈ [2],
Θ
b (k) (X; W) − Θ
m
b (k) (X; W0 )
m
2
= (k)
∇W fm (X; W)∇W fm
(k)
(X; W)⊤ (k)
− ∇ W fm (X; W0 )∇W fm
(k)
(X; W0 )⊤
2
2
(k)
≤ ∇W f m (X; W) − ∇W fm
(k)
(X; W0 ) (S.13)
2
+ 2 ∇W fm
(k)
(X; W0 ) (k)
∇W f m (X; W) − ∇W fm
(k)
(X; W0 ) .
2 2
The justification of the inequality from above is given below (which is an expanded version of the three
equations (364-366) in (Bartlett et al., 2021)): for all n-by-(pd) matrices A and B,
1 1
AA⊤ − BB ⊤ = (A − B)(A + B)⊤ + (A + B)(A − B)⊤
2 2 2 2
1
(A − B)(A + B)⊤ 2 + (A + B)(A − B)⊤ 2
≤
2
1
∥A − B∥2 × (A + B)⊤ 2 + ∥A + B∥2 × (A − B)⊤
≤
2 2
= ∥A − B∥2 × ∥A + B∥2
≤ ∥A − B∥2 × (∥A − B + B∥2 + ∥B∥2 )
≤ ∥A − B∥2 × (∥A − B∥2 + 2 ∥B∥2 ) .
(k)
Coming back to the inequality in Equation (S.13), we next bound the two terms ∇W fm (X; W0 ) and
2
(k) (k)
∇W fm (X; W) − ∇W fm (X; W0 ) there.
2
28
Published in Transactions on Machine Learning Research (02/2025)
2 2 n X
X m 2
(k)
∇W f m (X; W0 ) ≤ (k)
∇W fm (X; W0 ) = (k)
∇wj fm (xi ; W0 )
2 F
i=1 j=1
n X m 2
X (k) 2 ∥xi ∥
= λm,j |σ ′ (Zj (xi ; W0 ))|
i=1 j=1
d
m
n X (k) n
≤ λm,j ≤ γk (S.14)
d j=1
d
where γ1 = γ and γ2 = 1 − γ. The second inequality uses the assumption that |σ ′ (x)| ≤ 1 for all x ∈ R and
Pm (k) Pm
∥xi ∥ ≤ 1 for all i ∈ [n]. The third inequality follows from the fact that j=1 λm,j ≤ j=1 λm,j = 1.
For the second term, we recall that Zj (xi ; W) = √1 w⊤ xi .
d j
Using this fact, we derive an upper bound for the
second term as follows:
2
(k)
∇W f m (X; W) − ∇W fm
(k)
(X; W0 )
2
2
≤ (k)
∇W fm (X; W) − (k)
∇W fm (X; W0 )
F
n X
X m 2
= (k)
∇wj fm (xi ; W) − ∇wj fm
(k)
(xi ; W0 )
i=1 j=1
n X m 2
xi
q
(k)
X
= λm,j aj √ [σ ′ (Zj (xi ; W)) − σ ′ (Zj (xi ; W0 ))]
i=1 j=1
d
n m
1 XX 2 (k) 2
= ∥xi ∥ λm,j |σ ′ (Zj (xi ; W)) − σ ′ (Zj (xi ; W0 ))| . (S.15)
d i=1 j=1
In the rest of the proof, we will derive a probabilistic bound on the upper bound just obtained, and show the
conclusions claimed in the lemma.
For any ϵ > 0, i ∈ [n], and j ∈ [m], we define the event
⊤
w0j xi − wj⊤ xi ≤ ∥w0j − wj ∥ ∥xi ∥ ≤ ϵ.
If w0j
⊤
xi > 0, then wj⊤ xi ≤ 0 and thus
⊤
w0j xi ≤ ϵ + wj⊤ xi < ϵ.
Alternatively, if w0j
⊤
xi ≤ 0, then wj⊤ xi > 0 and thus
⊤
−w0j xi ≤ ϵ − wj⊤ xi ≤ ϵ.
29
Published in Transactions on Machine Learning Research (02/2025)
Using the observation that we have just explained and the fact that w0j
⊤
xi ∼ N (0, ∥xi ∥2 ), we obtain, for a
random variable N ∼ N (0, 1),
ϵ ϵ
Pr(Ai,j (ϵ)) ≤ Pr |N | ≤ = erf √
∥xi ∥ ∥xi ∥ 2
v
u 2 ! !
u ϵ
≤ 1 − exp − 4
t √ /π
∥xi ∥ 2
s
2ϵ2 ϵ
≤ 2
≤ , (S.16)
∥xi ∥ π ∥xi ∥
where the second inequality uses erf(x) ≤ 1 − exp(−(4x2 )/π). Let Ψ(W0 ) be the constraint on W =
p
(w1⊤ , . . . , wm ) defined by
⊤ ⊤
δ 2 cm,j ′
W ∈ Ψ(W0 ) ⇐⇒ ∥w0j ′ − wj ′ ∥ ≤ for all j ′ ∈ [m].
4
" #
2
E sup (k)
∇W fm (X; W) − (k)
∇ W fm (X; W0 )
W∈Ψ(W0 ) 2
m
n X
" #
1 X (k) 2
≤ ∥xi ∥2 λm,j E sup |σ ′ (Zj (xi ; W)) − σ ′ (Zj (xi ; W0 ))|
d i=1 j=1 W∈Ψ(W0 )
n X m
1 X (k)
≤ ∥xi ∥2 λm,j Pr (∃W ∈ Ψ(W0 ) s.t. σ ′ (Zj (xi ; W)) ̸= σ ′ (Zj (xi ; W0 )))
d i=1 j=1
n m
1 XX δ 2 cm,j
(k)
= ∥xi ∥2 λm,j Pr ∃wj s.t. ∥w0j − wj ∥ ≤ and σ (wj xi ) ̸= σ (w0j xi )
′ ⊤ ′ ⊤
d i=1 j=1
4
n m
1 XX (k)
∥xi ∥2 λm,j Pr Ai,j (δ 2 cm,j /4)
≤
d i=1 j=1
n X
m
(δ /4)
2 X (k)
≤ ∥xi ∥λm,j cm,j
d i=1 j=1
2 m
n(δ /4) X (k)
≤ λm,j cm,j .
d j=1
The first inequality uses the bound in Equation (S.15), and the fourth inequality uses the inequality derived
in Equation (S.16).
30
Published in Transactions on Machine Learning Research (02/2025)
We bring together the bound on the expectation just shown and also the bounds proved in Equations (S.13)
and (S.14). Recalling that γ1 = γ and γ2 = 1 − γ, we have
" #
E sup Θ
b (k) (X; W) − Θ
m
b (k) (X; W0 )
m
W∈Ψ(W0 ) 2
" #
2
≤E sup (k)
∇W f m (X; W) − (k)
∇ W fm (X; W0 )
W∈Ψ(W0 ) 2
" #
+ 2E sup (k)
∇ W fm (X; W0 ) (k)
∇W fm (X; W) − (k)
∇W fm (X; W0 )
W∈Ψ(W0 ) 2 2
" #
2
≤E sup (k)
∇W f m (X; W) − ∇W fm
(k)
(X; W0 )
W∈Ψ(W0 ) 2
r " #
n
+2 γk E sup (k)
∇W fm (X; W) − ∇W fm
(k)
(X; W0 )
d W∈Ψ(W0 ) 2
" #
2
≤E sup (k)
∇W fm (X; W) − ∇W fm
(k)
(X; W0 )
W∈Ψ(W0 ) 2
v " #
r u
n u 2
(k) (k)
+2 γk tE sup ∇W fm (X; W) − ∇W fm (X; W0 )
d W∈Ψ(W0 ) 2
v
m m
n(δ 2 /4) X (k)
r u
n u n(δ 2 /4) X (k)
≤ λm,j cm,j + 2 γk t λm,j cm,j
d j=1
d d j=1
v
m m
2n u
u X
δ n X (k) (k)
≤ λ cm,j + tγk λm,j cm,j .
2 d j=1 m,j d j=1
The third inequality uses Jensen’s inequality, and the last uses the fact that δ/2 ≥ (δ/2)2 . Hence, for each
k = 1, 2, by Markov inequality, we have, with probability at least 1 − (δ/2),
v
m m
2n √ u
uX
n X (k) (k)
sup Θ
b (k) (X; W) − Θ
m
b (k) (X; W0 )
m ≤ λm,j cm,j + γk t λm,j cm,j .
W∈Ψ(W0 ) 2 d j=1 d j=1
By union bound, the conjunction of the above inequalities for the k = 1 and k = 2 cases holds with probability
at least 1 − δ.
We prove the last remaining claim using the following lemma.
If A and B are real symmetric matrices, then
31
Published in Transactions on Machine Learning Research (02/2025)
Thus,
inf eigmin (Θ
b (1) (X; W))
m
W∈Ψ(W0 )
≥ eigmin (Θ
b (1) (X; W0 )) −
m sup Θ
b (1) (X; W) − Θ
m
b (1) (X; W0 )
m
W∈Ψ(W0 ) 2
v
m m
2n
u X
n X (1) (1)
≥ eigmin (Θ
b (1) (X; W0 )) − λ cm,j +
u
m
tγ λm,j cm,j
d j=1 m,j d j=1
v
m m
2nγ
u
nγ X X
= eigmin (Θ
b (1) (X; W0 )) − cm,j +
u
m
t cm,j
dm j=1 dm1/2 j=1
holds with probability at least 1 − δ. Equation (S.12) then follows from the fact that for all W,
eigmin (Θ
b m (X; W)) ≥ eig (Θ b (1)
m (X; W)).
min
We now bring together the results from Proposition D.1 and Lemmas E.1 and E.2, and identify a sufficient
condition for Theorem A.1, which corresponds to the condition in Lemma 3.4 in (Du et al., 2019b).
Lemma E.3. Consider δ ∈ (0, 1). Assume that Assumptions 3.1 and 3.3 hold, the activation function is
ReLU, and cm,j > 0 for all j ∈ [m]. Also, assume that γ > 0 and
m m
8n log 4n 8n 162 2 X
X n
m ≥ max δ
, cm,j , 2 2 cm,j .
dκn dκn j=1 d κn j=1
Define r
nλm,j 4 δ 2 cm,j
′
Rm,j = ∥y − u0 ∥ and Rm,j = .
d γκn 64
′
If Rm,j < Rm,j for all j ∈ [m] with probability at least 1 − 2δ , then on an event with probability at least 1 − δ,
we have that for all j ∈ [m], Rm,j
′
< Rm,j and the following properties also hold for all t ≥ 0:
(a) eigmin (Θ
b m (X; Wt )) ≥ γκn
4 ;
To see why doing so is sufficient, pick an arbitrary t0 ≥ 0, and assume the above inequality for all s ≥ 0.
Then, by event A and Lemma E.2, for all 0 ≤ s ≤ t0 , we have the following upper bound on the change of
32
Published in Transactions on Machine Learning Research (02/2025)
the Gram matrix from time 0 to s, and the following lower bound on the smallest eigenvalue of Θ
b m (X; Ws ):
2
X
Θ
b m (X; Ws ) − Θ
b m (X; W0 ) ≤ Θ
b (k) (X; Ws ) − Θ
m
b (k) (X; W0 )
m
2 2
k=1
v
m
2
2n u m (k)
uX
X n X (k)
≤ λ cm,j + t λm,j cm,j
d j=1 m,j d j=1
k=1
m √ v
m
2 2·nu
u
nX X
≤ λm,j cm,j + t λm,j cm,j
d j=1 d j=1
and
v
m m
2nγ t
uX
b (1) (X; W0 )) − nγ
X
eigmin (Θ
b m (X; Ws )) ≥ eig (Θ cm,j +
u
min m cm,j
dm j=1 dm1/2 j=1
v
m m
γκn 1 4n 1 8n t
uX
γκn X
cm,j + 1/2 ·
u
≥ − · · cm,j
2 4 m dκn j=1 m dκn j=1
γκn γκn γκn
≥ − = .
2 4 4
We now apply Lemma E.1 with ζ being set to 2 ,
γκn
which gives
and
r
nλm,j 4
∥wt0 j − w0j ∥ ≤ ∥y − u0 ∥ = Rm,j
′
for all j ∈ [m].
d γκn
We have just shown that all the four properties in the lemma hold for t0 .
It remains to prove Equation (S.17) under the event A and the assumption that Rm,j
′
< Rm,j for all j ∈ [m]
holds on this event. Suppose that Equation (S.17) fails for some j ∈ [m]. Let
E.4 Lemma bounding the NTK change and minimum eigenvalue - Smooth activation case
We now give a version of Lemma E.2 for the smooth activation case (that is, under Assumption 3.2). The
proof of this version is similar to the one for Lemma 5 in (Oymak & Soltanolkotabi, 2020), and uses the three
equations (364-366) in (Bartlett et al., 2021).
33
Published in Transactions on Machine Learning Research (02/2025)
Lemma E.4. Assume that Assumptions 3.1 to 3.3 hold. Let cm,j > 0 for every j ∈ [m]. Then, for any fixed
W = (w1⊤ , . . . , wm ) , if it satisfies
⊤ ⊤
cm,j
∥w0j − wj ∥ ≤ for all j ∈ [m],
2
we have
v
m m
nM 2 X (k) 2
uX
nM u (k)
Θ
b (k) (X; W) − Θ
b (k) (X; W0 ) ≤ λ c + λm,j c2m,j for all k ∈ [2]
4d2 j=1
m m m,j m,j
t
2 d3/2
j=1
and
v
2 Xm uXm
nM γ nM γ
eigmin (Θ
b m (X; W)) ≥ eig (Θ b (1) (X; W0 )) − c2 + (S.19)
u
c2 .
4d2 m j=1 m,j d3/2 m1/2 j=1 m,j
min m
t
Note that this lemma has a deterministic conclusion, although its original counterpart (Lemma E.2) has a
probabilistic one.
Proof. The beginning part of the proof is essentially an abbreviated version of the beginning part of the proof
of Lemma E.2. This repetition is intended to help the reader by not forcing her or him to look at the proof of
Lemma E.2 beforehand.
For k ∈ [2], let
m q
(k)
X
(k)
fm (−; W) : Rd → R, (k)
fm (x; W) = λm,j aj σ(Zj (x; W)),
j=1
(k)
and define ∇W fm (X; W) to be the n-by-(pd) matrix whose i-th row is the md-dimensional row vector
(k)
(∇W fm (xi ; W))⊤ .
For all k ∈ [2], we have
Θ
b (k) (X; W) − Θ
m
b (k) (X; W0 )
m
2
= (k)
∇W fm (X; W)∇W fm
(k)
(X; W)⊤ (k)
− ∇ W fm (X; W0 )∇W fm
(k)
(X; W0 )⊤
2
2
(k)
≤ ∇W f m (X; W) − ∇W fm
(k)
(X; W0 ) (S.20)
2
+ 2 ∇W fm
(k)
(X; W0 ) (k)
∇W f m (X; W) − ∇W fm
(k)
(X; W0 ) .
2 2
(k)
To see why this inequality holds, see the proof of Lemma E.2. We bound the two terms ∇W fm (X; W0 )
2
(k) (k)
and ∇W fm (X; W) − ∇W fm (X; W0 ) in Equation (S.20). We bound the first term as follows:
2
2 2 n X
X m 2
(k)
∇W f m (X; W0 ) (k)
≤ ∇W fm (X; W0 ) = (k)
∇wj fm (xi ; W0 )
2 F
i=1 j=1
n X
m 2
X (k) 2 ∥xi ∥
= λm,j |σ ′ (Zj (xi ; W0 ))|
i=1 j=1
d
m
n X (k) n
≤ λm,j ≤ γk
d j=1
d
34
Published in Transactions on Machine Learning Research (02/2025)
where γ1 = γ and γ2 = 1 − γ. The second inequality uses the assumption that |σ ′ (x)| ≤ 1 for all x ∈ R and
Pm (k) Pm
∥xi ∥ ≤ 1 for all i ∈ [n]. The third inequality holds because j=1 λm,j ≤ j=1 λm,j = 1. For the second
term, we recall that |σ ′′ (x)| ≤ M and so σ ′ is M -Lipschitz, and also that Zj (xi ; W) = √1d wj⊤ xi . Using these
facts, we derive an upper bound for the second term as follows:
2
(k)
∇W f m (X; W) − ∇W fm
(k)
(X; W0 )
2
2
(k)
≤ ∇W fm (X; W) − ∇W fm
(k)
(X; W0 )
F
n X
X m 2
= (k)
∇wj fm (xi ; W) − ∇wj fm
(k)
(xi ; W0 )
i=1 j=1
n X m 2
xi
q
(k)
X
= λm,j aj √ [σ ′ (Zj (xi ; W)) − σ ′ (Zj (xi ; W0 ))]
i=1 j=1
d
n m
1 X 2
X (k) 2
= ∥xi ∥ λm,j [σ ′ (Zj (xi ; W)) − σ ′ (Zj (xi ; W0 ))]
d i=1 j=1
n m
1 XX (k) 2
≤ λm,j [σ ′ (Zj (xi ; W)) − σ ′ (Zj (xi ; W0 ))]
d i=1 j=1
n X
2 X m 2
M (k)
⊤
≤ λm,j (wj − w0j ) xi
d2 i=1 j=1
2m
nM X (k) 2
≤ λ ∥wj − w0j ∥
d2 j=1 m,j
m
nM 2 X (k) 2
≤ λ c .
4d2 j=1 m,j m,j
The second to last step uses the Cauchy-Schwartz inequality, and the last step uses our assumption that
c
2 for all j ∈ [m]. From the derived bounds on the first and second terms in the last line of
∥wj − w0j ∥ ≤ m,j
Equation (S.20), it follows that
v
m 2
r u m
nM X (k) 2 n u nM 2 X (k) 2
Θ
b (k) (X; W) − Θ
b (k) (X; W0 ) ≤ λ c +2 γk t 2 λ c
m m
2 4d2 j=1 m,j m,j d 4d j=1 m,j m,j
v
m m
nM 2 X (k) 2
u X
nM u (k)
= λ c + 3/2 tγk λm,j c2m,j .
4d2 j=1 m,j m,j d j=1
eigmin (Θ
b m (X; W)) ≥ eig (Θ
min
b (1) (X; W))
m
≥ eigmin (Θ
b (1) (X; W0 )) − Θ
m
b (1) (X; W) − Θ
m
b (1) (X; W0 )
m .
2
Thus,
v
m 2
uXm
b (1) (X; W0 )) − nM γ nM γ t
X
eigmin (Θ
b m (X; W)) ≥ eig (Θ c2 +
u
c2 .
min m
4d2 m j=1 m,j d3/2 m1/2 j=1 m,j
35
Published in Transactions on Machine Learning Research (02/2025)
E.5 Lemma on a sufficient condition for Theorem 5.1 - Smooth activation case
We now give a version of Lemma E.3 for the smooth activation case (i.e., under Assumption 3.2). It brings
together the results from Proposition D.1 and Lemmas E.1 and E.4, and identifies a sufficient condition for
Theorem A.1, which corresponds to the condition in Lemma 3.4 in (Du et al., 2019b).
Lemma E.5. Assume that Assumptions 3.1 to 3.3 hold. Let δ ∈ (0, 1), and cm,j > 0 for all j ∈ [m]. Assume
that γ > 0 and
m m
8n log δ nM δ
2n 2 2 X 4n M δ
2 2 2 X
m ≥ max , c2 , c2 .
dκn 8d2 κn j=1 m,j d3 κ2n j=1 m,j
′
If Rm,j < Rm,j for all j ∈ [m] with probability at least 1 − 2δ , then on an event with probability at least 1 − δ,
we have that for all j ∈ [m], Rm,j
′
< Rm,j and the following properties also hold for all t ≥ 0:
(a) eigmin (Θ
b m (X; Wt )) ≥ γκn
4 ;
′
(c) ∥wtj − w0j ∥ ≤ Rm,j for all j ∈ [m]; and
qP
nM 2 δ 2
Pm m
(d) ∥Θ
b m (X; Wt ) − Θ
b m (X; W0 )∥2 ≤
82 d2 j=1 λm,j c2m,j + nM δ
23/2 d3/2 j=1 λm,j c2m,j .
Proof. The proof is very similar to that of Lemma E.3, although the concrete bounds in these proofs differ
due to the differences between Lemma E.2 and Lemma E.4.
Suppose Rm,j
′
< Rm,j for all j ∈ [m] on some event A′ having probability at least 1 − 2δ . Also, we would like
to instantiate Proposition D.1 with δ/2, so that its claim holds with probability at least 1 − 2δ . Let A be the
intersection of A′ with the event that the claim in Proposition D.1 holds with δ/2. By the union bound, A
has probability at least 1 − δ. We will show that on the event A, the four claimed properties of the lemma
hold.
It will be sufficient to show that
To see why doing so is sufficient, pick an arbitrary t0 ≥ 0, and assume the above inequality for all s ≥ 0.
Then, by the event A and Lemma E.4, for all 0 ≤ s ≤ t0 , we have the following upper bound on the change of
the Gram matrix from time 0 to s, and the following lower bound on the smallest eigenvalue of Θb m (X; Ws ):
Θ
b m (X; Ws ) − Θ
b m (X; W0 )
2
≤ Θ
b (1) (X; Ws )
m − Θ
b (1) (X; W0 )
m + Θ b (2) (X; Ws ) − Θ
m
b (2) (X; W0 )
m
2 2
v
2 2 m uX m
nM δ X nM δ u
≤ λm,j c2m,j + 3/2 3/2 t λm,j c2m,j
64d2 j=1
2 d j=1
36
Published in Transactions on Machine Learning Research (02/2025)
and
v
m
2 2
uX m
b (1) (X; W0 )) − nM δ γ nM δγ t
X
eigmin (Θ
b m (X; Ws )) ≥ eig (Θ c2 +
u
c2
min m
64d2 m j=1 m,j 4d3/2 m1/2 j=1 m,j
v
m m
γκn 1 nM δ 2 2 X
1
uX
γκn nM δ t
c2 +
u
> − · · c2
2 4 m 16d2 κn j=1 m,j m1/2 d3/2 κn j=1 m,j
γκn 1 1
γκn γκn
≥ − + = .
2 4 2 2 4
We now apply the version of Lemma E.1 for the analytic activation σ, with ζ being set to 2 .
γκn
This application
gives
Lm (Wt0 ) ≤ e−(γκn t0 )/2 Lm (W0 )
and
r
nλm,j 4
∥wt0 j − w0j ∥ ≤ ∥y − u0 ∥ = Rm,j
′
for all j ∈ [m].
d γκn
We have just shown that all the four properties in the lemma hold for t0 .
It remains to prove Equation (S.21) under the event A. Suppose that Equation (S.21) fails for some j ∈ [m].
Let
t1 = inf {t | ∥wtj − w0j ∥ > Rm,j for some j ∈ [m] } .
Then, by the continuity of wtj on t, we have
F Proof of Theorem A.1 on the global convergence of gradient flow (ReLU case)
The proof of Theorem A.1 essentially follows Lemma E.3, which itself follows from the secondary Proposi-
tion D.1 and Lemmas E.1 and E.2, derived in Appendices D and E. Pick δ ∈ (0, 1). Let
s
1 2 · 5122
D = n2 C 2 +
d γ 2 δ 5 κ2n d
where C is the assumed upper bound on the |yi |’s. Assume γ > 0 and
2 2 2 2 !
8n log 4n 8nD 16 n D
m ≥ max δ
, ,
κn d dκn d2 κ2n
37
Published in Transactions on Machine Learning Research (02/2025)
Note that
2 2
m 2 m 2 m
8n 8nD 8nD
X X X
cm,j =
p
· λm,j ≤ · λm,j · m
dκn j=1 dκn j=1
dκn j=1
2
8nD
= · m ≤ m2 ,
dκn
Thus,
m m
8n log 4n 8n 162 2 X
X n
m ≥ max δ
, cm,j , 2 2 cm,j .
dκn dκn j=1 d κn j=1
As a result, we can now employ Lemma E.3. Thus, if we find an event A′ such that the probability of A′ is at
least 1 − (δ/2) and under A′ , we have Rm,j
′
< Rm,j , then the conclusion of Lemma E.3 holds. In particular,
we may further calculate conclusions (c) and (d) of Lemma E.3 as
s
δ 2 cm,j δ2 p 1 2 · 5122
∥wtj − w0j ∥ ≤ ′
Rm,j < Rm,j = = · λm,j · n2 C 2 +
64 64 d γ 2 δ 5 κ2n d
s
8n 1 2 p
= · C +
2 · λm,j ,
κn d1/2 d γ2δ
and
m √ v
m
2 2·nu
uX
nX
∥Θ
b m (X; Wt ) − Θ
b m (X; W0 )∥2 ≤ λm,j cm,j + t λm,j cm,j
d j=1 d j=1
m √ v
2 2·n √ u m 3/2
uX
n X 3/2
= ·D· λm,j + · D·t λm,j
d j=1
d j=1
s
m
1 2 · 5122 X 3/2
n
= · n2 C 2 + · λ
d d γ 2 δ 5 κ2n d j=1 m,j
√ v
u m 3/2
1/4 uX
2 2·n 1 2 512 2
·
+ · n2 C 2 + · t λm,j
d d γ 2 δ 5 κ2n d j=1
s
m
512n2 1 2
3/2
X
= · C 2+ · λ
κn d3/2 d γ 2 δ 5 j=1 m,j
v
u m 3/2
1/4 uX
64n3/2 1 2
+ 1/2 · C2 + · t λm,j .
κn d5/4 d γ 2 δ5 j=1
38
Published in Transactions on Machine Learning Research (02/2025)
n m
X 1 X h i
= yi2 + λm,j E (wj⊤ xi )2 1{w⊤ xi ≥0}
i=1
d j=1
j
1
≤n C + 2
.
d
1 2
2
∥y − u0 ∥ < n C 2 + .
d δ
Let A′ be the corresponding event for the above inequality. Then, under A′ , we have
r
nλm,j 4
Rm,j =
′
∥y − u0 ∥
d γκn
r s
1 2 4
nλm,j
< · n C + 2 ·
d d δ γκn
s
1 2 · 42
= λm,j · n2 C 2 +
p
d γ 2 δκ2n d
δ 2 cm,j δ 2 cm,j
= < = Rm,j .
128 64
Thus, A′ is the desired event.
G Proof of Theorem 5.1 on the global convergence of gradient flow (smooth case)
The proof of the theorem is similar to that of Theorem A.1. It derives from Lemma E.5, which itself follows
from the secondary Proposition D.1 and Lemmas E.1 and E.4, derived in Appendices D and E. Recall that
C1 = sup E[σ(cz)2 ]
c∈(0,1]
where the expectation is taken over the real-valued random variable z with the distribution N (0, 1/d). To see
that C1 is finite, note that since |σ ′ (x)| ≤ 1 for all x ∈ R, we have
39
Published in Transactions on Machine Learning Research (02/2025)
8n 2n 2 n M C 2 + C1 2 n M C 2 + C1
10 3 2 15 4 2
m ≥ max · log , · , ·
κn d δ κ3n d3 γ2δ κ4n d4 γ2δ
s
2 · 642
= λm,j · n2 (C 2 + C1 )
p
cm,j
γ 2 δ 3 κ2n d
where C is the assumed upper bound on the |yi |’s. Note that
m m
nM 2 δ 2 X 2 nM 2 δ 2 X 2 · 642
c = λ m,j · n 2
C 2
+ C 1
8d2 κn j=1 m,j 8d2 κn j=1 γ 2 δ 3 κ2n d
m
nM 2 δ 2 2 2 2 · 642 X
= · n C + C 1 · λm,j
8d2 κn γ 2 δ 3 κ2n d j=1
210 n3 M 2 C 2 + C1
= ×
κ3n d3 γ2δ
and
m m
4n2 M 2 δ 2 X 2 4n2 M 2 δ 2 X 2 · 642
c = λm,j · n 2
C 2
+ C1
d3 κ2n j=1 m,j d3 κ2n j=1 γ 2 δ 3 κ2n d
m
4n2 M 2 δ 2 2 2 2 · 642 X
= · n C + C1 · λm,j
d3 κ2n γ 2 δ 3 κ2n d j=1
215 n4 M 2 C 2 + C1
= 4 4
× .
κn d γ2δ
Thus,
m m
8n log 2n nM 2 2 X
δ 4n 2
M 2 2 X
δ
m ≥ max δ
, c2 , c2 .
dκn 8d2 κn j=1 m,j d3 κ2n j=1 m,j
This allows us to employ Lemma E.5. Hence, it is sufficient to find an event A′ such that the probability of
A′ is at least 1 − (δ/2) and under A′ , we have Rm,j
′
< Rm,j . The desired conclusion then follows from the
conclusion of Lemma E.5, and the below calculations: if ∥wtj − w0j ∥ ≤ Rm,j
′
and Rm,j
′
< Rm,j , then
δcm,j
∥wtj − w0j ∥ < Rm,j =
8 s
δ p 2 · 642
= · λm,j · n2 (C 2 + C1 ) 2 3 2
8 γ δ κn d
s
n 128(C 2 + C1 )
= λm,j ×
p
,
κn d 1/2 γ2δ
40
Published in Transactions on Machine Learning Research (02/2025)
∥Θ
b m (X; Wt ) − Θ
b m (X; W0 )∥2
v
m
nM 2 δ 2 X um
uX
nM δ
≤ 2 2 λm,j cm,j + 3/2 3/2 t
2
λm,j c2m,j
8 d j=1 2 d j=1
m
nM 2 δ 2 X 2 (C + C1 )2 · 64
2 2
= λm,j λ m,j n
43 d2 j=1 γ 2 δ 3 κ2n d
v
m
(C 2 + C1 )2 · 642
uX
nM δ u
+ 3/2 3/2 t λm,j λm,j n 2
2 d j=1
γ 2 δ 3 κ2n d
v
m
3 2 X
27
(C 2
+ ) 2 u m 2 210 (C 2 + C1 )
uX
n M C1 n M
= λ 2
m,j
+ t λ .
κ2n d3 j=1 γ2δ κn d2 j=1 m,j γ2δ
Note that
n
X
2
E[∥y − u0 ∥ ] = yi2 − 2yi E[fm (xi ; W0 )] + E[fm (xi ; W0 )2 ]
i=1
n
X Xm
= yi2 − 2yi · 0 + E λm,j σ(Zj (xi ; W0 ))2
i=1 j=1
n
X m
X
= yi2 + λm,j E σ(Zj (xi ; W0 ))2
i=1 j=1
≤ n C + C1 .
2
2 2
∥y − u0 ∥ < n C 2 + C1 .
δ
Let A′ be the corresponding event for the above inequality. Then, under A′ , we have
r
nλm,j 4
Rm,j =
′
∥y − u0 ∥
d γκn
r
2 4
r
nλm,j
< · n (C 2 + C1 ) ·
d δ γκn
s
2 · 42
= λm,j · n2 (C 2 + C1 ) 2 2
p
γ δκn d
δcm,j δcm,j
= < = Rm,j .
16 8
Thus, A′ is the desired event.
Our convergence proof follows the structure of the convergence proof of (Du et al., 2019a, Theorem 5.1) with
necessary modifications, which in particular account for the changing weights and Gram matrices in our
setup.
41
Published in Transactions on Machine Learning Research (02/2025)
The proof is by induction on the number of gradient-update steps s. Here is a sketch of the proof for the
inductive case. We start by decomposing the error at step s + 1:
where I1 = η Θ
b m (X; Ws )(y − us ) and I2 = (us+1 − us − I1 ). We can then show that with high probability,
both the third and the fourth terms in Equation (S.23) are O(η 2 )∥y − us ∥2 , so that the sum of these terms
can be bounded from above by (ηγκn /4)∥y − us ∥2 if η is sufficiently small. On the other hand, the second
term can be bounded using the minimum eigenvalue of the positive definite Gram matrix:
−2(y − us )⊤ I1 = −2η(y − us )⊤ Θ b m (X; Ws )(y − us )
≤ −2η eigmin (Θ
b m (X; Ws ))∥y − us ∥2 .
We will show that if the network is large enough, with high probability, −2η eigmin (Θ
b m (X; Ws )) in the above
upper bound is at most −3ηγκn /4. Putting all these together gives the required bound: with high probability,
2
∥y − us+1 ∥2 ≤ ∥y − us ∥2 − 2(y − us )⊤ I1 − 2(y − us )⊤ I2 + ∥us+1 − us ∥
3ηγκn ηγκn
≤ ∥y − us ∥2 − ∥y − us ∥2 + ∥y − us ∥2
4 4
ηγκn
≤ 1− ∥y − us ∥2
2
ηγκn s+1
≤ 1− ∥y − u0 ∥2 .
2
Before proving the theorem, we show two useful facts. Let u(W) be the n-dimensional vector
which consists of the network outputs on the training inputs under the parameters W. Note that for each
gradient-update step s ∈ N ∪ {0}, the vector u(Ws ) is equal to us , the notation that we have been using in
the main text of the paper. We also define u′ (W) to be the following n-by-m matrix:
∂u
u′ (W) = .
∂W
42
Published in Transactions on Machine Learning Research (02/2025)
be the Euler discretisation of the gradient flow of the output. Here η > 0 is the learning rate.
Lemma H.1. For all W and j ∈ [m],
∂Lm (W)
p
λm,j n
≤ √ ∥y − u(W)∥.
∂wj d
Proof.
n
!
∂Lm (W) X wj⊤ xi xi
= (u(W)i − yi ) × λm,j aj × σ ′
p
√ ×√
∂wj i=1
d d
n
!
X wj⊤ xi xi
(u(W)i − yi ) × λm,j aj × σ ′
p
≤ √ ×√
i=1
d d
p n
λm,j X
≤ √ × |u(W)i − yi |
d i=1
p
λm,j n
≤ √ ∥y − u(W)∥.
d
The next lemma gives an upper bound on ∥y − us+1 ∥. As we will show shortly, this upper bound will play a
crucial role in the proof of Theorem 6.1.
Lemma H.2. Assume Assumptions 3.1 to 3.3. Then, for all s ∈ N ∪ {0}, we have
2η 2 M n3/2 η 2 n2
∥y − us+1 ∥2 ≤ 1 − 2η eigmin (Θ
b m (s)) +
2
∥y − us ∥ + 2 × ∥y − us ∥2 . (S.24)
d d
Proof. Write
us+1 − us = u
e s+1 − us + us+1 − ue s+1 .
| {z } | {z }
I1 I2
Then, we have
2 2
∥y − us+1 ∥ = ∥(y − us ) − (us+1 − us )∥
2 2
= ∥y − us ∥ − 2(y − us )⊤ (us+1 − us ) + ∥us+1 − us ∥
2 2
= ∥y − us ∥ − 2(y − us )⊤ I1 − 2(y − us )⊤ I2 + ∥us+1 − us ∥ .
(y − us )⊤ I1 = (y − us )⊤ (e
us+1 − us ) = η(y − us )⊤ Θ
b m (s)(y − us )
≥ η eig (Θ b m (s)) ∥y − us ∥2 .
min
n
∂Lm (W) X
L′m (W) = = (u(W)i − yi )u′ (W)i = u′ (W)⊤ (u(W) − y)
∂W i=1
43
Published in Transactions on Machine Learning Research (02/2025)
and
Then,
Also,
n
X n
X
∥L′m (s)∥ = (yi − usi )u′si ≤ |yi − usi | ∥u′si ∥
i=1 i=1
and
m
!!2
⊤
2
X wsj xi ∥xi ∥2 1
∥u′si ∥ = λm,j a2j σ ′
√ ≤ ,
j=1
d d d
n √
1 X n
∥L′m (s)∥ ≤ √ |yi − usi | ≤ √ ∥y − us ∥ .
d i=1 d
Let W(s,r) = Ws − rL′m (s). For j ∈ [m], write w(s,r)j for the part of W(s,r) going to the j-th node. Then,
for all i ∈ [n],
m
! ⊤
!!2
⊤ w(s,r)j xi 2
2 X wsj xi ∥xi ∥
u′si − u′ (W(s,r) )i = λm,j a2j σ′ √ − σ′ √
j=1
d d d
m 2 ∥x ∥2
X ⊤ i
≤ M2 λm,j a2j wsj − w(s,r)j xi
j=1
d2
2 m
M X 2
≤ λm,j wsj − w(s,r)j
d2 j=1
M2 2
≤ 2 Ws − W(s,r) .
d
√
n
Ws − W(s,r) = r ∥L′m (s)∥ ≤ η √ ∥y − us ∥ .
d
44
Published in Transactions on Machine Learning Research (02/2025)
Thus,
n Z
X η ⊤ 2
∥I2 ∥ =
2
u′si − u (W(s,r) )i
′
L′m (s) dr
i=1 r=0
Xn Z η ⊤ 2
≤ u′si − u (W(s,r) )i
′
L′m (s) dr
i=1 r=0
Xn Z η 2
≤ ∥u′si − u′ (W(s,r) )i ∥ × ∥L′m (s)∥ dr
i=1 r=0
n Z η √ √ 2
X ηM n n
≤ ∥y − us ∥ × √ ∥y − us ∥ dr
i=1 r=0 d3/2 d
4 2 3
η M n
= ∥y − us ∥4
d4
2 2
η M n3/2 2
= ∥y − us ∥ .
d2
As the upper bound depends quadratically on η, we can choose it small enough for gradient descent to
converge, as we will show in the proof of Theorem 6.1 in the next subsection.
Recall that ∥y − us+1 ∥2 can be expressed as the sum of four terms:
2 2 2
∥y − us+1 ∥ = ∥y − us ∥ − 2(y − us )⊤ I1 − 2(y − us )⊤ I2 + ∥us+1 − us ∥ . (S.25)
Thus far, we have bounded the second and third terms on the RHS of Equation (S.25):
These bounds lead to the first three terms in the claimed upper bound of Equation (S.24). It remains to get
an appropriate upper bound of the fourth term on the RHS of Equation (S.25).
45
Published in Transactions on Machine Learning Research (02/2025)
Using the bound on the derivative of the loss in Lemma H.1, we complete the proof:
n
X
2
∥us+1 − us ∥ = (u(s+1)i − usi )2
i=1
n
m ⊤
! !!2
w(s+1)j xi ⊤
X X wsj xi
=
p
λm,j aj σ √ −σ √
i=1 j=1
d d
n
m ⊤
! ! 2
w(s+1)j xi ⊤
X X p wsj xi
≤ λm,j aj σ √ −σ √
i=1 j=1
d d
2
n m ⊤ ⊤
X X p w(s+1)j xi wsj xi
≤ λm,j aj √ − √
i=1 j=1
d d
2
n m p
X X λm,j aj
≤ √ ∥w(s+1)j − wsj ∥∥xi ∥
i=1 j=1
d
! m p 2
Xn X λm,j aj
≤ ∥xi ∥2 × √ × ∥w(s+1)j − wsj ∥
i=1 j=1
d
2
m
(W )
p
X λm,j aj ∂L m s
≤n× √ × η
j=1
d ∂w sj
2
m p p
X λm,j aj η λm,j n
≤n× √ × √ ∥y − us ∥
j=1
d d
2
m
η 2 n2 X
≤ 2 ∥y − us ∥2 λm,j
d j=1
η 2 n2
= ∥y − us ∥2 .
d2
Using the lemmas we have just shown, we will prove global convergence of gradient descent. Recall the
assumed bound C on |yi | for every i ≥ 1 in Assumption 3.1, and also
C1 = sup E[σ(cz)2 ]
c∈(0,1]
where the expectation is taken over the real-valued random variable z distributed as N (0, 1/d). As shown in
Appendix G, C1 is finite.
By the argument in Appendix G again, there exists an event E1 such that E1 happens with probability at
least 1 − (δ/2) and conditioned on E1 , we have
2
r
∥y − u0 ∥ < n(C 2 + C1 ) . (S.26)
δ
Meanwhile, by Proposition D.1, there is an event E2 such that E2 happens with probability at least 1 − (δ/2)
and conditioned on E2 , we have
b m (0)) > γκn .
eigmin (Θ (S.27)
2
46
Published in Transactions on Machine Learning Research (02/2025)
Let E3 be the event that is the conjunction of E1 and E2 . This event happens with probability at least 1 − δ,
and under this event, Equations (S.26) and (S.27) both hold.
Condition on E3 . We prove the inequality in Equation (16) by induction on s. The base case of s = 0 is
immediate. To prove the inductive case, assume that s ≥ 1, and that the inequality in Equation (16) holds
for all s′ = 0, 1, . . . , s − 1.
Let α = ηγκn /2 and β = (1 − α)1/2 and
r
ηn 8λm,j (C 2 + C1 )
cm,j = .
1−β δd
Then,
m m
η n 8(C + C1 )
2 2 2
η n 8(C 2 + C1 )
X X 2 2
c2m,j = λ m,j = .
(1 − β)2 (1 − β)2
j=1
δd j=1
δd
s−1
X
∥wsj − w0j ∥ ≤ ∥w(s′ +1)j − ws′ j ∥
s′ =0
s−1
X ∂Lm (Ws′ )
≤ η
∂ws′ j
s′ =0
s−1 r
X λm,j n
≤ η ∥y − us′ ∥
d
s′ =0
r s−1
λm,j n X ′
≤η (1 − α)s /2 ∥y − u0 ∥
d
s′ =0
r
η λm,j n
≤ ∥y − u0 ∥
1−β d
r
2
r
η λm,j n
≤ n(C 2 + C1 )
1−β d δ
r
1 ηn 8λm,j (C + C1 )
2 cm,j
= × =
2 1−β δd 2
where the third inequality uses the bound shown in Lemma H.1, the fourth inequality follows from the
induction hypothesis, and the sixth inequality uses the bound in (S.26). Thus, by Lemma E.4 with cm,j from
above and the lower bound on the minimum eigenvalue in Equation (S.27), we have
eigmin (Θ
b m (s))
v
2 X m uX m
nM γ nM γ
≥ eigmin (Θ
b (1) (X; W0 )) − c2 +
u
c2
4d2 m j=1 m,j d3/2 m1/2 j=1 m,j
m
t
s !
nM 2 γ η n 8(C 2 + C1 ) η 2 n2 8(C 2 + C1 )
2 2
γκn nM γ
= − + 3/2 1/2
2 4d2 m (1 − β)2 δd d m (1 − β)2 δd
√
2η n M γ(C 2 + C1 ) 8ηn2 M γ(C 2 + C1 )1/2
2 3 2
γκn
= − + .
2 d3 m(1 − β)2 δ d2 m1/2 (1 − β)δ 1/2
47
Published in Transactions on Machine Learning Research (02/2025)
b m (s)) + 2η M n ∥y − us ∥ + η n ∥y − us ∥2
2 3/2 2 2
≤ 1 − 2η eigmin (Θ
d2 d2
b m (s)) + 2η M n (1 − α)s/2 ∥y − u0 ∥ + η n ∥y − us ∥2
2 3/2 2 2
≤ 1 − 2η eigmin (Θ
d2 d2
!
2η 2
r
2 3/2 2 2
M n η n
≤ 1 − 2η eigmin (Θ
b m (s)) + (1 − α) s/2
n(C + C1 ) + 2
2 ∥y − us ∥2 .
d2 δ d
Thus, we can complete the proof of this inductive case if we show that
!
2η 2
r
2 3/2 2 2
M n η n ηγκn
2η eigmin (Θ
b m (s)) − (1 − α)s/2
n(C 2 + C1 ) − 2 ≥
d2 δ d 2
which is equivalent to
!
2 ηn2
r
ηM n3/2 γκn
eigmin (Θ
b m (s)) ≥ (1 − α)s/2 n(C 2 + C1 ) + 2 + .
d2 δ 2d 4
We will show this sufficient condition by proving the following stronger inequality (stronger because of the
lower bound on eigmin (Θb m (s)) that we have derived above):
√
2η 2 n3 M 2 γ(C 2 + C1 ) 8ηn2 M γ(C 2 + C1 )1/2
γκn
− +
2 d3 m(1 − β)2 δ d2 m1/2 (1 − β)δ 1/2
!
2 ηn2
r
ηM n3/2 γκn
≥ (1 − α)s/2 n(C 2 + C1 ) + 2 + ,
d2 δ 2d 4
which is equivalent to
√
2η 2 n3 M 2 γ(C 2 + C1 ) 8ηn2 M γ(C 2 + C1 )1/2
γκn
≥ +
4 d3 m(1 − β)2 δ d2 m1/2 (1 − β)δ 1/2
2 ηn2
r
ηM n3/2
+ (1 − α) s/2
n(C 2+C ) +
1 .
d2 δ 2d2
But the four summands on the RHS of the above inequality are at most γκn /16 by the assumed upper bound
on η, the assumed lower bound on m, and the fact that (1 − α) ≤ 1. Thus, the inequality from above holds,
as desired.
48
Published in Transactions on Machine Learning Research (02/2025)
where ⊗ denotes the Kronecker product and B = ( λm,1 a1 . . . λm,m am ) ∈ R1×m . We sometimes view B
p p
as a row vector and write B⊤ to mean the corresponding m-dimensional (column) vector. Let
X = UDV⊤
be a reduced SVD of the data matrix X, where U is a n × k matrix with orthonormal columns, D is a
diagonal k × k matrix, V is a d × k matrix with orthonormal columns, and k ≤ min(n, d) is the rank of X.
Define
1
V ′ = qP (B⊤ ⊗ V) ∈ Rmd×k .
m
j=1 λm,j
Therefore,
qPm
j=1 λm,j
A = U √ D (V′ )⊤
d
The solution of Equation (S.28) under gradient flow or gradient descent with the initialisation W0 is given by
√
d
W∞ = A† y + V⊥ ′
(V⊥ ) W0 = qP
′ ⊤
V′ D−1 U⊤ y + V⊥′
(V⊥ ) W0
′ ⊤
j λ m,j
W0 = V′ (V′ )⊤ W0 + V⊥
′
(V⊥ ) W0 .
′ ⊤
From these facts, we can derive a formula that describes the changes in weights during the training based on
gradient flow or gradient descent:
√
d
W∞ − W0 = qP V′ D−1 U⊤ y + (V⊥′
(V⊥ ) W0 ) − W0
′ ⊤
j λm,j
√ !
d
= B ⊗P⊤
VD U y − (V′ (V′ )⊤ W0 )
−1 ⊤
j λm,j
1
=P B⊤ ⊗ β ∞ − VV⊤ β 0
j λm,j
49
Published in Transactions on Machine Learning Research (02/2025)
Pm p
where β 0 = j=1 λm,j aj w0j and
√ √
β∞ = dX† y = d VD−1 U⊤ y
almost surely as m → ∞; hence the denominator in (17) is of order 1. Similarly, under the mean-field scaling,
m
X 1 2 1 h ⊤ 2 i
m× σ(Z k (x; W 0 )) → E x w 01
m2 d
k=1
almost surely as m → ∞; hence the denominator in (17) is of order 1/m. For the numerator, from
Equation (19), we have
m Pm 2
2 1 λm,j ⊤ 2
Pmj=1
X
λm,j σ(Zj (x; W∞ )) − σ(Zj (x; W0 )) = x (β ∞ − VV ⊤
β ) .
j=1
d ( k=1 λm,k )2 0
Pm 2
λm,j Pm
Under the scaling (1), (Pj=1λ )2 = j=1 λ2m,j → (1 − γ)2 j≥1 λ e2 . Hence feature learning occurs if and
P
m,k
j
k Pm 2
λm,j
only if γ < 1. Under the mean-field scaling, Pj=1 2
= 1/m. Hence feature learning occurs. Additionally,
( λm,k )
k
as the (λm,j )j≥1 are ordered, we have
2 1 λ2 2
max λm,j σ(Zj (x; W∞ )) − σ(Zj (x; W0 )) = max P m,j 2 x⊤ (β ∞ − VV⊤ β 0 )
j=1,...,m d (
k λm,k )
j=1,...,m
2
λm,1 2
= P x⊤ (β ∞ − VV⊤ β 0 ) .
d( k λm,k ) 2
The right-hand side is the sum of two independent Gaussian random vectors, and is therefore a Gaus-
λ λ2
sian random vector, with mean P m,j β and covariance matrix (P m,j ( k λm,k − λm,j )(VV⊤ )2 +
P
λm,k ∞ λm,k )2
k k
2
λ λ
λm,j Id − P m,j
λ
VV⊤ = λm,j (Id − P m,j
λ
VV⊤ ). The distributional convergence in Equation (21)
m,k m,k
k k
then follows from Slutsky’s theorem.
50
Published in Transactions on Machine Learning Research (02/2025)
Meanwhile, we have
X p
λm,j aj w∞j
j>⌊ρm⌋
X Xp X p
λm,j β ∞ − VV⊤ λm,k ak w0k + λm,j Id − λm,j VV⊤ aj w0j
≤
j>⌊ρm⌋ k̸=j j>⌊ρm⌋
X Xp X p
λm,j ∥β ∞ ∥ + VV⊤ λm,k ak w0k + λm,j Id − λm,j VV⊤ aj w0j .
≤
j>⌊ρm⌋ k̸=j j>⌊ρm⌋
Also,
v
u
u 2
Xp u Xp
E VV⊤ λm,k ak w0k ≤u
tE VV
⊤ λm,k ak w0k
k̸=j k̸=j
q
= (1 − λm,j ) trace(VV⊤ )
√
≤ d,
and
v
u
u 2
X p u X p
λm,j Id − λm,j VV⊤ aj w0j λm,j (Id − λm,j VV⊤ ) aj w0j
E ≤u
tE
j>⌊ρm⌋ j>⌊ρm⌋
s X
2
= λm,j trace (Id − λm,j VV⊤ )
j>⌊ρm⌋
s X
= λm,j trace Id − 2λm,j VV⊤ + λ2m,j VV⊤
j>⌊ρm⌋
s X
≤ λm,j × d × (1 − λm,j )2
j>⌊ρm⌋
s X
≤ d λm,j .
j>⌊ρm⌋
51
Published in Transactions on Machine Learning Research (02/2025)
Our proof of Theorem 7.9 relies on the following observation on the linear combinations of continuous
independent real-valued random variables.
Lemma I.1. Let z1 , . . . , zn be independent continuous real-valued random variables. Let B ⊂ R be a finite
subset of the real numbers such that B ̸= {0}. Then, almost surely,
n
X
n
min bi zi > 0.
b∈B \{0,...,0}
i=1
Pn Pn
Proof. Denote S = B n \ {0, ..., 0}. For any b = (b1 , ..., bn ) ∈ S, i=1 bi zi , so that Pr( i=1 bi zi = 0) = 0.
Hence, since S is finite,
n
! ( n )!
X [ X
Pr min bi zi = 0 = Pr bi zi = 0
b∈S
i=1 b∈S i=1
n
!
X X
≤ Pr bi zi = 0
b∈S i=1
= 0.
The proof also uses the our globally-made standard assumption that for every random variable Z ∼ N (0, s2 )
for some s > 0, the expectation E[σ(Z)2 ] is finite and greater than 0.
Proof of Theorem 7.9. Since non-uniform feature Pm learning implies feature learning, we prove the former
only. We start by showing that the denominator j=1 λm,j (σ(Zj (xi ; W0 )))2 in the condition for non-uniform
feature learning converges to a positive finite value almost surely as m tends to ∞. To see this, note
m m
!
X X 1 λ
ej
lim λm,j (σ(Zj (xi ; W0 )))2 = lim γ· + (1 − γ) · Pm σ(Zj (xi ; W0 ))2
m→∞
j=1
m→∞
j=1
m ′
j =1 λ
ej ′
Pm e
limm→∞ j=1 λ j σ(Zj (xi ; W0 ))
m
!
X 1 2
= γ · lim σ(Zj (xi ; W0 )) + (1 − γ) ·
2
Pm e
m→∞
j=1
m limm→∞ ′ λj ′
j =1
∞
X
= γ · EZ∼N (0,∥xi ∥2 /d) σ(Z)2 + (1 − γ) · ej σ(Zj (xi ; W0 ))2 .
λ
j=1
The expectation in the first summand is positive and finite by our globally-made assumption on the activation
function σ. Also, the infinite sum in the second summand is positive almost surely because it is greater than
e1 σ(Z1 (xi ; W0 ))2 but λ
λ e1 σ(Z1 (xi ; W0 ))2 is almost surely positive; λ
e1 > 0 and σ(Z1 (xi ; W0 )) is almost surely
non-zero due to the injectivity of σ and the continuity of the random variable Z1 (x; W0 ). Furthermore, the
sum is almost surely finite as well, because its expectation is EZ∼N (0,∥xi ∥2 /d) [σ(Z)2 ] which is finite by our
globally-made assumption on the activation function σ. Thus, the limit of the denominator is positive and
finite almost surely.
Since the denominator in the condition of non-uniform feature learning converges to a positive finite value
almost surely, the condition holds if
2
lim inf max λm,j (σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))) > 0 almost surely. (S.29)
m→∞ j∈[m]
52
Published in Transactions on Machine Learning Research (02/2025)
Note that the limit here is not redundant since W1 depends on m. The new condition in Equation (S.29)
can be simplified further. It holds whenever
2
lim inf Z1 (xi ; W1 ) − Z1 (xi ; W0 ) > 0. (S.30)
m→∞
To see this, note that by the assumption of the theorem and the inverse function theorem, σ −1 is a well-
defined continuous function and also that Z1 (xi ; W0 ) does not depend on m. As a result, the inequality in
Equation (S.30) implies
2
lim inf σ(Z1 (xi ; W1 )) − σ(Z1 (xi ; W0 )) > 0, (S.31)
m→∞
because otherwise some subsequence of (σ(Z1 (xi ; W1 )))m would converge to σ(Z1 (xi ; W0 )) as m tends to
∞, but then by the continuity of σ −1 , the corresponding subsequence of (Z1 (xi ; W1 ))m would converge to
Z1 (xi ; W0 ), which contradicts Equation (S.30). Now using Equation (S.31), the assumption γ > 0, and the
fact that λ
e1 > 0, we can prove the condition in Equation (S.29) as follows:
2 2
lim inf max λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 )) ≥ lim inf λm,1 σ(Z1 (xi ; W1 )) − σ(Z1 (xi ; W0 ))
m→∞ j∈[m] m→∞
2
≥ lim inf (1 − γ)λ
e1 σ(Z1 (xi ; W1 )) − σ(Z1 (xi ; W0 ))
m→∞
> 0.
We now show that Equation (S.30) holds almost surely. Note that
2 w ⊤ x ⊤
2
11 i w01 xi
Z1 (xi ; W1 ) − Z1 (xi ; W0 ) = √ − √
d d
1 ⊤ 2
= η ∇wtj L(Wt ) t=0 xi
d
n
! !2
⊤
η2 X p w0j xi′ xi⊤′ xi
= yi′ λm,1 a1 σ ′
√ √
d d d
i′ =1
n
! !!2
⊤
η 2 λm,1 X w0j xi′
= yi ′ σ ′
√ xi⊤′ xi
d2 ′ d
i =1
n
! !!2
⊤
η (1 − γ)λ
2 e1 X w 0j x i′
≥ yi′ σ ′ √ xi⊤′ xi .
d2 ′ d
i =1
Since the lower bound from above does not depend on m, we have
n
! !!2
⊤
2 η 2 (1 − γ)λ
e1 X w0j xi′
lim inf Z1 (xi ; W1 ) − Z1 (xi ; W0 ) ≥ 2
y i′ σ ′
√ xi⊤′ xi .
m→∞ d d
i′ =1
Since η 2 (1 − γ)λ
e1 /d2 is positive, this lower bound is positive almost surely whenever the summation inside
the square is positive almost surely.
w⊤
0j xi
Conditioning on w0j and noting that σ ′ √
d
∥xi ∥2 > 0, we have by Lemma I.1 that almost surely
n
! !
⊤
X w0j xi′
yi′ σ ′
√ xi⊤′ xi > 0.
i′ =1
d
We may use this lemma since the yi′ ’s are independent from w0j and so their distributions are unaffected by
the conditioning. Now note that this almost-sure positivity of the summation holds regardless of which value
the conditioned w0j takes. Thus, the summation is positive almost surely without the conditioning. This
completes the proof.
53
Published in Transactions on Machine Learning Research (02/2025)
The proof of the theorem uses the following lemma on quadratic combinations of continuous independent
random variables.
Lemma I.2. Let z1 , . . . , zn be continuous independent real-valued random variables. Let B be an n-by-n
real-valued matrix such that Bii ̸= 0 for some i ∈ [n]. Then, almost surely,
n X
X n
zi zi′ Bii′ > 0.
i=1 i′ =1
Proof. Let i ∈ [n] such that Bii ̸= 0. Then, when viewed as a polynomial on zi ,
n X
X n
zi zi′ Bii′
i=1 i′ =1
is a quadratic polynomial with a non-zero coefficient for the term zi2 . As a result, the zero set of this
polynomial on zi has measure zero with respect to Lebesgue measure, that is, the Lebesgue measure of the set
( n X n
)
X
zi zi zi′ Bii′ = 0 ⊆ R
i=1 i′ =1
is zero (because the zero set of any analytic function has zero Lebesgue measure). Furthermore, zi is a
continuous random variable, and so we have
h i
E 1{Pn Pn zi z ′ B ′ =0} { zi′ | i′ ∈ [n], i′ ̸= i} = 0.
i=1 i′ =1 i ii
As a result,
n X
n
!
X h i
Pr zi zi′ Bii′ = 0 = E 1{Pn Pn zi z ′ B ′ =0}
i=1 i′ =1 i ii
i=1 i′ =1
h h ii
= E E 1{Pn Pn zi z ′ B ′ =0} { zi′ | i′ ∈ [n], i′ ̸= i}
i=1 i′ =1 i ii
= E[0] = 0.
Proof of Theorem 7.10. We first compute a lower bound of the squared norm of the gradient, which does
not depend on m.
n
! 2
⊤
2 X w0j xi xi
∇wtj L(Wt ) t=0 = ′
p
yi λm,j aj σ √ √
i=1
d d
n
! 2
⊤
λm,j X w0j xi
= yi σ ′ √ xi
d i=1
d
n
! 2
⊤
(1 − γ)λej X
′
w0j xi
≥ yi σ √ xi
d i=1
d
d X n X n
! !
⊤ ⊤
(1 − γ)λej X w0j xi w0j xi′
= yi yi′ σ ′
√ σ ′
√ xik xi′ k
d d d
k=1 i=1 i′ =1
n X n
! !!
⊤ ⊤
(1 − γ)λej X w 0j xi w 0j x i ′
= yi yi′ xi⊤ xi′ σ ′ √ σ′ √ .
d i=1 ′ d d
i =1
54
Published in Transactions on Machine Learning Research (02/2025)
Thus,
n X n
! !!
⊤ ⊤
2 (1 − γ)λ
ej X w0j xi w0j xi′
lim inf ∇wtj L(Wt ) t=0
≥ yi yi′ xi⊤ xi′ σ ′ √ σ ′
√ .
m→∞ d i=1 ′
d d
i =1
But by assumption, ((1 − γ)λ ej )/d is positive. The other factor in the lower bound is also positive with
probability one. To see this, note that the yi ’s in the factor are continuous independent random variables,
independent also from w0j , and
!2
⊤
w0j xi
∥xi ∥ σ2 ′
√ > 0 for all i ∈ [n],
d
due to Assumption 3.1 and the assumption that σ ′ > 0. As a result, conditioned on w0j , by Lemma I.2, the
factor is positive almost surely with respect to the conditional distributions of the yi ’s, which are the same
as the original unconditional distributions of them due to the independence of the yi ’s with respect to w0j .
Since this positivity holds regardless of which value w0j takes, it also holds without the conditioning on w0j .
This completes the proof.
Proof. The lemma follows from a straightforward calculation using the fact that fm (x; W0 ) = 0 for all
x ∈ Rd .
Lemma J.2. Assume Assumptions 3.1, 7.7 and 7.8. Then, we have that for all m, j ∈ [m], and i ∈ [n],
( )
2 η 2 c2 (1 − γ)2 λ
e2 (1 − γ)λ
j
ej (w⊤ xi )2
0j
λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 )) ≥ 1{w⊤ x ≥0} · min ,
0j i d2 d
where c depends only on the inputs/outputs (in particular, not depending on m) and is almost surely strictly
positive (almost surely, with respect to the input/output).
55
Published in Transactions on Machine Learning Research (02/2025)
Pn
Denoting δi = ((η λm,j aj )/d) i′ =1 (1{w⊤ x ≥0} xi′ xi )yi′ , we have that
⊤
p
0j i′
! !
⊤ ⊤
w0j xi w0j xi
|σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))| = σ √ + δi −σ √
d d
! !
⊤ ⊤
w0j xi w0j xi
≥ 1{w⊤ x ≥0} σ √ + δi −σ √
0j i
d d
!
⊤
w0j xi
≥ 1{w⊤ x ≥0}
0j i
1{w⊤ xi +√dδi ≥0} |δi | + 1{w⊤ xi +√dδi <0} √ .
0j 0j d
We then get
( )
2 (w0j
⊤
xi )2
λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 )) ≥ 1{w⊤ x ≥0} λm,j min δi2 , .
0j i d
n
!2
η2 X
≥ 2 λm,j 1{w⊤ x ≥0} min bi′ yi′
d 0j i b∈S
i′ =1
Note that the yi ’s are continuous and independent random variables by Assumption 7.8. Thus, by Lemma I.1,
with probability one,
Xn
c = min bi′ yi′ > 0.
b∈S
i′ =1
Proof of Theorem A.2. Note that the condition for non-uniform feature learning in Equation (S.2) implies
that for feature learning in Equation (S.1). Thus, we will prove only the former condition.
By our setup, we have that λm,1 ≥ . . . ≥ λm,k ≥ (1 − γ)λ
ek > 0 for all m. Also, by Lemma J.2, we have that
for 1 ≤ j ≤ k,
( )
2 η 2 c2 (1 − γ)2 λ
e2 (1 − γ)λ
j
ej (w⊤ xi )2
0j
λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 )) ≥ 1{w⊤ xi ≥0} · min , .
0j d2 d
56
Published in Transactions on Machine Learning Research (02/2025)
2
max λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))
j∈[m]
( )!
η 2 c2 (1 − γ)2 λ
e2 (1 − γ)λ
j
ej (w⊤ xi )2
0j
≥ max 1{w⊤ xi ≥0} · min , ,
j∈[k] 0j d2 d
which implies
2
lim inf max λm,j σ(Zj (xi ; W1 )) − σ(Zj (xi ; W0 ))
m→∞ j∈[m]
( )!
e2 (1 − γ)λ
η 2 c2 (1 − γ)2 λj
ej (w⊤ xi )2
0j
≥ max 1{w⊤ xi ≥0} · min , . (S.32)
j∈[k] 0j d2 d
We will show that for all δ ∈ (0, 1/2), with probability at least 1−(1/2+δ)k , the lower bound in Equation (S.32)
is positive and
∞
X
0< λm,j σ(Zj (xi ; W0 ))2 < ∞. (S.33)
j=1
Now condition on E ∩ Eδ′ . Then, there exists some j ≤ k such that w0j
⊤
xi ≥ ϵ. Thus, the lower bound in
Equation (S.32) is positive as shown below:
(
(1 (1 ej ′ (w⊤ ′ xi )2
)!
2 2 2 e2
η c − γ) λj ′ − γ)λ 0j
max 1w⊤ x ≥0 · min ,
j ′ ∈[k] 0j ′ i d2 d
( )
η 2 c2 (1 − γ)2 λ e2 (1 − γ)λ
j
ej (w⊤ xi )2
0j
≥ 1{w⊤ xi ≥0} · min ,
0j d2 d
( )
η 2 c2 (1 − γ)2 λe2 (1 − γ)λ
j
e j ϵ2
≥ min ,
d2 d
> 0.
57
Published in Transactions on Machine Learning Research (02/2025)
(1 − γ)λ ej
xi ≥0} · (wj xi )
⊤ 2
≥ · 1{w⊤
d j
(1 − γ)λ ej
≥ · ϵ2
d
> 0.
Furthermore, without any conditioning, we have
∞
X
λm,j ′ · σ(Zj ′ (xi ; W0 ))2 < ∞
j ′ =1
almost surely, because again without any conditioning, the usual expectation of the right-hand side of the
above inequality is finite as shown below:
∞ ∞ ∞
X X X ∥xi ∥2
λm,j ′ · σ(Zj ′ (xi ; W0 ))2 = λm,j ′ · E σ(Zj ′ (xi ; W0 ))2 =
E λm,j ′ ·
′ ′ ′
2d
j =1 j =1 j =1
∥xi ∥2
= < ∞.
2d
Thus, Equation (S.33) holds with probability at least 1 − (1/2 + δ)k . This completes the proof of the theorem.
We first compute a lower bound for the squared norm of the gradient, which does not depend on m.
n
! 2
⊤
2 X w 0j xi xi
∇wtj L(Wt ) t=0 = ′
p
yi λm,j aj σ √ √
i=1
d d
n
! 2
⊤
λm,j X w0j xi
= yi σ ′ √ xi
d i=1
d
n
! 2
⊤
(1 − γ)λ ej X
′
w0j xi
≥ yi σ √ xi
d i=1
d
d X n X n
! !
⊤ ⊤
(1 − γ)λ ej X w 0j xi w 0j x i′
= yi yi′ σ ′ √ σ′ √ xik xi′ k
d ′ d d
k=1 i=1 i =1
n X n
! !!
⊤ ⊤
(1 − γ)λ ej X w0j xi w0j xi′
= ⊤
yi yi′ xi xi′ σ ′
√ σ ′
√
d i=1 ′
d d
i =1
n X n
(1 − γ)λ
ej X
= yi yi′ xi⊤ xi′ 1{w⊤ xi ≥0} 1{w⊤ x ′ ≥0} .
d i=1 ′
0j 0j i
i =1
Thus,
n X n
2 (1 − γ)λ
ej X
lim inf ∇wtj L(Wt ) t=0
≥ yi yi′ xi⊤ xi′ 1{w⊤ 1 ⊤
x ≥0} {w0j xi′ ≥0} .
m→∞ d i=1 ′
0j i
i =1
58
Published in Transactions on Machine Learning Research (02/2025)
But by assumption, the factor ((1 − γ)λ ej )/d in the lower bound is always positive. The claim of the theorem
follows from the property that the other factor in the lower bound is also positive with probability at least
1/2. In the rest of the proof, we will show why this is so.
Note that
∥xi ∥2 1{w⊤ xi ≥0} > 0.
0j
K.1 Regression
In Figures S.2, S.3, S.4 and S.5 we respectively provide the detailed results for the datasets concrete, energy,
airfoil and plant.
K.2 Classification
We provide in Figure S.6 detailed results for the MNIST dataset, and in Figure S.7 results for the CIFAR–10
dataset. In Figure S.8, we provide further details on the individual impact of the parameter γ ∈ [0, 1]. Recall
that the smaller the value of γ, the more asymmetry is introduced, where γ = 1 recovers the iid model. We
can see from the experiments that pruning performance is improved as γ becomes smaller.
We provide here additional experimental results, as in Appendix K, but with a different activation function.
The experimental setting is the same as described in Section 8, except that the swish activation function is
replaced by the ReLU function. Although our theory does not cover the convergence of GD with the ReLU,
the experimental results obtained in this section are quantitatively similar to those obtained with the swish
function.
L.1 Regression
In Figures S.9, S.10, S.11 and S.12 we respectively provide detailed results for the datasets concrete, energy,
airfoil and plant.
L.2 Classification
We provide in Figures S.13, S.14 and S.15 detailed results for respectively the MNIST, CIFAR10 and
CIFAR100 experiments.
59
Published in Transactions on Machine Learning Research (02/2025)
0.20 0.24
γ = 1.0
α = 0.7, γ = 0.5 0.22
0.15 α = 0.5, γ = 0.2
0.20
Training risk
α = 0.4, γ = 0.0
Test risk
0.10 0.18
0.05 0.16
0.14
0.00
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
60
75
01
2.5
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
0.0
50
50
01
eig min (Θ̂ m (X; W t ))
0.0
2.0
25
40
kw tj − w 0j k
01
0.0
1.5
00
30
01
0.0
75
1.0
00
20
0.0
50
10 0.5 00
0.0
25
00
0
0.0
0.0
00
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
00
0.0
0.6
0.6 γ = 1.0
0.5 α = 0.7, γ = 0.5
0.4 0.5 α = 0.5, γ = 0.2
Training risk
α = 0.4, γ = 0.0
Test risk
0.3 0.4
0.2 0.3
0.1 0.2
0.0 0.1
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.2: Results for the concrete dataset (swish). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
risks for transfer learning, and 7) test risks for transfer learning.
60
Published in Transactions on Machine Learning Research (02/2025)
0.08
γ = 1.0 0.07
0.06 α = 0.7, γ = 0.5 0.06
α = 0.5, γ = 0.2 0.05
Training risk
α = 0.4, γ = 0.0
Test risk
0.04 0.04
0.03
0.02 0.02
0.01
0.00
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
60
3.5
5
00
0.0
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
50 3.0
4
eig min (Θ̂ m (X; W t ))
00
2.5
0.0
40
kw tj − w 0j k
2.0
3
00
30
0.0
1.5
20
2
00
0.0
1.0
10
0.5 1
00
0.0
0 0.0
0
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
00
0.0
0.25
0.2 α = 0.4, γ = 0.0
Test risk
0.20
0.15
0.1 0.10
0.05
0.0 0.00
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.3: Results for the energy dataset (swish). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
risks for transfer learning, and 7) test risks for transfer learning.
61
Published in Transactions on Machine Learning Research (02/2025)
0.40
γ = 1.0
0.30
α = 0.7, γ = 0.5 0.35
0.25 α = 0.5, γ = 0.2
Training risk
Test risk
0.20
0.15 0.25
0.10 0.20
0.05
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
100 1e 5
3.0
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
0
1.5
80
2.5
5
eig min (Θ̂ m (X; W t ))
1.2
kw tj − w 0j k
60 2.0
0
1.0
1.5
5
40
0.7
1.0
0
20 0.5
0.5 5
0.2
0
0
0.0
0.0
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6 Iterations 1e6
0.5 γ = 1.0
0.6
0.4
α = 0.7, γ = 0.5
0.5 α = 0.5, γ = 0.2
Training risk
0.4
0.2 0.3
0.1 0.2
0.0 0.1
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.4: Results for the airfoil dataset (swish). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
risks for transfer learning, and 7) test risks for transfer learning.
62
Published in Transactions on Machine Learning Research (02/2025)
0.062 0.066
γ = 1.0 0.065
0.060 α = 0.7, γ = 0.5 0.064
0.058 α = 0.5, γ = 0.2
0.063
Training risk
Test risk
α = 0.4, γ = 0.0
0.056 0.062
0.054 0.061
0.060
0.052
0.059
0.050
0.058
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
17.5 3.0 1e 5
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
15.0 2.5
1.2
12.5
10.0
0.8
1.5
7.5
0.6
5.0 1.0
0.4
2.5 0.5
0.2
0.0 0.0
0.0
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6 Iterations 1e6
0.18 0.18
γ = 1.0
0.16 0.16
α = 0.7, γ = 0.5
0.14 0.14 α = 0.5, γ = 0.2
Training risk
α = 0.4, γ = 0.0
Test risk
0.12 0.12
0.10 0.10
0.08 0.08
0.06 0.06
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.5: Results for the plant dataset (swish). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
risks for transfer learning, and 7) test risks for transfer learning.
63
Published in Transactions on Machine Learning Research (02/2025)
1.0
0.9
0.8 0.8
0.7
Training accuracy
Test accuracy
0.6 0.6
γ = 1.0 0.5
0.4 α = 0.7, γ = 0.5 0.4
α = 0.5, γ = 0.2 0.3
0.2 α = 0.4, γ = 0.0 0.2
0.1
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
6
40
00
0.7
0.0
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
0.6
5
00
0.5
0.0
kw tj − w 0j k
20 0.4
4
00
0.0
0.3
10 0.2
3
00
0.0
0.1
0
2
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
00
0.0
Iterations 1e6 Iterations 1e6 Iterations 1e6
1.0
0.9
0.8 0.8
0.7
Training accuracy
Test accuracy
0.6 0.6
0.5
0.4 0.4
0.3
0.2 0.2
0.1
500 1000 1500 2000 500 1000 1500 2000
Number of pruned nodes Number of pruned nodes
0.9
0.8 0.8
Training accuracy
0.7
Test accuracy
0.6 0.6
0.5
0.4 0.4
0.3
0.2 0.2
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.6: Results for the MNIST dataset (swish). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
accuracies for pruning, 7) test accuracies for pruning, 8) training accuracies for transfer learning, and 9) test
accuracies for transfer learning.
64
Published in Transactions on Machine Learning Research (02/2025)
Figure S.7: Results for the CIFAR–10 dataset (swish). From left to right and top to bottom, 1) test accuracies
through training, 2) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 3) test risks of the pruned models, and 4) test accuracies of the pruned
models.
65
Published in Transactions on Machine Learning Research (02/2025)
Figure S.8: Results for the CIFAR–10 dataset (swish). Impact of the parameter γ. From left to right and top
to bottom, 1) test accuracies through training, 2) differences in weight norms ∥wtj − w0j ∥ with j’s being the
neurons having the maximum difference at the end of the training, 3) test risks of the pruned models, and 4)
test accuracies of the pruned models.
66
Published in Transactions on Machine Learning Research (02/2025)
0.30
γ = 1.0 0.325
0.25
α = 0.7, γ = 0.5 0.300
0.20 α = 0.5, γ = 0.2 0.275
Training risk
α = 0.4, γ = 0.0
Test risk
0.15 0.250
0.10 0.225
0.200
0.05 0.175
0.00 0.150
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
3.0
60
10
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
0.0
2.5
50
40
06
1.5
0.0
30
04
20 1.0
10 0.0
0.5 02
0.0
0
00
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
0.0
0.6
0.5 γ = 1.0
0.5 α = 0.7, γ = 0.5
0.4 α = 0.5, γ = 0.2
Training risk
0.3
0.2 0.3
0.1 0.2
0.0 0.1
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.9: Results for the concrete dataset (ReLU). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
risks for transfer learning, and 7) test risks for transfer learning.
67
Published in Transactions on Machine Learning Research (02/2025)
0.07 0.08
γ = 1.0 0.07
0.06 α = 0.7, γ = 0.5
0.06
0.05 α = 0.5, γ = 0.2
Training risk
Test risk
0.04
0.03
0.03
0.02
0.02
0.01
0.01
0.00
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
80
0
00 .002 .005 .007 .010 .012 .015 .017 .020
70 3.5
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
0
5
60 3.0
0
eig min (Θ̂ m (X; W t ))
0
2.5
kw tj − w 0j k
50
0
5
40 2.0
0
0
30 1.5
0
5
20
0
1.0
0
10 0
0.5 5
0
0
0
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
0.0
α = 0.4, γ = 0.0
Test risk
0.075 0.08
0.06
0.050
0.04
0.025 0.02
0.000 0.00
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.10: Results for the energy dataset (ReLU). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
risks for transfer learning, and 7) test risks for transfer learning.
68
Published in Transactions on Machine Learning Research (02/2025)
0.40 0.45
γ = 1.0
0.35 0.40
α = 0.7, γ = 0.5
0.30 α = 0.5, γ = 0.2
0.35
Training risk
Test risk
0.20 0.30
0.15
0.25
0.10
0.05 0.20
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
5
02
3.5
0.0
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
80
3.0
0
eig min (Θ̂ m (X; W t ))
02
0.0
60 2.5
kw tj − w 0j k
5
01
2.0
0.0
40
1.5
0
01
0.0
20 1.0
5
00
0.5
0.0
0
0
00
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
0.0
0.6
γ = 1.0
0.4 α = 0.7, γ = 0.5
0.5
α = 0.5, γ = 0.2
Training risk
0.3
0.4 α = 0.4, γ = 0.0
Test risk
0.2 0.3
0.1 0.2
0.0 0.1
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.11: Results for the airfoil dataset (ReLU). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
risks for transfer learning, and 7) test risks for transfer learning.
69
Published in Transactions on Machine Learning Research (02/2025)
0.065
γ = 1.0 0.068
0.060 α = 0.7, γ = 0.5
α = 0.5, γ = 0.2 0.066
Training risk
Test risk
0.055 α = 0.4, γ = 0.0
0.064
0.050 0.062
0.045 0.060
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
10 3.0
0
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
01
0.0
8 2.5
6 2.0
6
00
4
0.0
1.5
4
00
2 1.0
0.0
2
0 0.5
0.0
00
2
0
00
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
0.0
α = 0.4, γ = 0.0
Test risk
0.08 0.08
0.07 0.07
0.06 0.06
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.12: Results for the plant dataset (ReLU). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
risks for transfer learning, and 7) test risks for transfer learning.
70
Published in Transactions on Machine Learning Research (02/2025)
1.0
Training accuracy
0.8 0.8
Test accuracy
0.6 0.6
0.4 0.4
0.2 0.2
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
Iterations 1e6 Iterations 1e6
40 1.0
06
0.0
kΘ̂ m (X; W t ) − Θ̂ m (X; W 0 )k 2
05
0.8
0.0
30
04
0.0
0.6
20
03
0.0
0.4
02
10
0.0
0.2
01
0.0
0
00
0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0 0.0 0.2 0.4 0.6 0.8 1.0
0.0
Iterations 1e6 Iterations 1e6 Iterations 1e6
1.0
0.8
0.8
Training accuracy
Test accuracy
0.6 0.6
0.4 0.4
0.2 0.2
0.9
0.8
0.8
Training accuracy
0.7
Test accuracy
0.6 0.6
0.5
0.4 0.4
0.3
0.2 0.2
20 40 60 80 100 20 40 60 80 100
Feature dimension Feature dimension
Figure S.13: Results for the MNIST dataset (ReLU). From left to right and top to bottom, 1) training risks,
2) test risks, 3) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 4) difference in NTG matrices, 5) minimum NTG eigenvalues, 6) training
accuracies for pruning, 7) test accuracies for pruning, 8) training accuracies for transfer learning, and 9) test
accuracies for transfer learning.
71
Published in Transactions on Machine Learning Research (02/2025)
Figure S.14: Results for the CIFAR–10 dataset (ReLU). From left to right and top to bottom, 1) test accuracies
through training, 2) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having the maximum
difference at the end of the training, 3) test risks of the pruned models, and 4) test accuracies of the pruned
models.
72
Published in Transactions on Machine Learning Research (02/2025)
Figure S.15: Results for the CIFAR–100 dataset (ReLU). From left to right and top to bottom, 1) test
accuracies through training, 2) differences in weight norms ∥wtj − w0j ∥ with j’s being the neurons having
the maximum difference at the end of the training, 3) test risks of the pruned models, and 4) test accuracies
of the pruned models.
73
Published in Transactions on Machine Learning Research (02/2025)
Figure S.16: Visualisation of features for MNIST data. We use the top two PCA components to plot the
points on a 2D space.
Figure S.17: Visualisation of features learnt for the Cifar10 experiment. The models are trained by taking as
input the hidden representation of a ResNet18 trained on ImageNet (first figure on the left). We use the top
two PCA components to plot the points on a 2D space.
This section aims at visualizing the main features learned in the MNIST and CIFAR experiments reported
in the main text. Inspired by (Yang & Hu, 2021), we plot the first two PCA components of the learned
features for MNIST (Figure S.16) and CIFAR10 (Figure S.17) datasets. For the MNIST dataset, as in (Yang
& Hu, 2021), the figures show that the features are quasi-random with the symmetric NTK setting, while
there is more separation under the asymmetric scaling. For the CIFAR10 experiment, which uses pre-trained
features on ImageNet, the features of the symmetric NTK are similar to those of the pre-trained features.
The features obtained by PCA better differentiates between the class.
N Hyper-parameter transfer.
When scaling-up neural networks, hyper-parameters tuning becomes prohibitively expensive. In practice, one
performs hyper-parameter optimization on a smaller version of the model, and uses (transfers) the found
values for training the larger model. However, this requires stability of the optimal parameters. As identified
in (Yang et al., 2022), the standard pytorch implementation is not stable as the width increases, which
can be a major challenge to scale-up models. In this section, we empirically show that the asymmetrical
parameterization enjoys stability of the optimal learning rate. We train FFNN with a single hidden layer on
Cifar10 for different width P = 1024, 2048, 4096. We compare the standard Pytorch parameterization with
the asymmetrical one (γ = 0.2, α = 0.5). The results are reported in S.18. As expected, in the standard
parameterization, the optimal learning rate shifts; as the width increases, the optimal learning rate becomes
smaller. On the other hand, with the asymmetrical scaling, the optimal learning rate remains stable as the
width increases
74
Published in Transactions on Machine Learning Research (02/2025)
Figure S.18: Stability of the optimal learning rate as the width increases. Training error in terms of (top)
accuracy (bottom) cross-entropy for (left) standard parameterisation and (right) asymmetrical parameterisa-
tion.
75