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

TL;DR

This paper provides a detailed, non-asymptotic analysis of gradient descent dynamics for finite-width multi-layer neural networks, capturing fluctuations and generalization beyond existing theories.

Contribution

It introduces the first finite-width, non-asymptotic state evolution theory for multi-layer neural networks, extending understanding beyond NTK, MF, and TP frameworks.

Findings

Captures Gaussian fluctuations in first-layer weights.

Allows weights to evolve from individual initializations.

Enables estimation of generalization error during training.

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.