Bayes-optimal learning of an extensive-width neural network from quadratically many samples

TL;DR

This work derives a closed-form Bayes-optimal generalization error for learning a single-hidden-layer neural network with a quadratic activation in the regime of extensive width and quadratically many samples, by reducing the problem to an extensive-rank matrix denoising framework and leveraging free probability techniques. It introduces GAMP-RIE, a polynomial-time algorithm whose state evolution matches the Bayes-optimal fixed point, thereby closing the gap between information-theoretic limits and computable methods in this setting. The analysis reveals a perfect-recovery threshold in the noiseless case and detailed asymptotics for various limits of the width-to-dimension ratio, clarifying the role of spectral properties of associated random matrices. Empirically, the authors show that noiseless gradient descent with random initialization can sample near-posterior interpolants, and that averaging over initializations can achieve Bayes-optimal performance, while in noisy settings GD behaves differently, highlighting a nuanced landscape of optimization versus information-theoretic limits. Overall, the work advances theoretical understanding of high-dimensional learning with extensive width and non-linear activations, and introduces techniques with potential applicability to other matrix-learning problems and phase-retrieval-like tasks.

Abstract

We consider the problem of learning a target function corresponding to a single hidden layer neural network, with a quadratic activation function after the first layer, and random weights. We consider the asymptotic limit where the input dimension and the network width are proportionally large. Recent work [Cui & al '23] established that linear regression provides Bayes-optimal test error to learn such a function when the number of available samples is only linear in the dimension. That work stressed the open challenge of theoretically analyzing the optimal test error in the more interesting regime where the number of samples is quadratic in the dimension. In this paper, we solve this challenge for quadratic activations and derive a closed-form expression for the Bayes-optimal test error. We also provide an algorithm, that we call GAMP-RIE, which combines approximate message passing with rotationally invariant matrix denoising, and that asymptotically achieves the optimal performance. Technically, our result is enabled by establishing a link with recent works on optimal denoising of extensive-rank matrices and on the ellipsoid fitting problem. We further show empirically that, in the absence of noise, randomly-initialized gradient descent seems to sample the space of weights, leading to zero training loss, and averaging over initialization leads to a test error equal to the Bayes-optimal one.

Figures (4)

• Figure 1: Left: The asymptotic MMSE of eq. \ref{['eq:MMSE_qhat']} for the noiseless $(\Delta=0)$ case, as a function of the sample complexity $\alpha$, for various width ratios $\kappa$. Right: Phase diagram representing the MMSE, brighter color indicates a higher value. The red curve is the perfect recovery transition line $\alpha_\PR$, see eq. \ref{['eq:alphaPR']}, and its origin is discussed in Section \ref{['sec:discussion']}.
• Figure 2: Left: Behavior of the asymptotic MMSE in the noiseless ($\Delta=0$) case as $\kappa$ gets increasingly small. The continuous lines are given by eq. \ref{['eq:MMSE_qhat']}, which we compare with the asymptotic $\kappa\to 0$ curve obtained by eq. \ref{['eq:small_kappa_MMSE_noiseless']}. We emphasize that the horizontal axis is $\alpha/\kappa$, which remains of order $\Theta(1)$ as $\kappa \to 0$: it corresponds to a number of samples $n$ of the same order as the number of parameters $dm$. Right: Comparison of the performance of GAMP-RIE with the asymptotic MMSE \ref{['eq:MMSE_qhat']} both in the noiseless ($\Delta = 0$) and in a noisy ($\sqrt{\Delta} = 0.25$) case, with $\kappa = 0.5$. Each dot is the average over $8$ runs of GAMP-RIE at a moderate size of either $d=100$ (circle dots) or $d=200$ (crosses). The error bars are the standard deviations of the MSE.
• Figure 3: Mean squared error (MSE) as a function of the sample complexity $\alpha$ for $\kappa\!=\!1/2$. Dots are simulations using GD with a single initialization averaged over $32$ realizations of the dataset, crosses are averages over $64$ initializations with $2$ realizations of the dataset. The continuous lines are the asymptotic MMSE given by \ref{['eq:MMSE_qhat']}. Left: noiseless $\Delta=0$ case. The colors indicate the size $d$. We can see how AGD appears to be well described by the theoretical MMSE. We used the learning rates $0.2$ for $d\!=\!200$ and $0.07$ for $d\!=\!100$. Right: Comparison of GD between the noisy $\sqrt{\Delta} \!=\! 0.25$ case ( red) and noiseless $\Delta \!=\! 0$ case ( blue). Adding noise makes AGD worse than the MMSE, and for sample complexity $\alpha \!\gtrsim\! 0.3$, all the initializations of GD converge to the same point, making the GD and AGD curves collapse.
• Figure 4: Left: Mean squared error as a function of the sample complexity $\alpha$, for $\kappa=1/2$ and $\Delta = 0.25^2$. Dots are simulations using GD with a single initialization averaged over $32$ realizations of the dataset, crosses are averages over $64$ initializations. The continuous line is the asymptotic MMSE given by \ref{['eq:MMSE_qhat']}. The colors indicate the strength of the regularization. Right: Trivialization threshold in the sample complexity $\alpha_T$ as a function of the noise level $\Delta$ in the teacher without regularization, $\lambda=0$. The measurement has a resolution of $0.1$ on the noise level and of $0.007$ on the sample complexity

Theorems & Definitions (7)

• theorem 1: Free entropy of matrix denoising
• theorem 2
• theorem 3: Stieltjes-Perron inversion formula
• lemma 1
• proof : Proof of Lemma \ref{['lemma:MMSEy_MMSES']}
• theorem 4: Theorem 1.1 of guionnet2002large
• theorem 5: State Evolution (informal) berthier2020stategerbelot2023graph