Over-parameterised Shallow Neural Networks with Asymmetrical Node
  Scaling: Global Convergence Guarantees and Feature Learning
We consider the optimisation of large and shallow neural networks via
gradient flow, where the output of each hidden node is scaled by some positive
parameter. We focus on the case where the node scalings are non-identical,
differing from the classical Neural Tangent Kernel (NTK) parameterisation. We
prove that, for large neural networks, with high probability, gradient flow
converges to a global minimum AND can learn features, unlike in the NTK regime.
We also provide experiments on synthetic and real-world datasets illustrating
our theoretical results and showing the benefit of such scaling in terms of
pruning and transfer learning.
Francois Caron, Fadhel Ayed, Paul Jung, Hoil Lee, Juho Lee, Hongseok Yang
# Over-parameterised Shallow Neural Networks

###### Abstract

We consider the optimisation of large and shallow neural networks via gradient flow, where the output of each hidden node is scaled by some positive parameter. We focus on the case where the node scalings are non-identical, differing from the classical Neural Tangent Kernel (NTK) parameterisation. We prove that, for large neural networks, with high probability, gradient flow converges to a global minimum AND can learn features, unlike in the NTK regime. We also provide experiments on synthetic and real-world datasets illustrating our theoretical results and showing the benefit of such scaling in terms of pruning 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 the properties of GD training of over-parameterised (or large-width) neural networks; that is, neural networks where the number \(m\) of hidden nodes is very large. In particular, under a "\(\sqrt{1/m}\)" scaling of the hidden nodes, Jacot et al. (2018) have shown that, as the number of nodes \(m\) tends to infinity, the solution obtained by GD coincides with that of kernel regression under a so-called limiting _Neural Tangent Kernel_ (NTK). Under the same node scaling, called _NTK scaling_, theoretical guarantees for the global convergence and generalisation properties have then been obtained for large (but finite) width neural networks (Du et al., 2019, 2019; Oymak and Soltanolkotabi, 2020; Arora et al., 2019; Bartlett et al., 2021).

However, it has been noted in a number of articles (Chizat et al., 2019; Yang, 2019; Arora et al., 2019) that in this large-width regime, feature learning does not occur; consequently, for large-width neural networks under the NTK scaling, GD training is performed in a _lazy training_ regime, in contrast with the typical feature-learning regime exhibited in deep neural networks.

### Summary of the main contributions

In this work, we investigate global convergence properties and feature learning in gradient-type training of neural networks under a more general asymmetrical node scaling. In particular, we consider that each node \(j=1,\ldots,m\) has a fixed node-specific scaling \(\sqrt{\lambda_{m,j}}\), where

\[\lambda_{m,j}=\gamma\cdot\frac{1}{m}+(1-\gamma)\cdot\frac{\widetilde{\lambda }_{j}}{\sum_{k=1}^{m}\widetilde{\lambda}_{k}} \tag{1}\]
where \(\gamma\in[0,1]\) and \(1\geq\widetilde{\lambda}_{1}\geq\widetilde{\lambda}_{2}\geq\ldots\geq 0\) are nonnegative fixed scalars with \(\sum_{j=1}^{\infty}\widetilde{\lambda}_{j}=1\). Note that \(\gamma=1\) corresponds to the NTK scaling \(\sqrt{1/m}\). If \(\gamma<1\), the node scaling is necessarily asymmetrical for large-width networks.1 For \((\widetilde{\lambda}_{j})_{j\geq 1}\), one can take for instance \(\widetilde{\lambda}_{j}=6\pi^{-2}j^{-2}\) for all \(j\geq 1\); another example is \(\widetilde{\lambda}_{1}=1\) and \(\widetilde{\lambda}_{j}=0\) for \(j\geq 2\).

Footnote 1: That is, for \(m\) large enough, there exists \(j\) and \(k\) such that \(\lambda_{m,j}\neq\lambda_{m,k}\).

We consider a shallow neural network with a ReLU or smooth activation function and without bias, where the first layer weights are trained via gradient flow and empirical risk minimisation under the \(\ell_{2}\) loss. We show that, under similar assumptions as in (Du et al., 2019, 2019) on the data, activation function, and initialisation, when the number of nodes \(m\) is sufficiently large:

* if \(\gamma>0\), then the training error goes to 0 at a linear rate with high probability; and
* feature learning arises if and only if \(\gamma<1\).

We also provide numerical experiments illustrating the theoretical results, and demonstrating empirically that such node-scaling is also useful for network pruning and transfer learning.

### Organisation of the paper

Section 2 introduces the neural network model with asymmetrical node scaling, the gradient flow updates, and the main assumptions on the data, activation function, and initialisation. Section 3 discusses the properties of the NTK at initialisation, and its infinite-width limit. Section 4 derives our main results on the convergence to a global minimum of gradient flow and sketches their proofs. Section 5 gives the main results regarding feature learning. Section 6 describes our experiments on simulated and real datasets, whose results illustrate our theoretical results and their potential applications. Missing detailed proofs can be found in the Appendix.

### Related work

**Large-width neural networks.** The analysis of large-width neural networks dates back to the work of Neal (1996) who showed the connection between Gaussian processes and neural network models in the infinite-width limit. Recent work has further explored this connection under different assumptions (Matthews et al., 2018; Lee et al., 2018; Yang, 2019; Favaro et al., 2020; Bracale et al., 2021; Jung et al., 2021; Lee et al., 2022).

**Large-width neural networks under the 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/\sqrt{m}\)" NTK scaling (Arora et al., 2019, 2019, 2019, 2019; Lee et al., 2019; Zou and Gu, 2019; Oymak and Soltanolkotabi, 2020; Zou et al., 2020). Crucially, when the size of the network is large enough with respect to the size of the training set, the training loss converges to a global minimum at a linear rate under gradient descent. However, under the NTK scaling, the hidden-layer features do not move significantly when the width is large, and the NTK scaling has been coined lazy training regime for this reason (Chizat et al., 2019; Woodworth et al., 2020).

**Large-width neural networks under the mean-field scaling.** An alternative scaling is the "\(1/m\)" mean field scaling (Rotskoff and Vanden-Eijnden, 2018; Mei et al., 2018, 2019; Sirignano and Spiliopoulos, 2020; Chen et al., 2021). The mean field scaling is also equivalent, up to symmetry, to the \(\mu P\) parameterisation of Yang and Hu (2021) in the case of shallow networks. While there is feature learning under the mean-field and \(\mu P\) scalings, the relative contribution of each individual activation to the network output becomes infinitesimal in the large-width limit. In contrast, under asymmetrical scaling, individual activations retain a nontrivial contribution even in the large-width limit.

**Asymmetrical scaling in neural networks.** The idea of using asymmetrical scaling parameters in the context of (stochastic) gradient descent optimisation of deep neural networks has been previously introduced by Wolinski et al. (2020). The focus of the work of Wolinski et al. (2020) was on the (empirical) usefulness in terms of pruning. Indeed, our experiments in Section 6 are also in line with their findings. Wolinski et al. (2020), however, only consider asymmetrical scaling with \(\gamma=0\) (no fixed part), and did not investigate the global convergence properties under such scaling. The property of random neural networks under (random) asymmetrical node scaling in the large-width limit has also been considered by Lee et al. (2022); but this paper did not investigate the training properties under gradient descent or gradient flow.

### Notations

For an integer \(n\geq 1\), let \([n]=\{1,\ldots,n\}\). For a multivariate real-valued function \(f:\mathbb{R}^{n}\rightarrow\mathbb{R}\), the gradient \(\nabla_{\mathbf{v}}f(\mathbf{v})\) is the \(n\)-dimensional column vector of partial derivatives \(\nabla_{\mathbf{v}}f(\mathbf{v})=\left(\frac{\partial f}{\partial v_{1}}( \mathbf{v}),\ldots,\frac{\partial f}{\partial v_{n}}(\mathbf{v})\right)^{\top}\) where \(\mathbf{v}=(v_{1},\ldots,v_{n})^{\top}\). For a square matrix \(B\), we denote its minimum and maximum eigenvalues by \(\mathrm{eig}_{\min}(B)\)
and \(\mathrm{eig}_{\max}(B)\), respectively. For a vector \(\mathbf{v}\in\mathbb{R}^{n}\), we denote by \(B=\mathrm{diag}(\mathbf{v})\) the \(n\)-by-\(n\) diagonal matrix with \(B_{ii}=v_{i}\) for \(i\in[n]\).

## 2 Problem setup

### Statistical model

We consider a shallow feedforward neural network (FFNN) with one hidden layer of \(m\geq 1\) hidden nodes and a scalar output. To simplify the analysis, we assume that there is no bias term. Let \(\mathbf{x}\in\mathbb{R}^{d}\) be some input vector, where \(d\geq 1\) is the input dimension. The model is defined as

\[\begin{split} Z_{j}(\mathbf{x};\mathbf{W})&=\frac{ 1}{\sqrt{d}}\mathbf{w}_{j}^{\top}\mathbf{x},\ \text{ for }j\in[m]\\ f_{m}(\mathbf{x};\mathbf{W})&=\sum_{j=1}^{m}\sqrt {\lambda_{m,j}}a_{j}\sigma(Z_{j}(\mathbf{x};\mathbf{W}))\end{split} \tag{2}\]

where \(f_{m}(\mathbf{x};\mathbf{W})\) is the scalar output of the neural network, \(Z_{j}(\mathbf{x};\mathbf{W})\) is the pre-activation of the \(j\)-th hidden node, \(\sigma:\mathbb{R}\rightarrow\mathbb{R}\) is the activation function, \(\mathbf{w}_{j}\in\mathbb{R}^{d}\) is the column vector of weights between node \(j\) at the hidden layer and the input nodes, and \(a_{j}\in\mathbb{R}\) is the weight between the hidden node \(j\) and the output node. The \(\lambda_{m,j}\)'s for \(j\in[m]\) are nonnegative scaling parameters for hidden nodes. Let \(\mathbf{W}=(\mathbf{w}_{1}^{\top},\dots,\mathbf{w}_{m}^{\top})^{\top}\) be a column vector of dimension \(md\) corresponding to the parameters to be optimised. We assume that \(\sigma\) admits a (weak) derivative \(\sigma^{\prime}\). For \(n\geq 1\), let \(\mathbf{\sigma}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}\) (resp. \(\mathbf{\sigma}^{\prime}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{n}\)) be the vector-valued multivariate function that applies the real-valued function \(\sigma\) (resp. \(\sigma^{\prime}\)) element-wise on each of the \(n\) input variables.

To simplify the analysis, we will assume throughout this article that the output weights \(a_{j}\) are randomly initialised and fixed afterwards: for all \(j\geq 1\),

\[a_{j}\mathop{\sim}\limits^{\text{iid}}\mathrm{Uniform}(\{-1,1\}). \tag{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 \(\lambda_{m,j}\) are fixed and parameterised as in Equation (1). By construction, we have \(\lambda_{m,1}>0\) and \(\sum_{j=1}^{m}\lambda_{m,j}=1\) for all \(m\geq 1\).

Recall that the case \(\gamma=1\) corresponds to the NTK scaling. The case that \(\gamma=0\) and \(\widetilde{\lambda}_{j}=\frac{1}{K}\) for \(j\in[K]\) for some \(K\leq m\) and \(0\) otherwise corresponds to a finite neural network with width \(K\). In the experiments, we will consider that \((\widetilde{\lambda}_{j})_{j\geq 1}\) is the probability mass of a Zipf law:

\[\widetilde{\lambda}_{j}=\frac{1}{\zeta(1/\alpha)}\frac{1}{j^{1/ \alpha}},\ \ j\geq 1 \tag{4}\]

for some \(\alpha\in(0,1)\), where \(\zeta\) is the Riemann zeta function. The parameter \(\alpha\) tunes how quickly \(\widetilde{\lambda}_{j}\) decreases with \(j\), smaller values corresponding to a more rapid decrease and a more asymmetrical node scaling.

### Training with gradient flow

Let \(\mathcal{D}_{n}=\{(\mathbf{x}_{i},y_{i})\}_{i\in[n]}\) be the training dataset, where \(n\geq 1\) is the number of observations. Let \(\mathbf{X}\) denote the \(n\)-by-\(d\) matrix whose \(i\)th row is \(\mathbf{x}_{i}^{\top}\). We want to minimise the empirical risk under the \(\ell_{2}\) loss. Consider the objective function

\[L_{m}(\mathbf{W})=\frac{1}{2}\sum_{i=1}^{n}(y_{i}-f_{m}(\mathbf{ x}_{i};\mathbf{W}))^{2} \tag{5}\]

which is non-convex in general. For a given dataset \(\mathcal{D}_{n}\), width \(m\geq 1\), output weights \(a_{j}\), and scaling parameters \((\lambda_{m,j})_{j\in[m]}\), we aim to estimate the trainable parameters \(\mathbf{W}\) by minimising \(L_{m}(\mathbf{W})\) using continuous-time gradient descent (a.k.a. gradient flow). Let \(\mathbf{W}_{0}\) be some initialisation and \((\mathbf{W}_{t})_{t>0}\) be the solution to the ordinary differential equation (ODE)

\[\frac{d\mathbf{W}_{t}}{dt}=-\nabla_{\mathbf{W}}L_{m}(\mathbf{W}_{t})\]
with \(\lim_{t\to 0}\mathbf{W}_{t}=\mathbf{W}_{0}\). Let \(\mathbf{w}_{tj}\) be the value of the parameter \(\mathbf{w}_{j}\) at time \(t\), and define \(Z_{tj}(\mathbf{x})=Z_{j}(\mathbf{x};\mathbf{W}_{t})\). Note that

\[\nabla_{\mathbf{w}_{j}}f_{m}(\mathbf{x};\mathbf{W}) =\frac{\partial f_{m}(\mathbf{x};\mathbf{W})}{\partial Z_{j}( \mathbf{x};\mathbf{W})}\nabla_{\mathbf{w}_{j}}Z_{j}(\mathbf{x};\mathbf{W})\] \[=\sqrt{\lambda_{m,j}}a_{j}\sigma^{\prime}(Z_{j}(\mathbf{x}; \mathbf{W}))\cdot\frac{1}{\sqrt{d}}\mathbf{x}.\]

Thus, under gradient flow, for \(j\in[m]\) and \(\mathbf{x}\in\mathbb{R}^{d}\),

\[\frac{d\mathbf{w}_{tj}}{dt} =\sum_{i=1}^{n}\nabla_{\mathbf{w}_{j}}f_{m}(\mathbf{x};\mathbf{W} _{t})\cdot(y_{i}-f_{m}(\mathbf{x}_{i};\mathbf{W}_{t}))\] \[=\sqrt{\lambda_{m,j}}\frac{a_{j}}{\sqrt{d}}\sum_{i=1}^{n}\sigma^ {\prime}(Z_{tj}(\mathbf{x}_{i}))\mathbf{x}_{i}\cdot(y_{i}-f_{m}(\mathbf{x}_{i };\mathbf{W}_{t})),\] \[\frac{dZ_{tj}(\mathbf{x})}{dt} =(\nabla_{\mathbf{w}_{j}}Z_{j}(\mathbf{x};\mathbf{W}_{t}))^{ \top}\frac{d\mathbf{w}_{tj}}{dt}\] \[=\sqrt{\lambda_{m,j}}\frac{a_{j}}{d}\mathbf{x}^{\top}\sum_{i=1}^ {n}\sigma^{\prime}(Z_{tj}(\mathbf{x}_{i}))\mathbf{x}_{i}\cdot(y_{i}-f_{m}( \mathbf{x}_{i};\mathbf{W}_{t})).\]

Note that the derivatives associated with each hidden node \(j\) are scaled by \(\sqrt{\lambda_{m,j}}\). We may view this per-node scaling as using different learning rates for different hidden nodes.

For an input \(\mathbf{x}\in\mathbb{R}^{d}\), the output of the neural network therefore satisfies the ODE

\[\frac{df_{m}(\mathbf{x};\mathbf{W}_{t})}{dt} =\nabla_{\mathbf{W}}f_{m}(\mathbf{x};\mathbf{W}_{t})^{\top}\frac {d\mathbf{W}_{t}}{dt}\] \[=\sum_{i=1}^{n}\Theta_{m}(\mathbf{x},\mathbf{x}_{i};\mathbf{W}_{t })\cdot(y_{i}-f_{m}(\mathbf{x}_{i};\mathbf{W}_{t})),\]

where \(\Theta_{m}:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow\mathbb{R}\) is the neural tangent kernel, defined by

\[\Theta_{m}(\mathbf{x},\mathbf{x}^{\prime};\mathbf{W}) =\nabla_{\mathbf{W}}f_{m}(\mathbf{x};\mathbf{W})^{\top}\nabla_{ \mathbf{W}}f_{m}(\mathbf{x}^{\prime};\mathbf{W})\] \[=\frac{1}{d}\mathbf{x}^{\top}\mathbf{x}^{\prime}\sum_{j=1}^{m} \lambda_{m,j}\sigma^{\prime}(Z_{j}(\mathbf{x};\mathbf{W}))\sigma^{\prime}(Z_{ j}(\mathbf{x}^{\prime};\mathbf{W})). \tag{6}\]

The associated neural tangent Gram (NTG) matrix \(\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W})\) is the \(n\)-by-\(n\) symmetric positive semi-definite matrix whose \((i,j)\)-th entry is \(\Theta_{m}(\mathbf{x}_{i},\mathbf{x}_{j};\mathbf{W})\). It takes the form

\[\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W})=\frac{1}{d}\sum_{j=1}^{m}\lambda_{ m,j}\operatorname{diag}\biggl{(}\sigma^{\prime}\biggl{(}\frac{\mathbf{X} \mathbf{w}_{j}}{\sqrt{d}}\biggr{)}\biggr{)}\mathbf{X}\mathbf{X}^{\top} \operatorname{diag}\biggl{(}\sigma^{\prime}\biggl{(}\frac{\mathbf{X}\mathbf{w} _{j}}{\sqrt{d}}\biggr{)}\biggr{)}. \tag{7}\]

### Main assumptions

In this section, we give the main assumptions used in this article. The first set of assumptions is on the training dataset \(\mathcal{D}_{n}=\{(\mathbf{x}_{i},y_{i})\}_{i\in[n]}\). The assumptions are rather mild, and similar to other assumptions used in the literature.

**Assumption 2.1** (Dataset).:
1. All inputs are non-zero and their norms are at most \(1\): \(0<\|\mathbf{x}_{i}\|\leq 1\) for all \(i\geq 1\).
2. For all \(i\neq i^{\prime}\) and all \(c\in\mathbb{R}\), \(\mathbf{x}_{i}\neq\mathbf{\alpha}\mathbf{x}_{i^{\prime}}\).
3. There exists \(C>0\) such that \(|y_{i}|\leq C\) for all \(i\geq 1\).

The next assumption concerns the activation function \(\sigma\). Standard activation functions ((leaky) ReLU, softplus, tanh, sigmoid, linear) satisfy this assumption.

**Assumption 2.2** (Activation function).: The activation function is either

1. the ReLU function \(\sigma(x)=\max(0,x)\), with weak derivative \(\sigma^{\prime}(x)=\mathbf{1}_{\{x>0\}}\); or
2. analytic, with \(|\sigma^{\prime}(x)|\leq 1\) and \(|\sigma^{\prime\prime}(x)|\leq M\) for some \(M>0\), and not a polynomial.
The last assumption, which is standard, is on the initialisation of the weights.

**Assumption 2.3** (Initialisation).: For \(j\in[m]\),

\[\mathbf{w}_{0j}\overset{\text{iid}}{\sim}\mathcal{N}\left(0,\mathrm{I}_{d}\right) \tag{8}\]

where \(\mathrm{I}_{d}\) is the \(d\)-by-\(d\) identity matrix.

## 3 Neural Tangent Kernel at initialisation and its limit

### Mean NTG at initialisation and its minimum eigenvalue

Let \(\mathbf{W}_{0}\) be a random initialisation from Equation (8). Consider the mean NTK at initialisation

\[\Theta^{*}(\mathbf{x},\mathbf{x}^{\prime}) =\mathbb{E}\left[\Theta_{m}(\mathbf{x},\mathbf{x}^{\prime}; \mathbf{W}_{0})\right]\] \[=\frac{1}{d}\mathbf{x}^{\top}\mathbf{x}^{\prime}\mathbb{E}\left[ \sigma^{\prime}\left(\frac{1}{\sqrt{d}}\mathbf{w}_{01}^{\top}\mathbf{x}\right) \sigma^{\prime}\left(\frac{1}{\sqrt{d}}\mathbf{w}_{01}^{\top}\mathbf{x}^{ \prime}\right)\right]. \tag{9}\]

Note that this mean NTK, which is also, by the law of large numbers, the limiting NTK under the \(1/\sqrt{m}\) scaling (Jacot et al., 2018), does not depend on \((\lambda_{m,j})_{j\geq 1}\) nor \(m\). Let \(\widehat{\Theta}^{*}(\mathbf{X})=\mathbb{E}\left[\widehat{\Theta}_{m}( \mathbf{X};\mathbf{W}_{0})\right]\) be the associated \(n\)-by-\(n\) mean NTG matrix at initialisation, whose \((i,i^{\prime})\)-th entry is \(\Theta^{*}(\mathbf{x}_{i},\mathbf{x}_{i^{\prime}})\). Let

\[\kappa_{n}=\mathrm{eig}_{\min}(\widehat{\Theta}^{*}(\mathbf{X})) \tag{10}\]

be the minimum eigenvalue of the mean NTG matrix at initialisation. This minimum eigenvalue plays an important role in the analysis of the global convergence properties in the symmetrical NTK 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, \(\kappa_{n}\) is strictly positive or bounded away from zero. We now state such a result, under the Assumptions of Section 2.3.

**Proposition 3.1** ((Du et al., 2019b, Theorem 3.1) and (Du et al., 2019a, Proposition F.1)).: _Under Assumptions 2.1 to 2.3, \(\kappa_{n}>0\)._

_Remark 3.2_.: Du et al. (2019b,a) make the assumption that the \(\mathbf{x}_{i}\) have unit norm. But their proof holds under the less strict Assumption 2.1(a).

### 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 B, follows from the triangle inequality and the law of large numbers, together with the facts that \(|\sigma^{\prime}(z)|\leq 1\) and \(\sum_{j\geq 1}\widetilde{\lambda}_{j}=1\).

**Proposition 3.3**.: _Consider a sequence \((\mathbf{w}_{0j})_{j\geq 1}\) of iid random vectors distributed as in Assumption 2.3. Assume Assumption 2.2 holds. Then,_

\[\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W}_{0})\to\widehat{\Theta}_{\infty}( \mathbf{X};\mathbf{W}_{0}) \tag{11}\]

_almost surely as \(m\to\infty\), where_

\[\widehat{\Theta}_{\infty}(\mathbf{X};\mathbf{W}_{0})=\gamma\widehat{\Theta}^{ *}(\mathbf{X})+(1-\gamma)\widehat{\Theta}_{\infty}^{(2)}(\mathbf{X};\mathbf{W }_{0})\]

_with \(\widehat{\Theta}_{\infty}^{(2)}(\mathbf{X};\mathbf{W}_{0})\) a random positive semi-definite matrix:_

\[\widehat{\Theta}_{\infty}^{(2)}(\mathbf{X};\mathbf{W}_{0})=\frac{1}{d}\sum_{ j=1}^{\infty}\widetilde{\lambda}_{j}\,\mathrm{diag}\bigg{(}\sigma^{\prime} \bigg{(}\frac{\mathbf{X}\mathbf{w}_{0j}}{\sqrt{d}}\bigg{)}\bigg{)}\mathbf{X} \mathbf{X}^{\top}\,\mathrm{diag}\bigg{(}\sigma^{\prime}\bigg{(}\frac{\mathbf{X }\mathbf{w}_{0j}}{\sqrt{d}}\bigg{)}\bigg{)}.\]

_Also, \(\mathbb{E}[\widehat{\Theta}_{\infty}(\mathbf{X};\mathbf{W}_{0})]=\mathbb{E}[ \widehat{\Theta}_{\infty}^{(2)}(\mathbf{X};\mathbf{W}_{0})]=\widehat{\Theta}^{ *}(\mathbf{X})\), and_

\[\mathbb{E}\left[||\widehat{\Theta}_{\infty}(\mathbf{X};\mathbf{W}_{0})-\widehat {\Theta}^{*}(\mathbf{X})||_{F}^{2}\right]=C_{0}(\mathbf{X})(1-\gamma)^{2}\sum_ {j\geq 1}\widetilde{\lambda}_{j}^{2} \tag{12}\]

_where \(||\cdot||_{F}\) denotes the Frobenius norm, and \(C_{0}(\mathbf{X})\geq 0\) is some positive constant equal to_

\[\sum_{1\leq i,i^{\prime}\leq n}\!\left(\frac{\mathbf{x}_{i}^{\top}\mathbf{x}_ {i^{\prime}}}{d}\right)^{2}\!\mathrm{Var}\bigg{(}\sigma^{\prime}\bigg{(}\frac{ 1}{\sqrt{d}}\mathbf{w}_{01}^{\top}\mathbf{x}_{i}\bigg{)}\sigma^{\prime}\bigg{(} \frac{1}{\sqrt{d}}\mathbf{w}_{01}^{\top}\mathbf{x}_{i^{\prime}}\bigg{)}\bigg{)}.\]
When \(\gamma=1\) (symmetric NTK scaling), the NTG converges to a constant matrix, and solutions obtained by GD coincide with that of kernel regression. Whenever \(\gamma<1\), Proposition 3.3 shows that the NTG is random at initialisation, even in the infinite-width limit, suggesting that we are not operating in the kernel regime asymptotically. As shown in Equation (12), the departure from the kernel regime, as measured by the total variance of the limiting random NTG, can be quantified by the nonnegative constant

\[(1-\gamma)^{2}\left(\sum_{j\geq 1}\widetilde{\lambda}_{j}^{2}\right)\in[0,1].\]

When this constant is close to 0, we approach the kernel regime; increasing this value leads to a departure from the kernel regime, and increases the amount of feature learning (see Section 5). The quantity \(\sum_{j\geq 1}\widetilde{\lambda}_{j}^{2}\in(0,1]\) is always strictly positive. More rapid decrease of the \(\widetilde{\lambda}_{j}\) as \(j\) increases will lead to higher values of \(\sum_{j\geq 1}\widetilde{\lambda}_{j}^{2}\). For example, when using the Zipf weights in Equation (4), we have \(\sum_{j\geq 1}\widetilde{\lambda}_{j}^{2}=\frac{\zeta(2/\alpha)}{\zeta(1/ \alpha)^{2}}\), which decreases with \(\alpha\), as shown in Figure 5 in the Appendix.

In this section, we described the behaviour of the NTG at initialisation in the infinite-width limit, and provided some intuition on the node scaling parameters. In the next two sections, we state results on the global convergence and feature learning properties of large, but finite, neural networks under such asymmetrical scaling.

## 4 Global convergence analysis

### Main results

**ReLU case.** Our main theorem for ReLU activation functions, which is given below, explains what happens during training via gradient flow. It says that with high probability, the loss decays exponentially fast with respect to \(\kappa_{n}\) and the training time \(t\), and the weights \(\mathbf{w}_{tj}\) and the NTG matrix change by \(O\left(\frac{n^{\lambda_{m,j}^{1/2}}}{\kappa_{n}d^{1/2}}\right)\) and \(O\left(\frac{n^{2}\sum_{j=1}^{\sum_{m=1}^{m}\lambda_{m,j}^{3/2}}}{\kappa_{n}d^ {3/2}}+\frac{n^{3/2}\sqrt{\sum_{j=1}^{\sum_{m=1}^{m}\lambda_{m,j}^{3/2}}}}{ \kappa_{n}^{\lambda_{m}^{\prime\prime}}d^{5/4}}\right)\), respectively.

**Theorem 4.1** (Global convergence, ReLU).: _Consider \(\delta\in(0,1)\). Let \(D_{0}=\sqrt{2C^{2}+(2/d)}\). Assume Assumption 2.1, Assumption 2.2(a), and Assumption 2.3. Assume \(\gamma>0\) and_

\[m\geq\max\left(\frac{2^{3}n\log\frac{4n}{\kappa_{n}d}}{\kappa_{n}d},\ \frac{2^{ 25}n^{4}D_{0}^{2}}{\kappa_{n}^{4}d^{3}\gamma^{2}\delta^{5}},\ \frac{2^{35}n^{6}D_{0}^{2}}{\kappa_{n}^{6}d^{5}\gamma^{2}\delta^{5}}\right).\]

_Then, with probability at least \(1-\delta\), the following properties hold for all \(t\geq 0\):_

1. \(\mathrm{eig}_{\min}(\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W}_{t}))\geq\frac {2\kappa_{n}}{4}\)_;_
2. \(L_{m}(\mathbf{W}_{t})\leq e^{-(\gamma\kappa_{n}t)/2}L_{m}(\mathbf{W}_{0})\)_;_
3. \(\|\mathbf{w}_{tj}-\mathbf{w}_{0j}\|\leq\frac{2^{3}nD_{0}}{\kappa_{n}d^{1/2} \gamma^{6/2}\sqrt{\lambda_{m,j}}}\) _for all_ \(j\in[m]\)_;_
4. \(\|\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W}_{t})-\widehat{\Theta}_{m}( \mathbf{X};\mathbf{W}_{0})\|_{2}\leq\left(\frac{2^{9}n^{2}D_{0}}{\kappa_{n}d^ {3/2}\gamma^{6/2}}\cdot\sum_{j=1}^{m}\lambda_{m,j}^{3/2}\right)+\left(\frac{2 ^{6}n^{3/2}D_{0}^{1/2}}{\kappa_{n}^{1/2}d^{5/4}\gamma^{1/2}d^{5/4}}\cdot \sqrt{\sum_{j=1}^{m}\lambda_{m,j}^{3/2}}\right)\)_._

The above theorem implies that, whenever \(\gamma>0\), the training error converges to 0 exponentially fast. Additionally, the weight change is bounded by a factor \(\sqrt{\lambda_{m,j}}\) and the NTG change is bounded by a factor \(\sqrt{\sum_{j=1}^{m}\lambda_{m,j}^{3/2}}\). We have (see Appendix A.2), as \(m\rightarrow\infty\), \(\lambda_{m,j}\rightarrow(1-\gamma)\widetilde{\lambda}_{j},\ \ j\geq 1\), and \(\sum_{j=1}^{m}\lambda_{m,j}^{3/2}\rightarrow(1-\gamma)^{3/2}\sum_{j=1}^{\infty} \widetilde{\lambda}_{j}^{3/2}\). If \(\widetilde{\lambda}_{j}>0\) (note that we necessarily have \(\widetilde{\lambda}_{1}>0\)) the upper bound in (c) is therefore vanishing in the infinite-width limit if and only if \(\gamma=1\) (NTK regime); similarly, the upper bound in (d) is vanishing if and only if \(\gamma=1\). Although we were not able to obtain matching lower bounds, we show in Section 5 that indeed feature learning arises whenever \(\gamma<1\).

**Smooth activation case.** We now state a similar result in the smooth activation case. First define

\[C_{1}=\sup_{c\in(0,1]}\mathbb{E}_{z\sim\mathcal{N}(0,1)}\left[\sigma\left(\frac {cz}{\sqrt{d}}\right)^{2}\right]. \tag{13}\]
**Theorem 4.2**.: _(Global convergence, smooth activation) Consider \(\delta\in(0,1)\). Assume Assumption 2.1, Assumption 2.2(b), and Assumption 2.3. Assume \(\gamma>0\) and_

\[m\geq\max\bigg{(}\frac{2^{3}n\log\frac{2n}{\delta}}{\kappa_{n}d},\ \frac{2^{10}n^{3}M^{2}(C^{2}+C_{1})}{ \kappa_{n}^{4}d^{3}\gamma^{2}\delta},\frac{2^{15}n^{4}M^{2}(C^{2}+C_{1})}{ \kappa_{n}^{4}d^{4}\gamma^{2}\delta}\bigg{)}.\]

_Then, with probability at least \(1-\delta\), the following properties hold for all \(t\geq 0\):_

1. \(\mathrm{eig}_{\min}(\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W}_{t}))\geq\frac {2\pi_{n}}{4}\)_;_
2. \(L_{m}(\mathbf{W}_{t})\leq e^{-(\gamma\kappa_{n},t)/2}L_{m}(\mathbf{W}_{0})\)_;_
3. \(\|\mathbf{w}_{tj}-\mathbf{w}_{0j}\|\leq\sqrt{\lambda_{m,j}}\times\frac{n}{ \kappa_{n}d^{2/3}}\sqrt{\frac{2^{7}(C^{2}+C_{1})}{7\delta}}\) _for all_ \(j\in[m]\)_;_
4. \(\|\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W}_{t})-\widehat{\Theta}_{m}( \mathbf{X};\mathbf{W}_{0})\|_{2}\leq\Big{(}\frac{2^{7}n^{3}M^{2}(C^{2}+C_{1})} {\kappa_{n}^{2}d^{3}\gamma^{2}\delta}\cdot\sum_{j=1}^{m}\lambda_{m,j}^{2} \Big{)}+\Big{(}\frac{2^{5}n^{2}M(C^{2}+C_{1})^{1/2}}{\kappa_{n}d^{2}\gamma^{ \delta/2}}\cdot\sqrt{\sum_{j=1}^{m}\lambda_{m,j}^{2}}\Big{)}.\)__

The same comments as in the previous section apply here. The proof is slightly different than in the ReLU case, and leads to a different factor \(\sqrt{\sum_{j=1}^{m}\lambda_{m,j}^{2}}\) for the change of the NTG in (d).

### Sketch of the proofs

We give here a sketch of the proofs of Theorems 4.1 and 4.2. The detailed proofs are given in Appendices E and F, with secondary lemmas and propositions given in Appendices C and D.

The structures of the proofs of Theorems 4.1 and 4.2 are similar to that of (Du et al., 2019, Theorem 3.2), which showed analogous results on the global convergence under the NTK scaling. There are some key differences however which we highlight below.

Gradient flow converges to a global minimum of the objective function if the minimum eigenvalue of the NTG matrix \(\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W}_{t})\) is bounded away from zero by some positive constant for all \(t\geq 0\). In the NTK scaling case, Du et al. (2019) showed that it is satisfied as, for \(m\) sufficiently large, with high probability:

1. the NTG matrix at initialisation is close to its mean, and the minimum eigenvalue is close to that of the mean NTG,
2. the weights \(\mathbf{w}_{tj}\) are almost constant over time, which implies
3. the NTG matrix is almost constant over time, hence
4. the minimum eigenvalue of the NTG matrix at time \(t\) is close to that at initialisation, which is bounded away from zero.

However, in the case of asymmetrical node scaling (\(\gamma<1\)), none of the points (i-iv) holds. At initialisation, the random NTG matrix may be significantly different from its mean. Additionally, as we will discuss in Section 5, both the weights and the NTG matrix substantially change over time. This therefore requires a somewhat different approach that we now describe.

Let \(\lambda_{m,j}^{(1)}=\frac{\gamma}{m}\) and \(\lambda_{m,j}^{(2)}=(1-\gamma)\frac{\bar{X}_{j}}{\sum_{k=1}^{n}\bar{\lambda}_ {k}}\), and note that \(\lambda_{m,j}^{(1)}+\lambda_{m,j}^{(2)}=\lambda_{m,j}\). For \(k\in\{1,2\}\), let \(\widehat{\Theta}_{m}^{(k)}\) be the \(n\)-by-\(n\) symmetric positive semi-definite matrices defined by Equation (7), with \(\lambda_{m,j}\) replaced by either \(\lambda_{m,j}^{(k)}\). Note that

\[\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W}_{t})=\widehat{\Theta}_{m}^{(1)}( \mathbf{X};\mathbf{W}_{t})+\widehat{\Theta}_{m}^{(2)}(\mathbf{X};\mathbf{W}_{ t}) \tag{14}\]

with \(\mathbb{E}[\widehat{\Theta}_{m}^{(1)}(\mathbf{X};\mathbf{W}_{0})]=\gamma \widehat{\Theta}^{*}(\mathbf{X})\).

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 over time. The important points of the proof are as follows. For large \(m\), with high probability,

1. \(\widehat{\Theta}_{m}^{(1)}(\mathbf{X};\mathbf{W}_{0})\) is close to its mean \(\gamma\widehat{\Theta}^{*}(\mathbf{X})\) and its minimum eigenvalue is therefore lower bounded by \(\frac{2\pi_{n}}{2}\);
2. while the weights \(\mathbf{W}_{t}\) may change significantly over time, \(\widehat{\Theta}_{m}^{(1)}(\mathbf{X};\mathbf{W}_{t})\) remains almost constant over time;
3. as a result, the minimum eigenvalue of \(\widehat{\Theta}_{m}^{(1)}(\mathbf{X};\mathbf{W}_{t})\) can be lower bounded by \(\frac{2\pi_{n}}{4}\);
4. this implies that the minimum eigenvalue of the overall NTG matrix \(\widehat{\Theta}_{m}(\mathbf{X};\mathbf{W}_{t})\) is lower bounded by \(\frac{2\pi_{n}}{4}\).

## 5 Feature learning analysis

In this section we provide some results about feature learning when \(\gamma<1\).
### Weight change by the first gradient update

We first show that on average, each individual weight in the network changes on the order of \(\lambda_{m,j}\) by the first gradient-update step. To state this formally, for \(j\in[m]\) and \(k\in[d]\), let \(w_{0jk}\) be the \(k\)-th component of the weight vector \(\mathbf{w}_{0j}\) at initialisation, and define a function \(g_{1}:\mathbb{R}^{d}\to\mathbb{R}\) by

\[g_{1}(\mathbf{x})=\mathbb{E}_{z\sim\mathcal{N}(0,1)}\left[\sigma\left(\frac{z \|\mathbf{x}\|}{\sqrt{d}}\right)\sigma^{\prime}\left(\frac{z\|\mathbf{x}\|}{ \sqrt{d}}\right)\right].\]

In the ReLU case, we have \(g_{1}(\mathbf{x})=\|\mathbf{x}\|/\sqrt{2\pi d}\). The next theorem describes our result on weight change, and its proof appears in Appendix G.

**Theorem 5.1**.: _Assume Assumption 2.2 holds. For all \(j\in[m]\) and \(k\in[d]\), we have_

\[\mathbb{E}\left[\left.\frac{dw_{ijk}}{dt}\right|_{t=0}\right]=-\frac{\lambda_ {m,j}}{\sqrt{d}}\sum_{i=1}^{n}x_{ik}g_{1}(\mathbf{x}_{i}).\]

_In particular, if \(\sigma\) is the ReLU function, this expectation is \(-\frac{\lambda_{m,j}}{d\sqrt{2\pi}}\sum_{i=1}^{n}x_{ik}\|\mathbf{x}_{i}\|\)._

Recall that \(\lambda_{m,j}\to(1-\gamma)\widetilde{\lambda}_{j}\) as \(m\to\infty\). Hence, if \(\gamma<1\) and \(\widetilde{\lambda}_{j}>0\), the expected change of the weight vector \(\mathbf{w}_{j}\) at the first GD update is non-vanishing in the infinite-width limit.

### NTK change by the first gradient update in the smooth activation case

We now consider a smooth activation function \(\sigma\) satisfying Assumption 2.2(b). The next result characterises how much the NTK kernel at time \(0\) changes on average, using the neural tangent hierarchy (Huang and Yau, 2020). Let \(g_{2}\) be a function of type \((\mathbb{R}^{d}\setminus\{0\})\times(\mathbb{R}^{d}\setminus\{0\})\times( \mathbb{R}^{d}\setminus\{0\})\to\mathbb{R}\) defined by

\[g_{2}(\mathbf{x}_{k},\mathbf{x}_{\ell},\mathbf{x}_{i})=\mathbb{E}_{(z_{1},z_{ 2},z_{3})}\left[\sigma^{\prime\prime}(z_{1})\sigma^{\prime}(z_{2})\sigma^{ \prime}(z_{3})\sigma(z_{3})\right],\]

where \((z_{1},z_{2},z_{3})\in\mathbb{R}^{3}\) is distributed as a multivariate zero-mean Gaussian with the following covariance matrix:

\[\Sigma=\frac{1}{d}\left(\begin{array}{ccc}\|\mathbf{x}_{k}\|^{2}&\mathbf{x }_{k}^{\top}\mathbf{x}_{\ell}&\mathbf{x}_{k}^{\top}\mathbf{x}_{i}\\ \mathbf{x}_{k}^{\top}\mathbf{x}_{\ell}&\|\mathbf{x}_{\ell}\|^{2}&\mathbf{x}_{ \ell}^{\top}\mathbf{x}_{i}\\ \mathbf{x}_{i}^{\top}\mathbf{x}_{k}&\mathbf{x}_{i}^{\top}\mathbf{x}_{\ell}&\| \mathbf{x}_{i}\|^{2}\end{array}\right).\]

The evolution equation for \(\Theta_{m}(\mathbf{x}_{k},\mathbf{x}_{\ell};\mathbf{W}_{t})\) is

\[\frac{d\Theta_{m}(\mathbf{x}_{k},\mathbf{x}_{\ell};\mathbf{W}_{ t})}{dt} =\Big{(}\nabla_{\mathbf{W}}\Theta_{m}(\mathbf{x}_{k},\mathbf{x}_{ \ell};\mathbf{W}_{t})\Big{)}^{\top}\,\frac{d\mathbf{W}_{t}}{dt}\] \[=-\sum_{i=1}^{n}(f_{m}(\mathbf{x}_{i};\mathbf{W}_{t})-y_{i})\times \Big{\langle}\nabla_{\mathbf{W}}\Theta_{m}(\mathbf{x}_{k},\mathbf{x}_{\ell}; \mathbf{W}_{t}),\,\nabla_{\mathbf{W}}f_{m}(\mathbf{x}_{i};\mathbf{W}_{t}) \Big{\rangle}.\]

**Theorem 5.2**.: _Assume Assumption 2.2(b) holds. For all \(\mathbf{x}_{k},\mathbf{x}_{\ell}\in(\mathbb{R}^{d}\setminus\{0\})\), we have_

\[\mathbb{E}\left[\left.\frac{d\Theta_{m}(\mathbf{x}_{k},\mathbf{x}_{\ell}; \mathbf{W}_{t})}{dt}\right|_{t=0}\right]=-\frac{\mathbf{x}_{k}^{\top}\mathbf{x }_{\ell}}{d^{3/2}}\times\left(\sum_{j=1}^{m}\lambda_{m,j}^{2}\right)\times \left[\sum_{i=1}^{n}\big{(}(\mathbf{x}_{k}^{\top}\mathbf{x}_{i})g_{2}(\mathbf{ x}_{k},\mathbf{x}_{\ell},\mathbf{x}_{i})+(\mathbf{x}_{\ell}^{\top}\mathbf{x}_{i})g_{2}( \mathbf{x}_{\ell},\mathbf{x}_{k},\mathbf{x}_{i})\big{)}\right].\]

The above theorem shows that the expected change to the NTK at initialisation is scaled by the factor \(\sum_{j=1}^{m}\lambda_{m,j}^{2}\), which converges to \((1-\gamma)^{2}\sum_{j}\widetilde{\lambda}_{j}^{2}\) as \(m\to\infty\). The expected change in the NTK at the first GD iteration is therefore bounded away from zero when \(\gamma<1\).

## 6 Experiments

The code to reproduce the experiments in this section can be found in this repository.
Figure 4: Results for Cifar10 data. From left to right, 1) test risk through training, 2) the differences in weight norms \(\|\mathbf{w}_{ij}-\mathbf{w}_{0j}\|\) with \(j\)s being the neurons having the maximum difference at the end of the training, 3) the test risks of the pruned models, and 4) test accuracies of the pruned models.

Figure 3: A subset of results for MNIST dataset. (Left) test accuracies of the pruned models, and (right) test accuracies of transferred models.

Figure 2: A subset of results for the regression experiments. From left to right, 1) the differences in weight norms \(\|\mathbf{w}_{tj}-\mathbf{w}_{0j}\|\) with \(j\)s being the neurons having the maximum difference at the end of the training for concrete dataset, 2) the differences in NTG matrices for energy dataset, 3) training risks of pruned models for airfoil dataset, and 4) test risks of transferred models for plant dataset.
### Simulated data

We first illustrate our theory on simulated data. We generate a dataset with \(n=100\) observations where, for \(i=1,\ldots,n\), \(\mathbf{x}_{i}\) is \(d=50\) dimensional and sampled uniformly on the unit sphere and \(y_{i}=\frac{5}{d}\sum_{j=1}^{d}\sin(\pi x_{i,j})+\varepsilon_{i}\) where \(\varepsilon_{i}\overset{\text{iid}}{\sim}\mathcal{N}(0,1)\). We use a wide shallow neural network as described in Section 2.1, with \(m=2000\) hidden nodes and scaling parameters \(\lambda_{m,j}\) defined as in Equation (1), where \(\gamma\in[0,1]\) and \((\widetilde{\lambda}_{j})_{j\geq 1}\) is the probability mass of a Zipf law with parameter \(\alpha\in(0,1)\), see Equation (4). We consider four values for the pair of parameters \((\gamma,\alpha)\): \(\gamma=1\), \((\gamma,\alpha)=(0.5,0.7)\), \((\gamma,\alpha)=(0.5,0.5)\), and \((\gamma,\alpha)=(0,0.4)\). For each setting, we run gradient descent 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 \(\gamma>0\), GD achieves near-zero training error. The minimum eigenvalue and the training rates increase with the value of \(\gamma\). For \(\gamma=1\), we have the highest minimum eigenvalue and the fastest training rate; however, there is no/very little feature learning: the weights and the NTG do not change significantly over the GD iterations. When \(\gamma<1\), there is clear evidence of feature learning: both the weights and the NTG change significantly over time; the smaller \(\gamma\) and \(\alpha\), the more feature learning arises.

### Regression

We also validate our model on four real-world regression datasets from the UCI repository2: concrete (concrete compressive strength, \((n,d)=(1030,9)\)), energy (energy efficiency, \((n,d)=(768,8)\)), airfoil (airfoil self-noise, \((n,d)=(1503,6)\)), and plant (combined cycle power plant, \((n,d)=(9568,4)\)). We split each dataset into training (40%), test (20%), and validation sets (40%), and the validation set is used to test transfer learning. We use the same parameters as in Section 6.1, and we now train our neural networks for 100 000 steps in each run.

Footnote 2: [https://archive.ics.uci.edu/ml/datasets.php](https://archive.ics.uci.edu/ml/datasets.php)

In addition, to further highlight the presence of feature learning in our model, we conduct two experiments. Firstly, we test the prunability of the networks. We gradually prune those hidden nodes which have small feature importance and measure risks after pruning. Here, feature importance is measured as \((\lambda_{m,j}\|\mathbf{w}_{t,j}\|^{2})_{j\in[m]}\), and our theory suggests that a model with smaller \(\gamma\) and/or \(\alpha\) values is likely to be more robust with respect to pruning, as long as \(\gamma<1\). Similar empirical findings were made by Wolinski et al. (2020) about the benefits of asymmetrical scaling for neural network pruning, in the case \(\gamma=0\). Secondly, we test the transferability of the 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 neural networks trained with the original training set. The features are taken to be the output of the hidden layers, so each data point in the validation set is represented with a \(m=2000\) dimensional vector. Then, we sort the feature dimensions with respect to feature importance as above 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 a ReLU activation function, and it is trained for 5,000 steps of gradient descent with a learning rate of 1.0. Again, our theory suggests that a model with smaller \(\gamma\) and \(\alpha\) values is likely to exhibit better transfer learning.

We summarise a subset of our results in Figure 2; additional results can be found in Appendix I. 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 \(\gamma\) and \(\alpha\). Also, we observe that models with smaller \(\gamma\) values are more robust to pruning, allowing them to retain relatively low training risks for larger numbers of pruned nodes. A similar phenomenon is seen in transfer learning, where models with smaller values of \(\gamma\) 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 a relatively small number of neurons.

### Classification

Finally, we apply our model on image classification tasks. We conduct two experiments, 1) a small-scale experiment under the same setting assumed in our theory, and 2) a larger-scale experiment under a more realistic setting.

**MNIST.** We take a subset of size 5,000 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 5,000 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 3. Other results can be found in Figure 10 in the Appendix.
**CIFAR10.** We consider a more challenging image classification task, CIFAR10. The dataset is composed of 60 000 images, among which 50 000 are used for training and the rest for testing. There are ten different classes. We illustrate the described benefits of asymmetrical node scaling and show they 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. This is the approach we implement here. We consider a Resnet-18 model, pre-trained on the ImageNet dataset. Using this model, we transform each one of the original images to a vector of dimension \(512\). We then train wide shallow neural networks as described in Section 2.1, with \(m=2000\), and output dimensions respectively \(10\). The experimental setup differs from previous results as 1) we use stochastic gradient descent with a mini-batch size of \(64\) instead of full batch gradient descent, here a step is a mini-batch step; 2) we use the cross-entropy loss instead of the MSE. All experiments are run five times, the learning rate is set to \(5.0\). In Figure 4, we report the results for the same four values of \((\gamma,\alpha)\) pairs. We can see that the main conclusions of the previous experiments hold in this setting as well, even though the theory does not apply directly. Additional results, investigating the effect of \(\gamma\), are given in Figure 11 in the Appendix.