A Functional-Space Mean-Field Theory of Partially-Trained Three-Layer Neural Networks

TL;DR

This paper develops a mean-field theory for a partially-trained three-layer neural network with a fixed first layer, showing linear decay of training loss and feature learning in the infinite-width limit.

Contribution

It generalizes mean-field theory to three-layer networks with a fixed first layer, introducing a functional-space perspective and analyzing training dynamics.

Findings

Training loss decays linearly to zero in the mean-field limit.

The theory captures feature learning in different scaling regimes.

Provides Rademacher complexity bounds for the solution space.

Abstract

To understand the training dynamics of neural networks (NNs), prior studies have considered the infinite-width mean-field (MF) limit of two-layer NN, establishing theoretical guarantees of its convergence under gradient flow training as well as its approximation and generalization capabilities. In this work, we study the infinite-width limit of a type of three-layer NN model whose first layer is random and fixed. To define the limiting model rigorously, we generalize the MF theory of two-layer NNs by treating the neurons as belonging to functional spaces. Then, by writing the MF training dynamics as a kernel gradient flow with a time-varying kernel that remains positive-definite, we prove that its training loss in $L_2$ regression decays to zero at a linear rate. Furthermore, we define function spaces that include the solutions obtainable through the MF training dynamics and prove Rademacher complexity bounds for these spaces. Our theory accommodates different scaling choices of the model, resulting in two regimes of the MF limit that demonstrate distinctive behaviors while both exhibiting feature learning.