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 f_*(boldsymbolx) propto sumj=1rλj σleft(langle boldsymbolθj, boldsymbolxrangleright), boldsymbolx ~ N(0,boldsymbolId), σ is the 2nd Hermite polynomial, and lbraceboldsymbolθj rbracej=1r subset ℝd are orthonormal signal directions. We consider the extensive-width regime r asymp dβ for β ∈ [0, 1), and assume a power-law decay on the (non-negative) second-layer coefficients λjasymp j-α for α ≥ 0. We present 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, sample size, and 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.