Emergence and scaling laws in SGD learning of shallow neural networks

TL;DR

The paper studies online SGD for learning a two-layer neural network with $P$ orthogonal teacher directions on Gaussian data, focusing on the extensive-width regime and activations with information exponent $k_*>2$. It develops a single-phase, online SGD analysis under a 2-homogeneous parameterization, proving polynomial sample and runtime guarantees and sharp per-direction recovery times through a greedy, automatic-deflation style mechanism. In the power-law setting $a_p\asymp p^{-\beta}$, it derives precise scaling laws for the MSE as a function of the student width $m$, SGD steps $t$, and data dimension $d$, showing that a superposition of many abrupt learning events yields a smooth, global decay. Simulations corroborate the theoretical predictions, illustrating the compute-optimal frontier and the alignment between idealized scaling curves and actual gradient-descent dynamics, with practical implications for understanding feature-learning and scaling in wide, shallow networks.

Abstract

We study the complexity of online stochastic gradient descent (SGD) for learning a two-layer neural network with $P$ neurons on isotropic Gaussian data: $f_*(\boldsymbol{x}) = \sum_{p=1}^P a_p\cdot σ(\langle\boldsymbol{x},\boldsymbol{v}_p^*\rangle)$, $\boldsymbol{x} \sim \mathcal{N}(0,\boldsymbol{I}_d)$, where the activation $σ:\mathbb{R}\to\mathbb{R}$ is an even function with information exponent $k_*>2$ (defined as the lowest degree in the Hermite expansion), $\{\boldsymbol{v}^*_p\}_{p\in[P]}\subset \mathbb{R}^d$ are orthonormal signal directions, and the non-negative second-layer coefficients satisfy $\sum_{p} a_p^2=1$. We focus on the challenging ``extensive-width'' regime $P\gg 1$ and permit diverging condition number in the second-layer, covering as a special case the power-law scaling $a_p\asymp p^{-β}$ where $β\in\mathbb{R}_{\ge 0}$. We provide a precise analysis of SGD dynamics for the training of a student two-layer network to minimize the mean squared error (MSE) objective, and explicitly identify sharp transition times to recover each signal direction. In the power-law setting, we characterize scaling law exponents for the MSE loss with respect to the number of training samples and SGD steps, as well as the number of parameters in the student neural network. Our analysis entails that while the learning of individual teacher neurons exhibits abrupt transitions, the juxtaposition of $P\gg 1$ emergent learning curves at different timescales leads to a smooth scaling law in the cumulative objective.

Figures (3)

• Figure 1: Power-law scaling of MSE loss as a result of superposition of emergent learning curves.
• Figure 2: Theoretical and empirical risk curves with $\beta=0.8$. $(a)$ Idealized scaling curves described in Section \ref{['sec: idealized dynamics and scaling laws']}. $(b)$ Empirical scaling curve of GD training on the population loss with $d=2048, P=1024$.
• Figure 3: The greedy maximum selection matrix. The red diagonal entries represent the relevant neurons that eventually achieve overlap close to $1$. The remaining irrelevant entries can be partitioned into three groups: the upper triangular entries $\bar{v}_{p, \pi(q)}$ with $p \in [P_*]$ and $p < q \in [P]$, the lower triangular entries, $\bar{v}_{k, \pi(p)}$ with $p \in [P_*]$ and $p < k \in [m]$, and the lower right block $\bar{v}_{k, \pi(q)}$ with $k > P_*, q > P_*$. We will control these blocks using the row gap (purple arrow), column gap (blue arrow), and the threshold gap (green arrows), respectively.

Theorems & Definitions (105)

• Theorem : (Informal) sample complexity
• Proposition : (Informal) scaling law
• Remark
• Remark
• Theorem 2.1: Main theorem for online SGD
• Proposition 2.2: Scaling laws
• Remark
• Corollary 2.3: Unstable scaling law
• Lemma 3.1: Initialization
• Lemma B.1