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

arXiv:2408.03733 · stat.ML, cond-mat.dis-nn, cs.IT, cs.LG, math.IT, math.PR · Submitted 2024-08-07 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Bayes-optimal learning of an extensive-width neural network from quadratically many samples".

Jane: The paper was written by Antoine Maillard, Emanuele Troiani, Simon Martin, Florent Krzakala and Lenka Zdeborová from ETH Zürich and EPFL and INRIA - École Normale Supérieure and ENS, Université PSL and CNRS and Sorbonne Université and Université de Paris.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Welcome back to the show, everyone! Today we're diving into a paper that's got a mouthful of a title: "Bayes-optimal learning of an extensive-width neural network from quadratically many samples." Jane, I gotta say, just reading that title makes my head spin a little.

Jane: It does sound intimidating, Tom, but the idea behind it is actually pretty beautiful. This paper is about figuring out the absolute best possible way to learn a certain type of neural network, and they've cracked it wide open. The authors are from ETH Zurich, EPFL, and ENS in Paris, a real powerhouse team.

Tom: So when you say "Bayes-optimal," you're talking about the gold standard, right? Like, the best any algorithm could possibly do, even with infinite computing power?

Jane: Exactly. It's the theoretical ceiling. And what they've done is found a closed-form formula for that ceiling when you're trying to learn a network with a quadratic activation function, and you have a very specific number of training examples.

Tom: And that's the "quadratically many samples" part. We're talking about needing a number of samples that scales with the square of the input dimension. So if your data has one hundred features, you might need ten thousand samples. That's a lot of data.

Jane: Right, and the reason that's interesting is because a previous paper showed that with just a linear number of samples, you can't do any better than simple linear regression. It's like trying to learn a complex curve, but you're only allowed to draw straight lines. This new paper says, "Okay, if you give us enough data, we can finally learn the actual curve."

Tom: So they're pushing into the regime where the network can actually start doing something interesting, and they've found the exact mathematical answer for how well you can do. That's a big deal.

Jane: It is a big deal, and it's the kind of result that gives us a benchmark. Now we know what the target is, we can start asking whether our practical algorithms are actually hitting it.

Tom: And that's exactly what we're going to dig into next. We've got the theoretical target, but how do we actually get there in practice? Stay tuned.

Summary: Jane: So Tom, we've established that this paper, "Bayes-optimal learning of an extensive-width neural network from quadratically many samples," gives us the theoretical best-case scenario. But what does that scenario actually look like?

Tom: Well, they've got this beautiful phase diagram. On one axis you have the sample complexity, which is basically how much data you're feeding it, and on the other you have the width of the network, which is how many hidden units it has. And there's this sharp line separating a region where you can perfectly recover the function from a region where you can't.

Jane: And that line is what they call the "perfect recovery threshold." Below it, you're stuck with some error. Above it, you can get the test error all the way down to zero. The formula they derived for that threshold is surprisingly simple, and it matches a naive counting of the degrees of freedom in the problem.

Tom: That's the part I find really elegant. It's like they're saying, "You need at least as many samples as there are independent knobs to turn in the network." And the math confirms it. For a narrow network, that threshold is lower, and for a wide network, it caps out at a specific value.

Jane: Right. And it's not just about the noiseless case. They also worked out what happens when there's noise in the data. In that case, the error decreases smoothly as you add more data, without that sharp transition. It's a more gradual improvement.

Tom: So we have this complete picture of the theoretical limits. But here's the million-dollar question, Jane: can any actual algorithm, running on a real computer, achieve this Bayes-optimal performance? Or is it just a beautiful mathematical fantasy?

Jane: That's the perfect question, and it's exactly what the paper tackles next. They didn't just stop at the theory; they built an algorithm to try to reach it. And that's where things get really interesting.

Tom: Alright, so we've got the target, and now we're going to hear about the arrow they shot to hit it. Let's get into it.

Improvements: Tom: So Jane, we've got this theoretical target from "Bayes-optimal learning of an extensive-width neural network from quadratically many samples." The natural question is, how do we actually hit it? And the paper's answer is an algorithm they call GAMP-RIE.

Jane: Right, and I love that name because it tells you exactly what it's doing. GAMP stands for Generalized Approximate Message Passing, which is a powerful technique for solving these high-dimensional inference problems. And RIE stands for Rotationally Invariant Estimator, which is a clever way to denoise matrices.

Tom: So they're combining two different tools. GAMP is really good at handling the data part, the measurements, and RIE is really good at handling the structure of the thing you're trying to find, which in this case is a matrix that represents the network's weights.

Jane: Exactly. And the key insight is that they can prove, in the high-dimensional limit, that this combined algorithm actually reaches the Bayes-optimal error. It's not just a heuristic that works well in practice; it's provably optimal.

Tom: That's huge. So we're not just saying "this is the best you can do theoretically." We're saying "here's a practical algorithm that gets you there." That's a rare combination.

Jane: It is. And they show it numerically, too. They ran the algorithm on simulated data, and the error it achieved matches the theoretical prediction almost perfectly, even for moderate-sized problems.

Tom: But I have to ask, because I know our engineer friend Meng will be wondering: is this just a theoretical curiosity, or is this something you could actually run on a real problem?

Jane: Well, the algorithm itself is polynomial time, so it's not computationally prohibitive. But the real-world impact is more about setting a benchmark. Now we have a standard to measure other, more practical algorithms against. For example, they compared it to gradient descent, which is the workhorse of deep learning.

Tom: And what did they find? Did gradient descent measure up?

Jane: That's the fascinating part, and it's a bit of a surprise. In the noiseless case, they found something really weird. We'll get into that next.

First Page: Jane: So Tom, we were just about to talk about what happens when you compare this optimal algorithm to plain old gradient descent. And the results from "Bayes-optimal learning of an extensive-width neural network from quadratically many samples" are honestly a bit wild.

Tom: Wild how? I'm all ears.

Jane: So, in the noiseless case, they found that a single run of gradient descent, starting from random weights, gives you an error that's almost exactly twice the Bayes-optimal error. But here's the kicker: if you run gradient descent multiple times from different random starting points and then average the results, you get down to the Bayes-optimal error.

Tom: Wait, so averaging over random initializations is the secret sauce? That's like saying if you throw a bunch of darts at a board and average where they land, you get the bullseye.

Jane: That's exactly the analogy. And it suggests something profound: that randomly-initialized gradient descent is actually sampling from the posterior distribution of the weights. It's not just finding one good solution; it's exploring the whole space of good solutions in a way that matches the Bayesian ideal.

Tom: That's a really surprising result. I would have thought gradient descent would get stuck in some local minimum and not explore at all.

Jane: Right, and that's why it's so interesting. It's a conjecture on their part, but the numerical evidence is strong. It's like the algorithm is doing something much smarter than we give it credit for.

Tom: But I'm guessing this doesn't hold up when you add noise to the problem, right?

Jane: You guessed right. When they add noise, that nice property breaks down. Averaging over initializations doesn't help you reach the Bayes-optimal error anymore. And they also see this "trivialization" phenomenon where, with enough data, all the different runs of gradient descent converge to the exact same solution, so averaging doesn't change anything.

Tom: So the story is different depending on whether you have noise or not. That's a really rich set of phenomena they've uncovered. I can't wait to hear what our guests Lu and Meng think about all this.

Jane: Me neither. We've got the theory, we've got the algorithm, and we've got these surprising empirical observations. There's a lot to unpack here.

Conclusion: Tom: Alright, let's wrap this up. We've been talking about "Bayes-optimal learning of an extensive-width neural network from quadratically many samples," and it's been a heck of a ride.

Jane: It really has. We started with the theoretical question: what's the absolute best you can do? And they found a closed-form answer for that, complete with a phase transition that tells you exactly how much data you need for perfect recovery.

Tom: Then they gave us a practical algorithm, GAMP-RIE, that actually reaches that theoretical limit. That's a rare and powerful combination.

Jane: And then they threw in that fascinating observation about gradient descent. The fact that averaging over random initializations can get you to the Bayes-optimal error in the noiseless case is a really deep and surprising result.

Tom: It makes you wonder what else gradient descent is secretly doing that we don't understand. This paper definitely opens up more questions than it answers.

Jane: For sure. And that's what great research does. It gives you a solid foundation and then points you toward the next set of challenges. The authors mention extending this to other activation functions, which is a huge open problem.

Tom: Well, we're excited to see what comes next. Thanks to everyone for listening, and we'll catch you on the next one.

Jane: Bye, everyone!

Antoine Maillard, Emanuele Troiani, Simon Martin, Florent Krzakala, Lenka Zdeborová

ETH Zürich · EPFL · INRIA - École Normale Supérieure · ENS, Université PSL · CNRS · Sorbonne Université · Université de Paris · EPFL

stat.ML, cond-mat.dis-nn, cs.IT, cs.LG, math.IT, math.PR

Submitted: 2024-08-07

Updated: 2026-08-10

Comments: 47 pages

Journal ref: Advances in Neural Information Processing Systems 37 (NeurIPS 2024)

DOI: 10.52202/079017-2609

Code: https://github.com/SPOC-group/ExtensiveWidthQuadraticSamples

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 85/100

Key concepts

Bayes-optimal learning
This refers to the theoretical ceiling—the absolute best possible performance an algorithm can achieve. It represents the ideal target for how well a learning process can perform given specific constraints, serving as a benchmark against which practical methods are measured.
Quadratically many samples
This describes the required amount of training data. The number of samples needed scales with the square of the input dimension. For example, if data has 100 features, approximately 10,000 samples are necessary to achieve optimal performance.
GAMP-RIE
This is a practical algorithm designed to reach the theoretical Bayes-optimal error. It combines Generalized Approximate Message Passing (GAMP) for handling data measurements and Rotationally Invariant Estimator (RIE) for denoising matrices.
Gradient Descent Averaging
In the absence of noise, running gradient descent multiple times from different random starting points and then averaging the results can achieve the Bayes-optimal error. This suggests that random initialization allows the algorithm to explore a wide space of good solutions.

Terminology

Summary

Summary

This paper addresses the problem of learning a target function corresponding to a single hidden layer neural network with a quadratic activation function and random weights, in the high-dimensional limit where the input dimension and network width are proportionally large. The authors consider the regime where the number of available samples is quadratic in the dimension, a setting previously identified as challenging and open. They derive a closed-form expression for the Bayes-optimal test error in this regime.

The setting is a teacher-student framework. The dataset consists of n samples D = yi, xi ni=1, where the input data are i.i.d. Gaussian vectors in dimension d: (xi)ni=1 ∼ N(0, Id). The teacher-weight vectors (wk∗)m k=1 are i.i.d. d-dimensional Gaussian vectors, and the noise (zi)ni=1 is i.i.d. Gaussian. The output labels are generated by a one-hidden layer teacher network with m hidden units and quadratic activation: yi = fW∗(xi):= (1/m) Σ k=1 m ((1/√d)(wk∗)⊤ xi + √∆ zi,k)2. The teacher weights W∗ and noise are unknown, but the form of the target function is known. The analysis is performed in the high-dimensional limit: d → ∞, with α:= n/d2 = Θ(1) and κ:= m/d = Θ(1).

The main contributions are summarized as follows:

  1. Closed-form expression for Bayes-optimal error: The authors provide the first closed-form expression for the Bayes-optimal test error for learning an extensive-width neural network from quadratically many samples. The main result, stated as Result 1, gives the asymptotic MMSE (minimum mean-squared error) as: MMSE = (2ακ/q̂) - (κ∆̃/2), where ∆̃:= 2∆(2+∆)/κ, and q̂ is a solution to the equation (1-2α) + (∆̃q̂)/2 = (4π2/(3q̂)) ∫ µ 1/q̂(y)3 dy. Here, µ t is the free convolution of the Marchenko-Pastur law and a scaled semicircular density.

  2. Perfect recovery threshold: In the noiseless case (∆=0), the formula implies that zero test error is achievable for a sample complexity α = n/d2 larger than a perfect-recovery threshold α PR, given by α PR = κ - κ2/2 if κ ≤ 1, and α PR = 1/2 if κ ≥ 1. This matches a naive counting of the number of degrees of freedom in the target function.

  3. GAMP-RIE algorithm: The authors introduce an algorithm, GAMP-RIE, which combines generalized approximate message passing (GAMP) with a rotationally-invariant matrix denoiser (RIE). They show that in the large size limit, this algorithm asymptotically achieves the Bayes-optimal error for all α, κ = Θ(1). The algorithm's state evolution is shown to match the fixed-point equations that provide the Bayes-optimal error.

  4. Technical approach: The derivation is enabled by establishing a link with recent works on optimal denoising of extensive-rank matrices and on the ellipsoid fitting problem. The approach involves a three-step proof sketch: (i) a universality conjecture that replaces the sensing matrices Zi with Gaussian GOE matrices, (ii) a conjecture on the free entropy of a matrix generalized linear model with a Wishart prior, and (iii) a theorem on the free entropy of extensive-rank matrix denoising, which is proven using results on HCIZ integrals and large deviations of Dyson Brownian motion.

  5. Empirical observations on gradient descent: The authors empirically compare the Bayes-optimal performance with that of gradient descent (GD). In the noiseless case, they observe that randomly-initialized GD seems to sample the space of interpolants, leading to twice the Bayes-optimal error. However, when averaged over initialization, the GD reaches an error very close to the Bayes-optimal one. In the noisy setting, they observe a gap between the MMSE and the performance of GD, even when averaged over initialization or regularized, and identify a landscape trivialization threshold where GD converges to the same function regardless of initialization.

The paper also discusses the small κ limit, recovering results from the finite hidden layer setting, and the large κ limit. It concludes by discussing limitations, particularly the restriction to quadratic activations, and outlines challenges for extending the analysis to generic activation functions, which would involve tensor models and open questions about their universality and denoising.

Improvements for AI systems

Based on the paper, here are specific improvements I can make to AI systems:

  • Implementation: I can implement the GAMP-RIE algorithm (Algorithm 1) as a new training method for single-hidden-layer neural networks with quadratic activations.

  • Capability: This algorithm provably achieves Bayes-optimal test error with quadratically many samples (n = αd2), whereas standard gradient descent requires more samples or achieves suboptimal error. It outperforms GD in noisy settings and matches the theoretical optimum in noiseless settings.

  • Implementation: I can integrate the analytical formula (Result 1, eqs. 8-9) into model selection tools.

  • Capability: Given the network width ratio κ = m/d and sample complexity α = n/d2, I can instantly predict the achievable test error before training, allowing practitioners to determine if they have sufficient data or need to adjust architecture. This eliminates expensive trial-and-error training runs.

  • Implementation: I can build a reusable module using the rotationally-invariant estimator (RIE) with the explicit form fRIE(R, Δ) = UfΔ(Λ)UT, where fΔ(λ) = λ - 2ΔhΔ(λ).

  • Capability: This module can denoise any symmetric matrix corrupted by Gaussian noise, achieving the theoretical minimum mean-squared error. It's useful for matrix completion, collaborative filtering, and covariance estimation tasks beyond neural networks.

  • Implementation: I can create a training scheduler that uses the perfect-recovery threshold αPR = κ - κ2/2 (for κ ≤ 1) or αPR = 1/2 (for κ ≥ 1) to automatically determine when to stop collecting data.

  • Capability: The system can tell users exactly how many samples are needed for perfect recovery given their network width, avoiding over-collection of data and reducing storage/computation costs.

  • Implementation: I can implement the averaged gradient descent (AGD) approach where multiple random initializations are averaged.

  • Capability: In noiseless settings, AGD achieves Bayes-optimal error (matching the theoretical MMSE), whereas single-run GD gives twice the optimal error. This is a simple drop-in improvement for existing training pipelines.

  • Implementation: I can use the theoretical MMSE formula to determine optimal regularization strength λ for noisy settings.

  • Capability: The paper shows that l2 regularization hurts performance in this setting. My system can automatically detect when regularization is counterproductive and suggest zero-regularization training, preventing common suboptimal choices.

  • Implementation: I can implement a detector that identifies when the loss landscape becomes trivial (all GD runs converge to same solution) using the threshold αT(Δ).

  • Capability: This tells practitioners when they can safely use single-initialization GD versus needing multiple runs, saving computational resources.

  • Implementation: I can generalize the GAMP-RIE algorithm to handle learnable second-layer weights using the framework in Appendix E (Claim 3).

  • Capability: This enables optimal learning for a broader class of neural networks where both layers are trained, with the same theoretical guarantees.

  • Implementation: I can build an efficient numerical solver for the free convolution equation (eq. 108) and the Stieltjes-Perron inversion (eq. 109) to compute µt densities.

  • Capability: This is a reusable computational tool for any problem involving free probability, including random matrix theory applications in wireless communications, finance, and quantum information.

  • Implementation: I can adapt the framework to handle non-Gaussian noise distributions using the universality conjecture (Conjecture 4.1).

  • Capability: The system can achieve optimal performance even when noise is not Gaussian, as long as the first two moments match, making it robust to real-world noise distributions.

These improvements collectively enable AI systems to achieve provably optimal performance in a specific but important class of neural network learning tasks, with clear theoretical guarantees that reduce uncertainty in deployment.

Sources

Related papers