How Two-Layer Neural Networks Learn, One (Giant) Step at a Time

Yatin DandiFlorent KrzakalaBruno LoureiroLuca PesceLudovic Stephan

article2024JMLR70 citations

Establishes exact sample complexity bounds showing how two-layer neural networks escape the lazy regime and learn multi-index target directions across single and multiple gradient descent steps.

Listen

Modern neural networks achieve remarkable practical success partly due to their ability to adapt internal feature representations to training data, escaping the restrictive performance ceilings of classical kernel methods and lazy training regimes. However, understanding exactly how, when, and with how much data networks learn these underlying structures in high-dimensional settings has remained an open theoretical challenge. The article evaluates how two-layer neural networks learn low-dimensional feature representations from high-dimensional isotropic Gaussian data during the early phase of training using large-batch gradient descent.

The authors analyze a supervised learning model where the target function depends only on a finite set of relevant directions within a high-dimensional space. The approach combines rigorous mathematical analysis—employing high-dimensional concentration inequalities, tensor decompositions, and conditional Gaussian equivalence techniques—with numerical simulations across various network widths, sample batch sizes, and iteration counts.

The analysis establishes three primary findings regarding sample complexity, iteration dynamics, and model capacity. First, for a single gradient descent step, a sample batch scaling linearly with the input dimension is necessary and sufficient to align network weights with the target function, but it only captures a single linear direction. Learning multiple relevant directions in one step requires a much larger batch scaling quadratically with the dimension, or even higher polynomial orders corresponding to the target function's leap index (its lowest non-zero polynomial order). Second, taking multiple gradient steps drastically reduces data requirements: a linear batch size per step is sufficient to learn multiple directions sequentially over time, provided the target directions satisfy a hierarchical "staircase" property where new directions link linearly to previously learned ones. Third, overparameterization in the hidden layer allows networks to fit nonlinear structures along learned directions in fewer iterations, whereas kernel methods and unlearned directions face strict polynomial sample barriers.

These findings provide actionable insights for practitioners designing large-scale distributed and federated training pipelines. They show that while single-sample methods require many sequential updates to escape saddle points, large-batch gradient updates provide a direct speed-up by capturing finite components of target directions in very few parallelizable steps. System designers face a clear trade-off: training with very large single batches parallelizes well but demands significantly more total data for complex targets, whereas multi-step procedures with moderately sized batches leverage the staircase structure to achieve substantially higher sample efficiency.

Organizations should structure training pipelines to balance batch size and step counts according to task complexity, employing multi-step schedules when data availability is constrained. The theoretical conclusions hold with high probability under the defined boundary conditions of isotropic Gaussian inputs, polynomial target functions, and clean layer-wise training splits. While extending exact generalization proofs to fully general activations remains an area of ongoing mathematical development, the fundamental mechanisms governing feature specialization and sample scaling provide a solid baseline for algorithmic optimization.

Département d’Informatique Ecole Normale Supérieure École Polytechnique Fédérale de Lausanne Information, Learning and Physics (IdePHICS) Laboratory Statistical Physics of Computation (SPOC) Laboratory Université de Rennes gaussians network architecture representation learning two-layer neural networks feature learning lazy training

arXiv: 2305.18270 lucpoisson/GiantStep

Paper: High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation, Jimmy Ba et al. (2022). Its analysis of how a single gradient step aligns two-layer-network features with a target direction provides the immediate one-step feature-learning framework that this paper extends to data requirements and multiple learning steps.

Paper: Revenge of Monosemanticity: Specialized Neurons Improve Data Efficiency in MLPs, Amirhesam Abedsoltan et al. (2026). It carries feature-learning theory into clustered regression, showing how specialized neurons can recover distinct local predictive directions even when no useful global low-dimensional structure exists.

Cover for How Two-Layer Neural Networks Learn, One (Giant) Step at a Time

Abstract

For high-dimensional Gaussian data, 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 an improvement in the approximation capacity with respect to the initialization. First, we compare the influence of batch size to that of multiple (but finitely many) steps. For a single gradient step, a batch of size n = O(d) is both necessary and sufficient to align with the target function, although only a single direction can be learned. In contrast, n = O(d²) is essential for neurons to specialize in multiple relevant directions of the target with a single gradient step. Even in this case, we show there might exist “hard” directions requiring n = O(d^ℓ) samples to be learned, where ℓ is known as the leap index of the target. Second, we show that the picture drastically improves over multiple gradient steps: a batch size of n = O(d) is indeed sufficient 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 allow for a drastic improvement in 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 intertwined role of the structure of the task to

Topics

gaussians network architecture representation learning technical analysis multiple targets recurrent neural networks feedforward neural networks layered

Keywords

two-layer neural networks feature learning lazy training gradient descent linear subspace learning hierarchical learning high-dimensional asymptotics gradient flow polynomial activations generalization error rates

1. Introduction
2. Summary of main results
2.1 A single Gradient step
2.2 Learning over many GD iterations
2.3 From features to generalization
3. Statement of main theoretical results
3.1 Preliminaries
3.2 Setting and assumptions
3.3 Single gradient step
3.4 Learning with many steps
3.5 From feature learning to generalization bounds
Conclusion
Acknowledgements
Appendix A. Numerical investigation
A.1 Learning with a single giant step
A.2 Learning with multiple steps
Appendix B. Gradient descent on the first layer
B.1 Technical assumptions
B.2 Preliminaries
B.3 Computing expectations
B.4 Concentration
B.5 Proof of Theorems 4 and 5
B.6 Spike+Bulk decomposition
B.7 Second step: Proof Sketch for Theorem 7
B.8 Proof of Theorem 7
B.9 Prediction of the alignment at the second step
B.10 Limitations of the Staircase Structure
Appendix C. Learning the second layer
C.1 Proof of Proposition 8
C.2 Full statement of Theorem 12
C.3 Conditional GET
C.4 Generalization Error Lower Bounds: Proof of Corollary 13
References

Knowls

Knowl 1 — Sequential Subspace Learning under Multi-Step Gradient Descent
theoretical result
Consider standard Gaussian data $z \sim \mathcal{N}(0, I_d)$ and a multi-index target function $y = f^*(z) = g^*(\langle w_1^*, z \rangle, \dots, \langle w_r^*, z \rangle)$ with $w_1^*, \dots, w_r^*$ orthonormal, where $g^*: \mathbb{R}^r \to \mathbb{R}$ and the student activation $\sigma: \mathbb{R} \to \mathbb{R}$ are fixed polynomials with $\deg(\sigma) \ge \deg(g^*)$ . Let $V^* = \text{span}(w_1^*, \dots, w_r^*)$ .

Define the sequence of nested subspaces $U_0^* \subseteq U_1^* \subseteq \dots \subseteq U_t^* \subseteq \dots$ by:

$U_0^* = \{0\}$

$U_{t+1}^* = U_t^* \oplus \text{span}\left(\{\mu_{U_t^*, x}(f^*) : x \in U_t^*\}\right)$

where $\mu_{U, x}(f) = \mathbb{E}_{x^\perp \sim \mathcal{N}(0, I_{U^\perp})}[\nabla_{x^\perp} f(x + x^\perp)]$ .

When training a two-layer neural network with width $p > \dim(U_t^*)$ , batch size $n = \Theta(d)$ , learning rate $\eta > 0$ , and symmetric random initialization over $t$ full-batch gradient descent steps on the first layer with fresh batches at each step, the weight matrix $W^t \in \mathbb{R}^{p \times d}$ satisfies the following with high probability as $n, d \to \infty$ :
1. There exists an almost surely positive random variable $X_{t, a}$ independent of $d, n$ such that:
$\inf_{v \in U_t^*, \|v\|=1} \|W^t v\| \ge X_{t, a} + O\left(\frac{\text{polylog}(d)}{\sqrt{d}}\right)$
1. For any direction $v \in (U_t^*)^\perp \cap V^*$ orthogonal to $U_t^*$ in the target subspace, all neuron overlaps remain vanishing:
$|\langle w_i^t, v \rangle| = O\left(\frac{\text{polylog}(d)}{\sqrt{d}}\right) \quad \text{for all } i \in [p]$

Thus, each gradient step expands the learned representation to strictly include all directions that are linearly coupled to the previously learned subspace through the conditional Hermite gradient.
Knowl 2 — Subspace Specialization via a Single Giant Gradient Step
theoretical result
Let data be drawn from $z \sim \mathcal{N}(0, I_d)$ with target $f^*(z) = g^*(\langle w_1^*, z \rangle, \dots, \langle w_r^*, z \rangle)$ having leap index $\ell \ge 1$ , where $V^* = \text{span}(w_1^*, \dots, w_r^*)$ . Assume the $\ell$ -th Hermite coefficient $\mu_\ell$ of the activation $\sigma$ is non-zero.

Let the first-layer weights $W^1 \in \mathbb{R}^{p \times d}$ be updated by a single gradient step from symmetric random initialization with batch size $n = \Omega(d^\ell)$ and learning rate $\eta = p d^{(\ell-1)/2}$ . Let $C_\ell^* = C_\ell(f^*)$ be the $\ell$ -th order Hermite tensor of $f^*$ , and let $u_1^*, \dots, u_{r_\ell}^*$ denote its higher-order singular vectors spanning $V_\ell^* = \text{span}(u_1^*, \dots, u_{r_\ell}^*)$ . With probability at least $1 - c e^{-c \log(d)^2}$ :
1. For every neuron $i \in [p]$ , the relative energy in the target subspace is strictly bounded away from zero:
$\frac{\|\Pi^* w_i^1\|^2}{\|w_i^1\|^2} \ge X_i$

where $X_1, \dots, X_p$ are i.i.d. non-negative random variables independent of $d$ with $\mathbb{E}[X_i] > 0$ , and $\Pi^*$ is the orthogonal projection onto $V^*$ .
1. The projected weights $\pi_i^1 = \Pi^* w_i^1$ concentrate asymptotically on $V_\ell^*$ :
$\|(I - \Pi_{V_\ell^*}) \pi_i^1\| \le c \frac{\text{polylog}(d)}{\sqrt{d}}$
1. If $p \ge r_\ell$ , the weights span the entire subspace $V_\ell^*$ : for any $\delta > 0$ , there exists $C_\delta > 0$ such that with probability $1 - \delta$ :
$\inf_{v \in V_\ell^*, \|v\|=1} \|W^1 v\| \ge C_\delta$

For $\ell = 1$ , $V_1^*$ is one-dimensional and only a single direction is learned; for $\ell > 1$ , neurons can specialize into multiple directions of $V_\ell^*$ simultaneously in a single step.
Knowl 3 — Impossibility of Feature Learning in Sub-Leap Sample Regimes
theoretical result
Consider standard Gaussian inputs $z \sim \mathcal{N}(0, I_d)$ and a multi-index target function $f^*(z)$ with target subspace $V^* = \text{span}(w_1^*, \dots, w_r^*)$ and leap index $\ell = \min\{j \in \mathbb{N} : \langle f^*, \mathcal{H}_j \rangle_\gamma \neq 0\} \ge 1$ .

If the first layer of a two-layer neural network is updated from random spherical initialization using a single gradient step with batch size $n = O(d^{\ell - \delta})$ for some $\delta > 0$ , then with probability at least $1 - c p e^{-c(\delta) \log(d)^2}$ over the data and initialization:

$\frac{\|\Pi^* w_i^1\|^2}{\|w_i^1\|^2} \le c \frac{\text{polylog}(d)}{d^{(1 \wedge \delta)/2}} \quad \text{for all } i \in [p]$

where $\Pi^*$ is the orthogonal projection matrix onto $V^*$ , and $c$ is a universal constant. Consequently, when $n = O(d^{\ell - \delta})$ , only a vanishing fraction of every neuron's weight aligns with the target subspace, proving that $n = \Omega(d^\ell)$ samples are necessary for feature learning with a single gradient step.
Knowl 4 — Conditional Gaussian Equivalence and Generalization in the Proportional Regime
theoretical result
In the proportional regime where sample size and network width scale linearly with dimension ( $n, p = \Theta(d)$ ), consider a target function $f^*(z)$ with non-zero first Hermite coefficient $v^* = C_1(f^*) \neq 0$ so that $V_1^* = \text{span}(v^*)$ . After one gradient step on the first layer with learning rate $\eta = \Theta(p)$ , let $W^1 \in \mathbb{R}^{p \times d}$ denote the updated weights.

For $z \sim \mathcal{N}(0, I_d)$ decomposed along a direction $v \in \mathbb{R}^d$ as $z = z_v v + z^\perp$ with $z_v = \langle z, v \rangle$ , define the conjugate kernel features $\phi_{\text{CK}}(z) = \sigma(W^1 z)$ and the conditional linear equivalent Gaussian features:

$\phi_{\text{CL}}(z; v) = \mu(z_v) + \Psi(z_v) z^\perp + \Phi(z_v) \xi$

where $\xi \sim \mathcal{N}(0, I_p)$ , and $\mu(z_v), \Psi(z_v), \Phi(z_v)$ match the first two conditional moments of $\phi_{\text{CK}}(z)$ given $z_v$ .

There exists a vector $v \in \mathbb{R}^d$ whose projection onto $V^*$ is proportional to $v^*$ such that the training and generalization errors of ridge regression on the second layer with features $\phi_{\text{CK}}$ converge in distribution to those with features $\phi_{\text{CL}}(z; v)$ . Consequently, the asymptotic generalization error of the trained two-layer network satisfies:

$\mathbb{E}\left[\left(f^*(z) - \hat{f}(z; W^1, \hat{a})\right)^2\right] \ge \|\mathcal{P}_{v^*, >1} f^*\|_\gamma^2 - o(1)$

where $\mathcal{P}_{v^*, >1} f^*$ denotes the component of $f^*$ whose conditional expansion orthogonal to $v^*$ has polynomial degree strictly greater than 1 in $z^\perp$ . Hence, a single step in the proportional regime can only fit non-linearities along the 1D spike direction $v^*$ .
Knowl 5 — Subspace Conditioning and Directional Polynomial Function Spaces
definition
Let $V$ be a vector space, $U \subseteq V$ a subspace, and $f: V \to \mathbb{R}$ a square-integrable function under the Gaussian measure $\gamma$ . For any $x \in U$ , the conditional function $f_{U, x}: U^\perp \to \mathbb{R}$ is defined by:

$f_{U, x}(x^\perp) = f(x + x^\perp)$

The first Hermite coefficient of this conditional function w.r.t. the standard Gaussian measure on $U^\perp$ is defined as:

$\mu_{U, x}(f) = \mathbb{E}_{x^\perp \sim \mathcal{N}(0, I_{U^\perp})}\left[\nabla_{x^\perp} f_{U, x}(x^\perp)\right] \in U^\perp$

For any integer $k \ge 0$ , $\mathcal{P}_{U, k}$ denotes the space of functions $f: V \to \mathbb{R}$ such that for every $x \in U$ , the conditional function $f_{U, x}$ is a polynomial in $x^\perp$ of degree at most $k$ . The orthogonal projection in $L^2(V, \gamma)$ onto $\mathcal{P}_{U, k}$ is denoted by $\mathcal{P}_{U, \le k}$ , and the projection onto its orthogonal complement is denoted by $\mathcal{P}_{U, >k}$ .
Knowl 6 — Generalization Error Lower Bound for Subspace Incomplete Features at Finite Width
theoretical result
Let $f^*(z) = g^*(\langle w_1^*, z \rangle, \dots, \langle w_r^*, z \rangle)$ be a multi-index target function with target subspace $V^* = \text{span}(w_1^*, \dots, w_r^*)$ on Gaussian inputs $z \sim \mathcal{N}(0, I_d)$ .

Assume that the hidden width $p$ remains bounded as $n, d \to \infty$ , and that the first-layer weights $W \in \mathbb{R}^{p \times d}$ only learn a proper subspace $U \subsetneq V^*$ , meaning that for every $v \in V^* \cap U^\perp$ , $|\langle w_i, v \rangle| = o(1)$ for all $i \in [p]$ . Then for any choice of second-layer weights $a \in \mathbb{R}^p$ satisfying $\|a\|_\infty \le O(1)$ , the generalization error of the two-layer network $\hat{f}(z; W, a) = \frac{1}{\sqrt{p}} \sum_{i=1}^p a_i \sigma(\langle w_i, z \rangle)$ satisfies:

$\mathbb{E}\left[\left(f^*(z) - \hat{f}(z; W, a)\right)^2\right] \ge \mathbb{E}_{P_U z}\left[\text{Var}\left(f^*(z) \mid P_U z\right)\right] - o(1)$

where $P_U$ is the orthogonal projection onto the subspace $U$ .
Knowl 7 — Leap Index of a Target Function
definition
Let $f^*: \mathbb{R}^d \to \mathbb{R}$ be a square-integrable function with respect to the standard Gaussian measure $\gamma = \mathcal{N}(0, I_d)$ , admitting the Hermite tensor expansion $f^*(z) = \sum_{j=0}^\infty \langle C_j(f^*), \mathcal{H}_j(z) \rangle$ , where $\mathcal{H}_j(z)$ is the $j$ -th order Hermite tensor and $C_j(f^*) = \langle f^*, \mathcal{H}_j \rangle_\gamma$ is the $j$ -th Hermite tensor coefficient.

The leap index $\ell$ of $f^*$ is defined as the lowest non-zero polynomial degree in its Hermite expansion:

$\ell = \min\left\{j \in \mathbb{N} : \langle f^*, \mathcal{H}_j \rangle_\gamma \neq 0\right\}$

For a single-index model $f^*(z) = g^*(\{w^*, z\rangle)$ , the leap index corresponds to the index of the first non-vanishing 1D Hermite coefficient of $g^*$ , also known as the information exponent.
Knowl 8 — Symmetry Barrier to Hierarchical Staircase Feature Learning
limitation
Consider an isotropic Gaussian covariate setting $z \sim \mathcal{N}(0, I_d)$ and a multi-index target function given by a symmetric additive combination of single-index functions:

$y = f^*(z) = \sum_{k=1}^r \sigma^*(\{w_k^*, z\rangle)$

where $w_1^*, \dots, w_r^*$ are orthonormal vectors and $\sigma^*$ has leap index 1 (i.e., its first Hermite coefficient is non-zero).

Under multi-step gradient descent with batch size $n = \Theta(d)$ , the subspace conditioning sequence $U_t^*$ satisfies:

$U_t^* = U_1^* = \text{span}\left(\sum_{k=1}^r w_k^*\right) \quad \text{for all } t \ge 1$

Because of the permutation symmetry among the coordinates $w_k^*$ , all directions orthogonal to $\sum_{k=1}^r w_k^*$ have zero conditional Hermite gradients ( $\mu_{U_1^*, x}(f^*) = 0$ for all $x \in U_1^*$ ). Consequently, multiple gradient steps with $n = \Theta(d)$ cannot learn any new orthogonal directions; escaping the initial 1D direction requires introducing asymmetry into the target function or increasing the batch size to $n = \Omega(d^2)$ .
Knowl 9 — Two-Phase Large-Batch Training Procedure for Two-Layer Neural Networks
algorithm
The two-phase training algorithm optimizes a two-layer neural network $\hat{f}(z; W, a) = \frac{1}{\sqrt{p}} \sum_{i=1}^p a_i \sigma(\langle w_i, z \rangle)$ on Gaussian inputs $z \sim \mathcal{N}(0, I_d)$ . In Phase 1, the first-layer weights $W \in \mathbb{R}^{p \times d}$ undergo $T$ steps of full-batch gradient descent with fresh batches of size $n$ and learning rate $\eta = O(p \sqrt{n/d})$ . In Phase 2, the second-layer weights $a \in \mathbb{R}^p$ are solved via ridge regression on an independent batch with regularization $\lambda$ .
Input: Data dimension $d$ , hidden width $p$ , number of first-layer GD steps $T$ , batch size $n$ , learning rate $\eta$ , ridge parameter $\lambda$
Output: Trained weights $(W_T, \hat{a})$
Initialize $W_0 \in \mathbb{R}^{p \times d}$ with rows $w_i^0 \sim \text{Unif}(\mathbb{S}^{d-1})$
Initialize $a_0 \in \mathbb{R}^p$ with $\sqrt{p} a_{0,i} \sim \text{Unif}([-1, 1])$ enforcing $w_{p-i+1}^0 = w_i^0$ and $a_{0, p-i+1} = -a_{0, i}$ for $i \in [p/2]$
for $t = 0$ to $T-1$ do
    Sample data matrix $Z_t \sim \mathcal{N}(0, I_{n \times d})$ and labels $Y_t = f^*(Z_t) \in \mathbb{R}^n$
    Compute gradient matrix $G_t \in \mathbb{R}^{p \times d}$ with rows:
         $G_{t, i} = \frac{a_{0, i}}{\sqrt{p}} \frac{1}{n} \sum_{\nu=1}^n z_t^\nu \sigma'(\langle w_i^t, z_t^\nu \rangle) (\hat{f}(z_t^\nu; W_t, a_0) - f^*(z_t^\nu))$
     $W_{t+1} = W_t - \eta G_t$
Sample independent batch $Z \sim \mathcal{N}(0, I_{n \times d})$ and labels $Y = f^*(Z) \in \mathbb{R}^n$
Compute feature activations $X = \sigma(W_T Z) \in \mathbb{R}^{n \times p}$
if $n < p$ then
     $\hat{a} = X^\top (X X^\top + \lambda I_n)^{-1} Y$
else
     $\hat{a} = (X^\top X + \lambda I_p)^{-1} X^\top Y$
return $(W_T, \hat{a})$

Knowl 10 — Iteration and Sample Complexities for Learning Target Directions

data/table

The table compares the required number of gradient steps $\tau$ and total sample complexity $n_T$ for two-layer neural networks to learn the relevant directions of multi-index targets $f^*$ of varying complexity across three learning algorithms: sequential one-sample SGD ( $n=1$ ), single-step large-batch GD ( $n = \Theta(d^\ell)$ ), and multi-step large-batch GD ( $n = \Theta(d)$ ).

Target Complexity	SGD ( $n=1$ )	One-step GD ( $n=\Theta(d^\ell)$ )	Multi-step GD ( $n=\Theta(d)$ )
Single-index, $\ell=1$	$\tau = n_T = d$	$\tau = 1, n_T = \Theta(d)$	$\tau = 1, n_T = \Theta(d)$
Single-index, $\ell=2$	$\tau = n_T = d \log d$	$\tau = 1, n_T = \Theta(d^2)$	$\tau = \Theta(\log d), n_T = \Theta(d \log d)$
Single-index, $\ell>2$	$\tau = n_T = d^{\ell-1}$	$\tau = 1, n_T = \Theta(d^\ell)$	$\tau = \Theta(d^{\ell-2}), n_T = \Theta(d^{\ell-1})$
Staircase	$\tau = n_T = d$	$\tau = 1, n_T = \Omega(d^2)$	$\tau = \Theta(1), n_T = \Theta(d)$

For staircase functions, multi-step large-batch GD with $n = \Theta(d)$ matches the optimal total sample complexity $n_T = \Theta(d)$ of one-pass SGD while reducing the iteration complexity from $\tau = d$ sequential steps to $\tau = \Theta(1)$ parallelizable steps. For hard single-index functions with leap $\ell > 1$ , a single step requires a batch size of $\Theta(d^\ell)$ , whereas multiple steps reduce the overall sample requirement to $\Theta(d^{\ell-1})$ .

Coverage note — Unproven conjectures (Conjectures 9 and 11) were excluded in favor of their proven special cases (Proposition 8, Theorem 12, and Corollary 13). Technical concentration lemmas used solely within proofs (e.g. Lemmas 14-37) were omitted.

References

1.Emmanuel Abbe, Enric Boix-Adsera, Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj. The staircase property: How hierarchical structure can guide deep learning. Advances in Neural Information Processing Systems, 34:26989–27002, 2021.
2.Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782–4887. PMLR, 2022.
3.Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics, 2023.
4.Luca Arnaboldi, Ludovic Stephan, Florent Krzakala, and Bruno Loureiro. From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks, February 2023. arXiv:2302.05882 [cond-mat, stat] type: article.
5.Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations, 2022.
6.Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 37932–37946. Curran Associates, Inc., 2022.
7.Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629–681, 2017.
8.Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learning Research, 22(106):1–51, 2021.
9.Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Information Processing Systems, 35:25349–25362, 2022.
10.Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks, 2023.
11.Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 9768–9783. Curran Associates, Inc., 2022.
12.Blake Bordelon and Cengiz Pehlevan. Dynamics of finite width kernel and prediction fluctuations in mean field neural networks, 2023.
13.Blake Bordelon, Abdulkadir Canatar, and Cengiz Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 1024–1034. PMLR, 13–18 Jul 2020.
14.David Bosch, Ashkan Panahi, and Babak Hassibi. Precise asymptotic analysis of deep random feature models, 2023.
15.Etienne Boursier, Loucas Pillaud-Vivien, and Nicolas Flammarion. Gradient flow dynamics of shallow relu networks for square loss and orthogonal inputs. Advances in Neural Information Processing Systems, 35:20105–20118, 2022.
16.S. Boyd, A. Ghosh, B. Prabhakar, and D. Shah. Randomized gossip algorithms. IEEE Transactions on Information Theory, 52(6):2508–2530, June 2006. ISSN 1557-9654. doi: 10.1109/TIT.2006.874516.
17.Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature Communications, 12(1):2914, May 2021. ISSN 2041-1723. doi: 10.1038/s41467-021-23103-1.
18.Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.
19.Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training in Differentiable Programming. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
20.Hugo Cui, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Generalization error rates in kernel regression: The crossover from the noiseless to noisy regime. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 10131–10143. Curran Associates, Inc., 2021.
21.Hugo Cui, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Error rates for kernel classification under source and capacity conditions, 2022.
22.Alex Damian, Eshaan Nichani, Rong Ge, and Jason D. Lee. Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models. Technical report, May 2023. arXiv:2305.10633 [cs, math, stat] type: article.
23.Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 5413–5452. PMLR, 02–05 Jul 2022.
24.Yatin Dandi, Ludovic Stephan, Florent Krzakala, Bruno Loureiro, and Lenka Zdeborová. Universality laws for gaussian mixtures in generalized linear models, 2023.
25.Lieven De Lathauwer, Bart De Moor, and Joos Vandewalle. A multilinear singular value decomposition. SIAM journal on Matrix Analysis and Applications, 21(4):1253–1278, 2000.
26.Luc Devroye, László Györfi, and Gábor Lugosi. A probabilistic theory of pattern recognition, volume 31. Springer Science & Business Media, 2013.
27.Oussama Dhifallah and Yue M. Lu. A precise performance analysis of learning with random features, 2020.
28.Rainer Dietrich, Manfred Opper, and Haim Sompolinsky. Statistical mechanics of support vector networks. Phys. Rev. Lett., 82:2975–2978, Apr 1999. doi: 10.1103/PhysRevLett.82.2975.
29.Konstantin Donhauser, Mingqi Wu, and Fanny Yang. How rotational invariance of common kernels prevents generalization in high dimensions. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2804–2814. PMLR, 18–24 Jul 2021.
30.Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet, editors, Proceedings of the 31st Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pages 1887–1930. PMLR, 06–09 Jul 2018. URL https://proceedings.mlr.press/v75/dudeja18a.html.
31.Federica Gerace, Bruno Loureiro, Florent Krzakala, Marc Mezard, and Lenka Zdeborova. Generalisation error in learning with random features and the hidden manifold model. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3452–3462. PMLR, 13–18 Jul 2020.
32.Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural network. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
33.Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 14820–14830. Curran Associates, Inc., 2020.
34.Sebastian Goldt, Bruno Loureiro, Galen Reeves, Florent Krzakala, Marc Mezard, and Lenka Zdeborova. The gaussian equivalence of generative models for learning with shallow neural networks. In Joan Bruna, Jan Hesthaven, and Lenka Zdeborova, editors, Proceedings of the 2nd Mathematical and Scientific Machine Learning Conference, volume 145 of Proceedings of Machine Learning Research, pages 426–471. PMLR, 16–19 Aug 2022.
35.Friedrich Gotze, Holger Sambale, and Arthur Sinulis. Concentration inequalities for polynomials in α-sub-exponential random variables. Electronic Journal of Probability, 2019.
36.Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyröla, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677, 2017.
37.Harold Grad. Note on N-dimensional hermite polynomials. Communications on Pure and Applied Mathematics, 2(4):325–330, 1949. ISSN 1097-0312. doi: 10.1002/cpa.3160020402.
38.W.H. Greub. Multilinear Algebra. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2012. ISBN 9783662007952.
39.Hong Hu and Yue M Lu. Universality laws for high-dimensional learning with random features. IEEE Transactions on Information Theory, 2022.
40.Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel. Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity. arXiv preprint arXiv:2106.15933, 2021.
41.Svante Janson. Gaussian hilbert spaces. Number 129. Cambridge university press, 1997.
42.Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. Sgd on neural networks learns functions of increasing complexity. Advances in neural information processing systems, 32, 2019.
43.Achim Klenke. Probability theory: a comprehensive course. Springer Science & Business Media, 2013.
44.Michel Ledoux and Michel Talagrand. Probability in Banach Spaces: Isoperimetry and Processes. Springer-Verlag, 1991. ISBN 9780387520131. Google-Books-ID: juC1QgAACAAJ.
45.Li Li, Yuxi Fan, Mike Tse, and Kuo-Yi Lin. A review of applications in federated learning. Computers & Industrial Engineering, 149:106854, 2020.
46.Bruno Loureiro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Marc Mezard, and Lenka Zdeborová. Learning curves of generic features maps for realistic datasets with a teacher-student model. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 18137–18151. Curran Associates, Inc., 2021.
47.Bruno Loureiro, Cedric Gerbelot, Maria Refinetti, Gabriele Sicuro, and Florent Krzakala. Fluctuations, bias, variance &; ensemble of learners: Exact asymptotics for convex losses in high-dimension. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 14283–14314. PMLR, 17–23 Jul 2022.
48.Song Mei and Andrea Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75(4):667–766, 2022.
49.Song Mei, Andrea Montanari, and Phan-Minh 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.
50.Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on learning theory, pages 2388–2464. PMLR, 2019.
51.Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Generalization error of random feature and kernel methods: Hypercontractivity and kernel matrix concentration. Applied and Computational Harmonic Analysis, 59:3–84, 2022. ISSN 1063-5203. doi: https://doi.org/10.1016/j.acha.2021.12.003. Special Issue on Harmonic Analysis and Machine Learning.
52.Boris Samuilovich Mityagin. The zero set of a real analytic function. Mathematical Notes, 107(3-4):529–530, 2020.
53.Andrea Montanari and Basil N. Saeed. Universality of empirical risk minimization. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 4310–4312. PMLR, 02–05 Jul 2022.
54.Gadi Naveh and Zohar Ringel. A self consistent theory of gaussian processes captures feature learning effects in finite cnns. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 21352–21364. Curran Associates, Inc., 2021.
55.M. Opper and R. Urbanczik. Universal learning curves of support vector machines. Phys. Rev. Lett., 86:4410–4413, May 2001. doi: 10.1103/PhysRevLett.86.4410.
56.Victor H. de la Pena and S. J. Montgomery-Smith. Decoupling Inequalities for the Tail Probabilities of Multivariate $U$ -Statistics. The Annals of Probability, 23(2):806–816, April 1995. ISSN 0091-1798, 2168-894X. doi: 10.1214/aop/1176988291.
57.Leonardo Petrini, Francesco Cagnetta, Eric Vanden-Eijnden, and Matthieu Wyart. Learning sparse features can lead to overfitting in neural networks, 2022.
58.Michael Polyak. Feynman diagrams for pedestrians and mathematicians. Graphs and patterns in mathematics and theoretical physics, 73:15–42, 2005.
59.Grant Rotskoff and Eric Vanden-Eijnden. Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 75(9):1889–1935, 2022. doi: https://doi.org/10.1002/cpa.22074.
60.Alessandro Rudi and Lorenzo Rosasco. Generalization properties of learning with random features. Advances in neural information processing systems, 30, 2017.
61.David Saad and Sara A. Solla. On-line learning in soft committee machines. Physical Review E, 52(4):4225–4243, October 1995. doi: 10.1103/PhysRevE.52.4225.
62.Dominik Schröder, Hugo Cui, Daniil Dmitriev, and Bruno Loureiro. Deterministic equivalent and error universality of deep random features learning, 2023.
63.Inbar Seroussi, Gadi Naveh, and Zohar Ringel. Separation of scales and a thermodynamic description of feature learning in some cnns. Nature Communications, 14(1):908, Feb 2023. ISSN 2041-1723. doi: 10.1038/s41467-023-36361-y.
64.James B. Simon, Madeline Dickens, Dhruva Karkada, and Michael R. DeWeese. The eigen-learning framework: A conservation law perspective on kernel regression and wide neural networks, 2022.
65.Stefano Spigler, Mario Geiger, and Matthieu Wyart. Asymptotic learning curves of kernel methods: empirical data versus teacher–student paradigm. Journal of Statistical Mechanics: Theory and Experiment, 2020(12):124001, dec 2020. doi: 10.1088/1742-5468/abc61d.
66.Aad van der Vaart and Jon Wellner. Weak Convergence and Empirical Processes: With Applications to Statistics. Springer Science & Business Media, March 1996. ISBN 9780387946405. Google-Books-ID: OCenCW9qmp4C.
67.Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint arXiv:1011.3027, 2010.
68.Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
69.Lechao Xiao, Hong Hu, Theodor Misiakiewicz, Yue Lu, and Jeffrey Pennington. Precise learning curves and higher-order scalings for dot-product kernel regression. Advances in Neural Information Processing Systems, 35:4558–4570, 2022.
70.Lenka Zdeborová. Understanding deep learning is also a job for physicists. Nature Physics, 16(6):602–604, 2020.

Citation

MLA

Dandi, Y., et al. “How Two-Layer Neural Networks Learn, One (Giant) Step at a Time”. Journal of Machine Learning Research, vol. 25, no. 349, 2024, pp. 1–5, https://www.jmlr.org/papers/v25/23-1543.html.

APA

Dandi, Y., Krzakala, F., Loureiro, B., Pesce, L., & Stephan, L. (2024). How Two-Layer Neural Networks Learn, One (Giant) Step at a Time. Journal of Machine Learning Research, 25(349), 1–65. https://www.jmlr.org/papers/v25/23-1543.html

Chicago

Dandi, Y., F. Krzakala, B. Loureiro, L. Pesce, and L. Stephan. 2024. “How Two-Layer Neural Networks Learn, One (Giant) Step at a Time”. Journal of Machine Learning Research 25 (349): 1–65. https://www.jmlr.org/papers/v25/23-1543.html.

Harvard

Dandi, Y. et al. (2024) “How Two-Layer Neural Networks Learn, One (Giant) Step at a Time”, Journal of Machine Learning Research, 25(349), pp. 1–65. Available at: https://www.jmlr.org/papers/v25/23-1543.html.

Vancouver

1. Dandi Y, Krzakala F, Loureiro B, Pesce L, Stephan L (2024) How Two-Layer Neural Networks Learn, One (Giant) Step at a Time. Journal of Machine Learning Research 25:1–65

BibTeX

@article{JMLR:v25:23-1543,
  author  = {Yatin Dandi and Florent Krzakala and Bruno Loureiro and Luca Pesce and Ludovic Stephan},
  title   = {How Two-Layer Neural Networks Learn, One (Giant) Step at a Time},
  journal = {Journal of Machine Learning Research},
  year    = {2024},
  volume  = {25},
  number  = {349},
  pages   = {1--65},
  url     = {http://jmlr.org/papers/v25/23-1543.html}
}

Metadata:DOI registry

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

How Two-Layer Neural Networks Learn, One (Giant) Step at a Time

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Sequential Subspace Learning under Multi-Step Gradient Descent

Knowl 2 — Subspace Specialization via a Single Giant Gradient Step

Knowl 3 — Impossibility of Feature Learning in Sub-Leap Sample Regimes

Knowl 4 — Conditional Gaussian Equivalence and Generalization in the Proportional Regime

Knowl 5 — Subspace Conditioning and Directional Polynomial Function Spaces

Knowl 6 — Generalization Error Lower Bound for Subspace Incomplete Features at Finite Width

Knowl 7 — Leap Index of a Target Function

Knowl 8 — Symmetry Barrier to Hierarchical Staircase Feature Learning