Wide neural networks with general weights: convergence rate and explicit dependence on the hyper-parameters

Abstract

Using Stein's method techniques introduced by Chatterjee (2008) and further extended by Kasprzak and Peccati (2022) and by Lachi\`eze-Rey and Peccati (2017), we derive novel quantitative bounds on the convergence in distribution of feed-forward fully connected neural networks (with Lipschitz activation functions) towards Gaussian processes, as the hidden layer width $n$ tends to infinity. We consider networks initialized with independent and identically distributed (i.i.d.) weights possessing sufficiently many finite moments, and i.i.d. Gaussian biases independent of the weights. Specifically, when the network is evaluated at a single input, we obtain convergence rates of order $O(n^{-1/2})$ in both total variation and Wasserstein distances. When evaluated at a general finite collection of inputs, we establish bounds of the same order in terms of the convex distance. All bounds are given in explicit and computable form. As a consequence of our estimates, we also deduce a novel convergence result in the regime where the depth of the neural network increases simultaneously with the width $n$, up to order $O\big((\log_2 n)^{1/3}\big)$. To the best of our knowledge, this is the first CLT in the infinite width/depth limit holding for general (nonlinear) Lipschitz activation functions and non-Gaussian weight distributions. Our analysis yields several results of independent interest, including: (i) an explicit lower bound on the determinant of the limiting covariance matrix and (ii) new advances in Stein's method, both for the one-dimensional Stein's equation associated with the square of a Lipschitz function and for the multivariate Stein's equation associated with the tensor product of a Lipschitz function with itself.