Learning quadratic neural networks in high dimensions: SGD dynamics and scaling laws

We study the optimization and sample complexity of gradient-based training of a two-layer neural network with quadratic activation function in the high-dimensional regime, where the data is generated as $y\propto {\sum }_{j=1}^r{\lambda }_j\sigma \left(⟨{\mathbf{\theta }}_{\mathbf{j}},\mathbf{x}⟩\right),\mathbf{x}\sim \mathcal{N}(0,{\mathbf{I}}_d)$, where $\sigma$ is the 2nd Hermite polynomial, and $\{{\mathbf{\theta }}_j{\}}_{j=1}^r\subset {\mathbb{R}}^d$ are orthonormal signal directions.

We consider the extensive-width regime $r\asymp d^{\beta }$ for $\beta \in (0,1)$, and assume a power-law decay on the (non-negative) second-layer coefficients ${\lambda }_j\asymp j^{-\alpha }$ for $\alpha \ge 0$.

We provide a sharp analysis of the SGD dynamics in the feature learning regime, for both the population limit and the finite-sample (online) discretization, and derive scaling laws for the prediction risk that highlight the power-law dependencies on the optimization time, the sample size, and the model width.

Our analysis combines a precise characterization of the associated matrix Riccati differential equation with novel matrix monotonicity arguments to establish convergence guarantees for the infinite-dimensional effective dynamics.