Bayes-optimal learning of an extensive-width neural network from quadratically many samples
Authors: Antoine Maillard
Emanuele Troiani
Simon Martin
Florent Krzakala
Lenka Zdeborova
Venue: NeurIPS 2024 poster (2024)
Keywords: Theory of neural networks
Bayes-optimal learning
non-convex optimization
statistical physics
high-dimensional statistics
Primary area: learning_theory
TL;DR: We show how to compute exact asymptotic predictions for the minimal generalization error reachable in learning an extensive-width shallow neural network, with a number of data samples quadratic in the dimension.

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

Bayes-optimal learning of an extensive-width neural network from quadratically many samples
Antoine Maillard
Emanuele Troiani
Department of Mathematics ETH Zürich, Switzerland
Statistical Physics Of Computation Laboratory EPFL, Switzerland
Simon Martin
Lenka Zdeborová
INRIA & Laboratoire de Physique ENS, Université PSL, France
Statistical Physics Of Computation Laboratory EPFL, Switzerland
Florent Krzakala
Information Learning and Physics Laboratory EPFL. Switzerland
Abstract
We consider the problem of learning a target function corresponding to a single hidden layer neural network, with 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 eтror to learn such 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 the more interesting regime where the number of samples quadratic the dimension. In this paper, we solve this challenge for quadratic activations and derive a closed-fore 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, 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-initaaizee gradient descent seems to sample the space weights, leading training loss, and averaging initialization leads to test error equal to the Bayes-optimal one.
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 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 detailed understanding of these fundamental limitations will eventually allow for 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
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

sample-come complexity in the so-called teacher-student setting, where the target function corresponds to "teacher" neural network. The architecture of this teacher neural network is chosen be fully connected feed-forward with 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/ E IR for each input data sample x; E RR. Given the architecture of the teacher networks (but not the values of the teacher-weights W*) and the training set of input-output pairs {v,, x; N== 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 benchmark.
In the so-called high-dimensional limit [Donoho, 2000], when the input training data are d-dimensional Gaussian vectors, in the limit d -> xo, the above research program has been carried out in detail over the last decades for small neural networks having only m = (1(1) hidden units, and learning from n ad data samples, where a (a(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 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- -ighndimensional limit (with d oo, m = (1) and n = Oa(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 target function corresponding to 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 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 multi-layer neural network (or even refinements like 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 be able to learn the target function with non-linear activations' to an infinitesimally small test error. This is perhaps intuitive as, with an extensive width, the number of parametersswwiihtts 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.
necessary in order to
Our contributions - In this paper, we step up to this challenge and derive a closed-form expression for the Bayes-optimal test error for target/teacher function corresponding neural network of extensive width, from quadratically many samples, for a particular case where the activation function (after the hidden layer) is quadratic. In partieular, our main contributions are:
one-hidden layer
We provide a closed-fome expression for the Bayes-optimal error of learning an extensive-widhe neural network from quadratically many samples, which the first type of such result to the best of our knowledge. Such form is enabled by the high- -igh-dimensional limit and corresponding concentration of quantities of interest. notably follows from our formula that, in the absence of noise in the target function, zero test error is achievable for sample complexity a n/dd larger than a perfect-recovery threshold a > apR where
x²
apR = K
2 if R
OPR = 2
if
K
2 1,
(I)
with x 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.
Note that for linear activations, the target functions reduces to linear regression and can be learned from linearly many samples.
2

---

We introduce the GAMP-RIE algorithm that combines generalized approximate m passing (GAMP) [Donoho et al., 2009, Rangan, 2011, Zdeborová and Krzakala, 2016] with a matrix denoiser that is based on so-called rotationally-invarian estimators (RIE) [Bun et al., 2016], and show that in the large size limit, this algorithm reaches the Bayes-optimal error for all o H = ((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-ranl matrix denoising [Maillard al., 2022b].
The derived formula involves the asymptotics of the Harish- Chandra-Itzykson -uber integral of random matrix theory -hhrish- 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 the one obtained by gradient descent. In the noiseless case we observe rather unusual and surprising scenario, as randomly -ommmiiiitializedd gra- dient descent seems to be sampling the space of interpolants, and leads twice the Bayes-optimal error. When averaged over initialization the gradient descent reaches an eгтor 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 phuse retrieval in the case of a single hidden unit (m = 1). Many works considered this problem in the high-dimensional limit d + oo, in the regime of Odd log d) samples;
c.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 nigh-dimensional limit has been established down to the constant in a n d. Authors of Mondelli and Montanari [2019] derived the weak recovery threshold for the noiseless case to be owR 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 informatiio-theorrtiialll 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 arr
and CAMP z 1.13 needed to achieve perfect learning for a Gaussian prior. Authors of Song et al. [2021] proposed a non non-robust polynomial algorithm capable of solving noiseless phase retrieval for a orr- 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-valiee valued case. Phase retrieval with generative priors was studied in Hand et al. [201 8), Aubin 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 (1),), the problem is special case of multi-index model that has been recently actively considered, e.g. in Aubin et al. [2019b], Bietti et al. [2023]. Damian et al. [2022, 2024], Collins-Woodfin not focused on the quadratic activations, as it does not bring particular simplification in this case.
Woodfin et al. [2023], Dandi et al. [2023, 2024b]. This line of work has
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 2 d (when the read-out layer is fixed as in our setting) in Du and Lee [2018 8J. Similar results were established in Soltanolkotabi et al. [2018], Venturi et al. [2019] for a slightly more general setting where the readout layer is Icarned.
Establishing results about sample complexity required for generalization in cases where m (or both m and m²) are ((d) is technically challenging, and so far, only a handful of works made progress in that direction. In particular, Gamarnik et al. [2019] considered m² 2 d and m 2 d, and have shown that a sample complexity n 2 d(d + 1)/2 is sufficient for perfect recovery of the target function. Sarao Mannelli et al. [2020b] considered the overparametrizee case with m = Oa(1) and m > d, and showed that gradient descent reaches exact recovery for sample complexity >> d(m 1) (m² + 1)m= /2, again considering the high-dimensional limit. Gradient descent of the population risk has been studied for general values of (m*, m) in Martin et al. [2024], along with discussion of the role of overparametrization.
3

---

2 Setting
As discussed above, we are studying the Bayes-optimal accuracy in the teacher-student setting. More concretely, we consider a dataset of n samples D = {1/, x; + where the input data is normal
Gaussian of dimension d: (x+)%=1 Lid. (O,,Ta). We then draw i..d. d-dimensional teacher-weight vectors (wi) ym i..d. (N La), and noise (2,)ī=1 1.dd (O, Im). Finally, the output labels (y)) ) obtained by one-hidden layer teacher network with m hidden units and quadratic activation:
y/ fw- (x.)
+ A 21,k
m
Crucially, we assume we know the form of the (stochastic) target function fw (.) (i.e. the value of m, A, and the form of eq. (2), including the fact that the activation function is quadratic) but we do not know the realization of neither the teacher weights W -(w;,
Win) nor the noise 2₁.
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 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 = Z(D) P prior ()) W, (x.
where Pprior (W) is prior distribution on the teacher weights W", and the likelihood P( X) can be seen probabilistic channel that generates the labels given the input data (x,) and the teacher weights W and (DD) is normalization constant. The Bayes-optimal (BO) estimator the labels for test sample Nerst not seen in the training set 'D then involves the average over the posterior distribution as follows (where F denotes the expectation over 21;
BO DD
(test) = IE Wyeest Ncest D] = Ey P(W dw.
We will evaluate the BO estimator in terms of its average generalization error, i.e. the mean squared error (MSE) achieved on a new sample. We define it in the following way:
(Messt
VD (xeest) ))'] A(2 + A)
MMSE2
Ew ,D.E your Steut
(4)
We denote it MMSE.t standing for minimum-MSE, as it is the minimum MSE achievable given the setting of the model, and we call MMSE lim MMSE3
Conventions for the MMSE Notice the peculiar multiplicative factor (m/2) and the additive term A2 + A) in eq. (4). As we detail in Appendix F. these factors ensure that MMSE - 1 for a a+ 0 (i.c. in the absence of data), and MMSE - 0 if the posterior concentrates around the true W (i.e. if (x)
[fw (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 high-dimensional limit, with an extensive-width architecture and quadratically many data samples:
d+00, a = n/6 = O11), K k =//a = O(1).
(5)
In all that follows, we consider the limit of eq. (5), so that n, d, all go to infinity together when we write c.g. lim4-+00- 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 "student" neural network
1 a n yv Fw(xz))² where fw(x)=
m m Σ
2
((W)-
Σ
Va
(wk) wx))Tx
(6)

---

2.0
1.00
Intensive width K = m/d
0.05
0.1
1.5
0.75
K 0.2
MMSE
0.50
0.5
1.0
< 1.0
5.0
0.25
0.5
0.00
0.0
0.0
0.2
0.4
0.6
0.0
0.2
0.4
0.6
Sample complexity a n dd
Sample complexity = n/d²
Figure I: Left: The asymptotic MMSE of eq. (7) for the noiseless (A 0) case, as a function of the sample complexity a, for various width ratios k. Right: Phase diagram representing the MMSE, brighter color indicates higher value. The red curve is the perfect recovery transition line aPR, see eq. (I), and its origin is discussed Section 5.
Note that this does not account for the noise, but activations neural networks are commonly considered deterministic, so consider this the most natural choice.
Minimization of loss over the weights W (w&) is commonly done using gradient descent
(GD): one initializes the weights as w( ~ Prior and then updates them to minimize the empirical loss, for an appropriately choice of learning rate, until convergence. Denoting the weights convergence as W(W(0) D) the estimator for test labels reads W D (Xtesst) S ,D) (Niest), )
As will see, it will be interesting consider also an estimator jAGD obtained by averaging the GD estimator on the labels over the initializations (' of the weights.
Main results
Notations We use tr(-) == (1/d)Ir[-] for the normalized trace. We denote GOE(d) the distribution
1.i.d. of symmetric matrices & E Rddd such that Eij ~ No (1 d)) for i S j. For kd with O, we denote Wm,d the Wishart distribution, and μAMP, the Marchenko-Pase 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-ddmmessoonl limit of eq. (5) by:
2ax RÂ
MMSE
9 2
where Á = (A(2+A)/4 and where q is a solution of the following equation:
[1/4(1))dy- &
4T²
(1 - 20x
(8)
2
3q
Here, 14 = PMP. R O...... (for 2 0) is the free convolution of the Marchenko-Pastue law and scaled semicircular density, see Appendix A for its precise definition.
Eq. (8) can be efficiently solved using a numerical scheme, which detailed in Appendix H.1. We present the results in Fig. 1. In what follows, detail our approach towards deriving Result 1, which a consequence of our main theoretical result stated Claim
Reduction to matrix estimation problem We first notice that by expanding the square in eq. (2), we can effectively reduce our learning task to an estimation problem in terms of S* = (L/m) Em k= (T We give analytical argument backing this observation in Appendix F.5.
5

---

Its conclusion is that, at leading order, the distribution of fw- (x) can be reduced to the following form, with ì == ddly-1a 1--11-1
フ - Tr AG. with ç (0, = 2A(2+ A)/K, Generalization error and MMSE on This equivalent problem gives us way to interpret the
(9)
/KK and where defined Z = (xxT - -L)//Vd
convention we chose for eq. (4). Indeed, if we denote S =[[S Z] the Bayes-optimal estimator
related to the problem of eq. (9), then MMSE <Etr((S* - S )²]. as proven in Lemma F.I.
The limit of the MMSE We now describe the general form of estimation problems covered by theoretical analysis, which encompasses the one described in eq. (9) (and thus the original eq. (2)).
The goal is recover the symmetric matrix S Rdxd, which was generated from the Wishart distribution Wm,dd from observations (y:), generated as
1/1 Pout (-ITrZz,"))
(10)
with Z, = (x - La)/va. The "channel" Pout accounts for possible non- non-linearities and noise, encompassing the case of additive Gaussian noise in eq. (9). We define the partition function as:
n
2( } Eg- II Pont (ya (r(SZ.))-
(1))
Notice that the averaged logarithm of Z is (up to an additive constant) equal to the mutual in- formation between the observations and the hidden variables: I(W"; ( n. log P (y1 Tr (l))). This links 2 to the optimal estimation of W, an important idea be- hind our study. We are now ready to state our main theoretical result. gives a sharp characterization of the Bayes-optimal error in any estimation problem of the type of eq. (10). By the reduction described above, can be directly applied to the original model of eq. (2), and will imply Result
[l Z
Claim 2. Assume that m = nd with k. > 0, and n = ad with a > 0. Let Qo = 1 + k-! Then:
The limit of the averaged log pog-partition function (sometimes called the free entropy) is given by
1
d= d ((u.z.) log 3 = sup gE[1.0o]
lim -
[(g) +
dyDE (y) Je (u (v.E)
(12)
RxR
where
(Qo - q)q
1 2 (12//)) log 4
1
((g)
= inf
920
4
2
8
(13)
v AT Qu - ②) exp {
dz
(z- v2g
[V(
} Pout (ulz).
4(Qu 9)
Here, (xx) == Ex.Y~ log |X - YY, and, for 2 0, μs == UMP H. O s.c..vvt is the free convolution of the Marchenko- Pastur distribution and a (scaled) semicircle law, see Appendix A for its definition.
For any a > O, except possibly in a countable set, the supremum in eq. (12) is reached in a unique q* E [1, Oo). Moreover, the asymptotic minimum mean- mean-squane squared error on the estimation of S achieved by the Bayes-optimal estimator SBO = F(S) x31,,; is equal to Qo -²":
lim Etr((S* - - 500 d-+00
) = Qu - 9²'
(14)
It is related to the MMSE of eq. (4) by MMSE = x(Qu 9*))
Specifying Claim 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 2 Notice that q" = lim E(tr(S''' BO) according to Claim 2. As the MMSE
decreases with a, it is clear that n² 2 q" (a = 0). When a = O. we have $ BO = F[S*] Ta. and thus q¹ (a = 0) = 1. We check in Appendix F.8 that the value 9² (a 0) 1 is recovered by eq. (12).

---

4 Derivation of the main results
We derive our main result (Claim 2) in two ways, First, we show how one can show Claim 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 growing body of work in high- nigh-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 ligh-dimmnsional models, a rigorous treatment of it remains elusive. It important, we feel, to present as well more mathemattcally 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 step sketch of 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., 20226, Pourkamali et al., 2024, Semerjian, 2024]. While 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 Da == (I/d)) log ((y,,, x:)),, ef. 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 nigh-dimensional limit of eq. (5).
Step 1: Universality with "Gaussian equivalent" problem - The first step of our approach is inspired by recent literature on the ellipsoid fitting problem [Maillard and Kunisky, 2024, Maillard and Bandeira, 2023]. It amounts to notice that, if Z, == (x - -Ia)/v by the central limit theorem, for any symmetric matrix S, Tr|Z,S] is (under mild boundedness conditions on the spectrum of S) approximately distributed as (O, 2tr(S?)) as d + co. 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 dinensiinnii 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, with G. GOE(d).
Conjecture 4.1 (Universality). We define
a Pout (v/(T)
(G) = d (())....) log Es II 1=1
(15)
where v/ ~ Pout (TrT( with St Wmm.d and G. 1.d.d. GOE(d). Then
lim b (G() = 0.
d
Conjecture 4.1 can be seen 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 Pout (which are satisfied by the Gaussian additive noise consider).
Step 2: A matrix generalized linear model with a Wishart prior- - By the first step above, we (G) and the corresponding estimation problem. A key observation is that one can view this problem as an instance of generalized linear model on S*, with a Gaussian data matrix whose i-th row is the flattening of the matrix G, 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* Wm,,. 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..dd [Aubin et al., 2019a, 2020]. In our setting, the structure
can focus on o (G())
7

---

of the Wishart prior raises several technical difficulties preventing to directly transpose the proof approaches of Barbier et al. [2019], so we state the following result as a conjecture. Conjecture 4.2 (The free entropy of a matrix generalized linear model). We have
lim (G) sup inf d-000
(Q0 q) (() + dyDD) log
NE[1.Qu] 920
4
RXR
where
1 ( =
lim Fy log Wsw. exp (- d T V25)3)
(16)
4 d+000 d²
the asymptotic free entropy of the matrix denoising problem Y = VQS* +. with & GOE(d). and S* ~ Wm.d, and we assume that the d + 30 limit in eq. (16) is well-defined.
Step 3: Extensive-rank matrix denoising - As a last step, we study the function (g)) 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 2 0, the limit in eq. (16) is well-defined,
and moreover (recall the definition (p)) and μt in Claim 2)
( - - (a))
1
log q
(17)
8
We provide very short and assumption free proof of Theorem 4. in Appendix F.2, which combines relation between (q) and HCTZ. 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 final remark, we notice that recent conjecture?
of Semerjian [2024] states that the free entropy of matrix denoising of S = (1/m) Ex km1 ( remains the same if one considers any i.i.d. prior for w 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 for any such teacher weight distribution.
The second part of Claim We briefly discuss the second part of Claim 2, concerning the large d
limit of IE tr -$B0, } The fact that the maximizer of eq. (12) is unique for almost all values of a can be seen by simple convexity arguments, see Appendix F.6. The relationship of q" with the
B0
asymptotic MMSE on the estimation of S is classical consequence of the I-MMSE theorem in generalized linear models of which eq. (9) is an instance, see e.g. Barbier et al. [2019] and Section D.5 of Maillard et al. [2020].
5 Discussion of the main results
Analysis of the Bayes-optimal estimator - We start by discussing the noiseless case (A = 0), which is deseribed by the phase diagram in Fig. 1. Since there is noise in the target function, we expect a sharp transition to zero MMSE at a critical sample complexity OPR- We analytically show in Appendix F.3 from eq. (8) that aPR 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 that there is perfect recovery as soon as a > 1/2. For all values of k we see the MMSE is a smooth curve going continuously from 1 at a 0 to apR. We derived the slope of the curve at apR to be (see Appendix F.4)
where the problem convex, where Gamarnik et al. [2019] shows
12
2
if 5
èMMSE a
K
1+
lopR
if 2 1.
K
"W mention here the "strong" conjecture of Semerjian [2024]. A weaker form of this conjecture is the universality of the best low-degree polynomial estimator for any i...d. prior.
8

---

Theory
(A
2.0
Theor (A 0.23")
GD (d 200)
GD
AGD (d 200)
Averaged (A 0)
1.5
252)
MSE
GD (d 100)
Averoged GD (A 0.26")
AGD (d 100)
1.0
0.5
0.0
0.0
0.1
0.2
0.3
0.4
0.0
0.1
0.2
0.3 0.4
Sample complexity a n/d
Sample complexity a n//d
Figure 2: Mean squared error (MSE) as function of the sample complexity a for n: 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 A O case. The colors indicate the size d. We can see how AGD appears to well described by the theoretical MMSE. We used the learning rates 0.2 for 200 and 0.07 for 100. Right: Comparison of GD between the noisy = case (red) and noiseless A 3 case (blue). Adding noise makes AGD worse than the MMSE, and for sample complexity a 2 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 observing concave dependence on a for small и it becomes convex when и increases and a is close to app- We also note that the smooth limit MMSE - quadratic number of samples being needed to learn better than linear regression.
- supports the result of Cui et al. [2023] about
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 a +
We considered analytically the limits и + 0 and к + oc, i.e. the limits of small and large (but still extensive 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 finite-size instances. In all the cases 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 fixed and equal to in Appendix G we generalize all our main results, thcoretical 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 gradient based approach to suboptimal, as the problem non-convex for k < 1. In Fig. 2, we compare (a) the MSE KS - ŠGD) 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, A 0, we very remarkably observe that the MSE reached by gradient descent very close to exactly twice larger than the asymptotic MMSE. Such relation is known in high- dimensional 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, aundoll-initializee 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,x² }/-1, plus a matrix in the orthogonal solution of the
space, that is S = Σ xxx,x Z. This means that gradient descent reaches 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- 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, observe 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 l2 regularization. We can characterize empirically another phase transition: for a sample complexity larger than at (A))).. GD converges to the same function independently of the initialization. In the noiseless A 0 case, this is simply the perfect recovery transition, and OpR = ar(A 0). while increasing the noise intensity makes the threshold lower until it reaches plateau, which for = 0.5 is at ary (A + oc) ~ 0.2. We display this numerical finding in Figure 5 (right) in Appendix H.3. A tight analytical study of the landscape-triviaaiaatee trivialization threshold ar A) as function of the noise variance A is left for future work.
6 Conclusion and limitations
In this work, we provide an explicit formula for the generalization MMSE when learning a target function in the form of a one-hidden layer neural network with quadratic activation in the limit of large dimensions, extensive width and a quadratic number of samples. The techniques deployed to obtain this result are novel and, we believe, of independent interest. There are many natural extensions of the present works. While we presented, additionally to the replica derivation, a mathematically sound derivation, a fully rigorous treatment, technical and lengthy task, is left for an extended version of this work. We analyzed the Bayes-optimal MMSE, presented the GAMP- 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 al. [2020], Stöger and Soltanolkotabi [2021]. Finally, we also presented the natural extension of our results and techniques to the case of learnable second layer.
algorithm
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 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.
Acknowledgements
We want to thank Giulio Biroli, Francis Bach, Guilhem Semerjian, Pierfrancesco Urbani, Vittorio Erba, Jason Lee and Afonso Bandeira for insightful discussions about this work. This work was supported by the Swiss National Science Foundation under grants SNSF SMArtNet (grant number 21 212049) and SNSF OperaGOST (grant number 200390).
10

---

References
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference Learning Theory, pages 2552-2623. PMLR, 2023.
Greg W Anderson, Alice Guionnet, and Ofer Zeitouni. An introduction random matrices. Cam- bridge university press, 2010.
Benjamin Aubin, Bruno Loureiro, Antoine Maillard, Florent Krzakala, and Lenka Zdeborová. The spiked matrix model with generative priors. Advances in Neural Information Processing Systems, 32, 2019a.
Benjamin Aubin, Antoine Maillard, Jean Barbier, Florent Krzakala, Nicolas Macris, and Lenka Zdeborová. The committee machine: computational to statistical gaps in learning a two-layers neural network. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):224023, jan 2019b.
Benjamin Aubin, Bruno Loureiro, Antoine Baker, Florent Krzakala, and Lenka Zdeborová. Ex- act asymptotics for phase retrieval and compressed sensing with random generative priors. In Mathematical and Scientific Machine Learning, pages 55-73. PMLR, 2020.
Jean Barbier, Mohamad Dia, Nicolas Macris, and Florent Krzakala. The mutual information in random linear estimation. In 2016 54th Anmual Allerton Conference Communication, Control, and Computing (Allerton), pages 625-632. IEEE, 2016.
Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Záeborová. Optimal errors and phase transitions high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, 1 (6(11):5451-5400 2019.
Gérard Ben Arous, Song Mei, Andrea Montanari, and Mihai Nica. The landscape of the spiked tensor model. Communications on Pure and Applied Mathematics, 72(11):222-23300, 2019.
Gérard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from nigh-dimensional inference. Journal of Machine Learning Research, 22 (106):1-51) 2021.
Florent Benaych-Georges and Raj Rao Nadakuditi. The eigenvalues and eigenvectors of finite, low rank perturbations of large random matrices. Advances Mathematics, 227(1)444--521 2011.
Raphael Berthier, Andrea Montanari, and Phan-Minh Nguyen. State evolution for approximate message passing with non-separable functions. Information and Inference: A Journal of the IMA, 91):333-79, 2020.
Alberto Bietti, Joan Bruna, and Loucas Vivien. On learning gaussian multi-inde index models with gradient flow. arXiv preprint vxXiv22310.19999 2023.
Joël Bun, Romain Allez, Jein-Philippe Bouchaud, and Marc Potters. Rotational invariant estimator for general noisy matrices. IEEE Transactions on Information Theory, 62(12))445--7490 2016.
Jian-Feng Cai, Meng Huang, Dong Li, and Yang Wang. Solving phase retrieval with random initial guess nearly as good by spectral initialization. Applied and Computational Harmonic Analysis, 58:60-84, 2022.
Francesco Camilli and Mare Mézard. Matrix factorization with neural networks. Physical Review E, 107(6))064308, 2023.
Francesco Camilli and Mare Mézard. The decimation scheme for symmetric matrix factorization. Journal of Physics A: Mathematical and Theoretical, 57(8)005000, 2024.
Emmanuel Candès and Xiaodong Li. Solving quadratic equations via phaselift when there are about as many equations as unknowns. Foundations of Computational Mathematics, 14:1017-1026.
2014.

---

Emmanuel J Candes, Thomas Strohmer, and Vladislav Voroninski, Phaselift: Exact and stable signal recovery from magnitude measurements via convex programming. Communications on Pure and Applied Mathematics, 66(8):1241-1274, 2013.
Patrick Charbonneau, Enzo Marinari, Giorgio Parisi, Federico Ricci-tersenght Gabriele Sicuro, Francesco Zamponi, and Mare Mezard. Spin Glass Theory and Far Beyond: Replica Symmetry Breaking after 40 Years. World Scientific, 2023.
Yuxin Chen and Emmanuel Candes. Solving random quadratic systems of equations is nearly as easy as solving linear systems. Advances in Neural Information Processing Systems, 28, 2015.
Yuxin Chen, Yuejie Chi, Jianqing Fan, and Cong Ma. Gradient descent with random initialization: Fast global convergence for nonconvex phase retrieval. Mathematical Programming, 176:5-37, 2019.
Elizabeth -ooodin, Courtney Paquette, Elliot Paquette, and Inbar Seroussi. Hitting the high-dimensional notes: An ode for sgd learning dynamics on glms and multi-index models. arXiv preprint arXiv:2308. 238,093300008777,, 2023.
Hugo Cui, Florent Krzakala, and Lenka Zdeborova. Bayes-optimal learning of deep random networks of extensive-widie In Proceedings of the 40th International Conference on Machine Learning. PMLR, 2023, URL https: /
http://eooceeeeeeee/ee///// mwl.m.rreeeeee press/v202) cui23b. html.
George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989.
Alex Damian, Loucas Pillaud- -ivien, Jason D Lee, and Joan Bruna. The computational complexity of learning gaussian single-index models. arXiv preprint arXiv:2403.05529, 2024.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. Conference on Learning Theory, pages 5413-5452. PMLR, 2022.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step time. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
Yatin Dandi, Ludovic Stephan, Florent Krzakala, Bruno Loureiro, and Lenka Zdeborová. Univer- sality laws for gaussian mixtures in generalized linear models. Advances in Neural Information Processing Systems, 36, 2024a.
Yatin Dandi, Emanuele Troiani, Luca Amaboldi, Luca Pesce, Lenka Zdeborova, and Florent Krzakala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. 2024b.
Laurent Demanet and Paul Hand. Stable optimizationless recovery from phaseless linear measure- ments. Journal of Fourier Analysis and Applications, 20:199-221 221,, 2014.
Jonathan Dong, Lorenzo Valzania, Antoine Maillard, Thanh-an Pham, Sylvain Gigan, and Michael Unser. Phase retrieval: From computational imaging to machine learning: tutorial. IEEE Signal Processing Magazine, 40(1):45-5)) 57, 2023.
David I. Donoho. High-dimensional data analysis: The curses and blessings of dimensionality. AMS math challenges lecture, 1(2000):32, 2000.
David L Donoho, Arian Maleki, and Andrea Montanari. pessage- algorithms for compressed sensing. Proceedings of the National Academy of Sciences, 106(45): 189 189199 2009.
Simon Du and Jason Lee. On the power of over- pee-parramttizztioon in neural networks with quadratic activation. In International conference on machine learning, pages 1329-13335 PMLR, 2018.
Andreas Engel. Statistical mechanics of learning. Cambridge University Press, 2001.
Marylou Gabrié. Mean-field inference methods for neural networks. Journal of Physics A: Mathe- matical and Theoretical, 53(22):223002, 2020.
12

---

David Gamarnik, Eren C Kızadda,, and Ilias Zadik. Stationary points of shallow neural networks with quadratic activation function. arXiv preprint arXiv:1912.01599 2019.
David Gamarnik, Cristopher Moore, and Lenka Zdeborová. Disordered systems insights on compu- tational hardness. Journal of Statistical Mechanics: Theory and Experiment, 2022(11): 11)1:1111:1,, 2022.
Elizabeth Gardner and Bernard Derrida. Three unfinished works on the optimal storage capacity of networks. Journal of Physics A: Mathematical and General, 22(12):1983 1989.
Cédric Gerbelot and Raphaël Berthier. Graph-based approximate message passing iterations. Infor- mation and Inference: A Journal of the IMA, 12(4):2562-2628, 2023.
Alice Guionnet and Ofer Zeitouni. Large deviations asymptotics for spherical integrals. Journal of functional analysis, 188(2):461-515, 2002.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in neural information processing systems, 30, 2017.
Dongning Guo, Shlomo Shamai, and Sergio Verdú. Mutual information and minimum mean- -quare error in gaussian channels. IEEE transactions on information theory, (1(4):12111299) 2005.
Géza Györgyi. First-order transition to perfect generalization in a neural network with binary synapses. Physical Review A, 41(12):7097, 1990.
Paul Hand, Oscar Leong, and Vlad Voroninski. Phase retrieval under generative prior. Advances in
Neural Information Processing Systems, 31, 2018.
Harish-Chandra. Differential operators semisimple lie algehra. American Journal of Mathematics, pages 87-120, 1957.
Uwe Helmke and John B Moore. Optimization and dynamical systems. Springer Science & Business Media, 2012.
Hong Hu and Yue M Lu. Universality laws for high-dimensional learning with random features. IEEE Transactions Information Theory, 69(3):1932- 1964, 2022.
Claude Itzykson and J-B Zuber. The planar approximation. ii Journal of Mathematical Physics, 21 (3):411-421 1980.
Adel Javanmard and Andrea Montanari. State evolution for general approximate message passing algorithms, with applications spatial coupling. Information and Inference: A Journal of the IMA, 2(2):115-144, 144, 2013.
Dmitriy Kunisky, Alexander S Wein, and Afonso S Bandeira. Notes computational hardness of hypothesis testing: Predictions using the low-degree likelihood ratio. In ISAAC Congress (International Society for Analysis, its Applications and Computation), pages 1-50. Springer, 2019.
Lucien Le Cam. Locally asymptotically normal families of distributions. certain approximations to families of distributions and their use in the theory of estimation and testing hypotheses. Univ. California Publ. Statist., 3:37, 1960.
Oon Lee and Kevin Schnelli. Tracy- distribution for the largest eigenvalue of real sample covariance matrices with general population. The Annals of Applied Probability, pages 3786-3839, 2016.
Thibault Lesieur, Léo Miolane, Marc Lelarge, Florent Krzakala, and Lenka Zdeborová. Statistical and computational phase transitions in spiked tensor estimation. In 2017 IEEE International Symposium on Information Theory (ISIT), pages 511-515. IEEE, 2017.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations, 2020.
13

---

Wangyu Luo, Wael Alghamdi, and Yue M Lu. Optimal spectral initialization for signal recovery with applications to phase retrieval. IEEE Transactions on Signal Processing. 67(9):2347-2366, 2019.
Antoine Maillard and Afonso S Bandeira. Exact threshold for approximate ellipsoid fitting of random points. arXiv preprint arXiv:2310.05777, 2023.
Antoine Maillard and Dmitriy Kunisky. Fitting an ellipsoid to random points: predictions using the replica method. IEEE Transactions on Information Theory, 2024.
Antoine Maillard, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase retrieval in high dimensions: Statistical and computational phase transitions. Advances Neural Information Processing Systems, 33: 13:11071-11088 2020.
Antoine Maillard, Florent Krzakala, Yue M Lu, and Lenka Zdeborová. Construction of optimal spectral methods in phase retrieval. In Mathematical and Scientific Machine Learning, pages 693-720. PMLR, 2022a.
Antoine Maillard, Florent Krzakala, Mare Mézard, and Lenka Zdeborová. Perturbative construction of mean-field equations in extensive-rank matrix factorization and denoising. Journal of Statistical Mechanics: Theory and Experiment, 2082(3,, 2022b.
Antoine Maillard, Afonso S Bandeira, David Belius, Ivan Dokmanić, and Shuta Nakajima. Injectivity of relu networks: perspectives from statistical physics. arXiv preprint wrXiv22302.14122 2023.
Antoine Maillard, Emanuele Troiani, Simon Martin, Florent Krzakala, and Zdeborová Lenka. Numerical code used for experimental results. https://githin httpggtttuu..oom/ttt Soppo///OOOo///////// srteessiieeitttlluuuarriaaaaaaap
2024.
Vladimir Alexandrovich Marchenko and Leonid Andreevich Pastur. Distribution of eigenvalues for some sets of random matrices. Mutematicheskii Sbormik, 114(4):507-536, 1967.
Simon Martin, Francis Bach, and Giulio Biroli. On the impact of overparamettrizzaion the training of shallow neural network in high dimensions. In International Conference on Artificial Intelligence and Statistics, pages 3655-3663. PMLR, 2024.
Mare Mezard and Andrea Montanari. Information, physics, and computation. Oxford University Press, 2009.
Marc Mézard, Giorgio Parisi, and Miguel Angel Virasoro. Spin glass theory and beyond: An Introduction to the Replica Method and Its Applications, volume 9. World Scientific Publishing Company, 1987.
Francesca Mignacco, Pierfrancesco Urbani, and Lenka Zdeborová. Stochasticity helps to navigate rough landscapes: comparing gradient-descent sesffenttddscent-bueeeee based algorithms in the phase retrieval problem. Machine Learning: Science and Technology, 2(3):035029, 2021.
Marco Mondelli and Andrea Montanari. Fundamental limits of weak recovery with applications to phase retrieval. Foundations of Computational Mathematics, 19(3):703-773, Jun 2019.
Andrea Montanari and Basil N Saeed. Universality of empirical risk minimization. In Conference on Learning Theory, pages 43 4310-4312. PMLR, 2022.
Andrea Montanari and Subhabrata Sen. A friendly tutorial on mean-field spin glass techniques for non-physicists. Foundations and Trends@ in Machine Learning, 17(1):1-173, 2024.
Opper and Haussler. Generalization performance of bayes optimal classification algorithm for learning perceptron. Physical review letters, 66 20:2677-2680. 1991.
Amelia Perry, Alexander S Wein, and Afonso S Bandeira. Statistical limits of spiked tensor models. In Annales de I' l'Institut Henri Pincaar-Probabilitts et Statistiques, volume 56, pages 230-264, 2020.
Farzad Pourkamali, Jean Barbier, and Nicolas Macris. Matrix inference in growing rank regimes.
IEEE Transactions Information Theory, 2024.
14

---

Sundeep Rangan. Generalized approximate message passing for estimation with random linear mixing. 2011 IEEE International Symposium on Information Theory Proceedings, pages 11722. IEEE, 2011.
Valentina Ros, Gérard Ben Arous, Giulio Biroli, and Chiara Cammarota. Complex energy landscapes in spiked-tensor and simple glassy models: Ruggedness, arrangements of local minima, and phase transitions. Physical Review X, (1):01 1003, 2019.
Stefano Sarao Mannelli, Giulio Biroli, Chiara Cammarota, Florent Krzakala, Pierfrancesco Urbani, and Lenka Zdeborová. Complex dynamics in simple neural networks: Understanding gradient flow in phase retrieval. Advances in Neural Information Processing Systems, 33:3265-3274, 2020a.
Stefano Sarao Mannelli, Eric Vanden- Eijnden, and Lenka Zdeborová. Optimization and generalization of shallow neural networks with quadratic activation functions. Advances in Neural Information 2020b.
Processing Systems,
Henry Schwarze. Learning a rule in a multilayer neural network. Journal of Physics Mathematical and General, 26(21):5781, 1993.
Guilhem Semerjian. Matrix denoising: Bayes-optimal estimators via low-degree polynomials. arXiv preprint 12X1::2402.1679, 2024.
Hyunjune Sebastian Seung, Haim Sompolinsky, and Naftali Tishby. Statistical mechanics of learning from examples. Physical review A, 45(8):6056, 1992.
Jack W Silverstein and Sang-Il Choi. Analysis of the limiting spectral distribution of large dimensional random matrices. Journal of Multivariate Analysis, 54(2):295 -09, 1995.
Mahdi Soltanolkotabi, Adel Javanmard, and Jason D Lee. Theoretical insights into the optimization landscape of over-parameterized shallow neural networks. IEEE Transactions on Information Theory, 65(2);742-769, 2018
Haim Sompolinsky, Naftali Tishby, and H Sebastian Seung. Learning from examples in large neural networks. Physical Review Letters, 65(13)16683, 1990.
Min Jae Song, Ilias Zadik, and Joan Bruna. On the cryptographic hardness of learning single periodic neurons. Advances in neural information processing systems, 34:29602-29615, 2021.
Roland Speicher. Free convolution and the random sum of matrices. Publications of the Research Institute for Mathematical Sciences, 29(5):731-744, 1993.
Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparametrized low-rank matrix reconstruction. Advances in Neural Information Processing Systems, 14:23831-238444 2021.
Antonia M Tulino and Sergio Verdú. Random matrix theory and wireless communications, Founda- tions and Trends in Communications and Information Theory, 1(1):1-182, 2004.
Luca Venturi, Afonso S Bandeira, and Joan Bruna. Spurious valleys in one-hidden-layer neural network optimization landscapes. Journal of Machine Learning Research, 20(133):1-34, 2019.
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018
Timothy LH Watkin, Albrecht Rau, and Michael Biehl. The statistical mechanics of learning rule. Reviews of Modern Physics, 65(2):499, 1993.
Eugene P Wigner. Characteristic vectors of bordered matrices with infinite dimensions. Annals of Mathematics, pages 548-5644 1955.
Lenka Zdeborová and Florent Krzakala. Statistical physics of inference: Thresholds and algorithms. Advances in Physics, 6555)2553-552, 2016.
15

---

A Additional definitions and conventions
Convention -Throughout this manuscript, use Ex denote the expectation solely the random variable X. We denote M (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 x d real matrices:
A matrix & is distributed according to the GOE(d) distribution (standing for Gaussian Orthogonal Ensemble) Exj Li.d. (N + 6:gl/d) for any $ i S S d.
For any m 2 1, matrix S is distributed according the Wishart distribution Wm,d WTW/m, where W e Rmxd with WAs i..d. N( 1) for k E [m], E [d].
For symmetric matrix M with eigenvalues (Ax) we denote μM == (1/d) Σ Di=1 8x, its empirical eigenvalue distribution (ESD). It is well known that for d - oo the ESD of GÓE(d) and Wm.d matrices converge to (respectively) the Wigner semicirele and the Marchenko-Pasee
density.
Theorem A.1. [Wigner [1955], Marchenko and Pastur [1967]] Let = ad for > 0, and let & GOF(d) and S Wm.d Then, as -> oc, the ESDs of & and S almost surely converge (in the sense of weak convergence) to the following probability distributions (respectively).
The semicircle law, with density
V4 x
(x)
1|11 S 2}.
(18)
/2T
We denote a ((.))((()) == 1-112210000 s.c. (x/ vt) the scaled semicircle law with variance
The Marchenko-Pen Pastur law, with density
(1 11-)))(xx+ H/(4, - x)( x-)
ifn S
2TT
(x)
(19)
if
2Tx
Here ==
= (1 K 11/222
Transforms of probability distributions- For any real probability measure μ. define its Stieltjes
transform (2) = [1 - 2)] for 2 E C. If C+ = {z C Im(2) > 0}, then gr(2) C+ for all 2 C_. Moreover, we have the Stieltjes- Perron inversion formula:
Theorem A.2 (Stieltjes-Perron inversion formula). For all < b, we have
p((a, b)) lim lim 1 (g + ic) - gulx ie)
640 <40 2im Jass
In particular, if μ has a continuous density with respect to the Lebesgue measure then:
dp lim 1 Im 9 (x + it). <0 π
Va E R, dx
We often use the logarithmic potential function E(p) = ((dx))((y)) log |x y). We further define the R- transform of μ
(s) = 9, (-s)
(20)
We refer to Tulíno and Verdú [2004] for more details on the definitions of this transform, e.g. concerning complete domain of definition. Informally, the R transform well-defined neighborhood of 0 for all measures which have bounded support. In particular, have for the semicircle and the Marchenko -asttu distributions:
Ra, Roppor 11,
(21)
16

---

Algorithm 1: GAMP-RIE
Result: The estimator S
Input:
Observations y E R" and "sensing vectors" Z, = (xxx² la)/vdd E Rdxd,
mitiaalize S ~ Wm,d and é, w, V randomly; while not converging do
Estimation of the variance and mean of [T;;;
v at and w -TriZ,'' gout (yis w vt-1)ve;
Variance and mean of estimated from the "channel" observations;
20
A'
Jout Your(yss wi v2t² and R S
Jour (yt, w/ Vz)Z,
dA
Update of the estimation of with the "prior" information;
S sit1 = SRIE R² 1
1
and ottl 2FFF
2.4²
2
t + 1;
end
Free additive convolution The main interest of the R-transform lies in its connection to the (additive) free convolution of measures. Informally, we can interpret the free convolution μ i v of two measures and y the limiting spectral measure of A B, where A and B are symmetric x d random matrices, with limiting spectral distributions μ and V, 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 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 Tulino and Verdú [2004] e.g.):
(s) = Rs((s) + (s).
(22)
Eq. (22) allows to efficiently compute the density of μ B v given the ones of μ and v. by relating the R transform to the Stieltjes transform, and then using the StieltjesPerroon inversion thcorem (Theorem A.2).
B
The GAMP-RIE algorithm
B.1 Polynomial-time optimal estimation with the GAMP-RIE algorithm
Let us recall crucial observation of Section the learning problem of eq. (2) can be effectively reduced to generalized linear model (GLM) on the matrix S (cf. eq. (10)):
Vi Pout (T Z,S"))
(23)
with Z, = (x.] -L)/Vi) ~ Wm,d, and Pout 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, (which can be viewed as "sensing vectors" applied to S*). Determining the optimal algorithm
in GLMs when the sensing vectors have arbitrary structure in general open. Anticipating on
universality argument for the MMSE (cf. Section 4), we "forget" momentarily about the structure of
{Z,), and assume that the optimal algorithm takes the form would have if the {Z.} were instead Gaussian matrices (i.e. GOE(d)). For generalized linear models with Gaussian sensing vectors, class of generalized approximate sassagg--assiigg (GAMP) algorithms have been extensively studied, and argued to reach optimal performance in the absence of a computational cetppputttionalto-ssttttitiiit to-statistical gap [Donoho et al., 2009, Rangan, 2011, Zdeborová and Krzakala, 2016) 6]. The GAMP algorithm includes a denoiser that is adjusted to the prior information about the signal S*, that is in our case Wishart distribution. Combining these two facts, we propose the GAMP-RIE algorithm in A Algorithm 1. An implementation of GAMP-RI RIE is accessible in the GitHub repository associated to this work [Maillard al., 2024].
17

---

The functions Jout, fRIE and FRIE appearing in Algorithm I are defined as follows. First, we let
Jout (v, w, V) dz (z w) (--2) Pout (ylz)
(24)
Jdze (2mp)2 Pont(yle
In particular, for the problem of eq. (9), we have
goutly, w; V)
A + V
The two functionn (fRIE, FRIE) are related the problem of matrix denoising, in which one aims at recovering a matrix So Wm.d from the observation of R So + Våg. with & GOE(d). We recall some important results on this problem, and how they relate to the definition of the functions (JRIE, FRIE).
ii)
The optimal estimator (in the sense of mean squared error) of So has been worked out in Bun et al. [2016], and belongs to the class of rotationallyiinvaiiant estimators (RIE). FRIE RR A) is this optimal estimator, and it admits the following explicit form. If R = UAUT is the spectral decomposition of R, and letting Ps |MP.* MD 0 s.c..VA / be its asymptotic eigenvalue distribution (see Appendix A for the definition of the free convolution μ and its relation to the sum of asymptotically free matrices), then fRIE (R, A) = Ufa(A)UT
where fa (A) x 2Ahs (A), with ha the Hilbert transform of Pa. More precisely:
(a = P.V. A.V.P//LAA/(. ps
Pa and ha can be evaluated numerically very efficiently, see Appendix A for details. (ii) FRIE(A) ()) 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]):
4T A? 3
FRIE A) =
= 4A2 dAPa()))a(()..
(25)
In Appendix B.2 we sketch the derivation of the state evolution of Algorithm 1, assuming 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 formal connection with the thformition-theorrtiie predictions of Claim 2. Notably, state-evolution of the GAMP-RIE algorithm (which turns out to follow from rigorous work on non-separable estimation with GAMP [Berthier al., 2020, Gerbelot and Berthier, 2023]), and show that its fixed points agree with the fixed point equations that provide the Bayes- Bayes-optimal error. In all regions of parameters that we investigated below we observed a unique fixed point, meaning that the GAMP-RIE algorithm asymptotically reaches the Bayes-optimal performance, see Appendix B.3
obtain a so-called
B.2 State evolution: a connection between Algorithm and Result
We briefly sketch here the crdistiicl--hyysice style derivation of the so-called state evolution of Algorithm I: this will draw a connection between the performance of the Bayes-optimal estimator, characterized by Result 1, and the estimator of Algorithm We define q² = tr((S) )T, and m² == tr/S's S* Thanks to sayes-optimallty, so-called Nishimori identities are preserved (see Zdeborová and Krzakala [2016] for more details), so that we have, at leading order d + oc, that g² m'.
show that, along the GAMP-RI trajectory, the
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:
The "sensing vectors" Z; are not Gaussian. We conjecture that the universality arguments discussed in Section 4 extend the analysis of the GAMP-RIE algorithm. This allows
to replace Z, by G; ~ i.idd GOE(d) when evaluating (g', m²) (i.e. when studying the high-dimessional performance of Algorithm 1). We are then able to make direct use of some results of Zdeborová and Krzakala [2016].
18

---

(ii) The prior over S* is not i.idd.: as we saw, this led to a non-trivval "denoising" part in Algorithm 1. The performance of this denoising procedure in the high-dimensional limit can however be characterized precislly, as the function FRIE admits a closed-form expression.
We now briefty expose the derivation, transposed to our setting under the universality assumption above. By definition of q², we have of 2(Qo q²). If wi z are centered and jointly Gaussian variables with E[w²]= = 24", E[>²] = 2Q2. and E[wz] = 2mm = 2g". and ~ Pou((11), we define
o' = 4a goue (y, w, )
(26)
so that =' in the n, d- + oc limit. For the "channel" part of the GAMP-RIE algorithm, the standard analysis for generalized linear model, alongside the universality phenomenon discussed above (which allows replacing Tr by Tr[G,S'] in the update of wf) shows that q' satisfies the equation:
a 1i 8q
[ (4 Jog
g' d== 4a
(27)
q=
where Jy(y.E) E) 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 After replacing Z, by Gaussian matrices G, in Algorithm 1, the argument is that at leading order as d + oc one has:
R
d
E
(28)
with & ~ GOE(d). Heuristic details on how to derive eq. (28) can be found again in Zdeborová and Krzakala [2016], see Section 6.4.1a there. By definition of FRIE, this implies g1+1 = tr[(S s++ Qo FRIE (()))())==) that by eq. (25):
)²]
1
4T
/ dApLig (A)²,
Qo
(29)
3(62)2 G'
with 14 = PMP V. 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 very interesting consequence. Indeed, the latter model admits 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, by GOE(d) matrices G;) Algorithm 1 can be written in the following form:
= Gi(u', Ex) - V Jout (y, (y,""=1, -V--1,,
(30)
n'
GT GT (y, Vt) + E-' Fx).
Let us clarify some notations used in eq. (30):
u' E RR, with p -2+1). can be seen the flattening of the symmetric matrix A'R' of Algorithm via the following canonical mapping. For S E Sa. we define vec(S) IRP by vec(S) = Su and vec(S).) = V2S, for i < j. This flattening is an isometry: (vec(S), vec = Tr(SR]. We have u' ('RR.
G E Raxp is a Gaussian i.i.d. matrix, whose elements have variance 2/d.
D-1 = -2x [gout (t, wt, V*) is related to A' by A = E,1 yt = 2/RIE (2..
Section VID.I1 in the arXiv version of Zdeborová and Krzakala [2016].
19

---

1.00
Theory (Noiseless)
Theory (A 0.25*)
0.75
AMP (d 100)
MMSE
AMP (d 200)
0.50
0.25
0.00
0.0
0.1
0.2
0.3
0.4
0.5
a n dP
Figure 3: Comparison of the performance of GAMP-RIE with the asymptotic MMSE (7) both in the noiseless (A 0) and in noisy (vā 0.25) case, with 0.5. Each dot is the average over 8 runs of GAMP-RIE at a moderate size of either d' = 100 (circle dots) or d 200 (crosses). The error bars are the standard deviations of the MSE.
(u', Ex) E R is the flattening of the RIE denoiser of Algorithm 1, i.e. if we denote R' R//E, the matrix such that u' =vec('* (///)/))):
vu E²) vec[,RIE Ex/2)].
Eq. (30) is the canonical form of the GAMP algorithm, 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 al. [2020]. Gerbelot and Berthier [2023]). Denote L' 7AMP == tr(S, and ÖAMP = n, Ei= 9out (yi, wt. V*)2 (recall the definition of these quantitiees in Algorithm 1). Assume that the "sensing matrices" (Z,) 1 in Algorithm I are replaced by (G.) which are i.i.d. GOE(d) matrices. Then for any 2 GAMP and d'AMP follow the state evolution equations (27) and (29) asymptotically as d, n cc.
Beyond the rigorous control of the GAMP-RIE algorithm, Theorem B. has an additional mathemati- cal consequence: it allows to leverage to prove results on the asymptotic MMSE and on the mutual information, as Theorem B. implies that they be used verbatim in setting. More precisely, the fact that the GAMP-RIE algorithm achieves an MSE with value given by Claim immediately yields that the latter is, at least, an upper bound the asymptotic MMSE (when assuming Gaussian GOE(d) "sensing vectors" G.). 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*)E log Z in Claim 2) a lower bound on the real one (see e.g. section 2.C in Barbier et al. [2016]).
of mathematical techniques that AMP algorithms
B.3 GAMP-RIE algorithm reaching the optimal error
In Fig. 3 we compare the asymptotic theoretical result for the Bayes-optimal error with the perfor- mance of the GAMP-RIE algorithm for 100 and d 200, in both the noiseless (bluc) 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 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. particular, unlike in the canonical phase retrieval problems (i.e. when m 1) [Barbier et al., 2019], we did not identify a computational-tt -ostttisttiall gap when learning this extensive-width quadratie-activation neural network.
20

---

C Generic activations
We give here a brief discussion on the extension of our work more generic activation function. While our derivation (ef. Section 4) heavily relies on the non- non-linearity being quadratic, a first natural extension would be to consider polynomial activations, with an output generated as (assuming a noiseless setting):
Xx
Vi
vd
k=1
for some integer p 2 3. One could also "linearize" this model, by writing it as (T", ,X)), in which X; are now p-tensors, defined as
T
m
X;
8p
dp/2
However, two main challenges arise when carrying out the program of Section 4 in this "tensor" model:
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 Saced, 2022, Dandi et al., 2024a, Maillard and Bandeira, 2023].
(ii) Secondly, the generalized form of Conjecture 4.2 would involve the free entropy of a tensor denoising problem. While rich literature has studied the fundamental limits of denoising low-rank tensors (see Lesieur al. [2017], Ben Arous et al. [2019], Ros al. [2019], Perry et al. [2020]. Gamarnik et al. [2022] and references therein), here T has rank m = ((d), and the optimal denoising of a large-rank tensor is, as far as we know, a completely open question.
These two challenges form the basis of an exciting but very challenging research program, which we leave for future work. Provided such program could be carried out for any polynomial activation, one might then hope to analyze generic activation functions, such as the ReLU or sigmoid, c.g. by decomposition over basis of orthogonal polynomials (such the Hermite basis), Ben Arous et al. 12021J, Abbe et al. [2023], Dandi et al. [2023, 2024b] for examples of such analyses in the case m (1).
D Derivation of Claim 2 from the replica method
In this section, we give non-rigorous derivation of eq. (12) using classical methods of statistical physics. We start from the definition of the partition function in eq. (1I). We denote D the standard Gaussian measure, and S(W) WTW/m /m.
(S (x
DW II Poue (D: TrZ,s(W))) 1=1
The replica method We make use of the heuristic replica trick [Mézard et al., 1987]. Letting t (L log Z, it consists in writing that lim 4 ba(r) = (1/d²) log E[ One then computes the d + oo limit of 42(r) for integer E N, before extending analytically the result to any O. 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 al. [1987] for an introduction to the replica method in the context of the statistical physics of disordered systems, Maillard et [2023], Montanari and Sen [2024] for mathematiiallly-frendd method, and to Mezard and Montanari [2009], Zdeborová and Krzakala [2016]. Gabrié [2020] for
(000rr) lim &a(r). with
descriptions of the
21

---

some of its applications in the context of theoretical computer science, high nigh-dimensional statistics, and machine learning.
The replicated free entropy- - We now compute the "replicated free entropy" Pa(r), for r E N. write as an average over + 1 replicas of the system, writing
Thanks to Bayes-optimality, S as the replica of index O. We write S = S(W") to simplify notations. We reach:
2r(r)
1 log de
DW a=0
f
II Pout (1) a=0
(31)
For fixed of matrices (8 Ja=o= by the central limit theorem the law of the variables 2" = Zr approach, as d + 0o, a correlated Gaussian distribution, with mean [[2*]] = 0, and covariance E[2*2*| = Ez (TT
2tr(S" as is easily checked from the fact that Z [xx - La)/ Vd, with x ~ N00 Ia). Since n = O(d)) the leading order of the term dy Ez II== Poin (y Tr(S*Z] will be the only one entering the leading order of &a(r). This means that we have, denoting the overlap matrix
trr=S)(SSSS)
(32)
that
1 Pa(r) =
p log
III DW
U dy 1 (4T) det Q
II
Pout (y) (5]2z)
n
+ od(1),
a=0
aQ /
e --
=
d²
log
DW a=0
dy (4m)
Pou (vl2a)
a=0
Qa dtr(S"(")) + oa(1). agh
II
(33)
Notice that the CLT-based argument above is made formal in Conjecture 4.1, and implies the univer- sality of o under the replacement of Z, by Gaussian GOE matrices Gg Since Q E R(++1)x(++1)
of finite size as + 0o, we can perform the Laplace method over Q in eq. (33), and we reach (omitting oa(1) terms as d + 0o, and recall nt/d2 + a):
fa(r) = sup
[J(Q) + a.Jom (Q)),
(34)
QEST
where S:+ is the set of positive semi-definite symmetric matrices of size + 1, and:
J(Q)
1
log
DW II
II asb
Qus - (
d²
a=0
(35)
dy dz - -2-
Jont (Q)
log
II a=0
Poun (v|=²).
(4m) / det
R1
Notice that can rewrite (Q) using Lagrange multipliers Q E S,+1 (or equivalently using the Fourier transform of the delta distribution, and the saddle point method on the Fourier parameters) as:
)(Q)= inf
Tr
de log / II
DW e ==0
e EE,b Tro
(36)
QES..1
The replica-symmetrri 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 matrices such that, for all a, E (0,... .r] with * 6:
Qua = =Q, Qus = Q
(37)
Qas 9, Qub = -
22

---

with 0 S Q, and Q.920.
The term Jont (Q) Under the ansatz of eq. (37), it is classical computation [Zdeborová and Krzakala, 2016] to reach:
{
dz
(= - 22
r+1
Junt ( = log
dy DE
exp
P Poun (ulz)
(38)
R2
(
4(Q 9)
The term ((Q) Using the replica-symmetric ansatz of eq. (37) in eq. (36), we get:
(r ()(QQ rqq)
log / /I +(++ Σ. I1((8 Th (E. st)
)(Q) = inf Q.4
+ d²
4
a=0
We now use the following Gaussian integration identity, for any symmetric matrix M:
(g~GOE(d) e [TT:MMe]
4 Tr(M²)
This allows to reach the following expression, which is analytic in
(r (+11)
QÔ r(r +
((() = inf
Q.4
4
1 log Le
(/DW>
}}
(39)
d
Recall that here S = S(W) = WT/
The limit - 0 From eqs. (34), (38) and (39), have:
1 QQ + de
log f DWe -Wwe-"Ppe 19 T-(59)
$ Ta(r = 0) = sup inf Q20 QER
(40)
This implies that Q = 0 and Q Qo
lim Es- (r((²) = =1 xー (recall m/d - k), and we correctly recover that 2r(r 0) O. Taking now the derivative with respect to r, followed by the r - limit, yields:
qu
lim P sup inf d-+00
dy Lo) Jo (v.
(41)
os45Qo90 920
JR2
lim
GE-GO (H4 Hg(E))
d
de
[ -Rmxd
Hq(E)=
DWe
+
(42)
de
V--22
2n(y,2)== /4T(Qo q)
exp
Pour (y)=).
(43)
4(Qn - 4)
In order to obtain from eq. (41) the prediction of eq. (12), it therefore suffices to show that, for any i 2 0:
lim E5~GOE(4) (H)(E) log H₂(E)) = Qoq E(m1/) log
1
(44)
d d²
2
We focus on deriving eq. (44) the remaining of this section. We note that we can rewrite the left-hand side as the free entropy of the following denoising problem:
== +
(45)
with & GOE(d), S
Wm.. and which we consider in the Bayes-optimal setting. Indeed, we can define the free entropy of this problem as
dg Tr - s)²]
Evss log da
DW exp
4
1
Ly log / DW exp (-4
dg
Tris
dg d4
I
4
de
2
1+ QQo
+
d² (aocor((HH4449 log H₂(E)].
(46)
4
23

---

Crucially, this auxiliary problem is again Bayes-optimal. which we will use in what follows.
Remark Eq. (45) defines a problem known as extensive-rank matrix denoising. The limit free en- tropy of this problem, as well as the analytical form of the Bayes-optimal estimator, for a rotationally- invariant prior S' and rotationally invariant noise È (which is here Gaussian) have both been understood and worked out completely [Bun et al., 2016, Maillard et al., 2022b, Pourkamali et al., 2024, Semerjian, 2024]. We will leverage these results (and partially re-derive them) in what follows.
We now use a change of variable to the singular values W, see e.g. Proposition 4.1.3 of Anderson et al. [2010]. We reach:
f DW exp (-2enss
+ de
Ca.m
dd e
3 EE
A
2
Rm
k=1
k=1
II A A k<k
EM X/
So(n)
DO
exp
TIOAAOY
(47)
in which S = WTW/m OAOT, with A = Diag((A) (A),,... Am; ..... 0)), and Ca.m > 0 is 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
1957, Itzykson and Zuber, 1980]:
2
log D0 { Ir [OSOTY]]
[ (od)
Od
HHCIZ (0, Y) =
= lim d-soc da
(48)
where S and Y are d x matrices with asymptotic eigenvalue distributions μss and jy. We can now apply the Laplace method in eq. (47) the eigenvalue distribution of S. As the problem of eq. (45) Bayes- optimal,, it is known that the typical eigenvalue distribution of S under the distribution of eq. (47) is μss uas- AMP ,,s as consequence of the so-called Nishimori identity, that AMP,, is the maximizer of the variational problem obtained by the use of Laplace's 's method, see Maillard et al. [2022b] for details. Since the asymptotic distribution of Y is (by eq. (45)) boy JMMP, M s.c.1//V we reach by eq. (46):
lim - do d
1
GE~GOE()) (E)(E) log Hg(E)] C(k) qQu 1
+ THCIZ( (M)
(49)
where ((k) is function of x m/d. dd. It can be easily seen that C(k) 0 by considering 0.
Fortunately, extensive-rank matrix denoising with Gaussian noise 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 > and any v, we have with μr == v H 0 s....v [Maillard et al., 2022b]:
1 S(m) [x2]
1
3 2 JHCIZ (t_ v)
1 E 0, 4t
logt t +
8
with I(u) = I u(dr)p(dy) log vl. Applying this formula with t 1/9 reach:
1 2uxx) + i [[tr((x2)]
log q + â
/HCIZ(Ā, PMP.N, =
Etr((S))))
8
4
- 1--2-mm E(ux) + 1 (2Qog + +1)
3 1
8 log
Combining it with eq. (49), we reach eq. (44) (recall that by M1y4 with the notations of eq. (44)).
E Large and small x limits
E.1 The small-s limit
We consider here the limit & - O, i.e. the limit of small (but still extensively large) hidden layer, and compute the limit of the MMSE curves shown in Fig. Since in the noiseless setting we have apR + O(k2) (cf. eq. (I)), we will work in the rescaled regime äx, with â remaining finite
24

---

1.00
R R-
0.01
0.75
R 0.05
0.1
MMSE
0.2
0.50
0.5
5.0
0.25
0.00
0.00
0.25
0.50
0.75
1.00
1.25
a/ti = n/dm
Figure 4: Behavior of the asymptotic MMSE in the noiseless (A 0) case gets increasingly small. The continuous lines are given by eq. (7), which we compare with the asymptotic к - 0 curve obtained by eq. (51). We emphasize that the horizontal axis is a /x,, which remains of order ((1) as - O: it corresponds to a number of samples n of the same order as the number of parameters dm.
as k. By analyzing eq. (8) in this regime (details are given in Appendix E.1.1), we reach that the MMSE satisfies, as K - 0:
1 A + A)
if <
MMSE
(50)
1 ( + A)
A(2 20 σ v 8)2 4(2+
if 2
In partieular, in the noiseless case (A 0), have:
1
MMSE
④ S + 45 2)
(51)
2
and we reach perfect recovery for ã = 1. This limit is illustrated Fig.
Remarkably, eq. (51) can be computed well by taking the limit + when assuming that m = (1) as + oo, a setting which was studied extensively in the literature (see Aubin al. [2019b] and references therein). We detail this computation in Appendix E.
E.1.1 Details of the small-k limit
Recall that by Claim 2, we have MMSE < (Qu 9²), with Qo = -- Since the MMSE remains finite as x + O, we consider the scaling = q/k, with < 7 < 1. We start again from eqs. (7) and (8). We denote A = A(2 + A). Eq. (7), combined with the scaling of implies that
á ~ x²/t for some finite > O, and since MMSE - 7 as K + O, we have
R²
1 1 A
2
Moreover, eq. (8) at order ((k) yields:
A
2n
O2(F (1., k=)
(52)
t
where
F(I, x)) = 4
(1/x221)² (y))
3x²
Notice that F(1, 0) since Mg ~ s.c.vv for E - oo, and f (ay) 3/{4=2g). the leading order of eq. (8) as K 0 is consistent but not informative.
Thus,
25

---

In what follows, we work out the small R. limit of F(t, k)), at first order in K. We denote 1/ (v) = (1/K)Ptyaa (y/n), so that the Stieltjes transform gx(2) == v(y)/(y 2)dy of satisfies the self- consistent equation (see Appendix A):
16
ig.
(53)
1 1
Moreover, we notice that Ve (## pMP,.. .x) EB 3 s.c.vtt so that the support of v remains bounded as K + O. We then proceed to expand in eq. (53). For any finite 2 E C, the leading order of the expansion easily given by -1/h th + 0₂(1), which gives that Us + s.c..vē However, as mentioned above, we need to go to the next order in this expansion compute eq- (52).
BBP-type transition We notice that KMP.PP.. (2) (1 x))(z) + kō(2 1) when k +
More precisely, it is composed of mass (1 k) in 0, and the rest of the mass is made up of continuous part supported between (1 - vR)²
1 2/K and (1 + vK)² 1 + 2VF. Vx can thus be seen the spectral density of the sum of GOE matrix (with variance f) and a small-rank perturbation matrix of rank m = kdd, with all non-zeru eigenvalues located close to We therefore expect by the so-called BBP transition phenomenon [Benaych- [Benaych-Georges and Nadakuditi, 201 I] that V. will possess set of m eigenvalues outside the semicircle bulk whenever the condition
1 2
(54)
yn.. Seee.v/ (2v()
satisfied, with Is.c.. (1()) = Exro [1/(X 2)] the Stieltjes transform of the semicirele. Since one can easily show that Is.c.. vz(2 = -t 1/2 eq. (54) equivalent N In this case, these "spiked" eigenvalues are located around the value (Benaych-Geeores and Nadakuditi, 2011]
ged ((-)) = R....v((1) + = 1+t
Moreover, as the width of the continuous part of K# is of size (VK,), we also expect this "spiked" part of the spectrum to have width (VVE).
Expansion of y/ Based on the remarks of the previous paragraph, we assume the following behavior for Ux as + 0. For any y E IR with y + 1 t, have
vv((y) = s.....(())) + kx)))()) + o()).
(55)
Furthermore, we also have, for all y R, when S 1:
VRuk y (1 +1+)
t p (2))
(56)
V
for finite density (()), with p = Eqs. (55) and (56) can be used expand the Stieltjes v(),p()) ,()), and then eq. (53) used find the values of these two functions. These computations are straightforward, and yield:
transform of Vx as function of D(1),
(y 2)
(V)
2T +t 4) 4t y² Hl s 221)).
(57)
(())
=Psc.vi-t-
Notice that the second equation eq. (57) only valid for t < 1, while the first one valid for all
t 2 O. One checks for instance that f (dy)) 1[t S 1). which implies that the normalization condition J (y)dy = 1 is well satisfied for all values of t 2 0. Using the expansion of eq. (57), obtain that
F(t,k) = 4
vx))))dy,
3
{
if
K
+ o(k).
1/t
if
So finally eq. (52) becomes
{ 22-1
if
2
1/t
1
26

---

And recall that MMSE 2āt A, so that
MMSE { 1(22-1 -) ift S 1,
{1
ift 2 1,
(MMSE + (2))) we reach that (1 + (2) (1 + A)/2, and 1 à + (1 - 2 +
Since
A otherwise. This yields eq. (50).
E.1.2
The small-x limit from large but finite hidden layer
We consider here the noiseless case:
1 m
(w]) Tx, vd
with 0(1) as n,d oc. We denote a n/d ām, and we assume that ā O(1) as m - oc (after n, d + oo). We can write the partition function (cf. eq. (11)) as:
Z
DW II
Pout
Ui
(58)
with Pout (ulz) = 5ly - ||z||/ /m) We can make a direct use of the results of Aubin et al. [2019h] to write:
lim
E log = extra.
Triqql + Ip mālo
(59)
d
d++20
1p
DE
Dw exp
( + + 2/2 w
Jpm
Rm
x log
Dw exp cxp
wTqw + + 1/2
Rm
{,00 0
00
= Ic
dy
DE
DZ Pout {ul(Im - 1/22 + +
Rm
Rm
log
DZP out ful(Tm //2 + //2
Here, q. q are symmetric m X m matrices, which satisfy moreover Im > q > 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 Importantly, the matrix can be interpreted as the "overlap matrix" of the model: if we denote (-) the average under the posterior measure in eq. (58), then we have
W
qet = E
(60)
where w, w' are two independent samples under (-). Moreover, thanks to the Bayes-optimality of the problem, it is known that the overlap concentrates [Zdeborová and Krzakala, 2016], in the sense that the random variable (w] wwww))/d concentrates on its average under (E() as d - co.
The "prior integral" Ip can be very easily computed with Gaussian integrals, and yields:
Ip = Tr log det + q).
(61)
We now focus on computing the leading order of le the large-m limit. can write
1c
= dy DE (y E) log (v.
(62)
Tq(,EE) = I 5 (Im 2/ 9/2
Rm
27

---

Let y = vmly tr(Im -q) (5*gE)/m)) We can change variables in eq. (62), and obtain:
= /
lc
dÿ dDDE DE (q E) log J (7. E) log m.
5 9
19(
DZd à
m
-r(I - q) (Tm 2 2
Rm
(63)
Notice that the additive term (1/2) log in Ic just amounts to renormalization of the partition function Z, so we remove this additional constant in what follows. We proceed to simplify (y in the large-m limit. We have
Jq(Ø.E) E)
z (Im - q)z
tr(Im - q) + ZT (Im -q) /) 1/26
DZ5
Vm
pm
m
du
e /miriTi -q) f
iuym
(T -al2 42 2T -y) 122
DZe
2tt
du iuz+ inttti/mm 4) dot Im +2 iullm 2)
ETq (Im 1/2 2u i
Im 2im -1
Tm 1/29112
2nt
du
ing (it-uutt((1n) - q)* 2n eTg (Tm 404mm-4444/2 (0000mmmm000//00000000000
2m
+ +0(1//m)
2TT
where
of = 2tr g) + m
=/ (2m 9)q 122.
Plugging it back into eq. (63) yields:
/ dy DE
1
(0)
y² log 2To
20
+0LI/m))
2
20
- 1 ʃ DE log 2
21
2 + (11/mm)
Since & ~ N'O0, Im). it follows from elementary concentration of measure that of concentrates on its average value o² given by:
d² 2tr((mm 4) 4tr[(Im q)ql 2tr[(Lm (Lm + +q)] 2tr q²].
All all we reach that (up to additive constants):
1 Ic 2
log tr[Im +
(64)
Combining eqs. (61) and (64) in eq. (59), we get at leading order in m, with + lim log Z:
1 - Φ extr
log tr o
m - マ traqq] + tr tr (m + Q) 2
(65)
Eq. (65) can be easily solved, and yields:
q
q (T 91
2ā
tr II
q2p
28

---

This implies that (recall 0 N q mm):
0if -
2
(66)
(2ā - 1) if 2
a 1
Im if 11
Now that we have obtained in eq. (66), we can compute the MMSE, or generalization emoo. Defining in eq. (4):
MMSE = Ew DD D Eree [[3test RO (Xiest)
2
the same arguments used in the proof of Lemma F.1 show that in the large m limit (but taken after oo), we have leading order
MMSE
d Etr[(S -B0))
m
d [tr((SBo)
with S = (1/m) Ex Wa W. Notice that SBO
(S), so that
1 E
w w
MMSE3
Σ
1Sk.Sm
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 = limg-+00 MMSEa
MMSE
LE m
Σ 15k.15m
Pe = tr(q2].
Combining with eq. (66), we reach:
lif S 2'
MMSE
45(1 ā) if مā
0if 1.
We have recovered eq. (51) from the limit m oc taken after d oo!
E.2 The large-R limit
We consider here oc, with remaining of order ( as Since the MMSE remains finite as well, we see from eq. (7) that we must have the scaling q = kt, with t remaining finite K co. very similar derivation the one of Appendix E.1.1 yields that eq. (8) in this limit becomes (with A = A(2+A) A)):
4
/
1
2a At lim
R-000 3rd (w)
1+t
Combining it with eq. (7) yields that
1 - 2nx
+ (11 2a + A)2 + San
MMSE
(67)
2
where we recall A(2 + A). In particular, for 0, we reach MMSE max(1 2ax 0),
coherently with the behavior shown in Fig. 1.
29

---

Other technicalities
F1 Properties of the MMSE of eq. (4)
Let S" := (1/m) k=1 (wi)), and S BO here the following lemma on the MMSE of eq. (4), under the high-dimensional limit of eq. (5):
[]] the Bayes-optimal estimator of S*. We show
Lemma F.1. For a constant C = (()) > 0:
(5)
27 B8)
C(s) n
MMSEa
KES ptr
<
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 8 D concen- trates on S' which implies that MMSEa -> 0. Conversely, for a = 0 (i.e. in the absence of data),
the Bayes-optimal estimator becomes S E(S*) Ia, so that Etr((S* $ BO, )²] x-1 Thus, we have MMSE, 1 for a = 0.
Proof of Lemma F.I. Notice that (cf. eq. (2)):
xtSx
[x]] = A +
d
with S = (1/m) E 1= W. Using this in eq. (3), and plugging it in eq. (4), we get (with
z Im) and NO0, Ia)):
Fs. DAN (
(
(1
|z|12
XT(S
BO
S
2vA
x
x
27
MMSF 2
2k
m
d
m
k=1
A(2 + A)
BO
|z|12)
xT
S*) )x}
44
ES DAN
A 1
tr((")
A(2 A)
2
m
d²
m
BO
(a) m
(S
X
S*xx
Es D.
2
dc
(b) m
Es D (1r((8
-UiSt-$BBO -SBO)
nF D Ir -B0-
(68)
2
where we used [||||||*] = m² + 2m and Etr(S*) = 1 in (a). and Ex-N (XTMX) = Tr[M22 M²² in (b). It remains to bound the first term of eq. (68) to conclude the proof of Lemma F.I. We notice that, by linearity of the trace, tB) is the Bayes-optimal estimator for tr(S*), i.c.
FstD (Ir(s: - -00) = min Fs. (tr((S) - r(D))²
(69)
r(D)
In particular, considering the estimator
r(D) = A),
1=1
1
x;S' x;
|
2
x W
d
+
A
1
+
Zi,k
Vd
30

---

we have using eq. (69):
( -s }}
n
)
|2:11
ED
< Es (xihtt)
s
x
1/)
la
+
+
2v4
x w
2
vd
i=1 1=1
(x)
< 3/11 12 Iz],
(70)
using the Schway-Schwarz inequality in (a), with
n
la 11)
11
=
tr
S
X
T -
TU f=1
27
12
A
21=1
nm
1
m
W
13
4AE
Zi,k
TUTZ
tedious but straightforward computation to compute [Ta]a as only involves the first moments of Gaussian random variables. We get (recall m kd):
2
11
(1
+ K )
nd
2A²
Ie knd
(71)
4A
I3
knd'
Combining eqs. (70) and (71), and plugging it back in eq. (68), we get
MMSE KES- tr (s - BO 27
>
C(k)
n
which ends the proof of Lemma F.1.
F.2 Proof of Theorem 4.1
First, we note that Theorem 1 of Pourkamali et al. [2024] implies that:
Qug
((() 2 /HCIZ( AMP,N, (1/g)
(72)
2
and we recall the definition of IhCIz in eq. (48). We recall then fundamental result proven in Guionnet and Zeitouni [2002]:
Theorem F2 (Theorem 1.1 of Guionnet and Zeitoumi [2002]). For any compactly supported proba- bility measures v and μ, and any > 0:
To (o, - -
3 logt 1 41. Eu|x
(73)
8
Moreover, the function J(: satisfies the following property. Let d be a distance on the space of probability measures on R that is compatible with the weak topology. Let X == R Viw, where W GOE(d), and R is a fixed (deterministic) matrix, with uniformly bounded spectral norm, and a compuctly supported limiting eigenvalue distribution Let px denote the empirical eigenvalue distribution X. Then, for any SM (R):
1
Tim lim log Pode( < -im lim inf 810 d=t00 d
log P(d(u, <)),
650 d de
-v;
(74)
31

---

In other words, the function (15 14) is the large deviations rate function (in the scale d²) for the empirical spectral measure of R + vłW, where R is fixed (deterministic) matrix with asymptotic spectral distribution μ., and W GOE(d). It is a well-known property of the free convolution (Speicher, 1993] that Pxx la E s.c..v as d - co, where the convergence is meant in the weak sense (and almost surely). Combining this result with eq. (74), have J E s.c.,v ps) 0. This yields by eq. (73):
1 IHCIZ(t 11, &...,vi)) Exe M &.c.v))) + log 7
(75)
Combining eqs. (72) and (75) yields eq. (17).
Remark- The proof above can be straightforwardly extended the free entropy of denoising any matrix S with iwvatiionaly-invariinn distribution and a compactly-supporte limiting eigenvalue dis- tribution (beyond the Wishart ensemble), as the results of Guionnet and Zeitouni [2002), Pourkamali et al. [2024] hold under these more general assumptions.
F.3 Perfect recovery threshold in the noiseless case
In this section, we give analytic argument derive the value of the perfect recovery threshold apR (see eq. (1)) in the noiseless setting. In the limit of perfect recovery the MMSE goes to 0, thus by eq. (7) (with A = this implies q co. Using eq. (8). can then write the equation satisfied by the perfect recovery threshold as
3(1 20PR)
lim / f ddy 1J0
(76)
4x
in which 14 MMP,, Mll m...y/ see Appendix A.
F.3.1 The case K <
Informal argument - Recall that in this case we can write xx) = (1 (1- (x))((()) + xM which iM compactly supported away from zero, see Appendix A. O, we thus expect ie to have a discontinuous support, made of two parts:
(a) A small semicircular density centered around 0, of radius ((VI), with mass (1 - x).
A smooth density, compactly supported away from zero, which has a well-defined limit as 0, and mass
Because of the factor t the right- hand side of eq. (76), only the part (a) will matter in the limit. Formal derivation - We first rewrite by a change of variable
ff ddyu(u)) f dv(Veu((Vi=))
It is clear that for all x * O, have (x) kUMP,x( that we can truncate the integral above all |a| S e/Vi, for any finite O. We will now show the following, for any R:
and f VMPP, (y)) xo, so
lim ViullaVtt (1 k) x.x.,,
(77)
We fix 2 C+ (where C_ C- = fz E Im(2) > 0)). Letting y = Vtz, we know from the Marchenko-Pasee theorem [Marchenko and Pastur, 1967] that gi (v) == E,, y)] is the unique solution in C4 to the equation
tg.
1 g/K g
Since y Vtz, is clear that g (t)/ and letting h Vtg. we casily get the expansion
1 1- Ø
h+ O(VE).
2
h
32

---

which can be inverted to
h = 22 4(1 x)
+ O(VB).
(78)
Notice that if denote (x)=) the Stieltjes transform of s....vTTxx. eq. (78) be written as (see e.g. Anderson et al. [2010]) h = (1 - (=),(=) + ((vt). By considering == x + is with x E R and
the limit e 0, we reach using the Stieltjes- Perron inversion theorem (Theorem A.2) that for any
x E R:
lim Vtu(LvV) (1 x))ac..T-x((3
(79)
Coming back to eq. (76) this implies:
3(1
20PR)
limm dy 1-0
4n
(1 k)² / dy
(1 x)) f /dyosc()))
4x (1 x2)2
Equivalently:
(1 k)²
x2
OPR
(80)
2
2
We notice that this critical value of n/d 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 x d symmetric matrices of rank m has, to leading order in d, p(x)d² degrees of freedom, where p(x)d is the dimension the Stiefel manifold orthonormal m- m-frames in Rd is well-known that p(k) = x² for d
[Helmke and Moore, 2012].
F.3.2 The case K 2
The case K > 1 is simpler to carry out. In this case, &MP, does not have a singular part at x 0. and Hr has smooth density as 0, and
/ dy dy
3
4x
1
so that
3(1 -2rrn)
/ dy (v) =0,
limm r10
4m²
and we reach aPR = 1/2, so that ap (asymptotically) coincides with the number d2/2 /2 of degrees of freedom of symmetric matrices. Since oPR is increasing with and has limit 1/2 both for k 1 1 and K. 1, we deduce that opn 1/2 for i = as well.
F.4 The derivative of the error at the perfect recovery threshold
Here, we extend the derivation of Section F.3 to compute the derivative of the MMSE with respect to the perfect recovery threshold. We start again from eqs. (7) and (8). Letting
1/4, we get,
with = orn:
OMMSE
0t Oa
20x
Oa
PR
PR
30m
mile)dda))
lim O.
(81)
72
(+0
We thus compute the next order of the expansion of 3y as 0.
33

---

F.4.1
The case K < 1
We extend the argument made in Section F.3.1. Notice that here the smooth part of the density, compactly supported away from zero, contributes at this order. Formally, for any small enough e >
ef
dy (1)) = txe² f [dyrmp.<())* + tt(),
3tx²
+ ot()))
(82)
2 - k)
On the other hand, we have around the singularity at 0:
- - IISE/VT
Misse dy dypp)
dz [Vtt)(Vt)))
(83)
We evaluate the next order of the right-hand side of eq. (83) using the same approach as in Sec- tion 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):
Zk²
Vim((Vt)= = (1 - -k)0sa_v1=x()) V7 2T(1 - /
4(1 - 13 2
+t 13 [3a - 61 - k) + 2(4 -n((1= x)?]
+ 0(t312)
(84)
2T(1 [1) x) =2/j3/2
for any |2| 2211 K. while Vina (Vtz) 013/12 if 2| > 2/1
This then yields:
of J
dy 3(1 - x)²
3tx
t
+ O(t3/2).
(85)
Se
2
4 x)
Combining eqs. (82) and (85) in eq. (81), we obtain (recall a = apR = K - k2 x²/12):
OMMSE
12
=
+
da
PR
K
1+
F.4.2 The case K 2
Again, we consider >> 1. The argument of Section F.4.1 generalizes immediately, removing the analysis of the singular part around We get directly
/ dy ((())))) =t t / dy PMP (3)) + oz)),
3tx?
2 11) + +c(l).
(86)
Plugging it in eq. (81), get in this case:
OMMSE
= 2
aa
Again, the specific case k = 1 can be tackled by continuity, as the derivative tends to 0 both as k f 1
and R
E55 Details the reduction to matrix estimation
We describe here how to effectively reduce the problem of eq. (2) to an estimation problem in terms of = (1/m) Σn I (w(
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 al., 2019]. Rather, analyze the leading order of eq. (2) and argue that (with high probability over
34

---

the distribution of the data and the teacher weights), the first non 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 backhone of a formal proof of this claim, which we leave for future work and would be carried e.g. by Gaussian interpolation techniques.
Let us define Z, = (x [aa/ Vd, and recall that X;- (N Ia). Expanding the square, we can rewrite the law of the output y/ fw- (x:) as
1 + Vd
| m
2VA A
1/4 A + tr(S*]
TrIZS*+ +
(87)
mVV
where (z,) i.d. (V Im). In what follows, we analyze the leading order of eq. (87). More specifically, denote V vdly: 1 - A), and we decompose
|z: m
+ 2V
A
m
y = TTZ Vd(tt((*) 1) -) Avd
2 x
w
(88)
772
Let us consider the leading order of the different terms of eq. (88). Since S ~ Wi.dd Tr(S"]
Σ ||w|ll//m strongly concentrates on its average. More precisely, by Bernstein's 's inequality (see Corollary 2.8.3 of Vershynin [2018]) we have, for all t 2 0:
P[[tr(S") - 1| 2 t| < 2 exp (-C& min(t,
where C >>> 0 depends only on x > 0. In particular,
P[ 2 d 1/41 s 2 exp(
so that we can replace In by 0 at leading order in eq. (88).
We now tackle 12, first for fixed (x; W Using that || strongly concentrates around its average, and the central limit theorem applied to the fluctuations of | one can see that for all i E [m]. we
have (with g, ~ (0, Im) independently of 2;, and _ denoting equality in distribution):
Z 2
2VA
m
d
24, x
W
m
myd
||2:12
112
A A
Vil
A
+
xi,xx. W
m
mvd d
g
2 x S'xx
x
4A
dd-+00 E:V
(89)
K
Lid. with E "N N(0, 1). independently of (x;, wi). The equivalence as d -> oo is given for a fixed i E [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 + co. 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 E [n]. as by the Hanson- Wright inequality and the union bound, we have (see Theorem 6.2.1 of Vershynin [2018]) for all t 2 0:
max X[S'x - tr((") 20 < 2n exp |C min ( d ve[n]
dt 115711, ||S* llop
dt
P(x))
(90)
for some constant C > O. Since the spectral norm of Wishart matrix [|St strongly concentrates on its average under the Wishart distribution (see Theorem 4.4.5 of Vershynin [2018]), we see that, uniformly over i E [n]]. the leading order of the variance in the right-hand side of eq. (89) is equal to Á = 2A(2 + A)/K. This ends our justification of eq. (9).
35

---

F.6 Unique maximizer q" in eq. (12)
Notice that if Jout = SRXR dyDe Jgly. E) log 1y E), then one can check that Jout is a strictly increasing function of y under mild regularity conditions 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 a >> 0 except possibly in a countable follows then from Proposition of Barbier et al. [2019], see also Appendix A.2 there.
Pout (namely assuming the presence of
F.7 Derivation of Result from Claim 2
In this section, we derive eqs. (7) and eq. (8) from Claim 2, in the case of Gaussian noise. More precisely, we assume Pout (yl22 2 exp[-(y - 2)² (22)] / V 2 in accordance with eq. (9). It is then an easy computation check (recall the defimition of Jg in eq. (13)):
J
1 log A 2(Qo q)}.
JRRR dyDE
We then reach that g = q" is characterized as the maximum of the following function:
F(a) -(4) logl A 2(Qo - 4)];
(91)
(Qo q)
I(q)
inf 420
Σ(m1/a)
log q
4
Recall that here μ4 = PMP.s a.c.. It known (see eqs. (77-78) of Semerjian [2024] e.g.) that
0E(44)
2n²
/
pely))dy.
Ōt
3
Thus, g =(()) can be characterized as the solution"
(Qu - 9)
+ x²
/ (1/g(1)
49 = 0.
(92)
4
30²
By eq. (91), q is solution in [1, Qa] to:
4a
((g)
(93)
A 2(Qo 4)
Recalling that MMSE = =<(Qu- - 4) by Claim 2, eq. (93) implies eq. (7). Combining eq. (93) with eq. (92). we reach eq. (8).
F88 The limit + 0
In this section, we check that the state evolution equations derived in Section F.7 yield indeed that q + 1 as - 0. Indeed, in this limit, Bo = [[S* = La. so that we must have q = [Br(s S*
Recall that q = 4a/ 4a/l2 [ + 2(Qo - 4)]. and that q is given by eq. (8). In particular, q - 0 as a - 0. Assuming the scaling an
O, we get
2
Ä
4
=
Q0
00
2
(94)
Ago
-2 +
F'))
where F(p) = (4T2/3) /3) 1/pp1/2p1ppll 1-pppp1 8dz. Letting =p(2)==pp = 1012-11112ppp0 pppp=zppppp-p p-1.pppp111111 11/2), we know by similar reasoning as the one of Section F.3 that the Stieltjes transform h = h1)2)) of Up satisfies the equation:
KVPP K K h
h.
'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 matrix denoising problem (Theorem 4. and using the I-MMSE theorem. We refer to Barbier et al. [2019]. Maillard al. [2020] for more details.
36

---

As - O, we can thus compute the expansion of (2) in powers of p. Applying then the Stieltjes- Perron inversion theorem (Theorem A.2), we get the expansion of 2,(2) in powers of p as:
v4 2²
32v
3p(2 - (4 + A 2²)
U2(2)
+ Opp)1))
2n
8T
& 2²
for |2| $ 2, and vp(2)= Opp2/2) for |2| 2 2. Plugging this expansion into F(p), we get:
F(p) 1 o(p).
Coming back to eq. (94), this gives d0 = 1w/[2+ Ax). and (recall Qo = + 1 then = s0 that our equations indeed consistent in the limit o O.
G Learning the second layer weights
We sketch here in a mathematically informal way the generalization of our results to the setting where the second layer weights are also learned. The second layer weights (at) are drawn i.i.d. from probability distribution Pa. and the student must learn (wi. of) from the observation of {x, ) and of
2
y/
Xx; + V A2ik
(95)
171
In the rest of this paper we focused the case Pa ô1. However, all our techniques and results can generalized to more generic choices of Pa, as we know show: particular, Claim the generalization of Claim 2 to this more general setting.
Throughout this section, we will assume for simplicity that Pa has bounded support, although expect our results to hold also for more general choices of Pa. We show how to extend Claim 2 to this case, by detailing the differences in the steps outlined in Section We eventually show that Algorithm can also be straightforwardly extended to this setting as well.
G.1 Generalizing the derivation
G.1.1 Reduction to matrix estimation
We first discuss the reduction matrix estimation problem, generalizing Section F.5 to this setting.
We define
m S'
a (w))
(96)
711
and we denote ma = Fr, [a] and Ca == Fr, [a²). We define the MMSE as (notice the additional factor ca with respect eq. (4)):
MMSE
m - Lw D Eress Necat (Mrest
BB (Xteat) )) A(2
(97)
2
By repeating the (mathematically informal) arguments of Section F.5 to this setting, we find that, at leading order m, d oc:
Vdly, tr((S)) THZZ VĀ.
(98)
with & lid. N(O, 1), and Ā 22(2+ + ACa)/K. We let
12
1/ = Vd
f/1
1/3
n1
j=1
Y
=
37

---

The observation of (y4)== is equivalent to the one of (ū.);- and Y. Notice that by eq. (98), we have
Y - A - (r((*) nyd
(TZ + vād)
(99)
1=1
Conditionally on S", the right-hand-side of eq. (99) is sum of n independent zero-mean random variables, which thus typically fluctuates in the scale" [[(nd) 11/1 = O(d Since 7/4 = vdly; - Y. this implies that at leading order we have
J T + V Vc.
(100)
The observer also has access to Y, alongside {ū 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.I, we reach that again we have MMSE =Etr[(S* S BO,, Moreover:
)d] as + oo.
MMSE KFs v. - $B0 )3}
= = [Es ( (Irl(S* - -S )
Conditioning on Y amounts condition on the value of tr(S*), , as detailed above. Let us make two important remarks:
As + xc, Y concentrates around its typical value [[[] = ma- Since the MMSE is bounded, we therefore have as - oo that MMSE kEs (T.) (tr((S* S RO
) m.).
( 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 the asymptotic spectral distribution of S*. In particular, at leading order:
MMSE xlss (ir((r((S* - =ma),
= xlss Mml (tr((S -S BO ) = ma),
(a) xlEs 1 (tr[(S* -
(101)
where in (a) we used that conditioning on tr(S*) m. 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 (Zi, yi generated by eq. (100).
G.1.2 Further steps of the derivation
Here, we notice that the arguments detailed in Section 4 on how to obtain an asymptotic expression of eq. (101) do not depend on the specific asymptotic spectral distribution of S More precisely:
A. Conjecture 4.1 can be directly extended to more general distributions 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 the main text, Conjecture 4.1 is an extension of Corollary 4.10 of Maillard and Bandeira [2023], which holds for generic choices 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 Vi (q) of eq. (16) is well-defined [Aubin al., 2019a, 2020].
C. Finally, the proof of Theorem 4.1 (see Appendix F.2) relies solely on the rotation invariance the distribution 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.. see e.g. Silverstein and Choi [1995], Lee and Schnelli [2016].
Recall that tr tr(((*))) (11) with high probability.
38

---

G.2 Conclusion: Claim 2 when learning the second layer
We are now ready to state the generalization of Claim learnable second layer. The effective problem we consider is the recovery of a symmetric matrix S E Raxd, which was generated as S = (1/m) Dk=1 o (WE)), from observations (y+)==1+ generated
ys Pout (+|T+
(102)
with Z, = [x Ia)/vd and x; Lid. (O(O,,)) Ta).
The asymptotic spectral distribution μx" of is called a generalized Marchenko-Pause distribution (or a free compound Poisson distribution: is also the free multiplicative convolution of the Marchenko- Pastur law and Pa. see Anderson 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]:
(x)
/
ka Pa(a)da.
(103)
K
Eq. (103) allows for an efficient numerical evaluation of μ² given Pa. Notice that Ex [X] = ma, and Ep- [X²] = m Ca/K.
The partition function for the learning problem of eq. (102) is again defined as:
(( | = Es II Poue (Mi Tr(SZ.))-
(104)
i=1
We then obtain the following generalization of Claim 2.
Claim 3. Assume that m = rd with i. > O, and n. = ad2 with a >> O. Recall that m. == E., [a] and
Co [p,, [. Let Qo [x m Co/ k.. Then:
-The limit of the averaged log- tog-partition function of eq. (104) is given by
1 lim
[(9) + [ dyDE
3 doo0 d ((i 1og
sup gE(m2.Qol
Jg(
(105)
JRXR
where
I(g)
= inf 420
(Qo - 2) 1
E(m11a)
1
log n
1
2
4
8)
(106)
/
(Qu 4) exp { (2- v2g
dz
(v
Pout (y]z).
4(Qo= - ①)
Here, M(p) == Ex [x YY, and, for t 2 0, pt == n² E o s.c.. V is the free convolution of n² and a (scaled) semicircle law (see Appendix A).
For any a > except possibly a countable set, the supremum eq. (105) is reached a unique g² E [må, Qo). Moreover, the asymptotic minimum mean-squared error on the estimation of S achieved by the Bayes-optimal estimator 3BO = =E(S' EESK xu,,, is equal to Qo g²:
-50 lim Eir((S* -$00 )2) Qo -
(107)
related the MMSE of eq. (97) by MMSE k(Qu - q²))
Therefore, generalizing Section F.7, Result holds well in this case, with Ā = 222(2 cc,A)/K. and = μ E s.c..V where μ² is characterized by eq. (103).
G.3 The GAMP-RIE algorithm
Finally, one can also generalize Algorithm this setting: the only change to perform is to adapt the functions FRIE and fRIE- Indeed, instead of denoising Wishart matrix (with an asymptotic spectrum given by the Marchenko-Pastue distribution), here one must denoise a matrix So with asymptotic spectral distribution given by n² defined in Appendix G.2. As mentioned, eq. (103) allows for an efficient numerical evaluation of μd given Pa- From there, one can adapt Algorithm to this case simply by replacing in the definitions of FRIE and fRIE the distribution PA by Ps p² E S...A.
39

---

H
Details on the numerics
3.0
MMSF
À 0
Trivialisation threshold aT
0.35
À 0.001
2.5
À 0.005
0.01
2.0
À 0.05
0.30
MSE
À 0.1
1.5
0.25
A 0.5
0.25
1.0
0.5
0.20
0.0
0.0
0.1
0.2
0.3
Sample complexity a n/d
Noisc level
Figure 5: Left: Mean squared error as function of the sample complexity a, for x 1/2 and A 0.252. Dots simulations using GD with single initialization averaged over 32 realizations of the dataset, crosses are averages over 64 initializations. The continuous line is the asymptotic MMSE given by (7). The colors indicate the strength of the regularization. Right: Trivialization threshold in the sample complexity ar as a function of the noise level A in the teacher without The measurement has resolution of 0.1 the noise level and of 0.007 on the sample complexity
regularization,
H.1 Solutions to the "state evolution" equations
We describe here how to solve eqs. (7).(). The first step to solve is obtain an analytical expression for Maa We refer to Appendix A for the definition of quantities used in this section. We recall that law and a scaled semicircular
14 == JMMP, x.. 0r s.c.. vi is the free convolution of the Marchenko-Pasee density. The R-transform of the scaled semicircle distribution is [Tulino and Verdú, 2004]:
Ro ao (2) 2t,
while for the Marchenko-Pane Pastur law we have
Ryser (2) R
We can now use (ef. Appendix A):
Rp, (2) R.pprar.. Ao no..v
Rr RR R Rexar. (3)
K
The Stieltjes transform g(=) 11(X =2)] of ue is the solution of the equation
+ g(2) (-9(=))
or equivalently
== -t(2) +
(108)
(Z))
g(2)
Among all the solutions to this equation, g(2) must be such that Im[g(=)] > if Im(x) > 0, and also satisfies g(2) 1/2 for co. Eq. (108) is third degree polynomial in g(2), and easily be
solved by algebraic solvers, and has single solution satisfying the constraints we described. Finally, (4(x) is given by the Stieltjes-Perror Perroon inversion theorem (see Appendix A):
Im + kc)
Pe (xx lim
(109)
2 +0
40

---

and we numerically choose E = 10 8. We now discuss the computation of the integral of (x)x) in (8). Notice that the integrand is only non-zere 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 @ 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 20/Qu This whole procedure is quite efficient and can be reproduced easily on any machine.
H.2 Gradient descent
In our experiments with gradient descent we are minimizing the objective R(W):
1 n 4 Σ (v= (ys fw(x,))
X
R(W)=
WRI
(110)
A
All the simulations are done in PyTorch with the student weights initialized in the prior. For "vanilla" gradient descent we iterate until convergence, and average over several repetitions. For averaged gradient descent (AGD) 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 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 train a networks with giant steps [Dandi et al., 2023].
In Figure the gradient descent run for zero regularization, À 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 the landscape of GD. We will expand this in Appendix H.3. All the error bars reported in Figure and Figure 5 (left) are standard deviations of the MSE measured on the samples. Figure 5 (right) has 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 used around 30 000 hours of computing time.
H.3 Additional experiments with GD
Here we study in more detail the phenomenology observed in Figure 2 (right) where in the presence of noise and at 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 trivialization of the landscape.
In Figure 5 (right) we study the trivialization threshold as a function of the noise level A. One needs to take care of the symmetries on W, so we first define S:
S(w()) (0) D)
1 m
(w(W(),,)) D)) W( DD).
where mean that for fixed dataset we run GD, then take matrix product to obtain S. This procedure allows us to define the dispersion
SGD = Ep tr (Ewrro) ((W .D) - S(w()), D) D
If the dispersion becomes zero it means that all the runs will converge to the same value. As we increase the sample complexity a the dispersion decreases, until it becomes zero. For each value of the noise level A 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 A.
In Figure 5 (left), where we studied the effect of la regularization on the weights, we can also see how even a relatively small X > 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 better error.
41

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract summarizes the motivation behind our analysis. its relation to previous literature, well as our main theoretical and experimental findings. We also outline the main techniques used in our theoretical derivations.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have included a paragraph related to discuss the scope of our results, as well their limitations, in Section 6.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
42

---

Answer: [Yes]
Justification: We presented in this paper exact claims and conjectures based on well-known analytical methods. Specifically, developed application of the so-called replica method in this new context of learning an extensive-width neural network with a number the dimension. Further, we presented mathematically
of samples which is quadratic sound derivation of our results, based on conjectures which are natural extensions of existing works, As mentioned in the text, fully rigorous derivation of results a very technical and lengthy avenue, and is left for 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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Our results are fully reproducible, as the algorithms we use in our experiments are clearly explained (and their hyperparameters given), and the code used to produce our results is provided.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, 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 provided or not.
If the contribution is 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 specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, 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 replicate the results, access to a hosted model (e.g., g., in the of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
43

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provide the numerical code used all the experiments presented in the paper, as well as the data obtained by running these experiments, in a public GitHub repository.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
wwssssssiiiiiyyyll for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: CCooeooooiisseeeeeeeeo www.ssooiccypyoeooooooooiosi for more details.
/ /nips. /twps.cc/ppubiiaoo wttpcc/pcpppppcudccccooooo
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The setting of our numerical experiments is fully detailed in the text alongside the related figures, and is clearly accessible in the numerical code provided.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We provide error bars for our experiments (ef. e.g. Figure 2), and clarify how the error bars were computed in Appendix H.2.
44

---

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, confi- dence intervals, or statistical significance tests, 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.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CL, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: In Appendix H.l and H.2 we provide an approximate description of the infrastructure needed to reproduce the figures.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported didn't 'tttt make it into the paper).
paper (e.g., preliminary or failed experiments that
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: https: //tpprr///////pp//// httpp/pllc//thcscuiiiiiieoo delnnesl
Answer: [Yes]
Justification: We checked that our paper and our research complies with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
45

---

Answer: [NA]
Justification: Our work is largely theoretical, and concerns the fundamental limits of learning with neural networks. Moreover, our numerical experiments limited to synthetic datasets. As such, we do not believe our work to have societal impact besides the long-term impact brought by better understanding of the theory of learning.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
As discussed above, this question is relevant the presented work.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: We only use synthetic data in our numerical experiments, and do not rely on any existing code.
Guidelines:
46

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide documented code alongside the paper, as well as details on how implemented the algorithmic procedures used in our experiments (e.g. gradient descent) in Appendix H.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
47

---

Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
48


Reviews:
[{"review_id": "QC1wsjFgGp", "reviewer": "Reviewer_bzi8", "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.", "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?\n\n2. In Lines 142–144, the term $P_{prior}(W)$ is mentioned as a prior on the teacher weights $W^*$. Shouldn’t this be $P_{prior}(W^*)$ instead? Could you clarify whether $W$ refers to teacher weights or student weights?\n\n3. How does the Bayes-optimal accuracy relate to the practical performance of gradient descent (GD)? Does high Bayes-optimal accuracy imply high accuracy for GD?\n\n4. Could you explain the origin and significance of the term $\\Delta(2+\\Delta)$ in Eq. (3)?\n\n5. In Line 128, there appears to be an extraneous 'm' following the definition of $D$.", "limitations": "Yes.", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper is robust in its theoretical analysis, providing clear mathematical derivations and a solid theoretical framework. The breadth of related works cited enriches the context and demonstrates the paper’s relevance to current research frontiers.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and dedication. Please refer to the common answer for most of their questions.\n\n- (Q2) -- $P_{prior}({\\bf W})$ is the distribution from which the teacher weights ${\\bf W}^\\star$ are sampled, which we denote ${\\bf W}^\\star \\sim P_{prior}$. Since this distribution is assumed to be known to the student, it appears naturally (as a consequence of Bayes' law) in the posterior distribution of the weights given the dataset.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Tt9UMLS34J", "reviewer": "Reviewer_xzmx", "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.", "questions": "As a non-specialist in this class of problems, my questions are simply why we are concerned with this class of problems? The authors mention connections to phase retrieval and matrix denoising problems; can we give a concrete example, however contrived, that motivates this work by its potential to eventually contribute to a concrete application in those domains,", "limitations": "The author's explanations of the limitations are absolutely clear as far as I can see.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This seems to be a simple and elegant example of learning for a well-controlled class of high-dimensional problems.", "weaknesses": "This is a rather technical result far removed from obvious useful applications; this is not a criticism in itself — we are all aware that Neurips is moving into the territory of COLT etc. However, in a generalist conference it is hard to communicate the significance of a technical result like this. This is not my area, so I am largely taking the author's assertions about the significance and novelty of this result at face value.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and dedication. Please refer to the common answer for the questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "u4IPgK3hFM", "reviewer": "Reviewer_NWbd", "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. \nHaving chosen the playground, the analysis proceeds with recasting the problem to a form that exploits the symmetries of the objective function, and a replica computation of the free energy. Such computations allow the author(s) to derive an asymptotic formula of the Minimum Mean Square Error. \nTo back the Physics-based derivation, they discuss a break-down of the steps required to formalize it completely. Of the three main points, two are conjectures. The final one instead has a mathematical proof. \nGiven that Gradient Descent is the go-to method, they explore its performance in relation to the theoretically optimal MMSE. Interestingly, they find that in the noiseless case it appears to sample weights from the posterior. The same picture breaks when the labels have noise. There, they discuss the appearance of a phase transition that trivializes the randomness over initialization of GD.", "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. \n- The regularization of Figure 3 is never discussed in text, why do you add it? Just to show that the distance from the MMSE becomes larger? What exactly counts to you as \"trivialize the gradient descent landscape almost completely\" in Figure 3 (Left), as you claim in lines 802-803? Please note that I understand the value of Figure 3 (right).\n- Related to the point above, why do you seek to show regularized objectives in the first place? I mean, what motivates your experiment? \n- Please, see the first weakness, where I planted some questions. \n### Typos\nI will write down in this subsection what I noticed by reading. Take it as a friendly contribution. All of them are largely not relevant to the evaluation of the paper.\n- The instruction block for the checklist was not deleted. \n- (line 161) you write \"student\", but I noticed that in other instances you used a different way of placing quotation marks, that is the TeX standard. \n- Equation (5) lacks punctuation.\n- (line 222) the end of the sentence \"using probabilistic technical amenable to rigorous treatment\" probably changed wording, maybe \"technical\"was meant to be \"techniques\"?\n- (line 779) the \"root\" quotation marks. \n- (line 792) is $30,\\,000$ really the right number? Why three zeros? Maybe I am just not used to this standard.", "limitations": "None.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Equation (1) and its natural link to the degrees of freedom of the target function is an appealing result. \n- The connections with ellipsoid fitting, computations in random matrix theory, formalized results on neural networks, and extensive-rank matrix denoising validate it. In simple words: it is not a made-up paper, but it is well positioned in between earlier works. \n- Understanding industry standard is crucial. Gradient Descent is the idealized version of what is actually used in practice. It is therefore valuable to understand it, as the author(s) did. \n- The result about GD opens up to continuations of the work into such direction of understanding formally what is happening there. This paper is not a dead-end. \n- Section 4 is an interesting take on trying to vindicate the replica method. I believe there is value in trying to break down what would be needed to make the Physics derivation a mathematically rigorous one. At the same time, it is of independent interest to use the replica method, as it tells us what we have to prove. Anyways, nice idea. \n- I was kind of surprised to see that as $\\kappa$ increases the MMSE worsens. \n- The techniques used are, in my honest opinion, beautiful and effective. \n- Even if the MMSE cannot be simulated (unless an algorithm achieving it was known), the author(s) bring to the table experiments for GD that are to the point and motivate further work. \n- The references are, to my limited single human reading stack, to the point. \n- The weaknesses I will discuss below are to be understood in the sense that I understand that this is a paper with lots of Physics-flavour. Personally, I agree with the modelling idea of finding a solvable model, hence the choice of a \"simplified setting\". Also, this simplified setting already requires non-trivial computations, and is therefore worth exploring. Nevertheless, I will outline what worries me to stimulate the discussion. \n\nOverall, I believe this was a good read. Thanks to the author(s)!", "weaknesses": "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. \n1. Perhaps the solvable one. You consider matching architectures, with $m = m^{\\star}$, I wonder, what is the technical challenge in analyzing the case in which the activations are the same but $m\\neq m^{\\star}$? \n2. Your target function is \"neuron agnostic\" in the sense that it takes a simple mean of the teacher neurons. How would your technique behave when we take a one hidden-neuron that also has random Gaussian weights $\\{a_k\\}_{k = 1}^m$ in the second layer? I have not thought this through extensively. \n3. The most challenging: a quadratic activation. While the Statistical Physics literature focuses a lot on quadratic activations/phase retrieval, which is thought to be an interesting problem, I am concerned about the extension in this setting to other activations. It appears to me that when recasting the task to a matrix estimation problem (lines 175-180), one heavily relies on \"expanding the square\" (lines 175). In other words, the symmetry of the activation function naturally builds the matrix $\\mathbf{S}^{\\star}$  (lines 176-177). Considering other (e.g. non-polynomial, or even non-square) activations, it looks to me as a very hard task. I do appreciate that the author(s) recognized this, but would be even happier if, in the spirit of section 4, the bonus page for publication were spent to discuss how a quadratic activation is not a dead-end. \n   More in depth, can you bring over an argument that can convince the reader that the method is worth pushing over the boundary of something that allows you to over-use Gaussian integrals (in the sense that square activations are perhaps the easiest non-linearity)?\n\nThe second concern is mostly structural and philosophical. I am convinced that all the results of the paper are of interest to the community. However, from my first reads, I had the impression that the \"Replica argument\", its \"formalization\" and the \"analysis of GD\" had little meaning if placed all inside a conference paper.\nFor example, while Section 4 is very useful for grounding the replica method, I only appreciate its value in the sense that it connects the dots with other works. In some sense, I find this counter-intuition that there are very few resources on Gaussian Equivalence and modern Statistical Physics of neural networks that discuss this, and then find a nice Section clearing out some aspects in a paper that should be 9 pages. I do wish at some point that these scattered nice sections will be condensated in some resource that will be accessible (e.g. a book). \nSimilarly, the analysis of GD could be justified for its industrial importance, and even simply because it is the quickest algorithm to compare with the MMSE prediction, but at the same time it poses the question: \"What is the soul of this work?\". \nPlease note that this weakness is more provocative than actually requiring action.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and dedication. Please refer to the common answer for the recurring questions.\n\n- (W1) -- It is indeed possible to consider that while the teacher has $m^\\star$ hidden units, the student has $m$ with $m > m^\\star$ (\"overparametrized\" regime) or $m < m^\\star$ (``underparametrized'' regime). While we have extended part of our analysis -- namely the reduction to matrix estimation and the universality conjecture -- to this case (provided both $m^\\star, m = \\Theta(d)$), new difficulties arise in the analysis of the resulting matrix generalized linear model, as a consequence of the lack of the Bayes-optimality. Nevertheless, based on so-called \"dilute\" limits of HCIZ integrals [E1], we were able to show that in the noiseless setting, in the overparametrized case $m > m^\\star$, the posterior mean estimator reaches perfect recovery for $\\alpha_{\\rm PR} = \\kappa - \\kappa^2/2$ (where $\\kappa = m/d$). This means that here the posterior mean estimator is far from being Bayes-optimal (as the Bayes-optimal estimator reaches perfect recovery at $\\kappa^\\star - (\\kappa^\\star)^2/2$ with $\\kappa^\\star = m^\\star/d$ as we saw).\n\n - (W2) -- It is indeed possible to consider learning the second layer weights. We verified that this effectively leads to consider the same matrix estimation problem, with ${\\bf S}^\\star$ now generated as\n    \n$$\n    {\\bf S}^\\star = \\frac{1}{m} \\sum_{k=1}^m a_k^\\star {\\bf w}^\\star_k ({\\bf w}^\\star_k)^\\top,\n$$\n \nand $(a^\\star_k)_{k=1}^m \\sim P_a$ are drawn i.i.d., for an arbitrary distribution $P_a$ (not necessarily Gaussian). The case studied in our work then corresponds to $P_a = \\delta_1$. We have checked that our analysis generalizes to this setting, and in particular we have derived the equivalent of Claim 1 for this setting. Remarkably, the conclusion is that the MMSE formulas (eqs.(14) and (15)) generalize directly, by replacing the Marchenko-Pastur law with a *generalized Marchenko-Pastur* (or free compound Poisson distribution), which is the free multiplicative convolution of the Marchenko-Pastur law and $P_a$. This distribution can be analytically characterized via its $R$-transform. We have added a detailed statement of these results in a new appendix of the revised version of the paper, but have not explored the phenomenology of this setting further.\n\n- Q1: we will make the legend of Fig.2 explicit. \n\n- (Q2 and Q3: On regularization and trivialization of the landscape) By \"trivialization\" of the landscape we mean that gradient descent from random initial conditions goes to a minimizer corresponding to the same function, independently of the initialization. As we discussed, this is not expected for $\\kappa < 1$, as the problem is then not convex. We study the effect of adding $\\ell_2$ regularisation on the weights, as we can see in Fig.2 (right) that the averaged GD algorithm seems to have an \"interpolation peak\" in noisy settings: it performs worse and worse as $\\alpha$ increases, until it becomes identical to non-averaged GD at $\\alpha = \\alpha_{\\rm triv}$. At this point $\\alpha_{\\rm triv}$ we therefore say that the landscape is trivialized. In simple networks with only one hidden unit, this interpolation peak is mitigated by adding such a a regularization [E2]. We can see in Fig.3 (left) that this does not happen for our network with an extensive number of hidden units: instead, the regularization decreases the trivialization threshold $\\alpha_{\\rm triv}$. Fig.3 (right) studies the effect of the noise level on this phenomenon.\n\n- We thank the reviewer for the list of typos that we will fix. \n\n- We thank the reviewer for their suggestion to write a book about the general approach. While we agree this is needed, and appreciate the comment, this is clearly beyond the scope of the paper. As for the \"soul of the work\", we are unsure what to answer.  \n\n\n-- [E1] \"Instanton approach to large-$n$ Harish-Chandra-Itzykson-Zuber integrals\", Bun&al, PRL (2014).\n\n-- [E2]  \"Optimal Regularization can Mitigate Double Descent\", Nakkiran&al, ICLR (2020).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "in2l2wJ4UU", "reviewer": "Reviewer_XJ6b", "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)).", "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. \n* Is Eqn (3) derived somewhere? If not, is it because it’s very elementary and can be found in a standard reference? Please give citation then. \n* Is it obvious that the GD estimator is the same as the ERM estimator in the discussion starting on Line 160?", "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", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* I found the paper very accessible despite having little pre-existing knowledge of this line of research. \n* The context of the current work is illustrated clearly. I did not have to do too much background reading to appreciate what has been done, and how the current submission contributes to advancing this line of research\n* Claim 1, which is central to the main result, is shown using the replica method as well as a more rigorous method. \n* For the particular setting described in Section 2, this work derives the MMSE under the regime described in Equation (4). This is novel as far as I am aware. \n* The empirical comparison between Bayes-optimal and (non-stochastic) gradient descent reveals an intriguing hypothesis that averaged over random initialisations, GD is close to being Bayes-optimal", "weaknesses": "* In terms of exposition, I found it a bit awkward that the closed-form asymptotic limit of MMSE_d is given in Eqn 14, which is outside of Section 3 “Main theoretical result”. I think the part of Section 5 titled “Results for the Bayer optimal estimator” is ready to be presented at the beginning of Section 3. \n* The notation and setup for Claim 1, to my eyes, belongs to Section 4 “Derivation of the main results”\n* It’s hard for me to imagine how this type of analysis can be generalised to a more complex architecture than the one-hidden layer network presented in Eqn (2). And I’m not talking about very exotic architectures. I mean even a small generalisation to multiple hidden layers. Is this really tenable using the theoretical tools employed here?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and dedication. Please refer to the common answer for the recurring questions.\n\n- The asymptotic limit (which we generically denote by $\\lim_{d \\to \\infty}$) is achieved by defining $m(d) = \\kappa d$, $n(d) = \\alpha d^2$, and then taking the limit of $d\\to\\infty$. This corresponds to take $m$ and $n$ to infinity while keeping the ratios $m/d$ and $n/d^2$ finite. We have added a clarification of this point in our revision.\n\n- Regarding the GD and ERM estimators, we want to clarify that they are in general not the same. The GD estimator is reached from a random initialization, and in non-convex scenarios (such as here when $\\kappa < 1$) is generally not a global minimizer of the empirical risk (which the ERM estimator is, by definition). The ERM estimator is generically not reachable by polynomial-time algorithms in non-convex problems. We study numerically the minimizer reached by GD on the stated loss from random initialization.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

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

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

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

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

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

$$
R = \frac{1}{2} E[(y- \hat{y}_{\mathcal{D}}^{\rm BO}(x) )^2].
$$

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

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

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

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

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

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

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

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

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

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

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

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

Meta-review: 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.
Decision: Accept (poster)