Genus expansion for non-linear random matrix ensembles with applications to neural networks

TL;DR

This paper introduces a unified framework using genus expansion and a novel series expansion to analyze non-linear random matrix ensembles and neural networks, providing new proofs and quantitative results on their convergence and spectral properties.

Contribution

It develops a generalized series expansion for neural networks that linearizes activation effects, enabling new analytical results on convergence and spectral distribution in neural networks.

Findings

Neural networks converge to Gaussian processes as width increases.

Quantified the convergence rate of Neural Tangent Kernel.

Computed moments of the spectral distribution of the Jacobian.

Abstract

We present a unified approach to studying certain non-linear random matrix ensembles and associated random neural networks at initialization. This begins with a novel series expansion for neural networks which generalizes Fa\'a di Bruno's formula to an arbitrary number of compositions. The role of monomials is played by random multilinear maps indexed by directed graphs, whose edges correspond to random matrices. Crucially, this expansion linearizes the effect of the activation functions, allowing for the direct application of Wick's principle and the genus expansion technique. As an application, we prove several results about neural networks with random weights. We first give a new proof of the fact that they converge to Gaussian processes as their width tends to infinity. Secondly, we quantify the rate of convergence of the Neural Tangent Kernel to its deterministic limit in Frobenius norm. Finally, we compute the moments of the limiting spectral distribution of the Jacobian (only the first two of which were previously known), expressing them as sums over non-crossing partitions. All of these results are then generalised to the case of neural networks with sparse and non-Gaussian weights, under moment assumptions.