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

Studies a teacher-student setting with very specific configurations of input data, teach weight distribution, noise distribution and teacher/student architecture. The minimum MSE achievable, Eqn (3), is subjected to a particular limit, Eqn (4), and shown in closed-form to be Eqn (14). From Eqn 14 and assuming the noiseless case., we can derive the "perfect-recovery" threshold, which is stated as one of the main contributions (Eqn (1)).

This paper explores the performance of an optimal (in the Bayes-Optimal sense) estimator in the teacher-student setting with aligned architectures being one-hidden layer neural networks with quadratic activations. The number of neurons is extensive in the dimension. Given previous works where the "proportional limit" regime was studied (i.e. asymptotically square dataset), finding that a vanilla estimator could attain optimal performance, they focus on a slightly different regime. Precisely, the author(s) build upon numerical observations of other works and let the number of samples scale quadratically in the dimension. Having chosen the playground, the analysis proceeds with recasting the problem to a form that exploits the symmetries of the objective function, and a replica computation of the free energy. Such computations allow the author(s) to derive an asymptotic formula of the Minimum Mean Square Error. To back the Physics-based derivation, they discuss a break-down of the steps required to formalize it completely. Of the three main points, two are conjectures. The final one instead has a mathematical proof. Given that Gradient Descent is the go-to method, they explore its performance in relation to the theoretically optimal MMSE. Interestingly, they find that in the noiseless case it appears to sample weights from the posterior. The same picture breaks when the labels have noise. There, they discuss the appearance of a phase transition that trivializes the randomness over initialization of GD.

The paper presents some learning theory for learning the large-dimension, large width perceptrons with a quadratic activation functions. The results seem to be twofold: Claim 1, and its specialisation in eq 14, which gives us the MSE test error. In addition they note that empirically the that the SGD solution averaged over random initialisations leads to a near-bayes-optimal learning error.

The paper explores Bayes-optimal learning of a neural network with extensive input dimensions and a single hidden layer using quadratic activations. It presents a closed-form expression for the optimal test error when the sample complexity is quadratic. This work connects to matrix denoising and ellipsoid fitting, providing both theoretical insights and empirical validations. Key findings include showing that randomly initialized gradient descent can approach the Bayes-optimal error.

We thank all referees for their interest in our work and their comments that will help to clarify our paper. We will answer here to the most common questions of the reviewers.

• Why are square-activated networks relevant? (Q1 of bzi8. Question of xzmx. W3 of NWbd. W3 of XJ6b). We want to stress that this is the first first work that studies non-linear two layer neural networks of extensive width, with enough samples to have feature learning. Furthermore, analogous models have been considered in the machine learning literature for analyzing the implicit bias of gradient descent [A1-A3] and as toy models for probing the advantage of depth in neural networks [B1, B2]. These existing works derive bounds, whereas our tight asymptotic result could help to settle the concerned questions about implicit regularization and advantage of depth more tightly. We thus anticipate that already the quadratic loss case will have broader impact along these lines. We will add these motivations in a revised version.

• Why is it important in the proof to use square activations? Q1 of bzi8. W3 of NWbd. As we emphasize in the paper, our approach relies on the square activation function. This allows to implement the main idea of our analysis, which is to reduce this problem of Bayes-optimal learning of a square-activated neural network of extensive width to the Bayes-optimal denoising of a matrix of extensive rank. Going beyond quadratic activations, e.g. to ReLU or sign, is a serious challenge: even if one were to consider a monomial activation of degree $k$ (i.e. $\phi(x) = x^k$), our analysis would map the learning problem to the Bayes-optimal denoising of a $k$-tensor of extensive rank. While denoising low-rank tensors has attracted a lot of attention (see for instance [C1-C3]), the extensive-rank case is still a widely open problem. We will add a discussion of this point in a revised conclusion.

• Does knowing the Bayes-Optimal performance inform us on Gradient Descent? Q3 of bzi8. The general answer is no. The gradient descent (GD) algorithm has no information on how the data is generated, so it is expected to be sub-optimal in general. Furthermore, for $\kappa < 1$ the problem is non-convex, so it is a priori reasonable to suspect that GD would require $\alpha>\alpha_{\rm PR}$ to achieve zero error (in the noiseless setting). Interestingly, we observe that, in the noiseless setting, $(i)$ a number $\alpha=\alpha_{\rm PR}$ of samples is sufficient for GD to reach perfect recovery, and $(ii)$ that averaging over different runs of GD gives Bayes-Optimal performance, for any value of $\alpha$. These are two surprising facts that our theory doesn't account for, and whose explanation is left for future work. We also emphasize that these two observations no longer hold in noisy scenarios (see Fig.3 right). However, it is worth pointing out that our Bayes-optimal analysis can inform us on the performance of an Approximate Message Passing algorithm [D1, D2] that will achieve the Bayes-optimal performance also in noisy settings. We will include an explicit description and analysis of this algorithm in the final version of the paper.

• Q4 of bzi8. Q2 of XJ6b. We thank the reviewers for suggesting a more detailed discussion of the origin of eq.(3). We study the Bayes-optimal performance of the network where the input data is corrupted by a Gaussian noise of variance $\Delta$, cf eq.(2). We clarify that the ``most natural'' definition of the MMSE would be (for $(x, y)$ new test samples):

Eq.(3) is the same quantity as $R$, up to a rescaling and an additive constant. As we show in Section 3, eq.(3) is equivalent to the matrix MMSE $\kappa E{\rm tr}[({\bf S}^\star - {\bf S})^2]$, in a problem with noiseless input data but noisy labels, with a variance $2\Delta(2+\Delta)/\kappa$. More generally, one should consider eq.(3) as the rescaling of $R$ that satisfies three properties: $(i)$ it is finite in the high-dimensional limit, $(ii)$ it is equal to $1$ for $\alpha = 0$ (in the absence of data), and $(iii)$ it goes to $0$ in the case of perfect recovery (if the posterior concentrates around $W^\star$). We will add in the revised version a detailed discussion of this point, and a mathematical proof of the equivalence of the MMSE of eq.(3) to the matrix MMSE $\kappa E {\rm tr}[({\bf S}^\star - {\bf S})^2]$.

• W1 and W2 of XJ6b. We thank the reviewer for suggesting a better organization of the main results of the paper. In the revised version, we will move eqs.(14) and (15) on the asymptotic MMSE to the beginning of the main results section. We will also clarify the relevance of the ``free entropy'' studied in Claim 1, recalling its relation to the mutual information, and thus to optimal estimation.

-- [A1] "Implicit Regularization in Matrix Factorization", Gunasekar&al, NeurIPS 2017.

-- [A2] "Towards Resolving the Implicit Bias of Gradient Descent for Matrix Factorization: Greedy Low-Rank Learning", Li&al, ICLR 2021.

-- [A3] "Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparametrized low-rank matrix reconstruction", Stöger&al, NeurIPS 2021.

-- [B1] "Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks", Nichani&al, NeurIPS 2023.

-- [B2] "Learning Hierarchical Polynomials with Three-Layer Neural Networks", Wang&al, ICLR 2024.

-- [C1] "Statistical and computational phase transitions in spiked tensor estimation", Lesieur&al, ISIT 2017.

-- [C2] "The landscape of the spiked tensor model", Ben Arous&al, Comm. Pure Appl. Math. (2019).

-- [C3] "Statistical limits of spiked tensor models", Perry &al, Ann. Inst. H. Poincaré (2020).

-- [D1] "Message-passing algorithms for compressed sensing", Donoho&al, PNAS (2009).

-- [D2] "Generalized approximate message passing for estimation with random linear mixing", Rangan, ISIT 2011.

This paper assumes a neural network teacher model and studies the performance of Bayes-optimal estimator. Authors assume a target function which is a one-hidden layer neural network with large width, quadratic activation, and random weights. They assume quadratically many samples, derive a closed-form expression for the Bayes-optimal test error. They further complement their results with empirical studies and compare the Bayes-optimal performance to the one obtained by gradient descent, supporting their story through numerical studies.

This paper was reviewed by four reviewers and received the following Scores/Confidence: 8/3, 4/2, 6/2, 8/2. I think the paper is studying an interesting topic and the results are relevant to NeurIPS community. The following concerns were brought up by the reviewers:

• poor motivation for broader audience. This should be addressed as it was a main concern of one of the reviewers.
• restricted setting. Although it is unfair to ask authors to generalize this setting, the limitations of the studied setting should be explicitly discussed in the paper.
• presentation can be improved in various places; see reviewer XJ6b's comments.

Authors should carefully go over reviewers' suggestions and address any remaining concerns in their final revision. Based on the reviewers' suggestion, as well as my own assessment of the paper, I recommend including this paper to the NeurIPS 2024 program.