How Two-Layer Neural Networks Learn, One (Giant) Step at a Time
We investigate theoretically how the features of a two-layer neural network
adapt to the structure of the target function through a few large batch
gradient descent steps, leading to improvement in the approximation capacity
with respect to the initialization. We compare the influence of batch size and
that of multiple (but finitely many) steps. For a single gradient step, a batch
of size $n = \mathcal{O}(d)$ is both necessary and sufficient to align with the
target function, although only a single direction can be learned. In contrast,
$n = \mathcal{O}(d^2)$ is essential for neurons to specialize to multiple
relevant directions of the target with a single gradient step. Even in this
case, we show there might exist ``hard'' directions requiring $n =
\mathcal{O}(d^\ell)$ samples to be learned, where $\ell$ is known as the leap
index of the target. The picture drastically improves over multiple gradient
steps: we show that a batch-size of $n = \mathcal{O}(d)$ is indeed enough to
learn multiple target directions satisfying a staircase property, where more
and more directions can be learned over time. Finally, we discuss how these
directions allows to drastically improve the approximation capacity and
generalization error over the initialization, illustrating a separation of
scale between the random features/lazy regime, and the feature learning regime.
Our technical analysis leverages a combination of techniques related to
concentration, projection-based conditioning, and Gaussian equivalence which we
believe are of independent interest. By pinning down the conditions necessary
for specialization and learning, our results highlight the interaction between
batch size and number of iterations, and lead to a hierarchical depiction where
learning performance exhibits a stairway to accuracy over time and batch size,
shedding new light on how neural networks adapt to features of the data.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, Ludovic Stephan
# Learning Two-Layer Neural Networks, One (Giant) Step at a Time

###### Abstract

We study the training dynamics of shallow neural networks, investigating the conditions under which a limited number of large batch gradient descent steps can facilitate feature learning beyond the kernel regime. We compare the influence of batch size and that of multiple (but finitely many) steps. Our analysis of a single-step process reveals that while a batch size of \(n=O(d)\) enables feature learning, it is only adequate for learning a single direction, or a single-index model. In contrast, \(n=O(d^{2})\) is essential for learning multiple directions and specialization. Moreover, we demonstrate that "hard" directions, which lack the first \(\ell\) Hermite coefficients, remain unobserved and require a batch size of \(n=O(d^{\ell})\) for being captured by gradient descent. Upon iterating a few steps, the scenario changes: a batch-size of \(n=O(d)\) is enough to learn new target directions spanning the subspace linearly connected in the Hermite basis to the previously learned directions, thereby a staircase property. Our analysis utilizes a blend of techniques related to concentration, projection-based conditioning, and Gaussian equivalence that are of independent interest. By determining the conditions necessary for learning and specialization, our results highlight the interaction between batch size and number of iterations, and lead to a hierarchical depiction where learning performance exhibits a stairway to accuracy over time and batch size, shedding new light on feature learning in neural networks.

## 1 Introduction

A number of theoretical efforts have been directed toward understanding the dynamics of training neural networks with gradient descent, aiming to uncover the underlying principles that govern their learning capabilities. By iteratively adjusting the network's weights based on the gradients of a chosen objective function, gradient descent seeks to minimize the discrepancy between predicted and true outputs. In this process, the network adapts to the data, learning which features in the input data are most important for explaining the labels. This adaptivity is in contrast to other machine learning pipelines such as random feature or kernel methods, and is behind the success of neural networks in efficiently solving problems from image classification to text generation.

In this paper, we delve into the training dynamics of shallow neural networks, focusing on the conditions that facilitate effective feature learning beyond the kernel or random features regime. We aim to answer the following fundamental question: can we precisely capture the presence of feature learning in the _early phase_ of gradient descent training? We shall thus investigate how the _very first few gradient steps_ on the first-layer parameters impact the representation abilities. Given we only perform a few steps, it will have to be giant ones, and we will use large, aggressive, learning rates. Our investigation highlights the crucial role of batch size in the training process in order to learn a meaningful representation in a few gradient steps, and the interplay between batch size, number of iterations and the complexity of the function one is attempting to learn. In fact, we shall establish the presence of a hierarchy in the learning of the different directions of the function to learn, in terms of their Hermite coefficients, as illustrated in Fig. 1 (for a single steps) and Fig. 2 (for multiple steps).

### Setting

We consider fully connected neural networks with \(p\) hidden neurons: given \(W\in\mathbb{R}^{p\times d}\) and \(\mathbf{a}\in\mathbb{R}^{p}\), the network function is defined as

\[\hat{f}(\mathbf{z};W,\mathbf{a})=\frac{1}{\sqrt{\rho}}\sum_{i=1}^{p}a_{i}\sigma( \langle\mathbf{w}_{i},\mathbf{z}\rangle)=\frac{1}{\sqrt{p}}a^{\top}\sigma(W\mathbf{z}), \tag{1}\]
where \(\sigma\) is the non-linear activation function applied elementwise. We study the performance of learning a distribution \(\rho\) on pairs \(\mathbf{z},y\), which corresponds to minimizing the _population loss_

\[\mathcal{R}(W,\mathbf{a})=\mathbb{E}_{(\mathbf{z},y)\sim\rho}\left[\mathcal{L}(\hat{f}( \mathbf{z};W,\mathbf{a}),y)\right], \tag{2}\]

where \(\mathcal{L}\) is the chosen loss function. In this work, we will focus on the case of the square loss

\[\mathcal{L}(\hat{y},y)=\frac{1}{2}\left(\hat{y}-y\right)^{2}\]

Data model --We assume the training data \((\mathbf{z},y)\in\mathbb{R}^{d+1}\) is sampled from the following generative model:

\[y=f^{*}(\mathbf{z})=\sum_{k=1}^{r}\sigma_{k}^{*}(\langle\mathbf{w}_{k}^{\star},\mathbf{z} \rangle), \mathbf{z}\sim\mathcal{N}(0,I_{d}). \tag{3}\]

where the weights \(W^{*}\in\mathbb{R}^{r\times d}\) are assumed normalized and orthogonal \(W^{*}{W^{*}}^{\top}=I_{r}\). This hypothesis class is a particular example of more general _multi-index functions_\(g(\langle\mathbf{w}_{1}^{\star},\mathbf{z}\rangle,\ldots,\langle\mathbf{w}_{r}^{\star}, \mathbf{z}\rangle)\) for which the target \(f^{*}\) only depends on the input data \(\mathbf{z}\) through \(r\) low-dimensional projections. Note that this can also be seens as fully-connected two-layer neural network with width \(r\). Learning this class with a two-layer network therefore boils down to learning the directions \(\mathbf{w}_{k}^{\star}\) on the top of the non-linear activations, therefore requiring _specialization_ of the hidden-units \(\mathbf{w}_{i}\) to the target directions [1]. Some of our results also hold for the more general multi-index model, where the dependency in the \(\langle\mathbf{w}_{k}^{\star},\mathbf{z}\rangle\) is arbitrary.

Training procedure --Similarly to [2, 3], we use a symmetric initialization to ensure that the network function is identically zero at the start of training: assuming that \(p\) is even,

\[a_{i}^{0}=-a_{p-i+1}^{0}\quad\text{and}\quad\mathbf{w}_{i}^{0}=\mathbf{w}_{p-i+1}^{0} \quad\text{for all}\;i\in[p/2]. \tag{4}\]

This condition is mainly necessary when \(p\) is small, as the initial network function might not concentrate around zero. Apart from this restriction, we draw

\[\sqrt{p}\cdot a_{i}^{0}\overset{i.i.d}{\sim}\operatorname{Unif}(\{-1,1\}) \quad\text{and}\quad w_{i}^{0}\overset{i.i.d}{\sim}\operatorname{Unif}( \mathcal{S}^{d-1}). \tag{5}\]

The distribution of the \(a_{i}\) can be replaced by any other distribution with positive variance, at least when \(p\) grows. After initialization, we split the training into two parts:

Figure 1: **Learning with a single gradient step: Illustration of the relationship between batch size and learning multi-index functions with a single giant steps with learning rate \(\mathcal{O}(p\sqrt{n/d})\), showcasing the stairway to generalization phenomenon: (a) a batch size of \(n=\mathcal{O}(d)\) results in learning only a single direction, approximating the target function with a single index function, (b) a larger batch size of \(n=\mathcal{O}(d^{2})\) facilitates specialization, allowing for learning multiple directions and coupling to all their Hermite coefficients, though directions with a null second Hermite coefficient remain unseen, (c) employing \(n=\mathcal{O}(d^{3})\) enables learning hard directions with the first nonzero coefficient of 3-degree (light blue neurons) and more generally, directions with \(\ell\)-Hermite coefficient zero require a batch size of \(n=\mathcal{O}(d^{\ell})\). (hard blue neurons)**
1. _First layer training_: we perform \(T\) steps of full-batch gradient descent \[\mathbf{w}_{i}^{t+1}=\mathbf{w}_{i}^{t}-\frac{\eta}{n}\sum_{\nu=1}^{n}\nabla_{\mathbf{w}_{j} }\mathcal{L}\left(\hat{f}(\mathbf{z}^{\nu};W^{t},\mathbf{a}),y^{\nu}\right),\] (6) where the training set \(\{(\mathbf{z}^{\nu},y^{\nu})\}_{\nu=1}^{n}\) is drawn i.i.d according to (3). We will assume that the training set for each gradient step is a fresh one, so that the total sample complexity for this step is \(Tn\).
2. _Second layer training_: once the first layer is trained, we perform a simple ridge regression with \(\ell_{2}\) penalty for the second layer: \[\hat{\mathbf{a}}=\operatorname*{arg\,min}_{\mathbf{a}\in\mathbb{R}^{p}}\frac{1}{n} \sum_{\nu=1}^{n}\mathcal{L}\left(\hat{f}(\mathbf{z}^{\nu};W^{t},\mathbf{a}),f^{\star} (\mathbf{z}^{\nu})\right)+\lambda\|\mathbf{a}\|^{2}.\] (7) Such a separation of the training between the first and second layer is a common setup for the theoretical study of training [3, 4, 5], and allows for a more tractable study of convergence.

Hermite expansion -We recall the definition of a Hermite expansion:

**Definition 1**.: _Let \(f:\mathbb{R}\to\mathbb{R}\) be a function that is square integrable w.r.t the Gaussian measure. There exists a unique sequence of real numbers \((c_{j})_{j\in\mathbb{N}}\), called the Hermite coefficients of \(f\), such that_

\[f(x)=\sum_{j=0}^{\infty}\frac{c_{j}}{j!}\operatorname{He}_{j}(x), \tag{8}\]

_with \(\operatorname{He}_{j}\) denoting the probabilist's Hermite polynomials [6]._

In the following, we will denote by \((\mu_{j})_{j\in\mathbb{N}}\) the Hermite coefficients of the activation function \(\sigma\), and \((\mu_{k,j}^{\star})_{j\in\mathbb{N}}\) those of the target function \(\sigma_{k}^{\star}\).

### Summary of main results

Our first set of results discuss the case of **a single gradient step**, which are illustrated in Fig. 1 (A particular numerical example is shown in Fig. 5). In this case:

Figure 2: **Learning with a few gradient steps:** Illustration of the staircase phenomenon when using few giants steps with learning rate \(\mathcal{O}(p\sqrt{n/d})\) in the linear regime \(n\approx d\). Upon iteration, one can learn more directions, one at a time, but _only_ in the subspace of directions “linearly connected” in the Hermite basis to the learned directions of the previous time-step (see text). In contrast to Fig. 1, this allows to generalize with multiple directions with \(\mathcal{O}(n)\) batches and does not require cancellation of previous directions, but this specialization is restricted to staircase functions (and “hard” neurons with zero first Hermite coefficients still require \(\mathcal{O}(d^{2})\) samples).
* For a large batch size \(n=\mathcal{O}(d)\) and learning rate \(\eta=\mathcal{O}(p\sqrt{n/d})\), meaningful feature learning can be obtained that surpasses random features. However, this is limited in that only a _single_ direction can be learned. Thus, if \(r>1\), only a single-neuron approximation of the target function will be obtained, specifically for the directions with non-zero Hermite coefficients.
* Surpassing the single direction approximation _requires_ utilizing a larger batch size of dimension \(n=\mathcal{O}(d^{2})\). Suppressing the main learned direction, we show that with a batch size of \(n=\mathcal{O}(d^{2})\) the gradient spans the space of many directions. This allows for the _specialization_ of the network weights, that is, learning multiple directions. In these directions, the learning process couples _all_ of the non-zero Hermite coefficients of the target function.
* Nonetheless, the "hard" directions (those for which the first \(\ell\) Hermite coefficients are zero), will continue to be disregarded. Learning these directions necessitates a batch size of \(n=\mathcal{O}(d^{\ell})\), as well as supressing the directions learned at \(n=\mathcal{O}(d^{\ell-1})\).

The situation drastically changes when allowing for **multiple gradient steps**. In this case, focusing on the linear batch size \(n=\mathcal{O}(d)\) regime, the phenomenology, as illustrated in Fig. 2, is as follows:

* While with a single step we can only learn a single direction in the teacher space, each additional gradient step allows for learning a new perpendicular direction, one at a time, upon the important condition that these directions are coupled through a staircase property (see sec.2.2 for details and definitions). Therefore, taking multiple steps allows to learn a multiple-index target with only \(n=\mathcal{O}(d)\) data, provided the staircase condition is met.
* However, directions that are not coupled through the staircase and with zero first Hermite coefficient cannot be learned in any finite number of steps. In fact, as discussed previously they require a batch size of \(n=\mathcal{O}(d^{2})\). In other words, while multiple steps help specialization, they do not allow for the identification of "hard" target directions.

We note that aside from the difference in the specialization capacity, taking a single or multiple steps have an important difference. Indeed, with one step learning with \(n=\mathcal{O}(d^{\ell})\) is only possible if the directions obtained from the gradient with \(n=\mathcal{O}(d^{\ell-1})\) are removed, for instance by learning the second layer or pre-processing the data. In contrast, each new step can add a new direction without the need to adapt the second layer weights.

The code to reproduce our figures is available on GitHub, and we refer to App. A for details on the numerical implementations.

### Related Works

The analysis of high-dimensional asymptotics of kernel regression has provided valuable insights into the advantages and limitations of kernel methods [7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]. In particular, a similar stairway picture as in Fig. 1 emerged from these works [19]. The key difference, however, is that at each regime \(n=\mathcal{O}(d^{\ell})\), kernels can only learn the \(\ell\)-th Hermite polynomial of the target. This should be contrasted with our feature learning regime where, once a direction is learned, all its Hermite coefficients are learned. Feature learning corrections to kernel methods have been investigated in [20, 21, 22, 23, 24, 25, 26]. On a complementary line of work, exact asymptotic results for the the random features models have been derived in the literature [27, 28, 29, 30, 31, 32, 33]. A large part of these results are enabled by the Gaussian equivalence property [34, 35, 36, 37].

In the context of neural network training dynamics, [38] established that the sample complexity of learning single-index targets with one-pass SGD is entirely determined by the first non-zero Hermite coefficient, which define an _information exponent_.

Closer to us are [39, 3] showed a sufficiently large single gradient step allows to beat a kernel method. However, their results are limited to single-index target and to a single (giant) gradient step. Here, we build on their results and consider multi-index targets and multiple steps. [3] proved that a batch size of \(n=d^{2}\) is _sufficient_ for learning multiple directions of a multi-index target. However, they did not establish whether this batch size is _necessary_, and were again limited to a single step. Our results reveal that a batch size of \(n=\mathcal{O}(d^{2})\) is indeed necessary for learning multiple directions in a single step, and we characterize which class of multi-index targets can be instead learned with \(n=O(d)\) with multiple steps.

In a related but different vein, [40], building upon the earlier work of [41] show how a "staircase" property of target functions characterizes the sample complexity arises for infinite width two-layer networks trained by one-pass SGD, a.k.a. the mean-field regime [42, 43, 44, 45]). Their focus, however, was on the peculiar case of _sparse boolean functions_ depending solely on linear coefficients, without considering direction-specific specialization. More recently, [4] provided partial results for the case of isotropic Gaussian data, always in the regime of one-pass SGD. Here we consider the more generic case of multi-index target functions with Gaussian data with large-batch SGD, and show that a "directional staircase" behavior arises when iterating a few giant gradient steps, while a related, but definitely different, picture arises with single steps depending on the batch size.

Finally, our results illustrate that during the early stage of training, neural networks first learn a single direction and require the learning of the linear part of the target function before specializing. This corroborates the findings of [46],
who showed that the neural network first fits the best linear classifier and subsequently learns functions of increasing complexity.

## 2 Main theoretical results

### Single gradient step

Our starting point is a _single_ giant gradient step, and the phenomenology described in Fig.1.

Feature learning in the first layer -Let us first focus on the effect of the batch learning of (6) on the first layer \(W\). For any \(i\in[p],k\in[r]\), we introduce the following quantities:

\[M\in\mathbb{R}^{p\times r}=\{m_{ik}\}_{i\in[p],k\in[r]},\quad m_{ik}=\langle \mathbf{w}_{i},\mathbf{w}_{k}^{\star}\rangle\quad\text{ and }\quad q_{i}=\|\mathbf{w}_{i}\|^{2}, \tag{9}\]

and the normalized correlation matrix

\[\Pi_{ik}=\frac{1}{\sqrt{p}}\frac{m_{ik}}{\sqrt{q_{i}}}. \tag{10}\]

We say that the first layer is (resp. isn't) _learned_ when \(\|\Pi\|_{F}=\Omega(1)\) (resp. \(o(1)\)). More specifically, let \(V^{\star}=\operatorname{span}\big{(}\{\mathbf{w}_{k}^{\star}\}_{k\in[r]}\big{)}\); we say that the subspace learned by \(W\) is \(U\subseteq V^{\star}\) if \(\|M\mathbf{u}\|=\Omega(1)\) for any \(\mathbf{u}\in U\), and \(\|M\mathbf{u}\|=o(1)\) for any \(\mathbf{u}\in U^{\perp}\).

Those definitions will be motivated in the next result section, which investigates the performance of the network when the first layer has/hasn't learned. For now, let us state our results. We define the _leap index_ of the learning problem as the first integer \(\ell\) such that \(\mu_{k,\ell}^{\star}\neq 0\) for some \(k\in[r]\).

**Theorem 1**.: _Let \(\ell\) be the leap index of the problem, and assume that \(\mu_{\ell}\neq 0\). Then, with high probability, after one step:_

1. _If_ \(\ell=1\) _and_ \(n=\Omega(d)\)_, let_ \[\mathbf{v}^{\star}=\sum_{k=1}^{r}\mu_{k,1}^{\star}\mathbf{w}_{k}^{\star}.\] (11)

_Then, when taking \(\eta=p\) the first layer learns, but it only learns the subspace \(\operatorname{span}(\mathbf{v}^{\star})\)._
2. _If_ \(\ell>1\) _and_ \(n=\Omega(d^{\ell})\)_, let_ \[V_{\ell}^{\star}=\operatorname{span}\big{(}\{\mathbf{w}_{k}^{\star}:\mu_{k,\ell}^ {\star}\neq 0\}\big{)}\,.\] (12)

_Then when taking \(\eta=p\sqrt{\nicefrac{{n}}{{d}}}\) the first layer learns, but it only learns \(V_{\ell}^{\star}\)._
3. _If_ \(n=\mathcal{O}(d^{\ell-\varepsilon})\) _for some_ \(\varepsilon>0\)_, then the first layer never learns._

Aspects of these results were already present in previous works: [39] shows the existence of a rank-one property of the gradient at initialization for \(n=\Theta(d)\), and [3] implies the positive part of our result for \(n=\Theta(d^{2})\), provided that \(V_{2}^{\star}=V^{\star}\) (which corresponds to their Assumption 2). We complement those results by providing matching lower bounds, as well as providing the picture for higher powers of \(d\).

### Few giant steps and the staircase phenomenon

We now move to the effect of multiple gradient steps, as described in Fig.2. As we shall see, while the effect of multiple steps is limited to a subclass of functions, it allows to learn them meaningfully with few iterations. For simplicity, we restrict ourselves to the case \(n=\mathcal{O}(d)\), although we expect a similar picture to hold in the other regimes. Informally, the picture is similar to a the "staircase" phenomenon observed in [4, 40]: once we learn a direction in the first-layer weights, the next gradient step can use the already learned subspace as a ladder to learn the next ones, which become linear when conditioning on the spike directions. To formalize this, we introduce the notion of _subspace conditioning_:

**Definition 2** (Subspace conditioning).: _Let \(V\) be a vector space, and \(U\subseteq V\) a subspace. For any function \(f:V\to\mathbb{R}\), and \(\mathbf{x}\in U\), we define the conditional function \(f_{U,\mathbf{x}}:U^{\perp}\to\mathbb{R}\) as_

\[f_{U,\mathbf{x}}(\mathbf{x}^{\perp})=f(\mathbf{x}+\mathbf{x}^{\perp}) \tag{13}\]

In short, the function \(f_{U,\mathbf{x}}\) corresponds to \(f\) "conditioned" on the projection of its argument in \(U\). Its first Hermite coefficient will be denoted as

\[\mu_{U,\mathbf{x}}(f)=\mathbb{E}_{\mathbf{x}^{\perp}\sim\mathcal{N}(0,I)}\left[\nabla _{\mathbf{x}^{\perp}}f_{U,\mathbf{x}}\right] \tag{14}\]

We are now equipped to state our conjecture:
**Conjecture 1**.: _Assume that \(n=\mathcal{O}(d)\), and \(\eta=\sqrt{p}\). Define a sequence of nested subspaces \(U_{0}^{\star}\subseteq U_{1}^{\star}\subseteq\dots\subseteq U_{t}^{\star} \subseteq\dots\) as_

* \(U_{0}^{\star}=\{0\}\)_,_
* _for any_ \(t\geq 0\)_,_ \(U_{t+1}^{\star}=U_{t}^{\star}\oplus\mathrm{span}\left(\{\mu_{U_{t}^{\star}, \mathbf{x}}(f^{\star}):\mathbf{x}\in U_{t}^{\star}\}\right)\)_._

_Then, after \(t\) gradient steps of the form (6), the first layer learns exactly \(U_{t}^{\star}\)._

It is easy to check that this conjecture is consistent with Theorem 1, since \(f_{\{0\},\mathbf{0}}=f\), and \(v^{\star}\) is exactly the first hermite of \(f^{\star}\).

Examples --As an illustration, we work out the subspaces \(U_{t}^{\star}\) for simplified versions of the examples in Figure 4. For simplicity, we will take \(\mathbf{w}_{k}^{\star}=e_{k}\), the \(k\)-th vector of the standard basis, so that \(\langle\mathbf{w}_{k}^{\star},\mathbf{z}\rangle=z_{k}\).

* \(f^{\star}(\mathbf{z})=z_{1}+z_{2}+z_{1}^{2}-z_{2}^{2}\): the normalized first hermite coefficient of \(f^{\star}\) is \(\mathbf{v}^{\star}=(e_{1}+e_{2})/\sqrt{2}\), so \(U_{1}^{\star}=\mathrm{span}(e_{1}+z_{2})\). Then, we can rewrite \(f^{\star}\) in the new basis \((\mathbf{v}^{\star},\mathbf{v}^{\perp})\), to get \[f^{\star}(\mathbf{z})=\sqrt{2}\langle\mathbf{v}^{\star},\mathbf{z}\rangle+2\langle\mathbf{v}^ {\star},\mathbf{z}\rangle\langle\mathbf{v}^{\perp},\mathbf{z}\rangle.\] Hence, if \(\mathbf{x}=\lambda\mathbf{v}^{\star}\), we have \(\mu_{U_{t}^{\star},\mathbf{x}}(f^{\star})=2\lambda\mathbf{v}^{\perp}\), and hence \(U_{2}^{\star}=V^{\star}\).
* \(f^{\star}(\mathbf{z})=z_{1}+z_{2}+z_{1}^{2}+z_{2}^{2}\): just as the above example, \(U_{1}^{\star}=\mathrm{span}(e_{1}+e_{2})\), and hence we perform the same change of basis. However, this time, \[f^{\star}(\mathbf{z})=\sqrt{2}\langle\mathbf{v}^{\star},\mathbf{z}\rangle+\langle\mathbf{v}^ {\star},\mathbf{z}\rangle^{2}+\langle\mathbf{v}^{\perp},\mathbf{z}\rangle^{2}.\] This implies that for any \(\mathbf{x}\in U_{1}^{\star}\), \(\mu_{U_{t}^{\star},\mathbf{x}}(f^{\star})=\mathbf{0}\), and the direction of \(v^{\perp}\) is never learned.

In Appendix B we provide a series of results and analytic predictions supporting Conjecture 1. In particular, the notion of subspace condition (Definition 2) arises through an inductive decomposition of the projections of the gradient along the teacher subspace. We are also able to verify the \(t=2\) case of Conjecture 1, under some technical conditions:

Figure 3: **Feature learning after a single step.** Specialization of hidden units in the \(n=\mathcal{O}(d^{\ell})\) regime (\(\ell=2,3\)). The plots show the cosine similarity with respect to the teacher vectors \((\mathbf{\theta}_{1},\mathbf{\theta}_{2})\) for the gradient of \(p=40\) different neurons, identified by different markers. The bisectrix of the first quadrant is shown as a continuous black line, the circle of unitary radius in black, and the circle of radius \(\nicefrac{{2}}{{\sqrt{d}}}\) in blue. In the upper panel, \((n,d)=(2^{18},2^{9})\), and \((n,d)=(2^{21},2^{7})\) in the lower one, and we use a symmetric \(2\)-index teacher, with matching student: \(\sigma=\sigma_{1}^{\star}=\sigma_{2}^{\star}\). **Left:**\(\sigma(z)=\mathrm{He}_{1}(z)\), **Center-Left:**\(\sigma(z)=\mathrm{He}_{2}(z)\), **Center-Right:**\(\sigma(z)=\mathrm{He}_{3}(z)\), **Right:**\(\sigma(z)=\mathrm{He}_{4}(z)\). We observe that for \(\ell=1\), we only learn a single direction, no matter the data quantity, while for \(\ell>1\) we learn every direction at once as soon as we reach \(n=d^{\ell}\). The small spread observed for \(\sigma=\mathrm{He}_{4}\) and \(n=d^{3}\) is likely caused by the small value of \(d\) used for the experiments. See details in App. A.
**Proposition 1** (Informal).: _In the setting of Conjecture 1, let \(\mathbf{v}^{\perp}\) be a vector in \(U_{2}^{\star}\cap{U_{1}^{\star}}^{\perp}\). Then, for any \(i\in[p]\) with \(a_{i}^{0}\) fixed to \(\frac{1}{\sqrt{p}}\) or \(-\frac{1}{\sqrt{p}}\),_

\[\left|\mathbb{E}\left[\langle\mathbf{w}_{i},\mathbf{v}^{\perp}\rangle\right]\right|= \Omega(1). \tag{15}\]

_On the other hand, for any \(\mathbf{v}\in U_{2}^{\star\perp}\) and \(i\in[p]\),_

\[\left|\mathbb{E}\left[\langle\mathbf{w}_{i},\mathbf{v}\rangle\right]\right|=o(1). \tag{16}\]

Furthermore, our theory leads to a precise prediction of the orientation of the gradient in the teacher subspace after a finite number of steps. This is illustrated in Figure 4.

### From feature learning to generalization bounds

We now investigate the consequence of Theorem 1 in the actual performance of the neural network. This is a fairly complex question: indeed, [12] shows that even in the absence of feature learning (i.e. when the first layer is fully random), it is possible to learn a polynomial approximation of \(f^{\star}\) of degree \(\ell\) as long as \(n,p=\Omega(d^{\ell+\delta})\) for some \(\delta>0\). Hence, our results and conjectures regarding generalization error need to depend on both the number of samples and the number of hidden neurons.

Finite \(p\) regime --This is the simplest case: in this setting, there is simply no way to fit anything beyond the learned directions, even with an infinite amount of data. This is formalized by the following proposition:

**Proposition 2**.: _Assume that \(p\) is bounded as \(n,d\to\infty\), and that the first layer \(W\) only learns a subspace \(U\subseteq V^{\star}\). For any choice of second layer \(\mathbf{a}\) such that \(\|\mathbf{a}\|_{\infty}\leq\mathbb{O}(1)\), we have_

\[\mathbb{E}\left[\left(f^{\star}(\mathbf{z})-\hat{f}(\mathbf{z};W,\mathbf{a})\right)^{2} \right]\geq\mathbb{E}\left[\operatorname{Var}\left(f^{\star}(\mathbf{z})\mid P_{ U}\mathbf{z}\right)\right]+o(1) \tag{17}\]

_where \(P_{U}\) is the orthogonal projection on \(U\)._

Figure 4: **Feature learning with multiple gradient steps. Top:** Generalization error as a function of \(n\) (\(d=512,p=256\)) after iterating the training procedure for six steps. **Bottom:** Cosine similarity of the projected gradient matrix \(G^{\perp}\) inside the teacher subspace for all the \(p\) neurons at a fixed ratio \(\nicefrac{{n}}{{d}}=4\), plotted at different stages of the training. The blue and purple lines are the theoretical prediction for the orientation of the gradient in the second step. We consider two different 2-index teacher functions: **Left:**\(\sigma_{1}^{\star}(z)=\sigma_{2}^{\star}(z)=\operatorname{He}_{1}(z)+\nicefrac{{ \operatorname{He}_{2}(z)}}{{2}}+\nicefrac{{\operatorname{He}_{4}(z)}}{{4!}}\). **Right:**\(\sigma_{1}^{\star}(z)=z-z^{2}\) and \(\sigma_{2}^{\star}(z)=z+z^{2}\). In accordance with Conjecture 1, the difference between the two cases is clear already after the first GD step: while on the left the gradient is stuck around the predicted rank-one spike after the first step (black line), on the right the gradient changes orientation in the second step, allowing to learn multiple features. See details in App. A.
The right-hand side of the above inequality corresponds to a predictor \(\hat{f}(\mathbf{z})=\mathbb{E}\left[f^{\star}(\mathbf{z})\mid P_{U}\mathbf{z}\right]\), which is the best predictor that depends only on \(P_{U}\mathbf{z}\). Conversely, we expect that with enough neurons, it is possible to approximate this predictor:

**Conjecture 2**.: _For any \(\delta>0\), there exists a \(p_{0}\in\mathbb{N}\) such that if \(p\geq p_{0}\), there exists a choice of second layer \(\mathbf{a}\) satisfying \(\|\mathbf{a}\|_{\infty}\leq\mathcal{O}(1)\), such that as \(n,p\to\infty\)_

\[\mathbb{E}\left[\left(f^{\star}(\mathbf{z})-\hat{f}(\mathbf{z};W,\mathbf{a})\right)^{2} \right]\leq\mathbb{E}\left[\mathrm{Var}\left(f^{\star}(\mathbf{z})\mid P_{U}\mathbf{z }\right)\right]+\delta+o(1) \tag{18}\]

Such a conjecture is proven in [3] for the specific case of \(U=V^{\star}\) (in which case the first term in the RHS is zero) and \(\sigma=\mathrm{relu}\), with additional bias terms.

Beyond the fixed regime --When the number of neurons is allowed to diverge, it is no longer impossible to learn the directions that are not present in \(M\), and the performance of the network may heavily depend on the training process for the second layer. We expect the following behavior:

1. On the directions present in \(W\), we are able to learn "everything".
2. On the orthogonal directions, we can do no better than a random feature model.

To formalize this conjecture, we introduce the following definition:

**Definition 3**.: _Let \(V\) be a vector space, and \(U\subseteq V\) a subspace. For any \(\ell\geq 0\), we define the space \(\mathcal{P}_{U,\ell}\) of functions \(f:V\to\mathbb{R}\) such that for any \(\mathbf{x}\in U\), the function \(f_{U,\mathbf{x}}\) introduced in Definition 2 is a polynomial of degree at most \(\ell\)._

Simply put, the space \(\mathcal{P}_{U,\ell}\) consists of polynomials in \(\mathbf{x}^{\perp}\) of degree at most \(\ell\), whose coefficients can be functions of \(\mathbf{x}\). We will denote by \(P_{U,\leq\ell}\) and \(P_{U,>\ell}\) the projections on \(\mathcal{P}_{U,\ell}\) and \(\mathcal{P}_{U,\ell}^{\perp}\) in \(\ell^{2}(\mathbb{R},\gamma)\), respectively, where \(\gamma\) is the Gaussian measure. Then, we conjecture the following:

**Conjecture 3**.: _Assume that \(p=\mathcal{O}(d^{\kappa_{1}})\) and \(n=\mathcal{O}(d^{\kappa_{2}})\), and that the first layer \(W\) only learns a subspace \(U^{\star}\subseteq V^{\star}\). Then, if \(\hat{\mathbf{a}}\) is obtained as in Eq. (7) for any value of \(\lambda\),_

\[\mathbb{E}\left[\left(f^{\star}(\mathbf{z})-\hat{f}(\mathbf{z};W,\hat{\mathbf{a}})\right) ^{2}\right]\geq\|P_{U^{\star},>\kappa}f^{\star}\|_{\gamma}^{2}+o(1), \tag{19}\]

_where \(\kappa=\min(\kappa_{1},\kappa_{2})\) and \(\|\cdot\|_{\gamma}\) is the norm in \(\ell^{2}(\mathbb{R},\gamma)\)._

It is easy to check that Proposition 2 corresponds to the \(\kappa=\kappa_{1}=0\) case of the above conjecture.

The linear case --We now show the \(\kappa_{1}=\kappa_{2}=1\) case of the above conjecture, often called the _proportional regime_. In this setting, the work [35] shows that for the random features \(W_{0}\), the minimization problem for the second layer (7) is equivalent to the one of a linear model; this result was extended in [39] to a trained first layer \(W_{1}\) for a "lazier" stepsize \(\eta=\mathcal{O}(1)\).

When \(\eta=\mathcal{O}(p)\), this equivalent model is not applicable anymore, due to the spike in the trained weight matrix \(W_{1}\); however, we show that this spike only allows the learning of non-linear functions parallel to the spike. Notice that the optimization problem (7) is equivalent to

\[\min_{\mathbf{a}\in\mathbb{R}^{p}}\frac{1}{n}\sum_{\nu=1}^{n}\left(\langle\mathbf{a}, \phi_{\text{CK}}(\mathbf{z}^{\nu})\rangle-f^{\star}(\mathbf{z}^{\nu})\right)^{2}+ \lambda\|\mathbf{a}\|^{2} \tag{20}\]

where \(\phi_{\text{CK}}(\mathbf{z})=\sigma(W\mathbf{z})/p\sqrt{n/d}\) is the feature map corresponding to the _conjugate kernel_ of the neural network. For a given direction \(\mathbf{v}\), we define the _conditional linear_ equivalent map \(\phi_{\text{CL}}(\mathbf{z};\mathbf{v})\) as follows: given the decomposition \(\mathbf{z}=z_{\mathbf{v}}\mathbf{v}+\mathbf{z}^{\perp}\),

\[\phi_{\text{CL}}(\mathbf{z};\mathbf{v})=\mu\left(z_{\mathbf{v}}\right)+\Psi(z_{\mathbf{v}})z^ {\perp}+\Phi(z_{\mathbf{v}})\mathbf{\xi} \tag{21}\]

where \(\mathbf{\xi}\sim\mathcal{N}(0,I_{p}-\mathbf{v}\mathbf{v}^{\perp})\), and \(\mu\in\mathbb{R}^{p},\Psi\in\mathbb{R}^{p\times d},\Phi\in\mathbb{R}^{p\times p}\) are chosen to match the first two conditional moments of \(\phi_{\text{CK}}\):

\[\mu(z_{\mathbf{v}}) =\mathbb{E}\left[\phi_{\text{CK}}(\mathbf{z})\mid z_{\mathbf{v}}\right], \quad\Psi(z_{\mathbf{v}})=\mathbb{E}\left[\phi_{\text{CK}}(\mathbf{z})(z^{\perp})^{ \top}\mid z_{\mathbf{v}}\right],\] \[\Phi(z_{\mathbf{v}}) =\mathrm{Cov}\left[\phi_{\text{CK}}(\mathbf{z})\mid z_{\mathbf{v}}\right]- \Psi(z_{\mathbf{v}})\Psi(z_{\mathbf{v}})^{\top}\]

Then, the following theorem holds:
**Theorem 2**.: _(informal) Assume that \(n,p=\Theta(d)\), and that the vector \(\mathbf{v}^{\star}\) defined in Theorem 1 is nonzero. Then, after one gradient step, there exists a vector \(\mathbf{v}\in\mathbb{R}^{p}\) such that:_

* _the projection of_ \(\mathbf{v}\) _on_ \(V^{\star}\) _is proportional to_ \(\mathbf{v}^{\star}\)_,_
* _the solutions_ \(\hat{\mathbf{a}}\)_,_ \(\tilde{\mathbf{a}}\) _to the optimization problem in (_20_) with the feature maps_ \(\phi_{\text{\sc cx}}(\mathbf{z})\) _and_ \(\phi_{\text{\sc cx}}(\mathbf{z};\mathbf{v})\) _yield the same generalization error._

The full statement of Theorem 2, with the relevant hypotheses, can be found in App. C. As we discuss in the Appendix, Theorem 2 also provides a proof of the case \(\kappa_{1}=\kappa_{2}=1\) of Conjecture 3:

**Corollary 1**.: _Conjecture 3 holds for \(\kappa_{1}=\kappa_{2}=1\)._

Informally, since anything orthogonal to \(\mathbf{v}\) maps to a linear model, we are only able to learn the linear part of \(f^{\star}\) on this subspace, which exactly corresponds to Conjecture 3.

## Conclusion

We provide a full picture of the learning dynamics when the first layer is trained using "large-batch" gradient steps. While some parts of those results are left as conjectures, we believe that they capture both the dynamics of multiple gradient steps, as well as the structure of the remaining directions. Proving those conjectures is likely to require fairly delicate concentration and random matrix arguments, which may be of independent mathematical interest.

## Acknowledgements

We thank Denny Wu, Theodor Misiakiewicz & Lenka Zdeborova for insightful discussions. We also acknowledge funding from the Swiss National Science Foundation grant SNFS OperaGOST, \(200021\_200390\) and the _Choose France - CNRS AI Rising Talents_ program.

Figure 5: **Learning with training of the second layer.** Simulation illustrating the different regimes in Fig.1, using \(d=512,p=1024\), a teacher function with two neurons and activation \(\sigma(z)=\operatorname{He}_{1}(x)+\operatorname{He}_{2}(x)/2!+ \operatorname{He}_{4}(x)/4!\), and a Relu student. (a) The first algorithm (green) applies a giant step and then learns the second layer. When \(n\gg d\), its performance goes beyond the linear predictor that would be obtained with a kernel method and reach the “single-index” regime in Fig.1. (b) To go beyond this regime, the second algorithm (blue) preprocesses the data to remove a plug-in estimate of the first Hermite coefficient. It reaches a lower plateau as \(n\approx d^{2}\), now beating the quadratic kernel. We contrast this behavior with the one of the random feature model (red). Details can be found in App. A.