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

Authors: antoine  maillard
emanuele  troiani
simon  martin
lenka  zdeborová
florent  krzakala
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 17:27:42

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.

Main paper: 
 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 activations1 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 whereEQUATIONwith κ = 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 Damian et al. [ , 2024]] , Collins-Woodfin et al. [2023] , Dandi et al. [2023 Dandi et al. [ , 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]  
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=1are obtained by a one-hidden layer teacher network with m hidden units and quadratic activation:y i = f W * (x i ) := 1 m m k=1 1 √ d (w * k ) ⊺ x i + √ ∆z i,k 2 .(2)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 weightsW * = (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:P(W|D) = 1 Z(D) P prior (W)P(y|W, {x i } n i=1 )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 overEQUATIONWe 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:EQUATIONWe 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 trueW ⋆ (i.e. if ŷ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:EQUATIONIn 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 7) for the noiseless (∆ = 0) case, as a function of the sample complexity α, for various width ratios κ. Right: Phase diagram representing the MMSE, brighter color indicates a higher value. The red curve is the perfect recovery transition line α PR , see eq. ( 1), and its origin is discussed in Section 5.EQUATIONNote 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 (0) , D) the estimator for test labels reads ŷGDW (0) ,D (x test ) := f Ŵ(W (0) ,D) (x test). As we will see, it will be interesting to consider also an estimator ŷAGD obtained by averaging the GD estimator on the labels over the initializations W (0) of the weights. ∼ 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:EQUATION)where ∆ := 2∆(2 + ∆)/κ, and where q is a solution of the following equation:(1 -2α) + ∆q 2 = 4π 2 3q µ 1/q (y) 3 dy. (8)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 ofS ⋆ := (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 :=EQUATIONwith ξ ∼ N (0, 1), ∆ := 2∆(2 + ∆)/κ, and where we definedZ := (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 Ŝopt = E[S| y, Z] the Bayes-optimal estimator related to the problem of eq. ( 9), then MMSE = κEtr[(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 asEQUATIONwithZ 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:EQUATION)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 + nE 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 byEQUATIONwhereEQUATIONHere, Σ(µ) := 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 uniqueq ⋆ ∈ [1, Q 0 ]. Moreover, the asymptotic minimum mean-squared error on the estimation of S ⋆ , achieved by the Bayes-optimal estimatorEQUATION)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 ⋆ ŜBO )] according to Claim 2. As the MMSE decreases with α, it is clear that q ⋆ ≥ q ⋆ (α = 0). When α = 0, we have Ŝ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). 
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, ifZ 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 defineEQUATIONwherey ′ i ∼ P out (•|Tr[G i S ⋆ ]), with S ⋆ ∼ W m,d and G i i.i.d.Thenlim d→∞ |Φ d -Φ (G) d | = 0.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 [Aubin et al., , 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 havelim d→∞ Φ (G) d = sup q∈[1,Q0] inf q≥0 (Q 0 -q)q 4 + Ψ(q) + α R×R dyDξJ q (y, ξ) log J q (y, ξ) ,whereEQUATION)is the asymptotic free entropy of the matrix denoising problem Y = √ qS ⋆ + ξ, 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)EQUATION)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 conjecture2 of Semerjian [2024] states that the free entropy of matrix denoising ofS ⋆ = (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 dlimit of E tr[(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] .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. 
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)∂MMSE ∂α αPR =      -2 - 4 κ + 12 1 + κ if κ ≤ 1, -2 + 2 κ if κ ≥ 1. 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 GD (d = 200) AGD (d = 200) GD (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 )Figure 2 : Mean squared error (MSE) as a function of the sample complexity α for κ = 1/2. Dots are simulations using GD with a single initialization averaged over 32 realizations of the dataset, crosses are averages over 64 initializations with 2 realizations of the dataset. The continuous lines are the asymptotic MMSE given by ( 7). Left: noiseless ∆ = 0 case. The colors indicate the size d. We can see how AGD appears to be well described by the theoretical MMSE. We used the learning rates 0.2 for d = 200 and 0.07 for d = 100. Right: Comparison of GD between the noisy √ ∆ = 0.25 case (red) and noiseless ∆ = 0 case (blue). Adding noise makes AGD worse than the MMSE, and for sample complexity α ≳ 0.3, all the initializations of GD converge to the same point, making the GD and AGD curves collapse.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.We also 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 ⋆ -Ŝ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. 
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. 
 
 Appendix: 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,dmatrices 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 densityEQUATION)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 densityEQUATION)Hereλ ± := (1 ± κ -1/2 ) 2 .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µ((a, b)) = lim δ↓0 lim ϵ↓0 1 2iπ b-δ a+δ [g µ (x + iϵ) -g µ (x -iϵ)]dx.In particular, if µ has a continuous density with respect to the Lebesgue measure then:∀x ∈ R, dµ dx = lim ϵ↓0 1 π Im g µ (x + iϵ).We often use the logarithmic potential function Σ(µ) := µ(dx)µ(dy) log |x -y|. We further define the R-transform of µ as:EQUATION)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:EQUATION)Algorithm 1: GAMP-RIEThe estimator Ŝ Input: Observations y ∈ R n and "sensing vectors"Z i := (x i x ⊺ i -I d )/ √ d ∈ R d×d ;Initialize Ŝ0 ∼ W m,d and ĉ, ω, V randomly; while not converging do• Estimation of the variance and mean of Tr[Z i Ŝ];V t = ĉt and ω t i = Tr[Z i Ŝt ] -g out (y i , ω t-1 i , V t-1 )V t ;• Variance and mean of S estimated from the "channel" observations;A t = 2α n n i=1 g out (y i , ω t i , V t ) 2 and R t = Ŝt + 1 dA t n i=1 g out (y i , ω t i , V t )Z i ;• Update of the estimation of S ⋆ with the "prior" information;Ŝt+1 = f RIE R t , 1 2A t and ĉt+1 = 2F RIE 1 2A t ; t = t + 1; endFree 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.):EQUATIONEq. ( 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).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)):EQUATIONwithZ 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] .The functions g out , f RIE and F RIE appearing in Algorithm 1 are defined as follows. First, we letEQUATION)In particular, for the problem of eq. ( 9), we haveg out (y, ω, V ) = y -ω ∆ + V .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), thenf RIE (R, ∆) = Uf ∆ (Λ)U ⊺ , where f ∆ (λ) = λ -2∆h ∆ (λ), with h ∆ the Hilbert transform of ρ ∆ . More precisely: h ∆ (λ) := P.V. 1 λ -t ρ ∆ (t)dt.ρ ∆ 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] ):EQUATIONIn 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.3We 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[( Ŝt ) 2 ], andm t := tr[ Ŝ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 ĉt = 2(Q 0 -q t ). If ω, z are centered and jointly Gaussian variables with E[ω 2 ] = 2q t , E[z 2 ] = 2Q 0 , and E[ωz] = 2m t = 2q t , and y ∼ P out (•|z), we defineEQUATIONso that A t = qt /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 replacingTr[Z i Ŝt ] by Tr[G i Ŝt ] in the update of ω t i )shows that qt satisfies the equation: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] .We have, 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:EQUATIONwith ξ ∼ GOE(d). Heuristic details on how to derive eq. ( 28) can be found again in Zdeborová andKrzakala [2016] , see Section 6.4.13 there. By definition of F RIE , this implies q t+1 := tr[( Ŝt+1) 2 ] = Q 0 -F RIE ((q t ) -1 ), so that by eq. ( 25):EQUATIONwith µ 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 GAMP-RIE 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:EQUATIONLet 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 = √ 2S 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. • vt (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• Σ -1 t := -2αEdiv[g out (t, ω t , V t )] is related to A t by A t = Σ -1 t . • V t := 2F RIE (Σ t /2).u t = vec(R t /Σ t ): vt (u t , Σ t ) := vec[f RIE (R t , Σ t /2)].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[ Ŝt S ⋆ ] and qt 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 qt 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] ).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.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):y i = 1 m m k=1 (w ⋆ k ) ⊺ x i √ d p ,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         T ⋆ := 1 m m k=1 (w ⋆ k ) ⊗p , X i := 1 d p/2 x ⊗p i .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] . 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, andS(W) = W ⊺ W/m. Z(S ⋆ , {x i } n i=1 ) = R m×d DW n i=1 P out (y i |Tr[Z i S(W)]) .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:EQUATION)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 covarianceE[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 dy 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 matrixEQUATIONthatEQUATIONNotice that the CLT-based argument above is made formal in Conjecture 4.1, and implies the universality of Φ d under the replacement of 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 → α):Z i by Gaussian GOE matrices G i . Since Q ∈ R (r+1)×(EQUATIONwhere S + r+1 is the set of positive semi-definite symmetric matrices of size r + 1, and:EQUATIONNotice 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:EQUATION)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:EQUATIONwith 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:EQUATION)The term J(Q) -Using the replica-symmetric ansatz of eq. ( 37) in eq. ( 36), we get:J(Q) = inf Q,q (r + 1)(Q Q -rq q) 4 + 1 d 2 log r a=0 DW a e - d( Q+ q) 4 a Tr[(S a ) 2 ]+ d q 4 Tr ( a S a ) 2 .We now use the following Gaussian integration identity, for any symmetric matrix M:E ξ∼GOE(d) e d 2 Tr[Mξ] = e d 4 Tr[M 2 ] .This allows to reach the following expression, which is analytic in r:EQUATION)Recall that here S = S(W) = W ⊺ W/m.The limit r → 0 -From eqs. ( 34), ( 38) and ( 39), we have:EQUATIONThis implies that Q = 0 andQ = 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:lim d→∞ Φ d = sup 0≤q≤Q0 inf q≥0 - q q 4 + α R 2dy DξJ q (y, ξ) log J q (y, ξ) (41)EQUATIONEQUATIONIn order to obtain from eq. ( 41) the prediction of eq. ( 12), it therefore suffices to show that, for any q ≥ 0:EQUATION)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:EQUATIONwith ξ ∼ 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 asEQUATIONCrucially, 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 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:EQUATIONin whichS = 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 DO to denote here the Haar measure over the orthogonal group O(d). The large-d limit of the last term is given by the HCIZ integral [Harish-Chandra, 1957, Itzykson and Zuber, 1980] :EQUATIONwhere 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):EQUATIONwhere 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] :- 1 2 Σ(µ t ) + 1 4t E µt [X 2 ] - 1 2 I HCIZ (t -1 , µ t , ν) - 3 8 + 1 4 log t + 1 4t E ν [X 2 ] = 0,with Σ(µ) := µ(dx)µ(dy) log |x -y|. Applying this formula with t = 1/q we reach:1 2 I HCIZ (q, µ MP,κ , µ Y ) = - 1 2 Σ(µ Y ) + q 4 E[tr(Y 2 )] - 3 8 - 1 4 log q + q 4 Etr[(S ⋆ ) 2 ], = - 1 2 Σ(µ Y ) + 1 4 (2Q 0 q + 1) - 3 8 - 1 4 log q.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-κ limitWe 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 7), which we compare with the asymptotic κ → 0 curve obtained by eq. ( 51). We emphasize that the horizontal axis is α/κ, which remains of order Θ(1) as κ → 0: it corresponds to a number of samples n of the same order as the number of parameters dm.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:EQUATIONIn particular, in the noiseless case (∆ = 0), we have:EQUATION)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.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 havet = κ 2 q = 1 -q + Λ 2 α .Moreover, eq. ( 8) at order O(κ) yields:EQUATIONwhereF (t, κ) := 4π 2 t 3κ 2 µ t/κ 2 (y) 3 dy.Notice that F (t, 0) = 1 since µ ξ ≃ σ s.c., √ ξ for ξ → ∞, and σ s.c., √ ξ (y) 3 dy = 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)/(yz)dy of ν satisfies the selfconsistent equation (see Appendix A):EQUATIONMoreover, 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 √ κ. ν κ canthus 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 conditionEQUATIONis 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., 54) is equivalent to t ≤ 1. In this case, these "spiked" eigenvalues are located around the value [Benaych-Georges and Nadakuditi, 2011 ]√ t (2 √ t) = -t -1/2 , eq. (g -1 s.c., √ t (-1) = R s.c., √ t (1) + 1 = 1 + t.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 haveEQUATION)Furthermore, we also have, for all y ∈ R, when t ≤ 1:EQUATIONfor a finite density ρ (1) , with ρ (1) (y)dy = 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:EQUATION)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) (dy) = -1{t ≤ 1}, which implies that the normalization condition ν κ (y)dy = 1 is well satisfied for all values of t ≥ 0. Using the expansion of eq. ( 57), we obtain thatF (t, κ) = 4π 2 t 3 ν κ (y) 3 dy, = 1 -κ 2 -t if t ≤ 1, 1/t if t ≥ 1 + o(κ).So finally eq. ( 52) becomes2 α - Λ t = 2 -t if t ≤ 1, 1/t if t ≥ 1,And recall that MMSE = 2 αt -Λ, so thatMMSE = t(2 -t) if t ≤ 1, 1 if t ≥ 1, 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 layerWe consider here the noiseless case:y i = 1 m m k=1 (w ⋆ k ) ⊺ x i √ d 2 ,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:EQUATIONwith P out (y|z) = δ(y -∥z∥ 2 /m). We can make a direct use of the results of Aubin et al. [2019b] to write:EQUATION                           I P := R m Dξ R m Dw 0 exp - 1 2 (w 0 ) ⊺ qw 0 + ξ ⊺ q1/2 w 0 × log R m Dw 0 exp - 1 2 w ⊺ qw + ξ ⊺ q1/2 w , I C := ∞ 0 dy R m Dξ R m DZ 0 P out y|(I m -q) 1/2 Z 0 + q 1/2 ξ × log R mDZP out y|(I m -q) 1/2 Z + q 1/2 ξ .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 haveEQUATIONwhere 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:EQUATION)We now focus on computing the leading order of I C in the large-m limit. We can write       I C= dy Dξ I q (y, ξ) log I q (y, ξ),EQUATION) Let y := √ m[y -tr(I m -q) -(ξ ⊺ qξ)/m].We can change variables in eq. ( 62), and obtain:EQUATION)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 large-m limit. We haveJ q ( y, ξ) = R m DZ δ y - √ m Z ⊺ (I m -q)Z m -tr(I m -q) + 2 Z ⊺ (I m -q) 1/2 q 1/2 ξ m , = du 2π e iu y+iu √ mtr(Im-q) DZe -iu √ m Z ⊺ (Im-q)Z m +2 Z ⊺ (Im-q) 1/2 q 1/2 ξ m , = du 2π e iu y+iu √ mtr(Im-q)-1 2 log det Im+2 iu(Im-q) √ m -2u 2 m ξ ⊺ q 1/2 (Im-q) 1/2 Im+ 2iu(Im-q) √ m -1 (Im-q) 1/2 q 1/2 ξ , = du 2π e iu y-u 2 tr[(Im-q) 2 ]-2u 2 m ξ ⊺ q 1/2 (Im-q)q 1/2 ξ+O(1/ √ m) , = 1 2πσ 2 ξ e - ( y) 2 2σ 2 ξ + O(1/ √ m),whereσ 2 ξ := 2tr[(I m -q) 2 ] + 4 m ξ ⊺ q 1/2 (I m -q)q 1/2 ξ.Plugging it back into eq. ( 63) yields:I C = dy Dξ 1 2πσ 2 ξ e - ( y) 2 2σ 2 ξ - 1 2 log 2πσ 2 ξ - y 2 2σ 2 ξ + O(1/ √ m), = - 1 2 Dξ log[2πσ 2 ξ ] - 1 2 + O(1/ √ m).Since ξ ∼ N (0, I m ), it follows from elementary concentration of measure that σ 2 ξ concentrates on its average value σ 2 given by:σ 2 := 2tr[(I m -q) 2 ] + 4tr[(I m -q)q] = 2tr[(I m -q)(I m + q)] = 2tr[I m -q 2 ].All in all we reach that (up to additive constants):EQUATION)Combining eqs. ( 61) and (64) in eq. ( 59), we get at leading order in m, with Φ := lim(1/d)E log Z:EQUATIONEq. ( 65) can be easily solved, and yields:   q = q(I m -q) -1 , q = 2 α tr[I m -q 2 ] q.This implies that (recall 0 ⪯ q ⪯ I m ):q =            0 if α ≤ 1 2 , (2 α -1)I m if 1 2 ≤ α ≤ 1, I m if α ≥ 1.(66)Now that we have obtained q in eq. ( 66), we can compute the MMSE, or generalization error. Defining it as in eq. ( 4):MMSE d := m 2 E W ⋆ ,D E ytest,xtest [(y test -ŷBO (x test )) 2 ],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 orderMMSE d = m d Etr[(S ⋆ -S BO ) 2 ] = 1 - m d Etr[(S BO ) 2 ],withS := (1/m) m k=1 w k w ⊺ k .Notice that S BO = ⟨S⟩, so thatMMSE d = 1 - 1 m E 1≤k,l≤m w ⊺ k w ′ l d 2 ,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 :MMSE = 1 - 1 m E 1≤k,l≤m q 2 kk ′ = 1 -tr[q 2 ].Combining it with eq. ( 66), we reach:MMSE =            1 if α ≤ 1 2 , 4 α(1 -α) if 1 2 ≤ α ≤ 1, 0 if α ≥ 1.We have recovered eq. ( 51) from the limit m → ∞ taken after d → ∞!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 + ∆)):1 -2α + Λt = lim κ→∞ 4π 2 3κt µ 1/[κt] (y) 3 dy, = 1 1 + t .Combining it with eq. ( 7) yields thatEQUATION)where we recall Λ = ∆(2 + ∆). In particular, for ∆ = 0, we reach MMSE = max(1 -2α, 0), coherently with the behavior shown in Fig. 1 .F.1 Properties of the MMSE of eq. ( 4)Let S ⋆ := (1/m) m k=1 w ⋆ k (w ⋆ k ) ⊺, and Ŝ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: MMSE d -κE S ⋆ ,D tr S ⋆ - ŜBO 2 ≤ C(κ) n .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 Ŝ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ŜBO = E[S ⋆ ] = I d , so that Etr[(S ⋆ - ŜBO ) 2 ] = κ -1 . Thus, we have MMSE d → 1 for α = 0.Proof of Lemma F.1. -Notice that (cf. eq. ( 2)):E z [f W (x)] = ∆ + x ⊺ Sx d , 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 )):EQUATIONwhere we used E[∥z∥ 4 ] = m 2 + 2m and Etr(S ⋆ ) = 1 in (a), andE 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( ŜBO ) is the Bayes-optimal estimator for tr(S ⋆ ), i.e.EQUATION)In particular, considering the estimatorr(D) := 1 n n i=1 (y i -∆), = 1 n n i=1 x i S ⋆ x i d + ∆ ∥z i ∥ 2 m -1 + 2 √ ∆ m m k=1 z i,k x ⊺ i w ⋆ k √ d ,we have using eq. ( 69):EQUATION≤ 3[EQUATIONusing the Cauchy-Schwarz inequality in (a), with                           I 1 := E   tr S ⋆ 1 n n i=1 x i x ⊺ i -I d 2   , I 2 := ∆ 2 E n i=1 ∥z i ∥ 2 nm -1 2 , I 3 := 4∆E   1 nm n i=1 m k=1 z i,k x ⊺ i w ⋆ k √ d 2   .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):EQUATIONCombining eqs. ( 70) and ( 71), and plugging it back in eq. ( 68), we getMMSE d -κE S ⋆ ,D tr S ⋆ - ŜBO 2 ≤ C(κ) n ,which ends the proof of Lemma F.1.First, we note that Theorem 1 of Pourkamali et al. [2024] implies that:EQUATIONand 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:EQUATIONMoreover, 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 + √ tW, 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):EQUATIONIn other words, the function J(ν; µ) is the large deviations rate function (in the scale d 2 ) for the empirical spectral measure of R + √ tW, 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):EQUATIONCombining 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.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 asEQUATIONin which µ t = µ MP,κ ⊞ σ s.c., √ t , see Appendix A.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 variablet dy µ t (y) 3 = dz [ √ tµ t ( √ tz)] 3 .It is clear that for all x ̸ = 0, we have µ t (x) → κν MP,κ (x) as t → 0, and ν MP,κ (y) 3 dy < ∞, 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:EQUATIONWe 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/(Xy)] is the unique solution in C + to the equationy = 1 1 + g/κ - 1 g -tg.Since y = √ tz, it is clear that g = O(1/ √ t), and letting h := √ tg, we easily get the expansionz = - 1 -κ h -h + O( √ t),which can be inverted toEQUATIONNotice 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:EQUATIONComing back to eq. ( 76) this implies:EQUATIONWe 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] .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, anddy µ MP,κ (y) 3 = 3 4π 2 κ 2 κ -1 , so that 3(1 -2α PR ) 4π 2 = lim t↓0 t dy µ t (y) 3 = 0,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 thresholdHere, 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 :EQUATIONWe thus compute the next order of the expansion of t µ t (y) 3 dy as t → 0.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:EQUATIONOn the other hand, we have around the singularity at y = 0:EQUATION)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):EQUATIONfor any |z| ≤ 2 √ 1 -κ, while √ tµ t ( √ tz) = O(t 3/2 ) if |z| > 2 √ 1 -κ.This then yields:EQUATIONCombining eqs. ( 82) and ( 85) in eq. ( 81), we obtain (recallα = α PR = κ -κ 2 /2): ∂MMSE ∂α PR = -2 - 4 κ + 12 1 + κ .F.4.2 The case κ ≥ 1Again, 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 t dy µ t (y) 3 = t dy µ MP,κ (yEQUATIONPlugging it in eq. ( 81), we get in this case:∂MMSE ∂α PR = -2 + 2 κ .Again, the specific case κ = 1 can be tackled by continuity, as the derivative tends to 0 both as κ ↑ 1 and κ ↓ 1.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.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 ) asEQUATIONwhere(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 decomposeEQUATION)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:P[|tr(S ⋆ ) -1| ≥ t] ≤ 2 exp -Cd 2 min(t, t 2 ) ,where C > 0 depends only on κ > 0. In particular,P[|I 1 | ≥ d -1/4 ] ≤ 2 exp(-C √ d),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):EQUATIONwithξ 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:EQUATIONfor 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 dyDξ 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 2In 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[-(yz) 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)):R×R dyDξ J q (y, ξ) log J q (y, ξ) = - 1 2 log[ ∆ + 2(Q 0 -q)].We then reach that q = q ⋆ is characterized as the maximum of the following function:EQUATIONRecall that here µ t := µ MP,κ ⊞ σ s.c., √ t . It is known (see eqs. (77-78) of Semerjian [2024] e.g.) that ∂Σ(µ t ) ∂t = 2π 2 3 µ t (y) 3 dy.Thus, q = q(q) can be characterized as the solution4 toEQUATIONBy eq. ( 91), q is a solution in [1, Q 0 ] to:EQUATION)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 α → 0In 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 = Etr[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 ∼ q0 α as α → 0, we getEQUATIONwhere F (p) := (4π 2 /3) [p -1/2 µ 1/p (z • p -1/2 )] 3 dz. 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:z = κ √ p κ + h √ p - 1 h -h.As p → 0, we can thus compute the expansion of h p (z) in powers of p. Applying then the Stieltjes-Perron inversion theorem (Theorem A.2), we get the expansion of ν p (z) in powers of p as:ν p (z) = √ 4 -z 2 2π + √ p 3z √ 4 -z 2 8π 3 - 3p(2 -z 2 )(4 + κ -z 2 ) 8π 3 κ √ 4 -z 2 + O(p 3/2 ),for |z| ≤ 2, and ν p (z) = O(p 3/2 ) for |z| ≥ 2. Plugging this expansion into F (p), we get:F (p) = 1 - p κ + o(p).Coming back to eq. ( 94), this gives q0 = 4κ/[2 + ∆κ], and (recall Q 0 = 1 + κ -1 ) then q = 1, so that our equations are indeed consistent in the limit α → 0.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 ofEQUATION)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 derivationWe first discuss the reduction to a matrix estimation problem, generalizing Section F.5 to this setting. We defineEQUATIONand we denote m a := E Pa [a] and c a := E Pa [a 2 ]. We define the MMSE as (notice the additional factor c a with respect to eq. ( 4)):EQUATION)By repeating the (mathematically informal) arguments of Section F.5 to this setting, we find that, at leading order as m, d → ∞:EQUATIONwith ξ i i.i.d.∼ N (0, 1), and ∆ := 2∆(2 + ∆c a )/κ. We let             y i := √ d   y i - 1 n n j=1 y i   , Y := 1 n n i=1 y i .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 haveEQUATIONConditionally 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 scale5 O[(nd) -1/2 ] = O(d -3/2 ). Since y i = √ d[y i -Y ],this implies that at leading order we haveEQUATIONThe 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 = κEtr[(S ⋆ -ŜBO ) 2 ] as d → ∞. Moreover:MMSE = κE S ⋆ ,Y,{ yi} tr[(S ⋆ - ŜBO ) 2 ], = κE Y [E S ⋆ ,{ yi} (tr[(S ⋆ - ŜBO ) 2 ]|Y )].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 ⋆ ,{ yi} (tr[(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:EQUATIONwhere 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).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:A. 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. B. 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 [Aubin et al., , 2020]] . C. 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] .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 asEQUATIONwithZ i := (x i x ⊺ i -I d )/ √ d and x i i.i.d.∼ N (0, 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 Marchenko-Pastur 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] :EQUATIONEq. ( 103) allows for an efficient numerical evaluation of µ ⋆ given P a . Notice that E µ ⋆ [X] = m a , andE µ ⋆ [X 2 ] = m 2 a + c a /κ.The partition function for the learning problem of eq. ( 102) is again defined as:EQUATIONWe 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 Pa [a] andc a := E Pa [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 byEQUATIONwhereEQUATIONHere, Σ(µ) := 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 ŜBO :EQUATION)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).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 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 ]: R σ s.c., √ t (z) = zt, while for the Marchenko-Pastur law we haveR µMP,κ (z) = κ κ -z .We can now use (cf. Appendix A):EQUATIONAmong 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):EQUATIONand 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.In our experiments with gradient descent we are minimizing the objective R(W):EQUATIONAll 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.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 Ŵ, so we first define Ŝ:Ŝ(W (0) , D) := 1 m Ŵ(W (0) , D) ⊺ Ŵ(W (0) , D),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 δ GD := E D tr E W (0) Ŝ(W (0) , D) -Ŝ(W (0) , D)2 .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.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.• 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. 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. 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?38th Conference on Neural Information Processing Systems (NeurIPS 2024).Note that for linear activations, the target functions reduces to linear regression and can be learned from linearly many samples.We mention here the "strong" conjecture ofSemerjian [2024]. A weaker form of this conjecture is the universality of the best low-degree polynomial estimator for any i.i.d. prior.Section VI.D.1 in the arXiv version ofZdeborová and Krzakala [2016].Notice that one can show that q is the minimizer of a convex function in eq. (91). This can be shown e.g. by recalling the relationship of this function to the free entropy of a matrix denoising problem (Theorem 4.1) and using the I-MMSE theorem. We refer toBarbier et al. [2019],Maillard et al. [2020] for more details.Recall that tr[(S ⋆ ) 2 ] = O(1) with high probability. 


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


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

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


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

Reviews: review: paper_summary: paper_summary: The paper explores Bayes-optimal learning of a neural network with extensive input dimensions and a single hidden layer using quadratic activations. It presents a closed-form expression for the optimal test error when the sample complexity is quadratic. This work connects to matrix denoising and ellipsoid fitting, providing both theoretical insights and empirical validations. Key findings include showing that randomly initialized gradient descent can approach the Bayes-optimal error.
strength_weakness: strength_weakness: The 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.
questions: questions: 1. Given the prevalence of activation functions like ReLU and sigmoid in practical applications, what prompted the choice of quadratic activation for this study?

2. 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?

3. 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?

4. Could you explain the origin and significance of the term $\Delta(2+\Delta)$ in Eq. (3)?

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

The 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.
For 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). 
Similarly, 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?". 
Please note that this weakness is more provocative than actually requiring action.
questions: questions: Can 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. 
- 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).
- Related to the point above, why do you seek to show regularized objectives in the first place? I mean, what motivates your experiment? 
- Please, see the first weakness, where I planted some questions. 
### Typos
I 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.
- The instruction block for the checklist was not deleted. 
- (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. 
- Equation (5) lacks punctuation.
- (line 222) the end of the sentence "using probabilistic technical amenable to rigorous treatment" probably changed wording, maybe "technical"was meant to be "techniques"?
- (line 779) the "root" quotation marks. 
- (line 792) is $30,\,000$ really the right number? Why three zeros? Maybe I am just not used to this standard.
limitations: limitations: None.
paper_summary: paper_summary: Studies a teacher-student setting with very specific configurations of input data, teach weight distribution, noise distribution and teacher/student architecture. The minimum MSE achievable, Eqn (3), is subjected to a particular limit, Eqn (4), and shown in closed-form to be Eqn (14). From Eqn 14 and assuming the noiseless case., we can derive the "perfect-recovery" threshold, which is stated as one of the main contributions (Eqn (1)).
strength_weakness: strength_weakness: * 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. 
* The notation and setup for Claim 1, to my eyes, belongs to Section 4 “Derivation of the main results”
* 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?
questions: questions: * 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. 
* 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. 
* Is it obvious that the GD estimator is the same as the ERM estimator in the discussion starting on Line 160?
limitations: limitations: Some 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

Comments: title: Thank you!
Comment to rebuttal
comment: Thank you for answering my questions. I keep my original score.
Thanks to authors for their detailed response. I maintain my original rating.
I am embarrassed at the quality of my review. I had intended to return and expand upon it before the deadline, but did not. I apologies to the authors and the AC for this low-quality feedback.

The main difficulty is that this paper is relatively far from my area of expertise and I do not feel qualified to assess its relationship to the broader field, which seems for this paper to be the crucial matter. I will keep the confidence of my review low in the hope that more-informed readers will provide more meaningful feedback.
Dear author(s), 
thank you for your explanatory comments. In this comment, I also acknowledge the general one you wrote for all reviewers. 

## General comments
Everything understood: as long as you state in your limitations the points you mention in the way you mention them here, we are done with these questions on my side. 

## Specific comments
- W1 & W2. Great to hear this. I am curious to see what the final version will be. 
- Q1-Q2-Q3. Thank you for this. For the trivialization, ok. 
- thank you for the typos. 
- as for the last comment, yes, it was more of a rant. Obviously out of scope. 
- as per the "soul of this work", I do stand with the idea that this is far more than a 9 page conference paper, but this is not a tragedy. 


Considering the points made above, I have raised my score according to the NEURIPS24 guidelines. I will keep an eye on the discussion with other reviewers and potentially update my score if other thresholds are surpassed.