Title: Bayes-optimal learning of an extensive-width neural network from quadratically many samples
Venue: NeurIPS 2024
Accepted: yes
Keywords: Theory of neural networks

Bayes-optimal learning

non-convex optimization

statistical physics

high-dimensional statistics
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 et al., 2023]
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.

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

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 et al., 2023] 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.

1 Introduction

Learning with multi-layer neural networks brought impressive progress and applications in many areas. It is well established that a large enough non-linear neural network can represent a large class of functions [Cybenko, 1989]. Yet the conditions under which the values of the weights can be found efficiently, and from how many samples of the data, remain theoretically elusive. While one may hope that a detailed understanding of these fundamental limitations will eventually allow for a more efficient training, answering such questions for general data and target function remains, however, beyond the reach of current theoretical methods.

In an early attempt to overcome the difficulty of the above generic question, a long line of work originating in Gardner and Derrida [1989], Sompolinsky et al. [1990] proposed to study the optimal

sample-complexity in the so-called teacher-student setting, where the target function corresponds to a 'teacher' neural network. The architecture of this teacher neural network is chosen to be fully connected feed-forward with a given number of layers, their widths, and activations. The values of each of the weights are generated independently, from a Gaussian distribution. This teacher neural network is then used to generate an output label y i ∈ R for each input data sample x i ∈ R d . Given the architecture of the teacher networks (but not the values of the teacher-weights W ∗ ) and the training set of input-output pairs { y i , x i } n i =1 , the smallest achievable test error can then be obtained by averaging the output of a student-neural network (with the same architecture as the teacher) over the values of weights drawn from the posterior distribution. We will refer to the accuracy reached this way as the Bayes-optimal one. It yields the fundamental limitations in learning such tasks, by any possible means, and can therefore serve as a benchmark.

In the so-called high-dimensional limit [Donoho, 2000], when the input training data are d -dimensional Gaussian vectors, in the limit d →∞ , the above research program has been carried out in detail over the last decades for small neural networks having only m = O d (1) hidden units, and learning from n = αd data samples, where α = O d (1) (see, e.g. Györgyi [1990], Opper and Haussler [1991], Seung et al. [1992], Watkin et al. [1993], Schwarze [1993], Barbier et al. [2019], Aubin et al. [2019b]). In the more recent literature, this setting is sometimes referred to as learning single-index and multi-index functions [Damian et al., 2022, Bietti et al., 2023, Collins-Woodfin et al., 2023, Dandi et al., 2023, 2024b, Damian et al., 2024]. While early works in this line originated in statistical physics and used the heuristic replica method [Mézard et al., 1987] to derive the closed-form expressions for quantities of interest in the high-dimensional limit (with d → ∞ , m = O d (1) and n = O d ( d ) ), a mathematical establishment followed using rigorous probabilistic methods [Barbier et al., 2019, Aubin et al., 2019b].

Reaching a closed-form expression for the Bayes-optimal sample complexity for target functions corresponding to multi-layer teacher neural networks is the next open and very challenging task. Among the recent work is Cui et al. [2023], that established (non-rigorously, using the replica method) the Bayes-optimal error for a target function corresponding to a multi-layer neural network of extensive width (i.e. linearly proportional to the dimension) from a number of samples also linear in the dimension. Interestingly, in this limit, the Bayes-optimal error resulted in a quite poor approximation of the function, which can be achieved as well by a simple linear regression on the input-output pairs. No method, be it a multi-layer neural network (or even refinements like a transformer), will be able to achieve better performance. [Cui et al., 2023] further argue, based on numerical evidence, that quadratically many samples in the dimension are necessary in order to be able to learn the target function with non-linear activations 1 to an infinitesimally small test error. This is perhaps intuitive as, with an extensive width, the number of parameters/weights in the teacher network is quadratic in dimension. However, such a regime is challenging for current theoretical tools. Reaching an analytical explicit expression for the Bayes-optimal performance in this regime, for the target function in the form of a neural network of extensive width, is an open, challenging, theoretical problem that has not yet been solved even for a single hidden layer architecture.

Our contributions In this paper, we step up to this challenge and derive a closed-form expression for the Bayes-optimal test error for a target/teacher function corresponding to a one-hidden layer neural network of extensive width, from quadratically many samples, for a particular case where the activation function (after the hidden layer) is quadratic. In particular, our main contributions are:

We provide a closed-form expression for the Bayes-optimal error of learning an extensive-width neural network from quadratically many samples, which is the first type of such result to the best of our knowledge. Such a form is enabled by the high-dimensional limit and corresponding concentration of quantities of interest. It notably follows from our formula that, in the absence of noise in the target function, zero test error is achievable for a sample complexity α = n/d 2 larger than a perfect-recovery threshold α > α PR where

with κ = m/d the ratio between the width m and the dimension d . We further notice that this matches a naive counting of the number of degrees of freedom in the target function.

We introduce the GAMP-RIE algorithm that combines generalized approximate message passing (GAMP) [Donoho et al., 2009, Rangan, 2011, Zdeborová and Krzakala, 2016] with a matrix denoiser that is based on so-called rotationally-invariant estimators (RIE) [Bun et al., 2016], and show that in the large size limit, this algorithm reaches the Bayes-optimal error for all α, κ = Θ(1) .

On the technical level, our result is enabled by combining results from the analysis of single-layer neural networks [Barbier et al., 2019] and extensive-rank matrix denoising [Maillard et al., 2022b]. The derived formula involves the asymptotics of the Harish-Chandra-Itzykson-Zuber integral of random matrix theory [Harish-Chandra, 1957, Itzykson and Zuber, 1980]. Our approach is notably inspired by recent results on the ellipsoid fitting problem [Maillard and Kunisky, 2024, Maillard and Bandeira, 2023]. These tools are of independent interest to the machine learning community, and we anticipate they will have other applications in the theory of learning.

We empirically compare the Bayes-optimal performance to the one obtained by gradient descent. In the noiseless case we observe a rather unusual and surprising scenario, as randomly-initialized gradient descent seems to be sampling the space of interpolants, and leads to twice the Bayes-optimal error. When averaged over initialization the gradient descent reaches an error that is very close to the Bayes-optimal. The rigorous establishment of these properties of gradient descent is left open.

Our experiments are reproducible, and accessible freely in a public repository [Maillard et al., 2024].

Further related works The problem studied in this work is known as phase retrieval in the case of a single hidden unit ( m = 1 ). Many works considered this problem in the high-dimensional limit d →∞ , in the regime of n = O ( d log d ) samples; see e.g. Candes et al. [2013], Chen et al. [2019], Demanet and Hand [2014]. A subsequent line of work established that the problem can be solved with only O ( d ) samples [Candès and Li, 2014, Chen and Candes, 2015, Cai et al., 2022].

Eventually, for Gaussian i.i.d. input data and i.i.d. teacher weights, the optimal sample complexity for learning phase retrieval in the high-dimensional limit has been established down to the constant in α = n/d . Authors of Mondelli and Montanari [2019] derived the weak recovery threshold for the noiseless case to be α WR = 1 / 2 for phase retrieval, and optimal spectral methods were shown to match this threshold in Luo et al. [2019], Maillard et al. [2022a]. The information-theoretically optimal accuracy and the one achieved by an approximate message passing algorithm were then derived in Barbier et al. [2019] for a general i.i.d. prior for the teacher weights. In the absence of noise, these results imply sample complexities α IT = 1 and α AMP ≈ 1 . 13 needed to achieve perfect learning for a Gaussian prior. Authors of Song et al. [2021] proposed a non-robust polynomial algorithm capable of solving noiseless phase retrieval for α ≥ α IT . Algorithms based on gradient descent were argued not to achieve the optimal sample complexity in Sarao Mannelli et al. [2020a], Mignacco et al. [2021]. Maillard et al. [2020] derived the MMSE for more general input data distributions, including the complex-valued case. Phase retrieval with generative priors was studied in Hand et al. [2018], Aubin et al. [2020]. We refer to a recent review [Dong et al., 2023] for an overview of the relations between these theoretical studies and practical applications of phase retrieval in imaging.

The case with different numbers of hidden units m ⋆ in the teacher and m in the student model, was also discussed in the literature. For m ∗ = O d (1) , the problem is a special case of a multi-index model that has been recently actively considered, e.g. in Aubin et al. [2019b], Bietti et al. [2023], Damian et al. [2022, 2024], Collins-Woodfin et al. [2023], Dandi et al. [2023, 2024b]. This line of work has not focused on the quadratic activations, as it does not bring particular simplification in this case.

The geometry of loss landscapes of one hidden-layer networks with quadratic activations was studied, and the absence of spurious local minima was established for m ≥ d (when the read-out layer is fixed as in our setting) in Du and Lee [2018]. Similar results were established in Soltanolkotabi et al. [2018], Venturi et al. [2019] for a slightly more general setting where the readout layer is learned.

Establishing results about sample complexity required for generalization in cases where m (or both m and m ∗ ) are Θ( d ) is technically challenging, and so far, only a handful of works made progress in that direction. In particular, Gamarnik et al. [2019] considered m ∗ ≥ d and m ≥ d , and have shown that a sample complexity n ≥ d ( d + 1) / 2 is sufficient for perfect recovery of the target function. Sarao Mannelli et al. [2020b] considered the overparametrized case with m ∗ = O d (1) and m > d , and showed that gradient descent reaches exact recovery for a sample complexity n > d ( m ∗ +1) -( m ∗ +1) m ∗ / 2 , again considering the high-dimensional limit. Gradient descent of the population risk has been studied for general values of ( m ∗ , m ) in Martin et al. [2024], along with a discussion of the role of overparametrization.

2 Setting

As discussed above, we are studying the Bayes-optimal accuracy in the teacher-student setting. More concretely, we consider a dataset of n samples D = { y i , x i } n i =1 where the input data is normal Gaussian of dimension d : ( x i ) n i =1 i . i . d . ∼ N (0 , I d ) . We then draw i.i.d. d -dimensional teacher-weight vectors ( w ∗ k ) m k =1 i . i . d . ∼ N (0 , I d ) , and noise ( z i ) n i =1 i . i . d . ∼ N (0 , I m ) . Finally, the output labels ( y i ) n i =1 are obtained by a one-hidden layer teacher network with m hidden units and quadratic activation:

Crucially, we assume we know the form of the (stochastic) target function f W ∗ ( · ) (i.e. the value of m , ∆ , and the form of eq. (2), including the fact that the activation function is quadratic) but we do not know the realization of neither the teacher weights W ∗ = ( w ⋆ 1 , · · · , w ⋆ m ) nor the noise z i .

Universality over the noise and weights distribution While we consider Gaussian distributions for the sake of our theoretical analysis, we expect our results to hold under more general i.i.d. models on both the noise and the teacher weights, under mild conditions of existence of moments. This is related to a recent conjecture of Semerjian [2024], see Sections 3 and 4.

Bayes-optimal test error Since we know the law of the dataset D , we can study the Bayes-optimal (BO) estimator , which minimizes the test error over all possible estimators. To do this, we use Bayes' theorem to obtain the posterior distribution P ( W |D ) of the weights W given the dataset:

where P prior ( W ) is a prior distribution on the teacher weights W ∗ , and the likelihood P ( y | W , X ) can be seen as a probabilistic channel that generates the labels given the input data ( x i ) n i =1 and the teacher weights W ∗ , and Z ( D ) is a normalization constant. The Bayes-optimal (BO) estimator of the labels for a test sample x test not seen in the training set D then involves the average over the posterior distribution as follows (where E z denotes the expectation over z 1 , · · · , z k )

We will evaluate the BO estimator in terms of its average generalization error, i.e. the mean squared error (MSE) achieved on a new sample. We define it in the following way:

We denote it MMSE d , standing for minimum-MSE, as it is the minimum MSE achievable given the setting of the model, and we call MMSE := lim d →∞ MMSE d .

Conventions for the MMSE Notice the peculiar multiplicative factor ( m/ 2) and the additive term -∆(2 + ∆) in eq. (4). As we detail in Appendix F.1, these factors ensure that MMSE → 1 for α → 0 (i.e. in the absence of data), and MMSE → 0 if the posterior concentrates around the true W ⋆ (i.e. if ˆ y BO D ( x ) = E z [ f W ⋆ ( x )] ). Moreover, as we also detail in Appendix F.1, eq. (4) matches the MMSE of a matrix estimation task to which we will reduce the original problem, see Section 3.

As motivated above, our goal is to analyze the MMSE in the high-dimensional limit, with an extensive-width architecture and quadratically many data samples:

In all that follows, we consider the limit of eq. (5), so that n, d, m all go to infinity together when we write e.g. lim d →∞ . As we will see, in this limit, the value of the MMSE for a given realization of the randomness concentrates on the averaged value defined in eq. (4).

Empirical risk minimization estimator A more standard way of learning the target function (2) is to minimize the empirical loss L corresponding to a 'student' neural network

0

.

0

0

.

2

0

.

4

0

.

6

Sample complexity

α

=

n/d

2

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

MMSE

κ

= 0

.

05

κ

= 0

.

1

κ

= 0

.

2

κ

= 0

.

3

κ

= 0

.

5

κ

= 1

.

0

κ

= 5

.

0

0

.

0

0

.

2

0

.

4

0

.

6

Sample complexity

α

=

n/d

2

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

Intensive width

κ

=

m/d

Note that this does not account for the noise, but activations in neural networks are commonly considered deterministic, so we consider this the most natural choice.

Minimization of the loss over the weights W = ( w k ) m k =1 is commonly done using gradient descent (GD): one initializes the weights as W (0) ∼ P prior and then updates them to minimize the empirical loss, for an appropriately choice of learning rate, until convergence. Denoting the weights at convergence as ˆ W ( W (0) , D ) the estimator for test labels reads ˆ y GD W (0) , D ( x test ) := ˜ f ˆ W ( W (0) , D ) ( x test ) . As we will see, it will be interesting to consider also an estimator ˆ y AGD obtained by averaging the GD estimator on the labels over the initializations W (0) of the weights.

3 Main results

Notations We use tr( · ) := (1 /d )Tr[ · ] for the normalized trace. We denote GOE( d ) the distribution of symmetric matrices ξ ∈ R d × d such that ξ ij i . i . d . ∼ N (0 , (1 + δ ij ) /d ) , for i ≤ j . For m = κd with κ > 0 , we denote W m,d the Wishart distribution, and µ MP ,κ the Marchenko-Pastur distribution with ratio κ . More details on classical definitions and notational conventions are given in Appendix A.

We start by stating the main result of our analysis, applied to the problem of eq. (2).

Result 1. The MMSE of eq. (4) is given in the high-dimensional limit of eq. (5) by:

where ˜ ∆ := 2∆(2 + ∆) /κ , and where ˆ q is a solution of the following equation:

Here, µ t := µ MP ,κ ⊞ σ s . c ., √ t (for t ≥ 0 ) is the free convolution of the Marchenko-Pastur law and a scaled semicircular density, see Appendix A for its precise definition.

Eq. (8) can be efficiently solved using a numerical scheme, which is detailed in Appendix H.1. We present the results in Fig. 1. In what follows, we detail our approach towards deriving Result 1, which is a consequence of our main theoretical result stated in Claim 2.

Reduction to a matrix estimation problem We first notice that by expanding the square in eq. (2), we can effectively reduce our learning task to an estimation problem in terms of S ⋆ := (1 /m ) ∑ m k =1 w ⋆ k ( w ⋆ k ) ⊺ . We give an analytical argument backing this observation in Appendix F.5.

Its conclusion is that, at leading order, the distribution of y = f W ⋆ ( x ) can be reduced to the following form, with ˜ y := √ d ( y -1 -∆) :

with ξ ∼ N (0 , 1) , ˜ ∆ := 2∆(2 + ∆) /κ , and where we defined Z := ( xx ⊺ -I d ) / √ d .

Generalization error and MMSE on S This equivalent problem gives us a way to interpret the convention we chose for eq. (4). Indeed, if we denote ˆ S opt = E [ S | ˜ y , Z ] the Bayes-optimal estimator related to the problem of eq. (9), then MMSE = κ E tr[( S ⋆ -ˆ S opt ) 2 ] , as proven in Lemma F.1.

The limit of the MMSE We now describe the general form of estimation problems covered by our theoretical analysis, which encompasses the one described in eq. (9) (and thus the original eq. (2)). The goal is to recover the symmetric matrix S ⋆ ∈ R d × d , which was generated from the Wishart distribution W m,d , from observations ( y i ) n i =1 , generated as with Z i := ( x i x ⊺ i -I d ) / √ d . The 'channel' P out accounts for possible non-linearities and noise, encompassing the case of additive Gaussian noise in eq. (9). We define the partition function as:

Notice that the averaged logarithm of Z is (up to an additive constant) equal to the mutual information between the observations and the hidden variables: I ( W ⋆ ; { y i }|{ x i } ) = E log Z + n E log P out ( y 1 | Tr[ Z 1 S ⋆ ]) . This links Z to the optimal estimation of W , an important idea behind our study. We are now ready to state our main theoretical result. It gives a sharp characterization of the Bayes-optimal error in any estimation problem of the type of eq. (10). By the reduction described above, it can be directly applied to the original model of eq. (2), and will imply Result 1.

Claim 2. Assume that m = κd with κ > 0 , and n = αd 2 with α > 0 . Let Q 0 := 1 + κ -1 . Then:

The limit of the averaged log-partition function (sometimes called the free entropy ) is given by

where

Here, Σ( µ ) := E X,Y ∼ µ log | X -Y | , and, for t ≥ 0 , µ t := µ MP ,κ ⊞ σ s . c ., √ t is the free convolution of the Marchenko-Pastur distribution and a (scaled) semicircle law, see Appendix A for its definition.

,

For any α > 0 , except possibly in a countable set, the supremum in eq. (12) is reached in a unique q ⋆ ∈ [1 , Q 0 ] . Moreover, the asymptotic minimum mean-squared error on the estimation of S ⋆ achieved by the Bayes-optimal estimator ˆ S BO := E [ S |{ y i , x i } ] , is equal to Q 0 -q ⋆ :

It is related to the MMSE of eq. (4) by MMSE = κ ( Q 0 -q ⋆ ) .

Specifying Claim 2 to the problem of eq. (9), we derive (details are given in Appendix F.7) Result 1, more precisely eqs. (7) and (8).

Polynomial-time optimal estimation with the GAMP-RIE algorithm In Appendix B, we motivate the definition of an algorithm (that we call GAMP-RIE ) to solve the problem of eq. (10). We further argue (based on a combination of theoretical results and numerical observations) that this algorithm is able to reach, in all regions of parameters we investigated, the optimal error described by Claim 2.

The condition q ≥ 1 -Notice that q ⋆ = lim d →∞ E [tr( S ⋆ ˆ S BO )] according to Claim 2. As the MMSE decreases with α , it is clear that q ⋆ ≥ q ⋆ ( α = 0) . When α = 0 , we have ˆ S BO = E [ S ⋆ ] = I d , and thus q ⋆ ( α = 0) = 1 . We check in Appendix F.8 that the value q ⋆ ( α = 0) = 1 is recovered by eq. (12).

4 Derivation of the main results

We derive our main result (Claim 2) in two ways. First, we show how one can show Claim 2 using the replica method , a heuristic but exact method (hence the word 'claim') which originated in statistical physics [Mézard et al., 1987], and has been used extensively in theoretical physics, as well as in a growing body of work in high-dimensional statistics, theoretical computer science, and theoretical machine learning [Mezard and Montanari, 2009, Zdeborová and Krzakala, 2016, Gabrié, 2020, Charbonneau et al., 2023]. The derivation, that has an interest on its own, is performed in detail in Appendix D and leverages recent progress on the problems of ellipsoid fitting [Maillard and Kunisky, 2024, Maillard and Bandeira, 2023] and extensive-rank matrix denoising [Maillard et al., 2022b, Pourkamali et al., 2024, Semerjian, 2024].

Despite the replica method being conjectured to yield exact results in a large class of high-dimensional models, a rigorous treatment of it remains elusive. It is important, we feel, to present as well a more mathematically sound derivation of our claims, and we thus give an alternative derivation of the Claim 2 using probabilistic techniques amenable to rigorous treatment. In what follows, we present a three-step sketch of a mathematical proof of Claim 2 that combines recent progress performed on the study of a problem known as the ellipsoid fitting conjecture [Maillard and Kunisky, 2024, Maillard and Bandeira, 2023] with the analysis of the fundamental limits of so-called generalized linear models [Barbier et al., 2019], as well as matrix denoising problems [Bun et al., 2016, Maillard et al., 2022b, Pourkamali et al., 2024, Semerjian, 2024]. While a complete mathematical treatment requires more work, we detail the main challenges arising in each of these steps, outlining a fully rigorous establishment of Claim 2.

We denote the free entropy Φ d := (1 /d 2 ) E log Z ( { y i , x i } ) , cf. eq. (11). We detail three precise results (two conjectures and a theorem), motivated by recent mathematical works, whose combination would rigorously establish the results of Claim 2. Recall that we consider the high-dimensional limit of eq. (5).

Step 1: Universality with a 'Gaussian equivalent' problem The first step of our approach is inspired by recent literature on the ellipsoid fitting problem [Maillard and Kunisky, 2024, Maillard and Bandeira, 2023]. It amounts to notice that, if Z i := ( x i x T i -I d ) / √ d , by the central limit theorem, for any symmetric matrix S , Tr[ Z i S ] is (under mild boundedness conditions on the spectrum of S ) approximately distributed as N (0 , 2tr[ S 2 ]) as d → ∞ . A large body of recent literature has established that the free entropy is universal for all data distributions sharing the same asymptotic distribution of their 'one-dimensional projections', see e.g. Hu and Lu [2022], Montanari and Saeed [2022], Dandi et al. [2024a], Maillard and Bandeira [2023]. This motivates the conjecture that the free entropy should remain identical (to leading order) if one replaces the matrices Z i with G i ∼ GOE( d ) .

Conjecture 4.1 (Universality) . We define

where y ′ i ∼ P out ( ·| Tr[ G i S ⋆ ]) , with S ⋆ ∼ W m,d and G i i . i . d . ∼ GOE( d ) . Then

Conjecture 4.1 can be seen as an extension of Corollary 4.10 of Maillard and Bandeira [2023], in the context of a teacher-student model. In particular, we expect it to hold under mild regularity conditions on the channel density P out (which are satisfied by the Gaussian additive noise we consider).

Step 2: A matrix generalized linear model with a Wishart prior By the first step above, we can focus on Φ ( G ) d , and the corresponding estimation problem. A key observation is that one can view this problem as an instance of a generalized linear model on S ⋆ , with a Gaussian data matrix whose i -th row is the flattening of the matrix G i . The limiting free entropy of such models has been worked out in Barbier et al. [2019], when the 'ground-truth vector' (here S ⋆ ) has i.i.d. elements. However, here the prior is far from being i.i.d. since S ⋆ ∼ W m,d . The results of Barbier et al. [2019] generalize naturally to other priors, but such extensions have only been rigorously analyzed in specific settings, e.g. for generative priors rather than i.i.d. [Aubin et al., 2019a, 2020]. In our setting, the structure

of the Wishart prior raises several technical difficulties preventing to directly transpose the proof approaches of Barbier et al. [2019], so we state the following result as a conjecture.

Conjecture 4.2 (The free entropy of a matrix generalized linear model) . We have

where

is the asymptotic free entropy of the matrix denoising problem Y = √ ˆ q S ⋆ + ξ , with ξ ∼ GOE( d ) , and S ⋆ ∼ W m,d , and we assume that the d →∞ limit in eq. (16) is well-defined.

Step 3: Extensive-rank matrix denoising As a last step, we study the function Ψ(ˆ q ) defined in eq. (16). The optimal estimators and limiting free entropy in matrix denoising have been worked out in Bun et al. [2016], Maillard et al. [2022b], and formally proven (under some assumptions) in Pourkamali et al. [2024], Semerjian [2024].

Theorem 4.1 (Free entropy of matrix denoising) . For any ˆ q ≥ 0 , the limit in eq. (16) is well-defined, and moreover (recall the definition of Σ( µ ) and µ t in Claim 2)

We provide a very short and assumption-free proof of Theorem 4.1 in Appendix F.2, which combines a relation between Ψ(ˆ q ) and HCIZ integrals of random matrix theory, proven in Pourkamali et al. [2024] (without any assumptions), and fundamental results on the large deviations of the Dyson Brownian motion [Guionnet and Zeitouni, 2002]. As a final remark, we notice that a recent conjecture 2 of Semerjian [2024] states that the free entropy of matrix denoising of S ⋆ = (1 /m ) ∑ m k =1 w ⋆ k ( w ⋆ k ) ⊺ remains the same if one considers any i.i.d. prior for w ⋆ k , under the matching of its first two moments with the Gaussian and the existence of all other moments. While the validity of this conjecture is subject to debate (see Section VII of Semerjian [2024], and the findings of Camilli and Mézard [2023, 2024]), in the present model it would imply universality of the generalization error given by Claim 2 for any such teacher weight distribution.

The second part of Claim 2 We briefly discuss the second part of Claim 2, concerning the large d limit of E tr[( S ⋆ -ˆ S BO ) 2 ] . The fact that the maximizer of eq. (12) is unique for almost all values of α can be seen by simple convexity arguments, see Appendix F.6. The relationship of q ⋆ with the asymptotic MMSE on the estimation of S ⋆ is a classical consequence of the I-MMSE theorem in generalized linear models of which eq. (9) is an instance, see e.g. Barbier et al. [2019] and Section D.5 of Maillard et al. [2020].

5 Discussion of the main results

Analysis of the Bayes-optimal estimator We start by discussing the noiseless case ( ∆ = 0 ), which is described by the phase diagram in Fig. 1. Since there is no noise in the target function, we expect a sharp transition to zero MMSE at a critical sample complexity α PR . We analytically show in Appendix F.3 from eq. (8) that α PR is given by the expression of eq. (1), and discuss how it is related to a naive counting argument of the 'degrees of freedom' of the target function. This transition was known for κ ≥ 1 where the problem is convex, where Gamarnik et al. [2019] shows that there is perfect recovery as soon as α > 1 / 2 . For all values of κ we see the MMSE is a smooth curve going continuously from 1 at α = 0 to 0 at α PR . We derived the slope of the curve at α PR to be (see Appendix F.4)

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

Sample complexity

α

=

n/d

2

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

MSE

Theory

d

= 200)

AGD (

d

= 200)

d

= 100)

AGD (

d

= 100)

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

Sample complexity

α

=

n/d

2

Theory (∆ = 0)

Theory (∆ = 0

.

25 2 )

GD (∆ = 0)

Averaged GD (∆ = 0)

GD (∆ = 0

.

25 2 )

Averaged GD (∆ = 0

.

25 2 )

It is interesting to observe that the convexity of the curve changes. While we are observing concave dependence on α for small κ it becomes convex when κ increases and α is close to α PR . We also note that the smooth limit MMSE → 1 as α → 0 supports the result of Cui et al. [2023] about a quadratic number of samples being needed to learn better than linear regression.

Wealso evaluated the MMSE in the presence of noise, where we observed it to decrease smoothly as α increases with no particular phase transition. We show an example of the theoretical prediction for the MMSE in this case in Fig. 2 right. As expected, in the presence of noise, it decreases monotonically and smoothly, and goes to zero as α →∞ .

We considered analytically the limits κ → 0 and κ →∞ , i.e. the limits of small and large (but still extensive in d ) hidden layer. The analysis of these limits are detailed in Appendix E.

Further, in Appendix B.3 we compare the asymptotic theoretical result for the Bayes-optimal error with the performance of the GAMP-RIE algorithm on finite-size instances. In all the cases we evaluated, we observed that GAMP-RIE reached the Bayes-optimal error characterized by Claim 2.

Finally, while we assumed in eq. (2) that the second layer weights are fixed and equal to 1 , in Appendix G we generalize all our main results, theoretical and algorithmic, to learnable second layer weights.

Comparison to the ERM estimator obtained by gradient descent The results discussed so far concern the Bayes-optimal MMSE, which requires evaluating the marginals of the posterior distribution. We now investigate numerically the performance of empirical risk minimization via gradient descent, which is the standard method of machine learning. It would be typical to expect a gradient based approach to be suboptimal, as the problem is non-convex for κ < 1 . In Fig. 2, we compare (a) the MSE κ tr[( S ⋆ -ˆ S GD ) 2 ] reached by gradient descent (GD) minimizing the loss (6) from random initialization, (b) the MSE reached by GD averaged over initializations, and (c) the MMSE derived from the theory.

In the noiseless case, ∆ = 0 , we very remarkably observe that the MSE reached by gradient descent is very close to exactly twice larger than the asymptotic MMSE. Such a relation is known in highdimensional generalized linear regression to hold between the Gibbs estimator, where test error is evaluated for weights that are sampled uniformly from the posterior, and the Bayes-optimal estimator that averages over the weights sampled from the posterior [Engel, 2001, Barbier et al., 2019]. In

general, there is no reason why the randomly initialized gradient descent should be able to sample the posterior measure. We nevertheless evaluate the average over the initialization of gradient descent and observe that, indeed, the MSE reached this way is consistent with the MMSE. This leads us to conjecture that in the noiseless one-hidden layer neural network with quadratic activation and a target function matching this architecture, randomly-initialized gradient descent samples the posterior despite the problem being non-convex, and hence its average achieves the MMSE.

Let us offer a heuristic argument for this perhaps intriguing phenomenon. It starts with the equivalent of the representer theorem: one can write S in the span of { x i x T i } n i =1 , plus a matrix in the orthogonal space, that is S = ∑ n i =1 β i x i x T i + Z . This means that gradient descent reaches one solution of the minimization with one additional spurious component. The Bayes optimal procedure would be to set this spurious reminder to zero since the data are not informative in this direction. It is reasonable (although non-trivial) to assume that this is what is achieved by averaging over initialization.

When comparing the MMSE to the performance of GD in the noisy setting, we observe a gap between the MMSE and the performance of gradient descent, even averaged over initialization or regularized (as shown in Appendix H.3, Figure 5 left). In particular, for the noisy case, we see that for small sample complexity, the averaged GD is close to matching the MMSE, but as the number of available samples increases, the error of the averaged and non-averaged versions of GD coincide. This is a sign of the trivialization of the landscape, in the sense that GD converges to the same function independently of the initialization: it can be quantified using the variances of the function reached by GD. This is investigated further in Appendix H.3, together with the effect of ℓ 2 regularization. We can characterize empirically another phase transition: for a sample complexity larger than α T (∆) , GD converges to the same function independently of the initialization. In the noiseless ∆ = 0 case, this is simply the perfect recovery transition, and α PR = α T (∆ = 0) , while increasing the noise intensity makes the threshold lower until it reaches a plateau, which for κ = 0 . 5 is at α T (∆ →∞ ) ≈ 0 . 2 . We display this numerical finding in Figure 5 (right) in Appendix H.3. A tight analytical study of the landscape-trivialization threshold α T (∆) as a function of the noise variance ∆ is left for future work.

6 Conclusion and limitations

In this work, we provide an explicit formula for the generalization MMSE when learning a target function in the form of a one-hidden layer neural network with quadratic activation in the limit of large dimensions, extensive width and a quadratic number of samples. The techniques deployed to obtain this result are novel and, we believe, of independent interest. There are many natural extensions of the present works. While we presented, additionally to the replica derivation, a mathematically sound derivation, a fully rigorous treatment, a technical and lengthy task, is left for an extended version of this work. We analyzed the Bayes-optimal MMSE, presented the GAMP-RIE algorithm that is able to reach it in polynomial time, and compared it to the performance of gradient descent numerically. We leave for future work the theoretical analysis of the properties of gradient descent that we discovered numerically. Of particular interest is the role played by the implicit nuclear norm regularization when starting from small initialization, as discussed for the matrix sensing problem e.g. in Gunasekar et al. [2017], Li et al. [2020], Stöger and Soltanolkotabi [2021]. Finally, we also presented the natural extension of our results and techniques to the case of a learnable second layer.

The main limitations of our setting are its restriction to Gaussian input data, random i.i.d. weights of the target/teacher neural network, quadratic activation, and a single hidden layer. Going beyond any of these limitations would be a compelling direction of research, in particular for more generic activation such as the ReLU or sigmoid function (we sketch this extension in Appendix C) and multiple layers, and we hope our work will spark interest in these directions.

References

Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learning Theory , pages 2552-2623. PMLR, 2023.

Greg W Anderson, Alice Guionnet, and Ofer Zeitouni. An introduction to random matrices . Cambridge university press, 2010.

Benjamin Aubin, Bruno Loureiro, Antoine Maillard, Florent Krzakala, and Lenka Zdeborová. The spiked matrix model with generative priors. Advances in Neural Information Processing Systems , 32, 2019a.

Benjamin Aubin, Antoine Maillard, Jean Barbier, Florent Krzakala, Nicolas Macris, and Lenka Zdeborová. The committee machine: computational to statistical gaps in learning a two-layers neural network. Journal of Statistical Mechanics: Theory and Experiment , 2019(12):124023, jan 2019b.

Benjamin Aubin, Bruno Loureiro, Antoine Baker, Florent Krzakala, and Lenka Zdeborová. Exact asymptotics for phase retrieval and compressed sensing with random generative priors. In Mathematical and Scientific Machine Learning , pages 55-73. PMLR, 2020.

Jean Barbier, Mohamad Dia, Nicolas Macris, and Florent Krzakala. The mutual information in random linear estimation. In 2016 54th Annual Allerton Conference on Communication, Control, and Computing (Allerton) , pages 625-632. IEEE, 2016.

Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences , 116(12):5451-5460, 2019.

Gérard Ben Arous, Song Mei, Andrea Montanari, and Mihai Nica. The landscape of the spiked tensor model. Communications on Pure and Applied Mathematics , 72(11):2282-2330, 2019.

Gérard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learning Research , 22 (106):1-51, 2021.

Florent Benaych-Georges and Raj Rao Nadakuditi. The eigenvalues and eigenvectors of finite, low rank perturbations of large random matrices. Advances in Mathematics , 227(1):494-521, 2011.

Raphael Berthier, Andrea Montanari, and Phan-Minh Nguyen. State evolution for approximate message passing with non-separable functions. Information and Inference: A Journal of the IMA , 9(1):33-79, 2020.

Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivien. On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310.19793 , 2023.

Joël Bun, Romain Allez, Jean-Philippe Bouchaud, and Marc Potters. Rotational invariant estimator for general noisy matrices. IEEE Transactions on Information Theory , 62(12):7475-7490, 2016.

Jian-Feng Cai, Meng Huang, Dong Li, and Yang Wang. Solving phase retrieval with random initial guess is nearly as good as by spectral initialization. Applied and Computational Harmonic Analysis , 58:60-84, 2022.

Francesco Camilli and Marc Mézard. Matrix factorization with neural networks. Physical Review E , 107(6):064308, 2023.

Francesco Camilli and Marc Mézard. The decimation scheme for symmetric matrix factorization. Journal of Physics A: Mathematical and Theoretical , 57(8):085002, 2024.

Emmanuel J Candès and Xiaodong Li. Solving quadratic equations via phaselift when there are about as many equations as unknowns. Foundations of Computational Mathematics , 14:1017-1026, 2014.

Emmanuel J Candes, Thomas Strohmer, and Vladislav Voroninski. Phaselift: Exact and stable signal recovery from magnitude measurements via convex programming. Communications on Pure and Applied Mathematics , 66(8):1241-1274, 2013.

Patrick Charbonneau, Enzo Marinari, Giorgio Parisi, Federico Ricci-tersenghi, Gabriele Sicuro, Francesco Zamponi, and Marc Mezard. Spin Glass Theory and Far Beyond: Replica Symmetry Breaking after 40 Years . World Scientific, 2023.

Yuxin Chen and Emmanuel Candes. Solving random quadratic systems of equations is nearly as easy as solving linear systems. Advances in Neural Information Processing Systems , 28, 2015.

Yuxin Chen, Yuejie Chi, Jianqing Fan, and Cong Ma. Gradient descent with random initialization: Fast global convergence for nonconvex phase retrieval. Mathematical Programming , 176:5-37, 2019.

Elizabeth Collins-Woodfin, Courtney Paquette, Elliot Paquette, and Inbar Seroussi. Hitting the high-dimensional notes: An ode for sgd learning dynamics on glms and multi-index models. arXiv preprint arXiv:2308.08977 , 2023.

Hugo Cui, Florent Krzakala, and Lenka Zdeborova. Bayes-optimal learning of deep random networks of extensive-width. In Proceedings of the 40th International Conference on Machine Learning . PMLR, 2023. URL https://proceedings.mlr.press/v202/cui23b.html .

George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems , 2(4):303-314, 1989.

Alex Damian, Loucas Pillaud-Vivien, Jason D Lee, and Joan Bruna. The computational complexity of learning gaussian single-index models. arXiv preprint arXiv:2403.05529 , 2024.

Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory , pages 5413-5452. PMLR, 2022.

Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning , 2023.

Yatin Dandi, Ludovic Stephan, Florent Krzakala, Bruno Loureiro, and Lenka Zdeborová. Universality laws for gaussian mixtures in generalized linear models. Advances in Neural Information Processing Systems , 36, 2024a.

Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborova, and Florent Krzakala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. 2024b.

Laurent Demanet and Paul Hand. Stable optimizationless recovery from phaseless linear measurements. Journal of Fourier Analysis and Applications , 20:199-221, 2014.

Jonathan Dong, Lorenzo Valzania, Antoine Maillard, Thanh-an Pham, Sylvain Gigan, and Michael Unser. Phase retrieval: From computational imaging to machine learning: A tutorial. IEEE Signal Processing Magazine , 40(1):45-57, 2023.

David L Donoho. High-dimensional data analysis: The curses and blessings of dimensionality. AMS math challenges lecture , 1(2000):32, 2000.

David L Donoho, Arian Maleki, and Andrea Montanari. Message-passing algorithms for compressed sensing. Proceedings of the National Academy of Sciences , 106(45):18914-18919, 2009.

Simon Du and Jason Lee. On the power of over-parametrization in neural networks with quadratic activation. In International conference on machine learning , pages 1329-1338. PMLR, 2018.

Andreas Engel. Statistical mechanics of learning . Cambridge University Press, 2001.

Marylou Gabrié. Mean-field inference methods for neural networks. Journal of Physics A: Mathematical and Theoretical , 53(22):223002, 2020.

David Gamarnik, Eren C Kızılda˘ g, and Ilias Zadik. Stationary points of shallow neural networks with quadratic activation function. arXiv preprint arXiv:1912.01599 , 2019.

David Gamarnik, Cristopher Moore, and Lenka Zdeborová. Disordered systems insights on computational hardness. Journal of Statistical Mechanics: Theory and Experiment , 2022(11):114015, 2022.

Elizabeth Gardner and Bernard Derrida. Three unfinished works on the optimal storage capacity of networks. Journal of Physics A: Mathematical and General , 22(12):1983, 1989.

Cédric Gerbelot and Raphaël Berthier. Graph-based approximate message passing iterations. Information and Inference: A Journal of the IMA , 12(4):2562-2628, 2023.

Alice Guionnet and Ofer Zeitouni. Large deviations asymptotics for spherical integrals. Journal of functional analysis , 188(2):461-515, 2002.

Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in neural information processing systems , 30, 2017.

Dongning Guo, Shlomo Shamai, and Sergio Verdú. Mutual information and minimum mean-square error in gaussian channels. IEEE transactions on information theory , 51(4):1261-1282, 2005.

Géza Györgyi. First-order transition to perfect generalization in a neural network with binary synapses. Physical Review A , 41(12):7097, 1990.

Paul Hand, Oscar Leong, and Vlad Voroninski. Phase retrieval under a generative prior. Advances in Neural Information Processing Systems , 31, 2018.

Harish-Chandra. Differential operators on a semisimple lie algebra. American Journal of Mathematics , pages 87-120, 1957.

Uwe Helmke and John B Moore. Optimization and dynamical systems . Springer Science & Business Media, 2012.

Hong Hu and Yue M Lu. Universality laws for high-dimensional learning with random features. IEEE Transactions on Information Theory , 69(3):1932-1964, 2022.

Claude Itzykson and J-B Zuber. The planar approximation. ii. Journal of Mathematical Physics , 21 (3):411-421, 1980.

Adel Javanmard and Andrea Montanari. State evolution for general approximate message passing algorithms, with applications to spatial coupling. Information and Inference: A Journal of the IMA , 2(2):115-144, 2013.

Dmitriy Kunisky, Alexander S Wein, and Afonso S Bandeira. Notes on computational hardness of hypothesis testing: Predictions using the low-degree likelihood ratio. In ISAAC Congress (International Society for Analysis, its Applications and Computation) , pages 1-50. Springer, 2019.

Lucien Le Cam. Locally asymptotically normal families of distributions. certain approximations to families of distributions and their use in the theory of estimation and testing hypotheses. Univ. California Publ. Statist. , 3:37, 1960.

Ji Oon Lee and Kevin Schnelli. Tracy-widom distribution for the largest eigenvalue of real sample covariance matrices with general population. The Annals of Applied Probability , pages 3786-3839, 2016.

Thibault Lesieur, Léo Miolane, Marc Lelarge, Florent Krzakala, and Lenka Zdeborová. Statistical and computational phase transitions in spiked tensor estimation. In 2017 IEEE International Symposium on Information Theory (ISIT) , pages 511-515. IEEE, 2017.

Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations , 2020.

Wangyu Luo, Wael Alghamdi, and Yue M Lu. Optimal spectral initialization for signal recovery with applications to phase retrieval. IEEE Transactions on Signal Processing , 67(9):2347-2356, 2019.

Antoine Maillard and Afonso S Bandeira. Exact threshold for approximate ellipsoid fitting of random points. arXiv preprint arXiv:2310.05787 , 2023.

Antoine Maillard and Dmitriy Kunisky. Fitting an ellipsoid to random points: predictions using the replica method. IEEE Transactions on Information Theory , 2024.

Antoine Maillard, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase retrieval in high dimensions: Statistical and computational phase transitions. Advances in Neural Information Processing Systems , 33:11071-11082, 2020.

Antoine Maillard, Florent Krzakala, Yue M Lu, and Lenka Zdeborová. Construction of optimal spectral methods in phase retrieval. In Mathematical and Scientific Machine Learning , pages 693-720. PMLR, 2022a.

Antoine Maillard, Florent Krzakala, Marc Mézard, and Lenka Zdeborová. Perturbative construction of mean-field equations in extensive-rank matrix factorization and denoising. Journal of Statistical Mechanics: Theory and Experiment , 2022(8):083301, 2022b.

Antoine Maillard, Afonso S Bandeira, David Belius, Ivan Dokmani´ c, and Shuta Nakajima. Injectivity of relu networks: perspectives from statistical physics. arXiv preprint arXiv:2302.14112 , 2023.

Antoine Maillard, Emanuele Troiani, Simon Martin, Florent Krzakala, and Zdeborová Lenka. Numerical code used for experimental results. https://github.com/SPOC-group/ ExtensiveWidthQuadraticSamples , 2024.

Vladimir Alexandrovich Marchenko and Leonid Andreevich Pastur. Distribution of eigenvalues for some sets of random matrices. Matematicheskii Sbornik , 114(4):507-536, 1967.

Simon Martin, Francis Bach, and Giulio Biroli. On the impact of overparameterization on the training of a shallow neural network in high dimensions. In International Conference on Artificial Intelligence and Statistics , pages 3655-3663. PMLR, 2024.

Marc Mezard and Andrea Montanari. Information, physics, and computation . Oxford University Press, 2009.

Marc Mézard, Giorgio Parisi, and Miguel Angel Virasoro. Spin glass theory and beyond: An Introduction to the Replica Method and Its Applications , volume 9. World Scientific Publishing Company, 1987.

Francesca Mignacco, Pierfrancesco Urbani, and Lenka Zdeborová. Stochasticity helps to navigate rough landscapes: comparing gradient-descent-based algorithms in the phase retrieval problem. Machine Learning: Science and Technology , 2(3):035029, 2021.

Marco Mondelli and Andrea Montanari. Fundamental limits of weak recovery with applications to phase retrieval. Foundations of Computational Mathematics , 19(3):703-773, Jun 2019.

Andrea Montanari and Basil N Saeed. Universality of empirical risk minimization. In Conference on Learning Theory , pages 4310-4312. PMLR, 2022.

Andrea Montanari and Subhabrata Sen. A friendly tutorial on mean-field spin glass techniques for non-physicists. Foundations and Trends® in Machine Learning , 17(1):1-173, 2024.

Opper and Haussler. Generalization performance of bayes optimal classification algorithm for learning a perceptron. Physical review letters , 66 20:2677-2680, 1991.

Amelia Perry, Alexander S Wein, and Afonso S Bandeira. Statistical limits of spiked tensor models. In Annales de l'Institut Henri Poincaré-Probabilités et Statistiques , volume 56, pages 230-264, 2020.

Farzad Pourkamali, Jean Barbier, and Nicolas Macris. Matrix inference in growing rank regimes. IEEE Transactions on Information Theory , 2024.

Sundeep Rangan. Generalized approximate message passing for estimation with random linear mixing. In 2011 IEEE International Symposium on Information Theory Proceedings , pages 2168-2172. IEEE, 2011.

Valentina Ros, Gérard Ben Arous, Giulio Biroli, and Chiara Cammarota. Complex energy landscapes in spiked-tensor and simple glassy models: Ruggedness, arrangements of local minima, and phase transitions. Physical Review X , 9(1):011003, 2019.

Stefano Sarao Mannelli, Giulio Biroli, Chiara Cammarota, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Complex dynamics in simple neural networks: Understanding gradient flow in phase retrieval. Advances in Neural Information Processing Systems , 33:3265-3274, 2020a.

Stefano Sarao Mannelli, Eric Vanden-Eijnden, and Lenka Zdeborová. Optimization and generalization of shallow neural networks with quadratic activation functions. Advances in Neural Information Processing Systems , 33:13445-13455, 2020b.

Henry Schwarze. Learning a rule in a multilayer neural network. Journal of Physics A: Mathematical and General , 26(21):5781, 1993.

Guilhem Semerjian. Matrix denoising: Bayes-optimal estimators via low-degree polynomials. arXiv preprint arXiv:2402.16719 , 2024.

Hyunjune Sebastian Seung, Haim Sompolinsky, and Naftali Tishby. Statistical mechanics of learning from examples. Physical review A , 45(8):6056, 1992.

Jack W Silverstein and Sang-Il Choi. Analysis of the limiting spectral distribution of large dimensional random matrices. Journal of Multivariate Analysis , 54(2):295-309, 1995.

Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the optimization landscape of over-parameterized shallow neural networks. IEEE Transactions on Information Theory , 65(2):742-769, 2018.

Haim Sompolinsky, Naftali Tishby, and H Sebastian Seung. Learning from examples in large neural networks. Physical Review Letters , 65(13):1683, 1990.

Min Jae Song, Ilias Zadik, and Joan Bruna. On the cryptographic hardness of learning single periodic neurons. Advances in neural information processing systems , 34:29602-29615, 2021.

Roland Speicher. Free convolution and the random sum of matrices. Publications of the Research Institute for Mathematical Sciences , 29(5):731-744, 1993.

Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparametrized low-rank matrix reconstruction. Advances in Neural Information Processing Systems , 34:23831-23843, 2021.

Antonia M Tulino and Sergio Verdú. Random matrix theory and wireless communications. Foundations and Trends® in Communications and Information Theory , 1(1):1-182, 2004.

Luca Venturi, Afonso S Bandeira, and Joan Bruna. Spurious valleys in one-hidden-layer neural network optimization landscapes. Journal of Machine Learning Research , 20(133):1-34, 2019.

Roman Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

Timothy LH Watkin, Albrecht Rau, and Michael Biehl. The statistical mechanics of learning a rule. Reviews of Modern Physics , 65(2):499, 1993.

Eugene P Wigner. Characteristic vectors of bordered matrices with infinite dimensions. Annals of Mathematics , pages 548-564, 1955.

Lenka Zdeborová and Florent Krzakala. Statistical physics of inference: Thresholds and algorithms. Advances in Physics , 65(5):453-552, 2016.

A Additional definitions and conventions

Convention Throughout this manuscript, we use E X to denote the expectation solely over the random variable X . We denote M + 1 ( R ) the set of real probability distributions.

Random matrix ensembles For any d ≥ 1 , we define two standard random matrix distributions over the space of symmetric d × d real matrices:

A matrix ξ is distributed according to the GOE( d ) distribution (standing for Gaussian Orthogonal Ensemble ) if ξ ij i . i . d . ∼ N (0 , [1 + δ ij ] /d ) for any 1 ≤ i ≤ j ≤ d .

For any m ≥ 1 , a matrix S is distributed according to the Wishart distribution W m,d if S = W ⊺ W /m , where W ∈ R m × d with W ki i . i . d . ∼ N (0 , 1) for k ∈ [ m ] , i ∈ [ d ] .

For a symmetric matrix M with eigenvalues ( λ i ) d i =1 , we denote µ M := (1 /d ) ∑ d i =1 δ λ i its empirical eigenvalue distribution (ESD). It is well known that for d → ∞ the ESD of GOE( d ) and W m,d matrices converge to (respectively) the Wigner semicircle and the Marchenko-Pastur density.

Theorem A.1. [Wigner [1955], Marchenko and Pastur [1967]] Let m = κd for κ > 0 , and let ξ ∼ GOE( d ) and S ∼ W m,d . Then, as d →∞ , the ESDs of ξ and S almost surely converge (in the sense of weak convergence) to the following probability distributions (respectively).

The semicircle law, with density

We denote σ s . c ., √ t ( x ) := t -1 / 2 σ s . c . ( x/ √ t ) the scaled semicircle law with variance t .

The Marchenko-Pastur law, with density

Transforms of probability distributions For any real probability measure µ , we define its Stieltjes transform g µ ( z ) := E µ [1 / ( X -z )] for z ∈ C . If C + := { z ∈ C : Im( z ) > 0 } , then g µ ( z ) ∈ C + for all z ∈ C + . Moreover, we have the Stieltjes-Perron inversion formula:

Theorem A.2 (Stieltjes-Perron inversion formula) . For all a < b , we have

In particular, if µ has a continuous density with respect to the Lebesgue measure then:

We often use the logarithmic potential function Σ( µ ) := ∫ µ (d x ) µ (d y ) log | x -y | . We further define the R -transform of µ as:

We refer to Tulino and Verdú [2004] for more details on the definitions of this transform, e.g. concerning its complete domain of definition. Informally, the R transform is well-defined in a neighborhood of 0 for all measures which have bounded support. In particular, we have for the semicircle and the Marchenko-Pastur distributions:

Algorithm 1: GAMP-RIE

Result: The estimator ˆ S

Input: Observations y ∈ R n and 'sensing vectors' Z i := ( x i x ⊺ i - I d ) / √ d ∈ R d × d

;

Initialize ˆ S 0 ∼ W m,d and ˆ c , ω , V randomly;

while not converging do

Estimation of the variance and mean of Tr[ Z i ˆ S ] ;

Variance and mean of S estimated from the 'channel' observations ;

Update of the estimation of S ⋆ with the 'prior' information ;

t

end

Free additive convolution The main interest of the R -transform lies in its connection to the (additive) free convolution of measures. Informally, we can interpret the free convolution µ ⊞ ν of two measures µ and ν as the limiting spectral measure of A + B , where A and B are symmetric d × d random matrices, with limiting spectral distributions µ and ν , and which are asymptotically free . While we refer to Anderson et al. [2010], Tulino and Verdú [2004] for mathematical discussions of asymptotic freeness, we recall that in particular if B is a GOE( d ) matrix independent of A , then A and B are asymptotically free. Crucially, the R transform is additive under free convolution (see Theorem 2.64 in Tulino and Verdú [2004] e.g.):

Eq. (22) allows to efficiently compute the density of µ ⊞ ν given the ones of µ and ν , by relating the R transform to the Stieltjes transform, and then using the Stieltjes-Perron inversion theorem (Theorem A.2).

B The GAMP-RIE algorithm

B.1 Polynomial-time optimal estimation with the GAMP-RIE algorithm

Let us recall a crucial observation of Section 3: the learning problem of eq. (2) can be effectively reduced to a generalized linear model (GLM) on the matrix S ⋆ (cf. eq. (10)):

with Z i := ( x i x ⊺ i -I d ) / √ d , S ⋆ ∼ W m,d , and P out a noise channel (which would be Gaussian in eq. (9)). An important difficulty in analyzing eq. (23) is the rather complex structure of the matrices Z i (which can be viewed as 'sensing vectors' applied to S ⋆ ). Determining the optimal algorithm in GLMs when the sensing vectors have arbitrary structure is in general open. Anticipating on a universality argument for the MMSE (cf. Section 4), we 'forget' momentarily about the structure of { Z i } , and assume that the optimal algorithm takes the form it would have if the { Z i } were instead Gaussian matrices (i.e. GOE( d ) ). For generalized linear models with Gaussian sensing vectors, a class of generalized approximate message-passing (GAMP) algorithms have been extensively studied, and argued to reach optimal performance in the absence of a computational-to-statistical gap [Donoho et al., 2009, Rangan, 2011, Zdeborová and Krzakala, 2016]. The GAMP algorithm includes a denoiser that is adjusted to the prior information about the signal S ⋆ , that is in our case a Wishart distribution. Combining these two facts, we propose the GAMP-RIE algorithm in Algorithm 1. An implementation of GAMP-RIE is accessible in the GitHub repository associated to this work [Maillard et al., 2024].

=

t

+1

;

The functions g out , f RIE and F RIE appearing in Algorithm 1 are defined as follows. First, we let

In particular, for the problem of eq. (9), we have

The two functions ( f RIE , F RIE ) are related to the problem of matrix denoising , in which one aims at recovering a matrix S 0 ∼ W m,d from the observation of R = S 0 + √ ∆ ξ , with ξ ∼ GOE( d ) . We recall some important results on this problem, and how they relate to the definition of the functions ( f RIE , F RIE ) .

( i ) The optimal estimator (in the sense of mean squared error) of S 0 has been worked out in Bun et al. [2016], and belongs to the class of rotationally-invariant estimators (RIE). f RIE ( R , ∆) is this optimal estimator, and it admits the following explicit form. If R = U Λ U ⊺ is the spectral decomposition of R , and letting ρ ∆ := µ MP ,κ ⊞ σ s . c ., √ ∆ be its asymptotic eigenvalue distribution (see Appendix A for the definition of the free convolution µ ⊞ ν and its relation to the sum of asymptotically free matrices), then f RIE ( R , ∆) = U f ∆ ( Λ ) U ⊺ , where f ∆ ( λ ) = λ -2∆ h ∆ ( λ ) , with h ∆ the Hilbert transform of ρ ∆ . More precisely:

ρ ∆ and h ∆ can be evaluated numerically very efficiently, see Appendix A for details.

( ii ) F RIE (∆) is defined as the asymptotic MMSE of the same matrix denoising problem. It can be written in the two equivalent forms (see Maillard et al. [2022b], Pourkamali et al. [2024], Semerjian [2024]):

In Appendix B.2 we sketch the derivation of the state evolution of Algorithm 1, assuming a universality result discussed in Section 4 holds as well for GAMP-RIE. Concretely, we show that one can analytically track the performance of its iterates in the high-dimensional limit, and we draw a formal connection with the information-theoretic predictions of Claim 2. Notably, we obtain a so-called state-evolution of the GAMP-RIE algorithm (which turns out to follow from rigorous work on non-separable estimation with GAMP [Berthier et al., 2020, Gerbelot and Berthier, 2023]), and show that its fixed points agree with the fixed point equations that provide the Bayes-optimal error. In all regions of parameters that we investigated below we observed a unique fixed point, meaning that the GAMP-RIE algorithm asymptotically reaches the Bayes-optimal performance, see Appendix B.3

B.2 State evolution: a connection between Algorithm 1 and Result 1

We briefly sketch here the statistical-physics style derivation of the so-called state evolution of Algorithm 1: this will draw a connection between the performance of the Bayes-optimal estimator, characterized by Result 1, and the estimator of Algorithm 1. We define q t := tr[( ˆ S t ) 2 ] , and m t := tr[ ˆ S t S ⋆ ] . Thanks to Bayes-optimality, one can show that, along the GAMP-RIE trajectory, the so-called Nishimori identities are preserved (see Zdeborová and Krzakala [2016] for more details), so that we have, at leading order as d →∞ , that q t = m t .

Up to some critical differences, we can transpose the derivation of Zdeborová and Krzakala [2016] of the state evolution of GAMP for generalized linear models with Gaussian sensing vectors, and i.i.d. priors, to our GAMP-RIE algorithm. The differences with our setting are twofold:

( i ) The 'sensing vectors' Z i are not Gaussian. We conjecture that the universality arguments discussed in Section 4 extend to the analysis of the GAMP-RIE algorithm. This allows us to replace Z i by G i i . i . d . ∼ GOE( d ) when evaluating ( q t , m t ) (i.e. when studying the high-dimensional performance of Algorithm 1). We are then able to make a direct use of some results of Zdeborová and Krzakala [2016].

( ii ) The prior over S ⋆ is not i.i.d.: as we saw, this led to a non-trivial 'denoising' part in Algorithm 1. The performance of this denoising procedure in the high-dimensional limit can however be characterized precisely, as the function F RIE admits a closed-form expression.

We now briefly expose the derivation, transposed to our setting under the universality assumption above. By definition of q t , we have ˆ c t = 2( Q 0 -q t ) . If ω, z are centered and jointly Gaussian variables with E [ ω 2 ] = 2 q t , E [ z 2 ] = 2 Q 0 , and E [ ωz ] = 2 m t = 2 q t , and y ∼ P out ( ·| z ) , we define

so that A t = ˆ q t / 2 in the n, d →∞ limit. For the 'channel' part of the GAMP-RIE algorithm, the standard analysis for generalized linear model, alongside the universality phenomenon discussed above (which allows replacing Tr[ Z i ˆ S t ] by Tr[ G i ˆ S t ] in the update of ω t i ) shows that ˆ q t satisfies the equation:

where J q ( y, ξ ) is defined in eq. (13). Eq. (27) is the very same as for the standard GAMP for generalized linear models [Zdeborová and Krzakala, 2016].

Wehave, however, a more structured prior. After replacing Z i by Gaussian matrices G i in Algorithm 1, the argument is that at leading order as d →∞ one has:

with ξ ∼ GOE( d ) . Heuristic details on how to derive eq. (28) can be found again in Zdeborová and Krzakala [2016], see Section 6.4.1 3 there. By definition of F RIE , this implies q t +1 := tr[( ˆ S t +1 ) 2 ] = Q 0 -F RIE ((ˆ q t ) -1 ) , so that by eq. (25):

with µ t := µ MP ,κ ⊞ σ s . c ., √ t . Notice that remarkably, eqs. (27) and (29) precisely match the extremization equations of the asymptotic free entropy, as given in Claim 2 and Result 1, exactly as for 'usual' generalized linear models with i.i.d. priors [Rangan, 2011, Javanmard and Montanari, 2013, Zdeborová and Krzakala, 2016].

Mathematical consequences The fact that, assuming the universality property above, our GAMPRIE algorithm can be seen as the usual GAMP algorithm in a generalized linear model with a non-separable prior has a very interesting consequence. Indeed, the latter model admits a rigorous state evolution thanks to the analysis of Berthier et al. [2020], Gerbelot and Berthier [2023]. To make this point clearer, we notice that (after replacing Z i by GOE( d ) matrices G i ) Algorithm 1 can be written in the following form:

Let us clarify some notations used in eq. (30):

u t ∈ R p , with p := ( d +1 2 ) , can be seen as the flattening of the symmetric matrix A t R t of Algorithm 1 via the following canonical mapping. For S ∈ S d , we define vec( S ) ∈ R p by vec( S ) ii = S ii and vec( S ) ij = √ 2 S ij for i < j . This flattening is an isometry: ⟨ vec( S ) , vec( R ) ⟩ = Tr[ SR ] . We have u t = vec( A t R t ) .

G ∈ R n × p is a Gaussian i.i.d. matrix, whose elements have variance 2 /d .

Σ -1 t := -2 α E div[ g out ( t , ω t , V t )] is related to A t by A t = Σ -1 t .

V t := 2 F RIE (Σ t / 2) .

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

0

.

5

α

=

n/d

2

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

MMSE

Theory (Noiseless)

Theory (∆ = 0

.

25

2

)

AMP (

d

= 100)

AMP (

d

= 200)

ˆ v t ( u t , Σ t ) ∈ R p is the flattening of the RIE denoiser of Algorithm 1, i.e. if we denote R t / Σ t the matrix such that u t = vec( R t / Σ t ) :

Eq. (30) is the canonical form of the GAMP algorithm, as written e.g. in Berthier et al. [2020], Gerbelot and Berthier [2023]. In particular, we can leverage their results to write:

Theorem B.1 (State Evolution (informal) Berthier et al. [2020], Gerbelot and Berthier [2023]) . Denote q t AMP := tr[ ˆ S t S ⋆ ] and ˆ q t AMP := 4 α n ∑ n i =1 g out ( y i , ω t i , V t ) 2 (recall the definition of these quantities in Algorithm 1). Assume that the 'sensing matrices' ( Z i ) n i =1 in Algorithm 1 are replaced by ( G i ) n i =1 , which are i.i.d. GOE( d ) matrices. Then for any t ≥ 0 , q t AMP and ˆ q t AMP follow the state evolution equations (27) and (29) asymptotically as d, n →∞ .

Beyond the rigorous control of the GAMP-RIE algorithm, Theorem B.1 has an additional mathematical consequence: it allows to leverage a set of mathematical techniques that use AMP algorithms to prove results on the asymptotic MMSE and on the mutual information, as Theorem B.1 implies that they can be used verbatim in our setting. More precisely, the fact that the GAMP-RIE algorithm achieves an MSE with value given by Claim 2 immediately yields that the latter is, at least, an upper bound on the asymptotic MMSE (when assuming Gaussian GOE( d ) 'sensing vectors' G i ). Additionally, the application of the I-MMSE theorem [Guo et al., 2005] shows that our claimed free entropy (i.e. the limit of (1 /d 2 ) E log Z in Claim 2) is a lower bound on the real one (see e.g. section 2.C in Barbier et al. [2016]).

B.3 GAMP-RIE algorithm reaching the optimal error

In Fig. 3 we compare the asymptotic theoretical result for the Bayes-optimal error with the performance of the GAMP-RIE algorithm for d = 100 and d = 200 , in both the noiseless (blue) and noisy (red) cases. We observe that even for such moderate sizes the agreement between the algorithmic performance and the theory is excellent.

We also stress that in all the cases we evaluated, the state evolution of the GAMP-RIE converges to the fixed point that corresponds to the Bayes-optimal performance. This means that the Bayes-optimal error discussed above is reachable efficiently with the GAMP-RIE algorithm. In particular, unlike in the canonical phase retrieval problems (i.e. when m = 1 ) [Barbier et al., 2019], we did not identify a computational-to-statistical gap when learning this extensive-width quadratic-activation neural network.

C Generic activations

We give here a brief discussion on the extension of our work to a more generic activation function. While our derivation (cf. Section 4) heavily relies on the non-linearity being quadratic, a first natural extension would be to consider polynomial activations, with an output generated as (assuming a noiseless setting):

for some integer p ≥ 3 . One could also 'linearize' this model, by writing it as y i = ⟨ T ⋆ , X i ⟩ , in which T ⋆ , X i are now p -tensors, defined as

However, two main challenges arise when carrying out the program of Section 4 in this 'tensor' model:

( i ) First, determining whether the universality Conjecture 4.1 holds for these models (and if yes, in which scaling of the number of samples n with d ) is a challenging open question that falls outside the scope of our results as well as of previous works on free entropy universality [Hu and Lu, 2022, Montanari and Saeed, 2022, Dandi et al., 2024a, Maillard and Bandeira, 2023].

( ii ) Secondly, the generalized form of Conjecture 4.2 would involve the free entropy of a tensor denoising problem. While a rich literature has studied the fundamental limits of denoising low-rank tensors (see Lesieur et al. [2017], Ben Arous et al. [2019], Ros et al. [2019], Perry et al. [2020], Gamarnik et al. [2022] and references therein), here T ⋆ has rank m = O ( d ) , and the optimal denoising of a large-rank tensor is, as far as we know, a completely open question.

These two challenges form the basis of an exciting but very challenging research program, which we leave for future work. Provided such a program could be carried out for any polynomial activation, one might then hope to analyze generic activation functions, such as the ReLU or sigmoid, e.g. by decomposition over a basis of orthogonal polynomials (such as the Hermite basis), see Ben Arous et al. [2021], Abbe et al. [2023], Dandi et al. [2023, 2024b] for examples of such analyses in the case m = O (1) .

D Derivation of Claim 2 from the replica method

In this section, we give a non-rigorous derivation of eq. (12) using classical methods of statistical physics. We start from the definition of the partition function in eq. (11). We denote D the standard Gaussian measure, and S ( W ) = W ⊺ W /m .

The replica method We make use of the heuristic replica trick [Mézard et al., 1987]. Letting Φ d := (1 /d 2 ) E log Z , it consists in writing that lim d →∞ Φ d = lim r → 0 ( ∂/∂r ) lim d →∞ Φ d ( r ) , with Φ d ( r ) := (1 /d 2 ) log E [ Z r ] . One then computes the d →∞ limit of Φ d ( r ) for integer r ∈ N , before extending analytically the result to any r ≥ 0 . While being non-rigorous, the replica method has achieved a great success in the study of both spin glasses and statistical learning models, and is widely conjectured to yield exact predictions. We refer the reader to Mézard et al. [1987] for an introduction to the replica method in the context of the statistical physics of disordered systems, Maillard et al. [2023], Montanari and Sen [2024] for mathematically-friendly descriptions of the method, and to Mezard and Montanari [2009], Zdeborová and Krzakala [2016], Gabrié [2020] for

some of its applications in the context of theoretical computer science, high-dimensional statistics, and machine learning.

The replicated free entropy We now compute the 'replicated free entropy' Φ d ( r ) , for r ∈ N . Thanks to Bayes-optimality, we can write it as an average over r +1 replicas of the system, writing S ⋆ as the replica of index 0 . We write S a := S ( W a ) to simplify notations. We reach:

For a fixed set of matrices { S a } r a =0 , by the central limit theorem the law of the variables z a := Tr[ S a Z ] approach, as d → ∞ , a correlated Gaussian distribution, with mean E [ z a ] = 0 , and covariance E [ z a z b ] = E Z [Tr[ S a Z ]Tr[ S b Z ]] = 2tr( S a S b ) , as is easily checked from the fact that Z d = ( xx ⊺ -I d ) / √ d , with x ∼ N (0 , I d ) . Since n = Θ( d 2 ) , the leading order of the term ∫ d y E Z ∏ r a =0 P out ( y | Tr[ S a Z ]) will be the only one entering the leading order of Φ d ( r ) . This means that we have, denoting the overlap matrix

that

Notice that the CLT-based argument above is made formal in Conjecture 4.1, and implies the universality of Φ d under the replacement of Z i by Gaussian GOE matrices G i . Since Q ∈ R ( r +1) × ( r +1) is of finite size as d → ∞ , we can perform the Laplace method over Q in eq. (33), and we reach (omitting o d (1) terms as d →∞ , and recall n/d 2 → α ):

where S + r +1 is the set of positive semi-definite symmetric matrices of size r +1 , and:

Notice that we can rewrite J ( Q ) using Lagrange multipliers ˆ Q ∈ S r +1 (or equivalently using the Fourier transform of the delta distribution, and the saddle point method on the Fourier parameters) as:

̸

The replica-symmetric ansatz An important assumption we make now is that there is a permutation symmetry between the different replicas in eq. (34), and we assume that this symmetry is not broken by the maximizer Q . This assumption is usually called replica symmetry , and is known to hold in generic statistical learning problems when they are in the Bayes-optimal setting [Zdeborová and Krzakala, 2016, Barbier et al., 2019]. Formally, we assume that the supremum over Q in eq. (34) (and the infimum over ˆ Q in eq. (36)) are reached in matrices such that, for all a, b ∈ { 0 , · · · , r } with a = b :

with 0 ≤ q ≤ Q , and ˆ Q, ˆ q ≥ 0 .

The term J out ( Q ) -Under the ansatz of eq. (37), it is a classical computation [Zdeborová and Krzakala, 2016] to reach:

The term J ( Q ) -Using the replica-symmetric ansatz of eq. (37) in eq. (36), we get:

We now use the following Gaussian integration identity, for any symmetric matrix M :

This allows to reach the following expression, which is analytic in r :

Recall that here S = S ( W ) = W W /m .

The limit r → 0 -From eqs. (34), (38) and (39), we have:

This implies that ˆ Q = 0 and Q = Q 0 = lim d →∞ E S ∼W m,d tr[ S 2 ] = 1 + κ -1 (recall m/d → κ ), and we correctly recover that Φ d ( r = 0) = 0 . Taking now the derivative with respect to r , followed by the r → 0 limit, yields:

In order to obtain from eq. (41) the prediction of eq. (12), it therefore suffices to show that, for any ˆ q ≥ 0 :

We focus on deriving eq. (44) in the remaining of this section. We note that we can rewrite the left-hand side as the free entropy of the following denoising problem:

with ξ ∼ GOE( d ) , S ⋆ ∼ W m,d , and which we consider in the Bayes-optimal setting. Indeed, we can define the free entropy of this problem as

Crucially, this auxiliary problem is again Bayes-optimal, which we will use in what follows.

Remark Eq. (45) defines a problem known as extensive-rank matrix denoising . The limit free entropy of this problem, as well as the analytical form of the Bayes-optimal estimator, for a rotationallyinvariant prior on S ⋆ and a rotationally invariant noise ξ (which is here Gaussian) have both been understood and worked out completely [Bun et al., 2016, Maillard et al., 2022b, Pourkamali et al., 2024, Semerjian, 2024]. We will leverage these results (and partially re-derive them) in what follows.

We now use a change of variable to the singular values of W , see e.g. Proposition 4.1.3 of Anderson et al. [2010]. We reach:

in which S = W ⊺ W /m = O Λ O ⊺ , with Λ = Diag(( λ 1 , · · · , λ m , 0 , · · · , 0)) , and C d,m > 0 is a constant depending only on m and d . Notice that we (slightly abusively) used the notation D O to denote here the Haar measure over the orthogonal group O ( d ) . The larged limit of the last term is given by the HCIZ integral [Harish-Chandra, 1957, Itzykson and Zuber, 1980]:

where S and Y are d × d matrices with asymptotic eigenvalue distributions µ S and µ Y . We can now apply the Laplace method in eq. (47) on the eigenvalue distribution of S . As the problem of eq. (45) is Bayes-optimal, it is known that the typical eigenvalue distribution of S under the distribution of eq. (47) is µ S = µ S ⋆ = µ MP ,κ , as a consequence of the so-called Nishimori identity, so that µ MP ,κ is the maximizer of the variational problem obtained by the use of Laplace's method, see Maillard et al. [2022b] for details. Since the asymptotic distribution of Y is (by eq. (45)) µ Y = µ MP ,κ ⊞ σ s . c ., 1 / √ ˆ q , we reach by eq. (46):

where C ( κ ) is a function of κ = m/d . It can be easily seen that C ( κ ) = 0 by considering ˆ q = 0 .

Fortunately, extensive-rank matrix denoising with Gaussian noise is one of the very few cases for which an easily tractable analytical form is known for the HCIZ integral. More specifically, we know that for any t > 0 and any ν , we have with µ t := ν ⊞ σ s . c ., √ t [Maillard et al., 2022b]:

with Σ( µ ) := ∫ µ (d x ) µ (d y ) log | x -y | . Applying this formula with t = 1 / ˆ q we reach:

Combining it with eq. (49), we reach eq. (44) (recall that µ Y = µ 1 / ˆ q with the notations of eq. (44)).

E Large and small κ limits

E.1 The smallκ limit

We consider here the limit κ → 0 , i.e. the limit of small (but still extensively large) hidden layer, and compute the limit of the MMSE curves shown in Fig. 1. Since in the noiseless setting we have α PR = κ + O ( κ 2 ) (cf. eq. (1)), we will work in the rescaled regime α = ˜ ακ , with ˜ α remaining finite

where

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

α/κ

=

n/dm

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

MMSE

κ

→

0

κ

= 0

.

01

κ

= 0

.

05

κ

= 0

.

1

κ

= 0

.

2

κ

= 0

.

5

κ

= 1

.

0

κ

= 5

.

0

as κ ↓ 0 . By analyzing eq. (8) in this regime (details are given in Appendix E.1.1), we reach that the MMSE satisfies, as κ → 0 :

In particular, in the noiseless case ( ∆ = 0 ), we have:

and we reach perfect recovery for ˜ α = 1 . This limit is illustrated in Fig. 4.

Remarkably, eq. (51) can be computed as well by taking the limit m → ∞ when assuming that m = O (1) as d → ∞ , a setting which was studied extensively in the literature (see Aubin et al. [2019b] and references therein). We detail this computation in Appendix E.1.2.

E.1.1 Details of the smallκ limit

Recall that by Claim 2, we have MMSE = κ ( Q 0 -q ⋆ ) , with Q 0 = 1 + κ -1 . Since the MMSE remains finite as κ → 0 , we consider the scaling q = ˜ q/κ , with 0 ≤ ˜ q ≤ 1 . We start again from eqs. (7) and (8). We denote Λ := ∆(2 + ∆) . Eq. (7), combined with the scaling of α , implies that ˆ q ∼ κ 2 /t for some finite t > 0 , and since MMSE = 1 -˜ q as κ → 0 , we have

Moreover, eq. (8) at order O ( κ ) yields:

Notice that F ( t, 0) = 1 since µ ξ ≃ σ s . c ., √ ξ for ξ → ∞ , and ∫ σ s . c ., √ ξ ( y ) 3 d y = 3 / [4 π 2 ξ ] . Thus, the leading order of eq. (8) as κ → 0 is consistent but not informative.

In what follows, we work out the small κ limit of F ( t, κ ) , at first order in κ . We denote ν κ ( y ) := (1 /κ ) µ t/κ 2 ( y/κ ) , so that the Stieltjes transform g κ ( z ) := ∫ ν ( y ) / ( y -z )d y of ν satisfies the selfconsistent equation (see Appendix A):

Moreover, we notice that ν κ = ( κ # µ MP ,κ ) ⊞ σ s . c ., √ t , so that the support of ν remains bounded as κ → 0 . We then proceed to expand in κ eq. (53). For any finite z ∈ C , the leading order of the expansion is easily given by z = -1 /h -th + o κ (1) , which gives that ν κ → σ s . c ., √ t . However, as mentioned above, we need to go to the next order in this expansion to compute eq. (52).

A BBP-type transition We notice that κ # µ MP ,κ ( x ) ≃ (1 -κ ) δ ( x ) + κδ ( x -1) when κ → 0 . More precisely, it is composed of a mass (1 -κ ) in 0 , and the rest of the mass κ is made up of a continuous part supported between (1 - √ κ ) 2 ≃ 1 -2 √ κ and (1 + √ κ ) 2 ≃ 1 + 2 √ κ . ν κ can thus be seen as the spectral density of the sum of a GOE matrix (with variance t ) and a small-rank perturbation matrix of rank m = κd , with all non-zero eigenvalues located close to 1 . We therefore expect by the so-called BBP transition phenomenon [Benaych-Georges and Nadakuditi, 2011] that ν κ will possess a set of m eigenvalues outside the semicircle bulk whenever the condition

is satisfied, with g s . c ., √ t ( z ) := E X ∼ σ s . c ., √ t [1 / ( X -z )] the Stieltjes transform of the semicircle. Since one can easily show that g s . c ., √ t (2 √ t ) = -t -1 / 2 , eq. (54) is equivalent to t ≤ 1 . In this case, these 'spiked' eigenvalues are located around the value [Benaych-Georges and Nadakuditi, 2011]

Moreover, as the width of the continuous part of κ # µ MP ,κ is of size O ( √ κ ) , we also expect this 'spiked' part of the spectrum to have a width O ( √ κ ) .

̸

Expansion of ν -Based on the remarks of the previous paragraph, we assume the following behavior for ν κ , as κ → 0 . For any y ∈ R with y = 1 + t , we have

Furthermore, we also have, for all y ∈ R , when t ≤ 1 :

for a finite density ρ (1) , with ∫ ρ (1) ( y )d y = 1 . Eqs. (55) and (56) can be used to expand the Stieltjes transform of ν κ as a function of ν (1) , ρ (1) , and then eq. (53) used to find the values of these two functions. These computations are straightforward, and yield:

Notice that the second equation of eq. (57) is only valid for t ≤ 1 , while the first one is valid for all t ≥ 0 . One checks for instance that ∫ ν (1) (d y ) = -1 { t ≤ 1 } , which implies that the normalization condition ∫ ν κ ( y )d y = 1 is well satisfied for all values of t ≥ 0 . Using the expansion of eq. (57), we obtain that

So finally eq. (52) becomes

And recall that MMSE = 2 αt -Λ , so that

Since t = (MMSE + Λ) / (2 ˜ α ) , we reach that t = (1 + Λ) / (2 ˜ α ) if ˜ α ≤ (1 + Λ) / 2 , and t = 1 -˜ α + √ (1 -tα ) 2 +Λ otherwise. This yields eq. (50).

E.1.2 The smallκ limit from a large but finite hidden layer

We consider here the noiseless case:

˜

with m = O (1) as n, d → ∞ . We denote α = n/d = ˜ αm , and we assume that ˜ α = Θ(1) as m →∞ ( after n, d →∞ ). We can write the partition function (cf. eq. (11)) as:

with P out ( y | z ) = δ ( y -∥ z ∥ 2 /m ) . We can make a direct use of the results of Aubin et al. [2019b] to write:

Here, q , ˆ q are symmetric m × m matrices, which satisfy moreover I m ⪰ q ⪰ 0 and ˆ q ⪰ 0 . The informal notation ' extr f ' in eq. (59) means that one should zero-out the gradient of the function f to compute the values of q , ˆ q .

The matrix q Importantly, the matrix q can be interpreted as the 'overlap matrix' of the model: if we denote ⟨·⟩ the average under the posterior measure in eq. (58), then we have

where w , w ′ are two independent samples under ⟨·⟩ . Moreover, thanks to the Bayes-optimality of the problem, it is known that the overlap concentrates [Zdeborová and Krzakala, 2016], in the sense that the random variable ( w ⊺ k w ′ l ) /d concentrates on its average under E ⟨·⟩ as d →∞ .

The 'prior integral' I P can be very easily computed with Gaussian integrals, and yields:

We now focus on computing the leading order of I C in the largem limit. We can write

Let ˜ y := √ m [ y -tr(I m -q ) -( ξ ⊺ q ξ ) /m ] . We can change variables in eq. (62), and obtain:

Notice that the additive term (1 / 2) log m in I C just amounts to a renormalization of the partition function Z , so we remove this additional constant in what follows. We proceed to simplify J q ( ˜ y, ξ ) in the largem limit. We have

where

Plugging it back into eq. (63) yields:

Since ξ ∼ N (0 , I m ) , it follows from elementary concentration of measure that σ 2 ξ concentrates on its average value σ 2 given by:

All in all we reach that (up to additive constants):

Combining eqs. (61) and (64) in eq. (59), we get at leading order in m , with Φ := lim(1 /d ) E log Z :

Eq. (65) can be easily solved, and yields:

This implies that (recall 0 ⪯ q ⪯ I m ):

Now that we have obtained q in eq. (66), we can compute the MMSE , or generalization error. Defining it as in eq. (4):

the same arguments used in the proof of Lemma F.1 show that in the large m limit (but taken after d →∞ ), we have at leading order

with S := (1 /m ) ∑ m k =1 w k w ⊺ k . Notice that S BO = ⟨ S ⟩ , so that

where w , w ′ are two independent samples under the posterior measure ⟨·⟩ . We know that the overlap concentrates (cf. the discussion around eq. (60)), so that at leading order, with MMSE := lim d →∞ MMSE d :

Combining it with eq. (66), we reach:

We have recovered eq. (51) from the limit m →∞ taken after d →∞ !

E.2 The largeκ limit

We consider here κ →∞ , with α remaining of order Θ(1) as κ →∞ . Since the MMSE remains finite as well, we see from eq. (7) that we must have the scaling ˆ q = κt , with t remaining finite as κ →∞ . A very similar derivation to the one of Appendix E.1.1 yields that eq. (8) in this limit becomes (with Λ := ∆(2 + ∆) ):

Combining it with eq. (7) yields that

where we recall Λ = ∆(2 + ∆) . In particular, for ∆ = 0 , we reach MMSE = max(1 -2 α, 0) , coherently with the behavior shown in Fig. 1.

F Other technicalities

F.1 Properties of the MMSE of eq. (4)

Let S ⋆ := (1 /m ) ∑ m k =1 w ⋆ k ( w ⋆ k ) ⊺ , and ˆ S BO := E [ S |D ] the Bayes-optimal estimator of S ⋆ . We show here the following lemma on the MMSE of eq. (4), under the high-dimensional limit of eq. (5):

Lemma F.1. For a constant C = C ( κ ) > 0 :

Lemma F.1 shows that we can consider the MMSE on S equivalently to the generalization MMSE of eq. (4).

Limits Notice that if the posterior concentrates around the true W ⋆ , then ˆ S BO = E [ S |D ] concentrates on S ⋆ , which implies that MMSE d → 0 . Conversely, for α = 0 (i.e. in the absence of data), the Bayes-optimal estimator becomes ˆ S BO = E [ S ⋆ ] = I d , so that E tr[( S ⋆ -ˆ S BO ) 2 ] = κ -1 . Thus, we have MMSE d → 1 for α = 0 .

Proof of Lemma F.1. - Notice that (cf. eq. (2)):

with S := (1 /m ) ∑ m k =1 w k w ⊺ k . Using this in eq. (3), and plugging it in eq. (4), we get (with z ∼ N (0 , I m ) and x ∼ N (0 , I d ) ):

where we used E [ ∥ z ∥ 4 ] = m 2 +2 m and E tr( S ⋆ ) = 1 in (a) , and E x ∼N (0 , I d ) [( x ⊺ Mx ) 2 ] = Tr[ M ] 2 + 2Tr[ M 2 ] in (b) . It remains to bound the first term of eq. (68) to conclude the proof of Lemma F.1. We notice that, by linearity of the trace, tr( ˆ S BO ) is the Bayes-optimal estimator for tr( S ⋆ ) , i.e.

In particular, considering the estimator

we have using eq. (69):

using the Cauchy-Schwarz inequality in (a) , with

It is a tedious but straightforward computation to compute { I a } 3 a =1 , as it only involves the first moments of Gaussian random variables. We get (recall m = κd ):

Combining eqs. (70) and (71), and plugging it back in eq. (68), we get

which ends the proof of Lemma F.1.

F.2 Proof of Theorem 4.1

First, we note that Theorem 1 of Pourkamali et al. [2024] implies that:

and we recall the definition of I HCIZ in eq. (48). We recall then a fundamental result proven in Guionnet and Zeitouni [2002]:

Theorem F.2 (Theorem 1.1 of Guionnet and Zeitouni [2002]) . For any compactly supported probability measures ν and µ , and any t > 0 :

Moreover, the function J ( ν ; µ ) satisfies the following property. Let d be a distance on the space of probability measures on R that is compatible with the weak topology. Let X := R + √ t W , where W ∼ GOE( d ) , and R is a fixed (deterministic) matrix, with uniformly bounded spectral norm, and a compactly supported limiting eigenvalue distribution µ . Let µ X denote the empirical eigenvalue distribution of X . Then, for any ν ∈ M + 1 ( R ) :

In other words, the function J ( ν ; µ ) is the large deviations rate function (in the scale d 2 ) for the empirical spectral measure of R + √ t W , where R is a fixed (deterministic) matrix with asymptotic spectral distribution µ , and W ∼ GOE( d ) . It is a well-known property of the free convolution [Speicher, 1993] that µ X → µ ⊞ σ s . c ., √ t as d →∞ , where the convergence is meant in the weak sense (and almost surely). Combining this result with eq. (74), we have J ( µ ⊞ σ s . c ., √ t ; µ ) = 0 . This yields by eq. (73):

Combining eqs. (72) and (75) yields eq. (17).

Remark The proof above can be straightforwardly extended to the free entropy of denoising any matrix S with a rotationally-invariant distribution and a compactly-supported limiting eigenvalue distribution (beyond the Wishart ensemble), as the results of Guionnet and Zeitouni [2002], Pourkamali et al. [2024] hold under these more general assumptions.

F.3 Perfect recovery threshold in the noiseless case

In this section, we give an analytic argument to derive the value of the perfect recovery threshold α PR (see eq. (1)) in the noiseless setting. In the limit of perfect recovery the MMSE goes to 0 , thus by eq. (7) (with ∆ = 0 ) this implies ˆ q →∞ . Using eq. (8), we can then write the equation satisfied by the perfect recovery threshold as

in which µ t = µ MP ,κ ⊞ σ s . c ., √ t , see Appendix A.

F.3.1 The case κ < 1

Informal argument Recall that in this case we can write µ MP ,κ ( x ) = (1 -κ ) δ ( x ) + κν MP ,κ ( x ) , in which ν MP ,κ is compactly supported away from zero, see Appendix A. As t → 0 , we thus expect µ t to have a discontinuous support, made of two parts:

( a ) A small semicircular density centered around 0 , of radius O ( √ t ) , with mass (1 -κ ) .

( b ) A smooth density, compactly supported away from zero, which has a well-defined limit as t → 0 , and a mass κ .

Because of the factor t in the right-hand side of eq. (76), only the part ( a ) will matter in the limit.

Formal derivation We first rewrite by a change of variable

It is clear that for all x = 0 , we have µ t ( x ) → κν MP ,κ ( x ) as t → 0 , and ∫ ν MP ,κ ( y ) 3 d y < ∞ , so that we can truncate the integral above to all | z | ≤ ε/ √ t , for any ε > 0 finite as t → 0 . We will now show the following, for any x ∈ R :

We fix z ∈ C + (where C + := { z ∈ C : Im( z ) > 0 } ). Letting y = √ tz , we know from the Marchenko-Pastur theorem [Marchenko and Pastur, 1967] that g t ( y ) := E µ t [1 / ( X -y )] is the unique solution in C + to the equation

Since y = √ tz , it is clear that g = O (1 / √ t ) , and letting h := √ tg , we easily get the expansion

̸

which can be inverted to

Equivalently:

so that

Notice that if we denote S κ ( z ) the Stieltjes transform of σ s . c ., √ 1 -κ , eq. (78) can be written as (see e.g. Anderson et al. [2010]) h = (1 -κ ) S κ ( z ) + O ( √ t ) . By considering z = x + iε with x ∈ R and the limit ε → 0 , we reach using the Stieltjes-Perron inversion theorem (Theorem A.2) that for any x ∈ R :

Coming back to eq. (76) this implies:

We notice that this critical value of n/d 2 coincides with a naive counting argument of degrees of freedom of S ⋆ . Indeed, as can be seen by the spectral decomposition, the set of d × d symmetric matrices of rank m has, to leading order in d , p ( κ ) d 2 degrees of freedom, where p ( κ ) d 2 is the dimension of the Stiefel manifold of orthonormal m -frames in R d . It is well-known that p ( κ ) = κ -κ 2 / 2 for d →∞ [Helmke and Moore, 2012].

F.3.2 The case κ ≥ 1

The case κ > 1 is simpler to carry out. In this case, µ MP ,κ does not have a singular part at x = 0 , and µ t has a smooth density as t → 0 , and

and we reach α PR = 1 / 2 , so that α PR d 2 (asymptotically) coincides with the number d 2 / 2 of degrees of freedom of symmetric matrices. Since α PR is increasing with κ , and has limit 1 / 2 both for κ ↑ 1 and κ ↓ 1 , we deduce that α PR = 1 / 2 for κ = 1 as well.

F.4 The derivative of the error at the perfect recovery threshold

Here, we extend the derivation of Section F.3 to compute the derivative of the MMSE with respect to α at the perfect recovery threshold. We start again from eqs. (7) and (8). Letting t := 1 / ˆ q , we get, with α = α PR :

We thus compute the next order of the expansion of t ∫ µ t ( y ) 3 d y as t → 0 .

F.4.1 The case κ < 1

We extend the argument made in Section F.3.1. Notice that here the smooth part of the density, compactly supported away from zero, contributes at this order. Formally, for any small enough ε > 0 :

On the other hand, we have around the singularity at y = 0 :

We evaluate the next order of the right-hand side of eq. (83) using the same approach as in Section F.3.1, going to next orders in the expansion as t → 0 of eq. (78). Using then again the Stieltjes-Perron inversion theorem, we reach with tedious but straightforward computations the generalization of eq. (79):

for any | z | ≤ 2 √ 1 -κ , while √ tµ t ( √ tz ) = O ( t 3 / 2 ) if | z | > 2 √ 1 -κ . This then yields:

Combining eqs. (82) and (85) in eq. (81), we obtain (recall α = α PR = κ -κ 2 / 2 ):

F.4.2 The case κ ≥ 1

Again, we consider κ > 1 . The argument of Section F.4.1 generalizes immediately, removing the analysis of the singular part around y = 0 . We get directly

Plugging it in eq. (81), we get in this case:

Again, the specific case κ = 1 can be tackled by continuity, as the derivative tends to 0 both as κ ↑ 1 and κ ↓ 1 .

F.5 Details on the reduction to matrix estimation

We describe here how to effectively reduce the problem of eq. (2) to an estimation problem in terms of S ⋆ := (1 /m ) ∑ m k =1 w ⋆ k ( w ⋆ k ) ⊺ .

Remark While our argument is backed by precise probabilistic concentration arguments, we notice that it is not a proof of the equivalence of the problems of eq. (2) and eq. (9) under all statistical tests, as would be implied e.g. by the contiguity of distributions [Le Cam, 1960, Kunisky et al., 2019]. Rather, we analyze the leading order of eq. (2) and argue that (with high probability over

the distribution of the data and the teacher weights), the first non-trivial order of the observations is characterized by the equivalent model of eq. (9). Notably, we do not claim the statistical equivalence of the problems of eq. (2) and eq. (9), but rather only that their asymptotic MMSEs coincide. While even this weaker statement is not formally implied by the arguments sketched below, we expect that they form the backbone of a formal proof of this claim, which we leave for future work and would be carried e.g. by Gaussian interpolation techniques.

Let us define Z i := ( x i x ⊺ i -I d ) / √ d , and recall that x i ∼ N (0 , I d ) . Expanding the square, we can rewrite the law of the output y i = f W ⋆ ( x i ) as

where ( z i ) n i =1 i . i . d . ∼ N (0 , I m ) . In what follows, we analyze the leading order of eq. (87). More specifically, we denote ˜ y i := √ d ( y i -1 -∆) , and we decompose

Let us consider the leading order of the different terms of eq. (88). Since S ⋆ ∼ W m,d , Tr[ S ⋆ ] = ∑ m k =1 ∥ w ⋆ k ∥ 2 /m strongly concentrates on its average. More precisely, by Bernstein's inequality (see Corollary 2.8.3 of Vershynin [2018]) we have, for all t ≥ 0 :

where C > 0 depends only on κ > 0 . In particular,

so that we can replace I 1 by 0 at leading order in eq. (88).

We now tackle I 2 , first for fixed ( x i , W ⋆ ) . Using that ∥ z i ∥ 2 strongly concentrates around its average, and the central limit theorem applied to the fluctuations of ∥ z i ∥ 2 , one can see that for all i ∈ [ n ] , we have (with g i ∼ N (0 , I m ) independently of z i , and d = denoting equality in distribution):

with ξ i i . i . d . ∼ N (0 , 1) , independently of ( x i , w ⋆ k ) . The equivalence as d → ∞ is given for a fixed i ∈ [ n ] : coherently with the remark above, we notice that a formal mathematical proof of equivalence of the two problems of eq. (2) and eq. (9) would rather need to tackle the joint law of all the observations, and to quantitatively control the deviation between the left and right-hand sides of eq. (89) as d →∞ . We leave such a proof for future work.

We finally note that the variance term on the right-hand side of eq. (89) strongly concentrates, uniformly in i ∈ [ n ] , as by the Hanson-Wright inequality and the union bound, we have (see Theorem 6.2.1 of Vershynin [2018]) for all t ≥ 0 :

for some constant C > 0 . Since the spectral norm of a Wishart matrix ∥ S ⋆ ∥ op strongly concentrates on its average under the Wishart distribution (see Theorem 4.4.5 of Vershynin [2018]), we see that, uniformly over i ∈ [ n ] , the leading order of the variance in the right-hand side of eq. (89) is equal to ˜ ∆ := 2∆(2 + ∆) /κ . This ends our justification of eq. (9).

F.6 Unique maximizer q ⋆ in eq. (12)

Notice that if J out ( q ) := ∫ R × R d y D ξ J q ( y, ξ ) log J q ( y, ξ ) , then one can check that J out is a strictly increasing function of q under mild regularity conditions on P out (namely assuming the presence of an additive Gaussian noise with arbitrarily small variance), see Proposition 21 of Barbier et al. [2019]. The fact that q ⋆ is uniquely defined for all values of α > 0 except possibly in a countable set follows then from Proposition 1 of Barbier et al. [2019], see also Appendix A.2 there.

F.7 Derivation of Result 1 from Claim 2

In this section, we derive eqs. (7) and eq. (8) from Claim 2, in the case of Gaussian noise. More precisely, we assume P out ( y | z ) = exp[ -( y -z ) 2 / (2 ˜ ∆)] / √ 2 π ˜ ∆ , in accordance with eq. (9). It is then an easy computation to check (recall the definition of J q in eq. (13)):

We then reach that q = q ⋆ is characterized as the maximum of the following function:

Recall that here µ t := µ MP ,κ ⊞ σ s . c ., √ t . It is known (see eqs. (77-78) of Semerjian [2024] e.g.) that

Thus, ˆ q = ˆ q ( q ) can be characterized as the solution 4 to

By eq. (91), q is a solution in [1 , Q 0 ] to:

Recalling that MMSE = κ ( Q 0 -q ) by Claim 2, eq. (93) implies eq. (7). Combining eq. (93) with eq. (92), we reach eq. (8).

F.8 The limit α → 0

In this section, we check that the state evolution equations derived in Section F.7 yield indeed that q → 1 as α → 0 . Indeed, in this limit, S BO = E [ S ⋆ ] = I d , so that we must have q = E tr[ S BO S ⋆ ] = 1 .

Recall that ˆ q = 4 α/ [ ˜ ∆+2( Q 0 -q )] , and that ˆ q is given by eq. (8). In particular, ˆ q → 0 as α → 0 . Assuming the scaling ˆ q ∼ ˆ q 0 α as α → 0 , we get

where F ( p ) := (4 π 2 / 3) ∫ [ p -1 / 2 µ 1 /p ( z · p -1 / 2 )] 3 d z . Letting ν p ( z ) := p -1 / 2 µ 1 /p ( z · p -1 / 2 ) , we know by a similar reasoning as the one of Section F.3 that the Stieltjes transform h = h p ( z ) of ν p satisfies the equation:

As p → 0 , we can thus compute the expansion of h p ( z ) in powers of p . Applying then the StieltjesPerron inversion theorem (Theorem A.2), we get the expansion of ν p ( z ) in powers of p as:

for | z | ≤ 2 , and ν p ( z ) = O ( p 3 / 2 ) for | z | ≥ 2 . Plugging this expansion into F ( p ) , we get:

Coming back to eq. (94), this gives ˆ q 0 = 4 κ/ [2 + ˜ ∆ κ ] , and (recall Q 0 = 1 + κ -1 ) then q = 1 , so that our equations are indeed consistent in the limit α → 0 .

G Learning the second layer weights

We sketch here in a mathematically informal way the generalization of our results to the setting where the second layer weights are also learned. The second layer weights ( a ⋆ k ) m k =1 are drawn i.i.d. from a probability distribution P a , and the student must learn ( w ⋆ k , a ⋆ k ) m k =1 from the observation of { x i } n i =1 and of

In the rest of this paper we focused on the case P a = δ 1 . However, all our techniques and results can be generalized to more generic choices of P a , as we know show: in particular, Claim 3 is the generalization of Claim 2 to this more general setting.

Throughout this section, we will assume for simplicity that P a has bounded support, although we expect our results to hold also for more general choices of P a . We show how to extend Claim 2 to this case, by detailing the differences in the steps outlined in Section 4. We eventually show that Algorithm 1 can also be straightforwardly extended to this setting as well.

G.1 Generalizing the derivation

G.1.1 Reduction to matrix estimation

We first discuss the reduction to a matrix estimation problem, generalizing Section F.5 to this setting. We define

and we denote m a := E P a [ a ] and c a := E P a [ a 2 ] . We define the MMSE as (notice the additional factor c a with respect to eq. (4)):

By repeating the (mathematically informal) arguments of Section F.5 to this setting, we find that, at leading order as m,d →∞ :

with ξ i i . i . d . ∼ N (0 , 1) , and ˜ ∆ := 2∆(2 + ∆ c a ) /κ . We let

The observation of ( y i ) n i =1 is equivalent to the one of ( ˜ y i ) n i =1 and Y . Notice that by eq. (98), we have

Conditionally on S ⋆ , the right-hand-side of eq. (99) is a sum of n independent zero-mean random variables, which thus typically fluctuates in the scale 5 O [( nd ) -1 / 2 ] = O ( d -3 / 2 ) . Since ˜ y i = √ d [ y i -Y ] , this implies that at leading order we have

The observer also has access to Y , alongside { ˜ y i } n i =1 . Notice that by the argument above, Y is (up to order d -3 / 2 ) a deterministic observation of tr[ S ⋆ ] . By eq. (97), and repeating the arguments of the proof of Lemma F.1, we reach that again we have MMSE = κ E tr[( S ⋆ -ˆ S BO ) 2 ] as d → ∞ . Moreover:

Conditioning on Y amounts to condition on the value of tr( S ⋆ ) , as detailed above. Let us make two important remarks:

( i ) As d → ∞ , Y concentrates around its typical value E [ Y ] = m a . Since the MMSE is bounded, we therefore have as d →∞ that MMSE = κ E S ⋆ , { ˜ y i } (tr[( S ⋆ -ˆ S BO ) 2 ] | Y = m a ) .

( ii ) As we will see in what follows (and exactly like in the case of fixed second layer), the leading order of the MMSE of the inference problem of eq. (100) only depends on the asymptotic spectral distribution of S ⋆ . In particular, at leading order:

where in (a) we used that conditioning on tr( S ⋆ ) = m a does not change the asymptotic spectral distribution of S ⋆ .

All in all, we focus on characterizing the MMSE given in eq. (101), for the inference problem of recovering S ⋆ from the knowledge of { Z i , y i } generated by eq. (100).

G.1.2 Further steps of the derivation

Here, we notice that the arguments detailed in Section 4 on how to obtain an asymptotic expression of eq. (101) do not depend on the specific asymptotic spectral distribution of S ⋆ . More precisely:

Conjecture 4.1 can be directly extended to more general distributions of S ⋆ than the Wishart distribution. Indeed, the heuristic argument explaining this universality phenomenon does not depend on the distribution of S ⋆ , and on a technical level, as mentioned in the main text, Conjecture 4.1 is an extension of Corollary 4.10 of Maillard and Bandeira [2023], which holds for generic choices of distributions of matrices.

Conjecture 4.2 is also straightforwardly extended here, simply replacing the Wishart prior by the more generic prior of eq. (96). More generally, we expect it to hold for any prior such that the function Ψ(ˆ q ) of eq. (16) is well-defined [Aubin et al., 2019a, 2020].

Finally, the proof of Theorem 4.1 (see Appendix F.2) relies solely on the rotation invariance of the distribution of S ⋆ , as well as the fact that S ⋆ admits a compactly supported asymptotic eigenvalue distribution. These two facts hold for the distribution of eq. (96) for compactly supported P a , see e.g. Silverstein and Choi [1995], Lee and Schnelli [2016].

G.2 Conclusion: Claim 2 when learning the second layer

We are now ready to state the generalization of Claim 2 to a learnable second layer. The effective problem we consider is the recovery of a symmetric matrix S ⋆ ∈ R d × d , which was generated as S ⋆ = (1 /m ) ∑ m k =1 a ⋆ k w ⋆ k ( w ⋆ k ) ⊺ , from observations ( y i ) n i =1 , generated as

i i i d i d

The asymptotic spectral distribution µ ⋆ of S ⋆ is called a generalized Marchenko-Pastur distribution (or a free compound Poisson distribution: it is also the free multiplicative convolution of the MarchenkoPastur law and P a , see Anderson et al. [2010]). µ ⋆ is compactly supported, and can be characterized by its R transform [Marchenko and Pastur, 1967, Silverstein and Choi, 1995, Tulino and Verdú, 2004]:

Eq. (103) allows for an efficient numerical evaluation of µ ⋆ given P a . Notice that E µ ⋆ [ X ] = m a , and E µ ⋆ [ X 2 ] = m 2 a + c a /κ .

The partition function for the learning problem of eq. (102) is again defined as:

We then obtain the following generalization of Claim 2.

Claim 3. Assume that m = κd with κ > 0 , and n = αd 2 with α > 0 . Recall that m a := E P a [ a ] and c a := E P a [ a 2 ] . Let Q 0 := E µ ⋆ [ X 2 ] = m 2 a + c a /κ . Then:

The limit of the averaged log-partition function of eq. (104) is given by

where

Here, Σ( µ ) := E X,Y ∼ µ log | X -Y | , and, for t ≥ 0 , µ t := µ ⋆ ⊞ σ s . c ., √ t is the free convolution of µ ⋆ and a (scaled) semicircle law (see Appendix A).

For any α > 0 , except possibly in a countable set, the supremum in eq. (105) is reached in a unique q ⋆ ∈ [ m 2 a , Q 0 ] . Moreover, the asymptotic minimum mean-squared error on the estimation of S ⋆ , achieved by the Bayes-optimal estimator ˆ S BO := E [ S |{ y i , x i } ] , is equal to Q 0 -q ⋆ :

It is related to the MMSE of eq. (97) by MMSE = κ ( Q 0 -q ⋆ ) .

Therefore, generalizing Section F.7, Result 1 holds as well in this case, with ˜ ∆ = 2∆(2 + c a ∆) /κ , and µ t := µ ⋆ ⊞ σ s . c ., √ t , where µ ⋆ is characterized by eq. (103).

G.3 The GAMP-RIE algorithm

Finally, one can also generalize Algorithm 1 to this setting: the only change to perform is to adapt the functions F RIE and f RIE . Indeed, instead of denoising a Wishart matrix (with an asymptotic spectrum given by the Marchenko-Pastur distribution), here one must denoise a matrix S 0 with asymptotic spectral distribution given by µ ⋆ defined in Appendix G.2. As mentioned, eq. (103) allows for an efficient numerical evaluation of µ ⋆ given P a . From there, one can adapt Algorithm 1 to this case simply by replacing in the definitions of F RIE and f RIE the distribution ρ ∆ by ρ ∆ = µ ⋆ ⊞ σ s . c ., √ ∆ .

H Details on the numerics

0

.

0

0

.

1

0

.

2

0

.

3

Sample complexity

α

=

n/d

2

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

3

.

0

MSE

MMSE

λ

= 0

λ

= 0

.

001

λ

= 0

.

005

λ

= 0

.

01

λ

= 0

.

05

λ

= 0

.

1

λ

= 0

.

25

λ

= 0

.

5

0

1

2

Noise level ∆

0

.

20

0

.

25

0

.

30

0

.

35

Trivialisation threshold

α

T

H.1 Solutions to the 'state evolution' equations

We describe here how to solve eqs. (7),(8). The first step to solve is to obtain an analytical expression for µ t . We refer to Appendix A for the definition of quantities used in this section. We recall that µ t := µ MP ,κ ⊞ σ s . c ., √ t is the free convolution of the Marchenko-Pastur law and a scaled semicircular density. The R -transform of the scaled semicircle distribution is [Tulino and Verdú, 2004]:

while for the Marchenko-Pastur law we have

We can now use (cf. Appendix A):

The Stieltjes transform g ( z ) = E µ t [1 / ( X -z )] of µ t is the solution of the equation

Among all the solutions to this equation, g ( z ) must be such that Im[ g ( z )] > 0 if Im( z ) > 0 , and also satisfies g ( z ) ∼ 1 /z for z →∞ . Eq. (108) is a third degree polynomial in g ( z ) , and can easily be solved by algebraic solvers, and has a single solution satisfying the constraints we described. Finally, µ t ( x ) is given by the Stieltjes-Perron inversion theorem (see Appendix A):

or equivalently

and we numerically choose ε = 10 -8 . We now discuss the computation of the integral of µ t ( x ) 3 in (8). Notice that the integrand is only non-zero over at most two finite intervals. Exact values of the edges are given by setting the discriminant of equation (108) to zero. The last step is finding a solution in ˆ q to equation (8). We find the function 'root' in Scipy, which uses a variant of the Powell hybrid method, to be performing quite well when initialized in the value 2 α/Q 0 . This whole procedure is quite efficient and can be reproduced easily on any machine.

H.2 Gradient descent

In our experiments with gradient descent we are minimizing the objective R ( W ) :

All the simulations are done in PyTorch with the student weights initialized in the prior. For 'vanilla' gradient descent we iterate until convergence, and average over several repetitions. For averaged gradient descent (AGD) we first generate the dataset, then train the student several times with starting weights independently sampled in the prior, and "average the weights" at the end of training. By this we mean that for each run we train until convergence, then obtain the matrix S and average it. Finally, we average this procedure over several repetitions. The learning rate is chosen to be suitably large, as it's typically better to train a networks with giant steps [Dandi et al., 2023].

In Figure 2 the gradient descent is run for zero regularization, λ = 0 . In Figure 5 (left) we then study the effect of regularization to check whether regularization helps to achieve the Bayes-optimal error, but conclude that it does not and in fact it hurts the performance. In Figure 5 (right) we study the effect of the noise on the landscape of GD. We will expand on this in Appendix H.3. All the error bars reported in Figure 2 and Figure 5 (left) are standard deviations of the MSE measured on the samples. Figure 5 (right) has a finite resolution indicated in the caption. A single run of vanilla GD for the models we display can be completed in at most 30 minutes on an average machine without using GPUs. For producing our figures we used around 30 000 hours of computing time.

H.3 Additional experiments with GD

Here we study in more detail the phenomenology observed in Figure 2 (right) where in the presence of noise and at a large sample complexity all the runs of GD seem to converge to the same prediction. In the figure we noticed that above certain sample complexity the averaged and non-averaged GD errors are identical. This suggests that GD will eventually lead the weights of the network to the same configuration up to the symmetries of the problem independently of the initial state. We call this a trivialization of the landscape.

In Figure 5 (right) we study the trivialization threshold as a function of the noise level ∆ . One needs to take care of the symmetries on ˆ W , so we first define ˆ S :

where we mean that for a fixed dataset we run GD, then take a matrix product to obtain S . This procedure allows us to define the dispersion

If the dispersion becomes zero it means that all the runs will converge to the same value. As we increase the sample complexity α the dispersion decreases, until it becomes zero. For each value of the noise level ∆ we indicate the minimum sample complexity for which the dispersion is either less than 10 -2 , or less than 10 -3 of the maximum dispersion at fixed ∆ .

In Figure 5 (left), where we studied the effect of ℓ 2 regularization on the weights, we can also see how even a relatively small λ > 0 regularization leads to a trivialization of the landscape again in the sense that different initializations of GD provide the same prediction and averaging does not lead to a better error.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract summarizes the motivation behind our analysis, its relation to previous literature, as well as our main theoretical and experimental findings. We also outline the main techniques used in our theoretical derivations.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We have included a paragraph related to discuss the scope of our results, as well as their limitations, in Section 6.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We presented in this paper exact claims and conjectures based on well-known analytical methods. Specifically, we developed an application of the so-called replica method in this new context of learning an extensive-width neural network with a number of samples which is quadratic in the dimension. Further, we presented a mathematically sound derivation of our results, based on conjectures which are natural extensions of existing works. As mentioned in the text, a fully rigorous derivation of our results is a very technical and lengthy avenue, and is left for a more suitable mathematical venue.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Our results are fully reproducible, as the algorithms we use in our experiments are clearly explained (and their hyperparameters given), and the code used to produce our results is provided.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide the numerical code used to run all the experiments presented in the paper, as well as the data obtained by running these experiments, in a public GitHub repository.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The setting of our numerical experiments is fully detailed in the text alongside the related figures, and is clearly accessible in the numerical code provided.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provide error bars for our experiments (cf. e.g. Figure 2), and clarify how the error bars were computed in Appendix H.2.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: In Appendix H.1 and H.2 we provide an approximate description of the infrastructure needed to reproduce the figures.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We checked that our paper and our research complies with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our work is largely theoretical, and concerns the fundamental limits of learning with neural networks. Moreover, our numerical experiments are limited to synthetic datasets. As such, we do not believe our work to have societal impact besides the long-term impact brought by a better understanding of the theory of learning.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: As discussed above, this question is not relevant to the presented work.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: We only use synthetic data in our numerical experiments, and do not rely on any existing code.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide a documented code alongside the paper, as well as details on how we implemented the algorithmic procedures used in our experiments (e.g. gradient descent) in Appendix H.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nStudies 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)).\n\nSTRENGTHS:\n* I found the paper very accessible despite having little pre-existing knowledge of this line of research. \n* The context of the current work is illustrated clearly. I did not have to do too much background reading to appreciate what has been done, and how the current submission contributes to advancing this line of research\n* Claim 1, which is central to the main result, is shown using the replica method as well as a more rigorous method. \n* For the particular setting described in Section 2, this work derives the MMSE under the regime described in Equation (4). This is novel as far as I am aware. \n* The empirical comparison between Bayes-optimal and (non-stochastic) gradient descent reveals an intriguing hypothesis that averaged over random initialisations, GD is close to being Bayes-optimal\n\nWEAKNESSES:\n* In terms of exposition, I found it a bit awkward that the closed-form asymptotic limit of MMSE_d is given in Eqn 14, which is outside of Section 3 “Main theoretical result”. I think the part of Section 5 titled “Results for the Bayer optimal estimator” is ready to be presented at the beginning of Section 3. \n* The notation and setup for Claim 1, to my eyes, belongs to Section 4 “Derivation of the main results”\n* It’s hard for me to imagine how this type of analysis can be generalised to a more complex architecture than the one-hidden layer network presented in Eqn (2). And I’m not talking about very exotic architectures. I mean even a small generalisation to multiple hidden layers. Is this really tenable using the theoretical tools employed here?\n\nQUESTIONS:\n* Line 151 defines $MMSE=\\lim_{d \\to \\infty} MMSE_d$. But from Eqn (4), it seems that the asymptotic regime is more than just taking $d \\to \\infty$. I found this confusing. \n* Is Eqn (3) derived somewhere? If not, is it because it’s very elementary and can be found in a standard reference? Please give citation then. \n* Is it obvious that the GD estimator is the same as the ERM estimator in the discussion starting on Line 160?\n\nLIMITATIONS:\nSome of the limitations that come to mind have already been acknowledged in the submission. Specifically, there is a need to go beyond Gaussian input data, etc, see Pg. 9 Line 350", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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. \nHaving 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. \nTo 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. \nGiven 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.\n\nSTRENGTHS:\n- Equation (1) and its natural link to the degrees of freedom of the target function is an appealing result. \n- The connections with ellipsoid fitting, computations in random matrix theory, formalized results on neural networks, and extensive-rank matrix denoising validate it. In simple words: it is not a made-up paper, but it is well positioned in between earlier works. \n- Understanding industry standard is crucial. Gradient Descent is the idealized version of what is actually used in practice. It is therefore valuable to understand it, as the author(s) did. \n- The result about GD opens up to continuations of the work into such direction of understanding formally what is happening there. This paper is not a dead-end. \n- Section 4 is an interesting take on trying to vindicate the replica method. I believe there is value in trying to break down what would be needed to make the Physics derivation a mathematically rigorous one. At the same time, it is of independent interest to use the replica method, as it tells us what we have to prove. Anyways, nice idea. \n- I was kind of surprised to see that as $\\kappa$ increases the MMSE worsens. \n- The techniques used are, in my honest opinion, beautiful and effective. \n- Even if the MMSE cannot be simulated (unless an algorithm achieving it was known), the author(s) bring to the table experiments for GD that are to the point and motivate further work. \n- The references are, to my limited single human reading stack, to the point. \n- The weaknesses I will discuss below are to be understood in the sense that I understand that this is a paper with lots of Physics-flavour. Personally, I agree with the modelling idea of finding a solvable model, hence the choice of a \"simplified setting\". Also, this simplified setting already requires non-trivial computations, and is therefore worth exploring. Nevertheless, I will outline what worries me to stimulate the discussion. \n\nOverall, I believe this was a good read. Thanks to the author(s)!\n\nWEAKNESSES:\nAs mentioned in lines 350-353, the setting is very restricted. While universality might come in help, at least for conjecturing that some aspects extend nicely in the high dimensional limit, there are at least three features that in my opinion are more complicated. Below, I will elaborate on them. \n1. Perhaps the solvable one. You consider matching architectures, with $m = m^{\\star}$, I wonder, what is the technical challenge in analyzing the case in which the activations are the same but $m\\neq m^{\\star}$? \n2. Your target function is \"neuron agnostic\" in the sense that it takes a simple mean of the teacher neurons. How would your technique behave when we take a one hidden-neuron that also has random Gaussian weights $\\{a_k\\}_{k = 1}^m$ in the second layer? I have not thought this through extensively. \n3. The most challenging: a quadratic activation. While the Statistical Physics literature focuses a lot on quadratic activations/phase retrieval, which is thought to be an interesting problem, I am concerned about the extension in this setting to other activations. It appears to me that when recasting the task to a matrix estimation problem (lines 175-180), one heavily relies on \"expanding the square\" (lines 175). In other words, the symmetry of the activation function naturally builds the matrix $\\mathbf{S}^{\\star}$  (lines 176-177). Considering other (e.g. non-polynomial, or even non-square) activations, it looks to me as a very hard task. I do appreciate that the author(s) recognized this, but would be even happier if, in the spirit of section 4, the bonus page for publication were spent to discuss how a quadratic activation is not a dead-end. \n   More in depth, can you bring over an argument that can convince the reader that the method is worth pushing over the boundary of something that allows you to over-use Gaussian integrals (in the sense that square activations are perhaps the easiest non-linearity)?\n\nThe second concern is mostly structural and philosophical. I am convinced that all the results of the paper are of interest to the community. However, from my first reads, I had the impression that the \"Replica argument\", its \"formalization\" and the \"analysis of GD\" had little meaning if placed all inside a conference paper.\nFor example, while Section 4 is very useful for grounding the replica method, I only appreciate its value in the sense that it connects the dots with other works. In some sense, I find this counter-intuition that there are very few resources on Gaussian Equivalence and modern Statistical Physics of neural networks that discuss this, and then find a nice Section clearing out some aspects in a paper that should be 9 pages. I do wish at some point that these scattered nice sections will be condensated in some resource that will be accessible (e.g. a book). \nSimilarly, the analysis of GD could be justified for its industrial importance, and even simply because it is the quickest algorithm to compare with the MMSE prediction, but at the same time it poses the question: \"What is the soul of this work?\". \nPlease note that this weakness is more provocative than actually requiring action.\n\nQUESTIONS:\nCan you please make the legend in Figure 2 (right) more explicative? I obviously get it, but do not understand why you chose to \"condensate it\". As a reader, the others are easier to parse, as they label everything. \n- The regularization of Figure 3 is never discussed in text, why do you add it? Just to show that the distance from the MMSE becomes larger? What exactly counts to you as \"trivialize the gradient descent landscape almost completely\" in Figure 3 (Left), as you claim in lines 802-803? Please note that I understand the value of Figure 3 (right).\n- Related to the point above, why do you seek to show regularized objectives in the first place? I mean, what motivates your experiment? \n- Please, see the first weakness, where I planted some questions. \n### Typos\nI will write down in this subsection what I noticed by reading. Take it as a friendly contribution. All of them are largely not relevant to the evaluation of the paper.\n- The instruction block for the checklist was not deleted. \n- (line 161) you write \"student\", but I noticed that in other instances you used a different way of placing quotation marks, that is the TeX standard. \n- Equation (5) lacks punctuation.\n- (line 222) the end of the sentence \"using probabilistic technical amenable to rigorous treatment\" probably changed wording, maybe \"technical\"was meant to be \"techniques\"?\n- (line 779) the \"root\" quotation marks. \n- (line 792) is $30,\\,000$ really the right number? Why three zeros? Maybe I am just not used to this standard.\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-03"}

comments: SUMMARY:
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.

STRENGTHS:
This seems to be a simple and elegant example of learning for a well-controlled class of high-dimensional problems.

WEAKNESSES:
This is a rather technical result far removed from obvious useful applications; this is not a criticism in itself — we are all aware that Neurips is moving into the territory of COLT etc. However, in a generalist conference it is hard to communicate the significance of a technical result like this. This is not my area, so I am largely taking the author's assertions about the significance and novelty of this result at face value.

QUESTIONS:
As a non-specialist in this class of problems, my questions are simply why we are concerned with this class of problems? The authors mention connections to phase retrieval and matrix denoising problems; can we give a concrete example, however contrived, that motivates this work by its potential to eventually contribute to a concrete application in those domains,

LIMITATIONS:
The author's explanations of the limitations are absolutely clear as far as I can see.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe paper is robust in its theoretical analysis, providing clear mathematical derivations and a solid theoretical framework. The breadth of related works cited enriches the context and demonstrates the paper’s relevance to current research frontiers.\n\nWEAKNESSES:\nThe motivation for focusing on Bayes-optimal accuracy with quadratic activation remains under-explained, which may leave readers uncertain about the broader applicability of the results. Furthermore, the implications of these findings for common neural network training techniques using ReLU or sigmoid activations are not discussed, which could limit the paper’s appeal to a broader audience.\n\nQUESTIONS:\n1. Given the prevalence of activation functions like ReLU and sigmoid in practical applications, what prompted the choice of quadratic activation for this study?\n\n2. In Lines 142–144, the term $P_{prior}(W)$ is mentioned as a prior on the teacher weights $W^*$. Shouldn’t this be $P_{prior}(W^*)$ instead? Could you clarify whether $W$ refers to teacher weights or student weights?\n\n3. How does the Bayes-optimal accuracy relate to the practical performance of gradient descent (GD)? Does high Bayes-optimal accuracy imply high accuracy for GD?\n\n4. Could you explain the origin and significance of the term $\\Delta(2+\\Delta)$ in Eq. (3)?\n\n5. In Line 128, there appears to be an extraneous 'm' following the definition of $D$.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}