Weak Convergence of L2-Regularized Two-Layer Neural Networks under SGD via Mean Field Theory

Abstract

Extending the work of Sirignano and Spiliopoulos (2020), we use mean field theory to study two-layer neural networks with l2-regularization trained under single-sample online stochastic gradient descent. We prove that in the asymptotic regime of both infinite training steps and infinite hidden layer width, such a neural network weakly converges to a deterministic and unique solution that satisfies a partial differential equation of the gradient flow form seen elsewhere in optimal transport and physics. Moreover, we show that the parameters of said neural network, despite being interdependent throughout training, asymptotically become independent. These results are only subject to loose moment bounds at initialization. Our proofs utilize a probabilistic approach on the network's training evolution instead of studying the geometry of the loss surface. We also provide numerical simulation results consistent with our theoretical guarantees.