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

TL;DR

This paper establishes explicit convergence rates of wide neural networks with general weights to Gaussian processes, using Stein's method, and explores the effects of increasing depth alongside width.

Contribution

It provides the first CLT for neural networks with Lipschitz activations and non-Gaussian weights in the simultaneous infinite width and depth regime.

Findings

Convergence rate of O(n^{-1/2}) in total variation and Wasserstein distances.

Bounds are explicit and computable for finite collections of inputs.

Novel CLT results for networks with non-Gaussian weights and increasing depth.

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.