Large-width asymptotics for ReLU neural networks with $\alpha$-Stable initializations

TL;DR

This paper extends large-width asymptotic analysis of neural networks to $ ext{alpha}$-Stable initializations, revealing how non-Gaussian stable distributions influence the limiting processes and training dynamics.

Contribution

It introduces the analysis of $ ext{alpha}$-Stable neural networks, showing their convergence to $ ext{alpha}$-Stable processes and characterizing their training dynamics with a new $ ext{alpha}$-Stable NTK.

Findings

Rescaled $ ext{alpha}$-Stable NNs converge to $ ext{alpha}$-Stable processes as width increases.

Activation functions influence the scaling law for $ ext{alpha}$-Stable processes.

Gradient descent achieves zero training error with a linear rate in large-width $ ext{alpha}$-Stable NNs.

Abstract

There is a recent and growing literature on large-width asymptotic properties of Gaussian neural networks (NNs), namely NNs whose weights are initialized as Gaussian distributions. Two popular problems are: i) the study of the large-width distributions of NNs, which characterizes the infinitely wide limit of a rescaled NN in terms of a Gaussian stochastic process; ii) the study of the large-width training dynamics of NNs, which characterizes the infinitely wide dynamics in terms of a deterministic kernel, referred to as the neural tangent kernel (NTK), and shows that, for a sufficiently large width, the gradient descent achieves zero training error at a linear rate. In this paper, we consider these problems for $\alpha$-Stable NNs, namely NNs whose weights are initialized as $\alpha$-Stable distributions with $\alpha\in(0,2]$. First, for $\alpha$-Stable NNs with a ReLU activation function, we show that if the NN's width goes to infinity then a rescaled NN converges weakly to an $\alpha$-Stable stochastic process. As a difference with respect to the Gaussian setting, our result shows that the choice of the activation function affects the scaling of the NN, that is: to achieve the infinitely wide $\alpha$-Stable process, the ReLU activation requires an additional logarithmic term in the scaling with respect to sub-linear activations. Then, we study the large-width training dynamics of $\alpha$-Stable ReLU-NNs, characterizing the infinitely wide dynamics in terms of a random kernel, referred to as the $\alpha$-Stable NTK, and showing that, for a sufficiently large width, the gradient descent achieves zero training error at a linear rate. The randomness of the $\alpha$-Stable NTK is a further difference with respect to the Gaussian setting, that is: within the $\alpha$-Stable setting, the randomness of the NN at initialization does not vanish in the large-width regime of the training.