Table of Contents
Fetching ...

High-Dimensional Analysis of Gradient Flow for Extensive-Width Quadratic Neural Networks

Simon Martin, Giulio Biroli, Francis Bach

TL;DR

This work develops a high-dimensional DMFT-like framework to analyze gradient flow and Langevin dynamics for shallow quadratic networks in an extensive-width regime, where data dimension and network width grow proportionally. By introducing a Gaussian surrogate for the rank-one quadratic measurements, the authors derive self-consistent stochastic equations describing the evolution of the student weights and a typical label, and they further reduce the long-time behavior to a denoising/Oja-flow dynamic. They uncover regimes of overparameterization that guarantee convergence to global minima, reveal double-descent phenomena under label noise, and provide explicit interpolation and perfect-recovery thresholds in the small-regularization limit, with connections to Bayes-optimal predictions and ERM. The results illuminate how overparameterization influences learning and generalization in nonlinear high-dimensional models and offer a principled link between dynamical learning trajectories and static optimality predictions, with broad implications for understanding real-world neural networks.

Abstract

We study the high-dimensional training dynamics of a shallow neural network with quadratic activation in a teacher-student setup. We focus on the extensive-width regime, where the teacher and student network widths scale proportionally with the input dimension, and the sample size grows quadratically. This scaling aims to describe overparameterized neural networks in which feature learning still plays a central role. In the high-dimensional limit, we derive a dynamical characterization of the gradient flow, in the spirit of dynamical mean-field theory (DMFT). Under l2-regularization, we analyze these equations at long times and characterize the performance and spectral properties of the resulting estimator. This result provides a quantitative understanding of the effect of overparameterization on learning and generalization, and reveals a double descent phenomenon in the presence of label noise, where generalization improves beyond interpolation. In the small regularization limit, we obtain an exact expression for the perfect recovery threshold as a function of the network widths, providing a precise characterization of how overparameterization influences recovery.

High-Dimensional Analysis of Gradient Flow for Extensive-Width Quadratic Neural Networks

TL;DR

This work develops a high-dimensional DMFT-like framework to analyze gradient flow and Langevin dynamics for shallow quadratic networks in an extensive-width regime, where data dimension and network width grow proportionally. By introducing a Gaussian surrogate for the rank-one quadratic measurements, the authors derive self-consistent stochastic equations describing the evolution of the student weights and a typical label, and they further reduce the long-time behavior to a denoising/Oja-flow dynamic. They uncover regimes of overparameterization that guarantee convergence to global minima, reveal double-descent phenomena under label noise, and provide explicit interpolation and perfect-recovery thresholds in the small-regularization limit, with connections to Bayes-optimal predictions and ERM. The results illuminate how overparameterization influences learning and generalization in nonlinear high-dimensional models and offer a principled link between dynamical learning trajectories and static optimality predictions, with broad implications for understanding real-world neural networks.

Abstract

We study the high-dimensional training dynamics of a shallow neural network with quadratic activation in a teacher-student setup. We focus on the extensive-width regime, where the teacher and student network widths scale proportionally with the input dimension, and the sample size grows quadratically. This scaling aims to describe overparameterized neural networks in which feature learning still plays a central role. In the high-dimensional limit, we derive a dynamical characterization of the gradient flow, in the spirit of dynamical mean-field theory (DMFT). Under l2-regularization, we analyze these equations at long times and characterize the performance and spectral properties of the resulting estimator. This result provides a quantitative understanding of the effect of overparameterization on learning and generalization, and reveals a double descent phenomenon in the presence of label noise, where generalization improves beyond interpolation. In the small regularization limit, we obtain an exact expression for the perfect recovery threshold as a function of the network widths, providing a precise characterization of how overparameterization influences recovery.
Paper Structure (193 sections, 37 theorems, 704 equations, 21 figures)

This paper contains 193 sections, 37 theorems, 704 equations, 21 figures.

Key Result

Proposition 4

Consider the case $\kappa \geq \min(\kappa^*, 1)$ and the variables $q, \xi$, solutions of the system of equations eq:SystemResult2. Then, as $\alpha \to \infty$, we have $q \to \lambda$ and $\xi = \Theta(\alpha^{-1})$. The limit of the gradient flow is then given by: and the MSE and the training loss write:

Figures (21)

  • Figure 1: Empirical distribution of the student labels $y_k(t) = \mathrm{Tr} ( W(t)W(t)^\top X_k )$ during optimization with gradient descent, defined in equation \ref{['eq:GDdynamics']}, with parameters $\kappa = 0.4, \kappa^* = 0.3$, $d = 150$, quadratic cost and no regularization, for two values of $\alpha$ and three values of time $t$. The black curve corresponds to the Gaussian density with zero mean and variance equal to the empirical variance of the labels.
  • Figure 2: Comparison between simulations of gradient descent, defined in equation \ref{['eq:GDdynamics']}, and numerical integration of the system of equations \ref{['eq:SystemResult2']}, for $\kappa = 0.7, \kappa^* = 0.5$ and zero label noise. Gradient descent results are averaged over 10 (left) and 50 (right) realizations of the initialization, teacher and data. Left: MSE and empirical loss value as a function of $\alpha$, for different values of $\lambda$. Dots correspond to GD simulations and full lines to the solution of equations \ref{['eq:SystemResult2']}. Right: eigenvalue distribution of $Z = WW^\top$ reached by GD (blue), restricted to its nonzero eigenvalues, and associated density computed with the asymptotic spectral distribution of the matrix in equation \ref{['eq:ZinfinityReg']} (black line) for three values of $\lambda$ and two values of $\alpha$.
  • Figure 3: Rank of the solution $\kappa_{\min}$ obtained with $\kappa = 1$, as a function of $\alpha$ and for different values of $\lambda$. $\kappa^* = 0.3$ (horizontal dashed line), $\Delta = 0$ (left) and $\Delta = 0.5$ (right). Curves obtained by simulating the system \ref{['eq:SystemResult2']} and using equation \ref{['eq:kappamin']} to compute $\kappa_{\min}$. Above these curves, the solution found by gradient flow does not depend on $\kappa$.
  • Figure 4: Comparison between simulations of gradient descent, defined in equation \ref{['eq:GDdynamics']}, and numerical integration of the system of equations \ref{['eq:SystemResult2']}, for $\kappa = 0.35, \kappa^* = 0.3$ and zero label noise. Gradient descent results are averaged over 10 (left) and 50 (right) realizations of the initialization, teacher and data. Left: MSE and empirical loss value as a function of $\alpha$ for different values of $\lambda$. Dots correspond to GD simulations and full lines to the solution of the system \ref{['eq:SystemResult2']}. Right: eigenvalue distribution of $Z = WW^\top$ reached by GD (blue), restricted to its nonzero eigenvalues, and associated density computed with the asymptotic spectral distribution of the matrix in equation \ref{['eq:ZinfinityReg']} (black line) for three values of $\lambda$ and $\alpha$.
  • Figure 5: MSE as a function of time for gradient descent trajectories (see equation \ref{['eq:GDdynamics']}) averaged over 10 realizations of the initialization, teacher and data. Parameters $\kappa = 0.4, \kappa^* = 0.3, \lambda = 0.01$, $\Delta = 0.25$ (left) and $0.5$ (right). Overfitting is characterized by portions where the MSE increases with time.
  • ...and 16 more figures

Theorems & Definitions (65)

  • Claim 1
  • Claim 2
  • Claim 3
  • Proposition 4
  • Claim 5
  • Proposition 6
  • Corollary 1
  • Proposition 7
  • Proposition 8
  • Proposition 9
  • ...and 55 more