Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Precise gradient descent training dynamics for finite-width multi-layer neural networks

Qiyang Han, Masaaki Imaizumi

TL;DR

This work delivers a first precise distributional description of gradient descent training dynamics for general multi-layer neural networks in the finite-width proportional regime, where the sample size and feature dimension grow proportionally while network width and depth stay bounded. By introducing a state-evolution framework built on an iterative reduction to auxiliary GFOM iterates, the authors obtain per-iteration Gaussian characterizations for the first-layer weights and concentration results for deeper layers, applicable to non-Gaussian features. They connect these dynamics to both training and generalization errors via low-dimensional Gaussian integrals, and propose an augmented gradient-descent procedure that yields consistent estimates of the generalization error at every iteration, enabling practical early stopping and hyperparameter tuning. A key theoretical implication is that, even under model misspecification, the learned predictor retains a single-index structure with an effective signal combining the true signal and initialization. The framework extends beyond two-layer networks, accommodates non-lazy training, and provides concrete algorithmic tools and simulation evidence, with rigorous proofs anchored in matrix-variate GFOM theory and non-asymptotic concentration analysis.

Abstract

In this paper, we provide the first precise distributional characterization of gradient descent iterates for general multi-layer neural networks under the canonical single-index regression model, in the `finite-width proportional regime' where the sample size and feature dimension grow proportionally while the network width and depth remain bounded. Our non-asymptotic state evolution theory captures Gaussian fluctuations in first-layer weights and concentration in deeper-layer weights, and remains valid for non-Gaussian features. Our theory differs from existing neural tangent kernel (NTK), mean-field (MF) theories and tensor program (TP) in several key aspects. First, our theory operates in the finite-width regime whereas these existing theories are fundamentally infinite-width. Second, our theory allows weights to evolve from individual initializations beyond the lazy training regime, whereas NTK and MF are either frozen at or only weakly sensitive to initialization, and TP relies on special initialization schemes. Third, our theory characterizes both training and generalization errors for general multi-layer neural networks beyond the uniform convergence regime, whereas existing theories study generalization almost exclusively in two-layer settings. As a statistical application, we show that vanilla gradient descent can be augmented to yield consistent estimates of the generalization error at each iteration, which can be used to guide early stopping and hyperparameter tuning. As a further theoretical implication, we show that despite model misspecification, the model learned by gradient descent retains the structure of a single-index function with an effective signal determined by a linear combination of the true signal and the initialization.

Precise gradient descent training dynamics for finite-width multi-layer neural networks

TL;DR

This work delivers a first precise distributional description of gradient descent training dynamics for general multi-layer neural networks in the finite-width proportional regime, where the sample size and feature dimension grow proportionally while network width and depth stay bounded. By introducing a state-evolution framework built on an iterative reduction to auxiliary GFOM iterates, the authors obtain per-iteration Gaussian characterizations for the first-layer weights and concentration results for deeper layers, applicable to non-Gaussian features. They connect these dynamics to both training and generalization errors via low-dimensional Gaussian integrals, and propose an augmented gradient-descent procedure that yields consistent estimates of the generalization error at every iteration, enabling practical early stopping and hyperparameter tuning. A key theoretical implication is that, even under model misspecification, the learned predictor retains a single-index structure with an effective signal combining the true signal and initialization. The framework extends beyond two-layer networks, accommodates non-lazy training, and provides concrete algorithmic tools and simulation evidence, with rigorous proofs anchored in matrix-variate GFOM theory and non-asymptotic concentration analysis.

Abstract

In this paper, we provide the first precise distributional characterization of gradient descent iterates for general multi-layer neural networks under the canonical single-index regression model, in the `finite-width proportional regime' where the sample size and feature dimension grow proportionally while the network width and depth remain bounded. Our non-asymptotic state evolution theory captures Gaussian fluctuations in first-layer weights and concentration in deeper-layer weights, and remains valid for non-Gaussian features. Our theory differs from existing neural tangent kernel (NTK), mean-field (MF) theories and tensor program (TP) in several key aspects. First, our theory operates in the finite-width regime whereas these existing theories are fundamentally infinite-width. Second, our theory allows weights to evolve from individual initializations beyond the lazy training regime, whereas NTK and MF are either frozen at or only weakly sensitive to initialization, and TP relies on special initialization schemes. Third, our theory characterizes both training and generalization errors for general multi-layer neural networks beyond the uniform convergence regime, whereas existing theories study generalization almost exclusively in two-layer settings. As a statistical application, we show that vanilla gradient descent can be augmented to yield consistent estimates of the generalization error at each iteration, which can be used to guide early stopping and hyperparameter tuning. As a further theoretical implication, we show that despite model misspecification, the model learned by gradient descent retains the structure of a single-index function with an effective signal determined by a linear combination of the true signal and the initialization.
Paper Structure (71 sections, 36 theorems, 326 equations, 6 figures, 2 tables)

This paper contains 71 sections, 36 theorems, 326 equations, 6 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Overview
  3. Our new theoretical framework
  4. A statistical application: Algorithmic estimation of generalization error
  5. A further theoretical implication: Structure of the learned model $f_{\textbf{W}^{(t)}}$
  6. Proof techniques
  7. Further related literature
  8. Other precise characterizations of the generalization error
  9. Sample complexity for learning single-index functions
  10. Organization
  11. Preliminaries: Gradient mappings
  12. Algorithmic gradient mappings
  13. Theoretical gradient mappings
  14. Meta theory: Precise gradient descent training dynamics
  15. Assumptions and some further notation
  16. ...and 56 more sections

Key Result

Proposition 2.2

The following gradient formula holds for $\alpha \in [1:L]$: Here $\mathsf{P}_{X,\bm{W}}^{(\alpha:L]}:\mathbb{R}^{m\times q}\to \mathbb{R}^{m\times q}$ is a back-propagation mapping defined via

Figures (6)

  • Figure 1: Algorithmic estimation of the generalization error.
  • Figure 2: Algorithmic estimation of the generalization error with a multi-index model.
  • Figure 3: Effect of wide networks.
  • Figure 4: Effect of non-regular activations.
  • Figure 5: Algorithmic estimation of the generalization error (noiseless).
  • ...and 1 more figures

Theorems & Definitions (75)

  • Definition 2.1
  • Proposition 2.2
  • Definition 2.3
  • Definition 2.4
  • Remark 1
  • Remark 2
  • Definition 3.1
  • Remark 3
  • Theorem 3.2
  • Remark 4
  • ...and 65 more