Precise characterization of the prior predictive distribution of deep ReLU networks

Introduction

It is well known that standard neural networks initialized with Gaussian weights tend to Gaussian processes (Rasmussen,, 2003) in the infinite width limit (Neal,, 1996; Lee et al.,, 2018; de G. Matthews et al.,, 2018), coined neural network Gaussian process (NNGP) in the literature. Although the NNGP has been derived for a number of architectures, such as convolutional (Novak et al.,, 2019; Garriga-Alonso et al.,, 2019), recurrent (Yang,, 2019) and attention mechanisms (Hron et al.,, 2020), little is known about the finite width case.

The reason why the infinite width limit is relatively tractable to study is that uncorrelated but dependent units of intermediate layers become normally distributed due to the central limit theorem (CLT) and as a result, independent. In the finite width case, zero correlation does not imply independence, rendering the analysis far more involved as we will outline in this paper.

One of the main motivations for our work is to better understand the implications of using Gaussian priors in combination with the compositional structure of the network architecture. As argued by Wilson and Izmailov, (2020); Wilson, (2020), the prior over parameters does not carry a meaningful interpretation; the prior that ultimately matters is the prior predictive distribution that is induced when a prior over parameters is combined with a neural architecture (Wilson and Izmailov,, 2020; Wilson,, 2020).

Studying the properties of this prior predictive distribution is not an easy task, the main reason being the compositional structure of a neural network, which ultimately boils down to products of random matrices with a given (non-linear) activation function. The main tools to study such products are the Mellin transform and the Meijer-G function (Meijer,, 1936; Springer and Thompson,, 1970; Mathai,, 1993; Stojanac et al.,, 2017), both of which will be leveraged in this work to gain theoretical insights into the inner workings of BNN priors.

Our results provide an important step towards understanding the interplay between architectural choices and the distributional properties of the prior predictive distribution, in particular:

We characterize the prior predictive density of finite-width ReLU networks of any depth through the framework of Meijer-G functions.

We draw analytic insights about the shape of the distribution by studying its moments and the resulting heavy-tailedness. We disentangle the roles of width and depth, demonstrating how deeper networks become more and more heavy-tailed, while wider networks induce more Gaussian-like distributions.

We connect our work to the infinite width setting by recovering and extending prior results (Lee et al.,, 2018; Matthews et al.,, 2018) to the infinite depth limit. We describe the resulting distribution in terms of its moments and match it to a normal log-normal mixture (Yang,, 2008), empirically providing an excellent fit even in the non-asymptotic regime.

Finally, we introduce generalized He priors, where a desired variance can be directly specified in the function space. This allows the practitioner to make an interpretable choice for the variance, instead of implicitly tuning it through the specification of each layer variance.

The rest of the paper is organized as follows: in Section 3.1, we introduce the relevant notation for the neural network that will be analyzed. We describe the prior works of Lee et al., (2018); Matthews et al., (2018) in more detail in Section 3.2 . Then, in Section 3.3, we introduce the Meijer-G function and the necessary mathematical tools. In Section 4 we derive the probability density function for a linear network of any depth and extend these results to ReLU networks, which represents the key contribution of our work. In Section 5, we present several consequences of our analysis, including an extension of the infinite width setting to infinite depth as well as precise characterizations of the heavy-tailedness in the finite regime. Finally, in Section 6, we show how one can design an architecture and a prior over the weights to achieve a desired prior predictive variance.

Related Work

Although Bayesian inference in deep learning has recently risen in popularity, surprisingly little work has been devoted to investigating standard priors and their implied implicit biases in neural architectures. Only through the lens of infinite width, progress has been made (Neal,, 1996), establishing a Gaussian process behaviour at the output of the network. More recently, Lee et al., (2018) and Matthews et al., (2018) extended this result to arbitrary depth. We give a brief introduction in Section 3.2. Due to their appealing Gaussian process formulation, infinite width networks have been extensively studied theoretically, leading to novel insights into the training dynamics under gradient descent (Jacot et al.,, 2018) and generalization (Arora et al., 2019b, ).

Although theoretically very attractive, the usage of infinite width has been severely limited by its inferior empirical performance (Arora et al., 2019a, ). While first insights into this gap have been obtained (Aitchison,, 2020; Aitchison et al.,, 2021), the picture is far from complete and a better understanding of finite width networks is still highly relevant for practical applications. In the finite regime, however, such precise characterizations in function space have been elusive so far and largely limited to empirical insights (Flam-Shepherd et al.,, 2017, 2018) and investigations of the heavy-tailedness of layers (Vladimirova et al.,, 2019; Fortuin et al.,, 2021). The field of finite-width corrections has recently gained a lot of attraction. Hanin and Nica, (2020) studies the simultaneous limit of width and depth of the Jacobian of a ReLU net. More recently, a number of concurrent works appearing shortly before or after ours, such as Zavatone-Veth and Pehlevan, (2021); Roberts et al., (2021); Li et al., (2021) study properties of large-but-finite neural nets. Notably, Zavatone-Veth and Pehlevan, (2021) concurrently derived similar results on the characterization of the prior predictive distribution of finite width networks, however, we did not discover them until after we had completed this work. Note also that our novel limiting behaviours (cf. Sec. 5.1) and our analytical insights into heavy-tailedness (cf. Sec. 5.2) clearly distinguish our work from theirs. In addition, our results also offer valuable guidance on prior design for ML practitioners (cf. Sec. 6). Li et al., (2021) derive a similar limiting result in the infinite-width and depth setting in the case of Resnets (ours if for fully-connected, cf. Sec. 5.1), albeit with a completely different proof technique. Moreover, we also give precise insights into the prior predictive distribution for finite width (cf. Sec. 4), whereas Li et al., (2021) only work with the limits.

Finally, this work is also related to the studies of signal propagation into finite-width random networks (Poole et al.,, 2016; Schoenholz et al.,, 2017) and initialization (He et al.,, 2015; Hanin and Rolnick,, 2018). In particular, He et al., (2015) uses a second moment analysis to specify the variance of the weights. In this sense, our approach extends it by deriving all the moments of the distribution.

Background

where $\bm{\theta}=(\bm{W}^{(1)},\dots,\bm{W}^{(L)})$ denotes the collection of all weights. Throughout this work, we assume standard initialization, i.e. $W_{ij}^{(l)}\sim\mathcal{N}(0,\sigma_{l}^{2})$, where weights in each layer can have a different variance $\sigma_{l}^{2}$. Often, it will be more convenient to work with the corresponding unit-level formulation, expressed through the recursive equations:

2 Prior Predictive Distribution and Infinite Width

3 Meijer-G function

The Meijer-G function is the central tool to our analysis of the predictive prior distribution in the finite width regime. The Meijer-G function is a ubiquitous tool, appearing in a variety of scientific fields ranging from mathematical physics (Pishkoo and Darus,, 2013) to symbolic integration software (Adamchik and Marichev,, 1990) to electrical engineering (Ansari et al.,, 2011). Despite its high popularity in many technical fields, there have only been a handful of works in ML leveraging this elegant and convenient theoretical framework (Alaa and van der Schaar,, 2019; Crabbe et al.,, 2020). In the following, we will introduce the Meijer-G function along with the relevant mathematical tools to develop our theory.

The defining property of Meijer-G functions is their closure under integration, i.e. the convolution of two Meijer-G functions is again a Meijer-G function. Combined with the fact that most elementary functions can be written as a Meijer-G function, this property becomes extremely powerful at expressing complicated integrals neatly. Our proofs leverage this result extensively by expressing the integrands encountered in the prior predictive function as Meijer-G functions. Throughout this text, we will only encounter Meijer-G functions of a simpler signature $G^{l,0}_{0,l}\left(\cdot|\bm{b}\right)$. For completeness, we show its functional form here:

Predictive Priors for Neural Networks

In this section we detail our theoretical results on the predictive prior distribution implied by a fully-connected neural network with Gaussian weights, both with and without ReLU non-linearities.

First, we consider linear networks, i.e. fully-connected networks where the post-activations coincide with the pre-activations. They can be characterized as the product of Gaussian random matrices, for which the result in terms of Meijer-G functions is known (consider for instance Ipsen, (2015)). To highlight the differences between the linear and the non-linear approach, we re-prove the linear case leveraging our proof technique and notation. For simplicity, we assume w.l.o.g. that the input is normalized, i.e $||\bm{x}{||}=1$. We now present the resulting distribution of the predictive prior:

Theorem 4.1. Suppose $l\geq 1$, and the input has dimension $d_{0}$. Then, the joint marginal density of the random vector $\bm{f}^{(l)}$ (i.e. the density of the $l$-th layer pre-activations) is proportional to: p(\bm{f}^{(l)})\propto G^{l,0}_{0,l}\left(\frac{||\bm{f}^{(l)}||^{2}}{2^{l}\sigma^{2}}\bigg{\rvert}0,\frac{1}{2}\left(d_{1}-d_{l}\right),\dots,\frac{1}{2}\left(d_{l-1}-d_{l}\right)\right), (9) where $\sigma^{2}=\prod_{i=1}^{l}\sigma_{i}^{2}$ . Our proof is based on an inductive technique by conditioning on the pre-activations $\bm{f}^{(l)}$ of the previous layer while analyzing the pre-activations of the next layer. Due to space constraints, we defer the full proof of Thm 4.1 to the Appendix B.2. In the following, to give a flavor of our technique and highlight the great utility of the Meijer-G function, we present the base case as well as the key technical result (Lemma 4.2) to perform the inductive step to obtain the final statement.

Here we restrict our attention to the first layer pre-activations which, given an input $\bm{x}$, are defined as:

Conditioned on the input, this is a sum of $d_{0}$ i.i.d. Gaussian random variables $W^{(1)}_{jk}\sim\mathcal{N}(0,\sigma^{2}_{1})$, which is again Gaussian with mean zero and variance $\sigma_{1}^{2}||\bm{x}||^{2}$, i.e. $f^{(1)}_{k}(\bm{x})\sim\mathcal{N}(0,\sigma_{1}^{2}||\bm{x}||^{2})\stackrel{{\scriptstyle(d)}}{{=}}\mathcal{N}(0,\sigma_{1}^{2})$. where the last equality follows from $||{\mathbf{x}}{||}=1$. As we are conditioning on the input, the joint distribution of the first layer units is composed of $d_{1}$ independent Gaussians and hence

Indeed, as anticipated from Thm. 4.1, the corresponding Meijer-G function encodes the Gaussian density:

Conditioning on the previous pre-activations $\bm{f}^{(l-1)}$ brings us into a similar setting as in the base case because the resulting conditional distribution is again a Gaussian:

In contrast to the fixed input $\bm{x}$, we now need to apply the law of total probability to integrate out the dependence on $\bm{f}^{(l-1)}$, leveraging the induction hypothesis for $p(\bm{f}^{(l-1)})$, to obtain the marginal distribution $p(\bm{f}^{(l)})$. This is where the Meijer-G function comes in handy as it can easily express such an integral. In combination with the closedness of the family under integration, this enables us to perform an inductive proof. We summarize this in the following Lemma:

In the previous paragraph we have computed the prior predictive distribution for a linear network with Gaussian weights. Now, we extend these results to ReLU networks. The proof technique used is very similar to its linear counterpart, the main difference stems from the need to decompose the distribution over active and inactive ReLU cells. As a consequence, the resulting density is a superposition of different Meijer-G functions, each associated with a different active (linear) subnetwork. This is presented in the following:

Analytic Insights into the Prior Predictive Distribution

Here we highlight how one can use the mathematical machinery of Meijer-G functions to derive interesting insights, relying on numerous mathematical results provided in the literature (Gradshteyn and Ryzik,, 2013; Brychkov,, 2008; Andrews,, 2011). With this rich line of work at our disposal, we can easily move from the rather abstract but mathematically convenient world of Meijer-G functions to very concrete results. We demonstrate this by recovering and extending the NNGP results provided in Lee et al., (2018); Matthews et al., (2018). In particular, our analysis allows for simultaneous width and depth limits, showing how different limiting distributions emerge as a consequence of the growth with respect to $L$. Finally, we characterize the heavy-tailedness of the prior-predictive for any width, providing further evidence that deeper models induce distributions with heavier tails, as observed in Vladimirova et al., (2019). We hope that this work paves the way for further progress in understanding priors, leveraging this novel connection through the rich literature on Meijer-G functions.

We will start by giving an alternative proof in the linear case for the Gaussian behaviour emerging as the width of the network tends to infinity, recovering the results of Lee et al., (2018); Matthews et al., (2018) in the restricted setting of having just one fixed input $\bm{x}$. We extend their results in the following ways:

We provide a convergence proof that is independent of the ordering of limits.

We characterize the distributions arising from a simultaneous infinite width and depth limit, considering different growth rates for depth $L$.

For ease of exposition, we focus on the equal width case, i.e. where $d_{1}=\dots=d_{l-1}=m$ with one output $d_{L}=1$. To have well-defined limits, one has to resort to the so-called NTK parametrization (Jacot et al.,, 2018), which is achieved by setting the variances as $\sigma_{1}^{2}=1$ and $\sigma_{i}^{2}=\frac{2}{m}$ for $i=2,\dots,l$. We summarize the result in the following:

2 Heavy-tailedness Increases with Depth

From the moments analysis, it is simple to recover a known fact about the prior distribution of neural networks, namely that deeper layers are increasingly heavy-tailed (Vladimirova et al.,, 2019). To see this, we can derive the kurtosis, a standard measure of tailedness of the distribution (Westfall,, 2014), defined as:

where $X$ is a univariate random variable with finite fourth moment. We can calculate the kurtosis of a ReLU network analytically at any width, relying on closed-form results for the lower order moments of the Binomial distribution. We outline the exact calculation in the Appendix C and state the resulting expression here:

Note how the kurtosis increases with depth $L$ and decreases with the width $m$, highlighting once again the opposite roles those two parameters take regarding the shape of the prior. As expected, for fixed depth $L$, the kurtosis converges to $3$, i.e. $\lim_{m\xrightarrow[]{}\infty}\kappa_{\text{ReLU}}(m,L)=3$, which is the kurtosis of a standard Gaussian variable $\mathcal{N}(0,1)$. For the simultaneous limit $L=\gamma m$, we find a value of $3e^{5\gamma}$, which exactly matches with the expression derived in Thm. 5.1:

Prior Design

In Section 5.2, we have already outlined how the choice of architecture influences the heavy-tailedness of the distribution. If a Gaussian-like output is desired, the architecture should be designed in such a way that the width $m$ significantly exceeds the depth $L$ (small $\gamma)$, while heavy-tailed predictive priors can be achieved by considering regimes where $L$ exceeds $m$ (big $\gamma)$. As we will see shortly after in this section, another consequence of our analysis is that we can now directly work with the variance in function space, instead of implicitly tuning it by changing the variances at each layer. In this way, the variance over the weights has a clear interpretation in terms of predictive variance, making it easier for the deep learning practitioner to take an informed decision when designing a prior for BNNs. This type of variance analysis has previously been used to devise better initialization schemes for neural networks coined "He-initialization" (He et al.,, 2015) in the Gaussian case and "Xavier initialization" (Glorot and Bengio,, 2010) for uniform weight initializations. Therefore, we coin the resulting prior Generalized He-prior.

We consider the parametrization $\sigma_{1}^{2}=1$ and $\sigma_{i}^{2}=\frac{2t_{i}^{2}}{m}$, which, as we show in Appendix C, leads to the following predictive variance:

Suppose we want a desired output variance $\sigma^{2}$. This can be achieved as follows: let $a_{i}$, with $i=2,\dots,l$, be $l-1$ coefficients such that $\sum_{i=1}^{l-1}a_{i}=l-1$. Then choose $t_{i}^{2}=\left(\sigma^{2}\right)^{\frac{a_{i}}{l-1}}$, implying that the layer variances $\sigma_{i}^{2}$ are given as

He-priors correspond to the special case where $\sigma^{2}=1$ and $a_{i}=1$, $i=1,\dots,l-1$, while a standard Gaussian $N(0,1)$ prior results in $\sigma^{2}=\left(\frac{m}{2}\right)^{l-1}$, $\sigma_{1}^{2}=1$, and $a_{i}=1$, $i=1,\dots,l-1$. Note how "far" He-priors can be from standard Gaussian for relatively deep and wide neural networks. While rather innocent-looking, using a standard Gaussian $\mathcal{N}(0,1)$ prior can lead to an extremely high output variance.

Combining with the previous insights, practitioners can now choose a desired output variance along with a desired level of heavy-tailedness in a controlled manner by specifying the architecture, i.e. the width $m$ and the depth $L$.

Discussion

Our work sheds light on the shape of the prior predictive distribution arising from imposing a Gaussian distribution on the weights. Leveraging the machinery of Meijer-G functions, we characterized the density of the output in the finite width regime and derived analytic insights into its properties such as moments and heavy-tailedness. An extension to the stochastic process setting in the spirit of Lee et al., (2018); Matthews et al., (2018) as well as to convolutional architectures (Novak et al.,, 2019) could bring theory even closer to practice and we expect similar results to also hold in those cases. This is however beyond the scope of our work and we leave it as future work.

Our technique enabled us to extend the NNGP framework to infinite depth, discovering how in the more general case, the resulting distribution shares the same moments as a normal log-normal distribution. This allowed us to disentangle the roles of width and depth, where the former induces a Gaussian-like distribution while the latter encourages heavier tails. Empirically, we found that the normal log-normal mixture provides an excellent fit to the true distribution even in the non-asymptotic setting, capturing both the cumulative and probability density function to a very high degree of accuracy. This surprising observation begs further theoretical and empirical investigations. In particular, discovering a suitable stochastic process incorporating the normal log-normal mixture, could lend further insights into the inner workings of neural networks. Moreover, the role of heavy-tailedness regarding generalization is very intriguing, potentially being an important reason underlying the gap between infinite width networks and their finite counterparts. This is also in-line with recent empirical works on priors, suggesting that heavy-tailed distributions can increase the performance significantly (Fortuin et al.,, 2021).

Using these insights, we described how one can choose a fixed prior variance directly in the output space along with a desired level of heavy-tailedness resulting from the choice of the architecture. We leave it as future work to also consider higher moments to give more nuanced control over the resulting prior in the function space. Finally, we hope that the introduction of the Meijer-G function sparks more theoretical research on BNN priors and their implied inductive biases in function space.

References

Appendix

Here we describe the properties of Meijer-G functions which we will use extensively in the following.

The first result concerns the Mellin transform of the Meijer-G function, which will be the key to solve the integrals that we will face later.

See Chapter 3.2 and 2.3 of Mathai and Saxena, (2006). Conditions of validity: for the class of Meijer-G functions that we consider here ($p$, $n=0$, $m=q=l$ and the coefficients are all real) the Mellin transform exists (see Mathai and Saxena, (2006), Section 2.3.1). ∎

To establish the base case $m=1$, we need the following results.

\exp(z)=G^{1,0}_{0,1}\left(-z\bigg{\rvert}0\right) .

multiplication by power property: $z^{d}G^{m,n}_{p,q}\left(z|\begin{smallmatrix}a_{1}&\dots&a_{p}\\ b_{1}&\dots&b_{q}\end{smallmatrix}\right)=G^{m,n}_{p,q}\left(z|\begin{smallmatrix}a_{1}+d&\dots&a_{p}+d\\ b_{1}+d&\dots&b_{q}+d\end{smallmatrix}\right)$ .

See Chapter 2.6 of Mathai and Saxena, (2006) for the first identity. The last property follows directly from the definition of Meijer-G function. ∎

To perform the inductive step, we will encounter the following integral, that can be expressed in terms of the Meijer-G function.

The conditions of validity for the class of Meijer-G that we consider here are again satisfied (see Appendix B of Stojanac et al., (2017)).

Appendix B Proof for Linear Networks and Derivation of their Moments

Here we collect all the results regarding linear networks, establishing the relevant technical Lemmas to derive the density and calculate the moments of the resulting distribution.

The units of any layer are uncorrelated, i.e.

for all layers $l$, and for all $k$, $k^{\prime}$ $\in[d_{l}]$

However, they are not independent, but only conditionally independent given the previous later’s units. As a remark, note that as $d_{1}\rightarrow\infty$, the units $f_{k}^{(2)}$ approach a Gaussian distribution, for which uncorrelation implies independence.

Here we prove the main technical Lemma (Lemma 4.2) that allows us to perform the inductive step.

, where it can be shown that its determinant is:

By noting that the integral we are trying to solve depends only on $||\bm{f}^{(l-1)}||^{2}=r^{2}$, we have that the density is, up to a normalization constant independent of $\bm{f}^{(l)}$:

where we can change the order of integration due to the fact that the integrand is positive in the integration region (Tonelli’s theorem). Now by using Proposition A.1, the inner integral has the following solution:

where we have simply applied the definition of the Meijer-G function.

B.2 Probability Density Function for linear networks

We proof the result on the probability density function for a linear network in the following.

Theorem. Suppose $l\geq 1$, and the input has dimension $d_{0}$. Then, the joint marginal density of the random vector $\bm{f}^{(l)}$ (i.e. the density of the $l$-th layer pre-activations) is proportional to: p(\bm{f}^{(l)})\propto G^{l,0}_{0,l}\left(\frac{||\bm{f}^{(l)}||^{2}}{2^{l}\sigma^{2}}\bigg{\rvert}0,\frac{1}{2}\left(d_{1}-d_{l}\right),\dots,\frac{1}{2}\left(d_{l-1}-d_{l}\right)\right), (54) where $\sigma^{2}=\prod_{i=1}^{l}\sigma_{i}^{2}$ . Proof. We proof by induction. For the base case, consider $l=1$. We have shown that

Therefore we can re-write its density as:

where we have used the identity between the exponential function and the Meijer-G function (Proposition A.2).

Now we can use the fact that the the units of the $l$-th layer are conditionally independent given the previous’ layer units. Furthermore the conditional distribution is Gaussian due to the fact that the weights are i.i.d Gaussian. Therefore we can write:

In the first step we have marginalized out the units of the $l-1$ layer, and applied the product rule of probabilities. In the second step we have applied the induction hypothesis.

The integral is in the form of Lemma 4.2. For the coefficients of the Meijer-G function $b_{2}=\frac{1}{2}\left(d_{1}-d_{l-1}\right)$, $\dots$, $b_{l-1}=\frac{1}{2}\left(d_{l-2}-d_{l-1}\right)$, note that:

holds for all $i\in[d_{l-2}]$ and clearly $b_{1}+\frac{1}{2}\left(d_{l-1}-d_{l}\right)=\frac{1}{2}\left(d_{l-1}-d_{l}\right)$ as $b_{1}=0$ in our case. Therefore by Lemma 4.2 we can conclude that:

B.3 CDF of prior predictive

We also derive the CDF of the linear network in the following theorem and proceed to prove it.

Theorem B.2 (CDF of prior predictive). Let $f^{l}$ be the output of a of a linear network of $l$ layers. We assume the final layer is one dimensional. Then the the cdf is $F_{l}(t):=1-P(f^{l}>t)$, $t>0$. We have that P(f^{l}>t)=\frac{t}{2C}G^{l+1,0}_{1,l+1}\left(\omega t^{2}\bigg{\rvert}\begin{smallmatrix}\frac{1}{2}&&&&\\ -\frac{1}{2}&0&b_{1}&\dots&b_{l-1}\end{smallmatrix}\right), (64) where $b_{i}=\frac{1}{2}(d_{i}-1)$, $i\in[l-1]$, $C$ is the normalization constant and $\omega=\frac{1}{2^{l}\sigma^{2}}$. Proof. Let $X=f^{l}$.

where in the first step we have used the result of 4.1, in the second step we have applied the substitution $y=\frac{x^{2}}{t^{2}}$, and in the last step we have used Equation 27 with $\rho=\frac{1}{2}$, $\sigma=1$ and $\alpha=\omega t^{2}$ ∎

B.4 Resulting Moments for Linear Networks

We are interested in the k-th moment of $Z$. Using spherical coordinates and the properties of the Meijer-G function in a similar way as the proofs above, we get:

Note that it can be equivalently written as:

If $d_{1}=\dots=d_{l-1}=m$, and $d_{l}=1$. For instance the variance ($k=1$) isBy symmetry, all the odd moments are zero:

B.5 Infinite width and depth limit

We also present the infinite-width and infinite-depth result for the linear case. Due to the linear nature, the proof simplifies significantly compared to the ReLU case.

where $\sigma^{2}=\prod_{i=1}^{l}\sigma_{i}^{2}$. Assuming $d_{1}=\dots d_{l-1}=m$ and $d_{l}=1$ and the NTK parametrization $\sigma_{1}^{2}=1$ and $\sigma_{2}^{2}=\dots=\sigma_{l}^{2}=\frac{1}{m}$ simplifies this to

Define the $k$-th order polynomial $p(m)=\left(\frac{m}{2}+k-1\right)\dots\left(\frac{m}{2}+1\right)\frac{m}{2}$. Denote its coefficients by $\alpha_{i}$ for $i=1,\dots,k$. We know that $\alpha_{k}=\frac{1}{2^{k}}$ and from Lemma C.3 that

Assuming constant depth, performing the division by $m^{k}$ thus leads to

Choosing $s=-\frac{\gamma}{2}$ and $t^{2}=\frac{\gamma}{2}$ hence recovers the moments exactly. ∎

B.6 Normalization Constant and Angular Constant

We complete the picture by calculating the normalization constant of the resulting distribution.

B.7 Angular constant

Now we can apply the Legendre duplication formula:

to the numerator term $2z=d_{l}-k-1$ and get:

from which we can conclude, after some elementary algebraic manipulations:

Finally we prove the technical Lemma that we used in the previous proof.

Lemma B.6. $\int_{0}^{\pi}\sin^{k}(x)dx=\frac{\Gamma(k)}{2^{k}\Gamma\left(\frac{k}{2}\right)\Gamma\left(\frac{k}{2}+1\right)}(2\pi).$ (104) Proof. By integrating by parts, and using some algebraic manipulation, it is easy to see that:

Evaluating the integral between and $\pi$, we get:

The following expression includes both the even and the odd case:

where $k!!$ is the double factorial. Following a very similar procedure, we get the desired result. ∎

In the ReLU case, we will see that that the integral is from to $\frac{\pi}{2}$. In that case, we get:

B.8 Kurtosis of Linear Networks

Using the closed form expressions from Section B.4, we can describe the kurtosis of the output $f^{(l)}$ as

In particular, the distribution is always more heavy-tailed than a Gaussian, for which $\kappa=3$ (i.e. the distribution is leptokurtic). The second obvious conclusion is that depth increases the heavy-tailedness exponentially, which is in-line with the theoretical results of (Vladimirova et al.,, 2019). On the contrary, the width has the effect of "normalizing" the distribution, in particular in the limit of large width we have that:

which is the kurtosis of the Gaussian distribution, as anticipated from Lemma B.3.

Appendix C Proofs for ReLU networks and Derivation of their Moments

where in the second to last step we have used the well known property of the Delta function $\int_{-\infty}^{\infty}f(x)\delta(x-x_{0})dx=f(x_{0})$ and in the last step we used the fact that the density $p$ is symmetric around 0. ∎

C.2 Proof of Theorem 4.3

The proof is again by induction. The base case ($l=2$) is stated in Lemma C.2. For the general case, we use again an identical approach as in Theorem 4.1. We expand the coefficients and write $c^{l}_{r_{l}}:=\binom{d_{l}}{r_{l}}\frac{1}{2^{d_{l}}\Gamma\left(\frac{d_{l}-r_{l}}{2}\right)}$, where $l>0$ is the layer index. Induction step: assume that the pre-nonlinearities have the following form:

where $d_{A}:=|\mathcal{S}\setminus A|=d_{l-1}-|A|$. Also, here we abuse the notation and consider that $\mathcal{S}$ is not in the power set, i.e., $\mathcal{S}\not\in\Omega$. This is to be consistent with the fact that we are handling the case in which at least one unit is active after the ReLU activation is applied. Now following a similar procedure as in Lemma C.2, we have

For each set $A$ we have a $(d_{l-1}-|A|)$- dimensional integral that can be solved using once again Lemma 4.2 see proof of Lemma C.2 for a small but important detail of this integral. Note that for the new Meijer-G coefficients of Lemma 4.2:

holds for all $i\in[d_{l-2}]$. Therefore the solution of each integral is equal to

The new coefficient for every set $A$ is:

Therefore, because the dependence on $A$ is only through its cardinality $r^{l-1}_{l-1}:=|A|$, we define:

The final form of this equation stated in the theorem is obtained by grouping all the coefficients not involving the Meijer-G function, and substituting $r_{i}\leftarrow d_{i}-r_{i}$ and use the property $\binom{d_{i}}{r_{i}}=\binom{d_{i}}{d_{i}-r_{i}}$.

If at least in one layer it happens that all post-activations are zero, then the distribution of the network is a point mass at 0. Let’s call this event $E$, and its probability $q_{0}$. The probability of its complement $\bar{E}$ is the probability that for all the intermediate layers, at least one unit is active. These are $l-1$ independent events, the probability of each being $\frac{2^{d_{i}}-1}{2^{d_{i}}}$ (one unit is active in $2^{d_{i}}-1$ cases out of the all possible combinations of units). Therefore we can conclude that:

C.3 Base case for ReLU nets

Note that the above integral is $(d_{1}-|A|)$ dimensional due to the effect of the delta. Now the integral(s) above can be solved in an equivalent manner as in the previous section using Lemma 4.2Note that the integral is only for the positive reals. Lemma 4.2 can still be used because when switching to spherical coordinates, we are interested in the radius part, while the angular constant can still be calculated, but now we the angles are all from 0 to $\frac{\pi}{2}$, and they are equal to

where we have used the fact that the expression depends on the set $A\in\Omega$ only through $|A|$, and therefore we can use the fact that the number of subsets with $r$ elements is given by the binomial coefficient $\binom{d_{1}}{r}$. Define:

Finally, there is the special case where all the units are inactive (set to zero). This happens with probability $q_{0}=\frac{1-2^{d_{i}-1}}{2^{d_{i}}}$. ∎

Any non empty subset of $d<d_{2}$ units has the same distribution (with terms involving $d_{2}$ replaced by $d$).

C.4 Resulting moments

Let $d_{l}=1$, $d_{1},\dots,d_{l-1}=m$, and $b_{i}=\frac{1}{2}(r_{i}-1)$, $i=1,\dots,l-1$.

For instance, for the variance ($k=1$) the sum becomes:

Now each sum can be solved independently:

Note how the variance of a ReLU net is significantly reduced if compared with the variance of a linear network of the same depth (compare with Eq. 78). Similarly, one can get the fourth moment:

Note how ReLU nets are more heavy-tailed than linear nets.

To calculate the asymptotic moments we need three technical Lemmas that express the quantities encountered in a better form. First we describe the coefficients of a factorized polynomial:

Then it holds that $\alpha_{m}=1$ and $\alpha_{m-1}=\sum_{i=1}^{m}a_{i}$.

Next we use Lemma C.3 to write the ratio of Gamma functions as a polynomial:

where $P_{k}$ is a $k$-th order polynomial with coefficients $\alpha_{k}=2^{-k}$ and $\alpha_{k-1}=\frac{k^{2}-k}{2^{k}}$.

The leading coefficient can easily be obtained from multiplying together the terms $\frac{m}{2}$. From Lemma C.3 we conclude that

Next we need to control the sums involving the factorials. Since we just expressed the ratio of Gamma functions as a polynomial, we essentially need to know how to control sums of the type

which amounts to controlling the moments of a binomial distribution with fault probability $p=\frac{1}{2}$. We do this as follows:

where $Q_{k-1}$ is a $k-1$-th order polynomial. Moreover, writing $Q_{l}$ in monomial basis

For a proof of the recursion, we refer to Boros and Moll, (2004); Benyi, (2005). Moreover the polynomials satisfy the recursion

Denote by $\alpha^{(k)}$ the coefficients of $Q_{k}$, so $\alpha^{(k)}_{0},\dots,\alpha^{(k)}_{k}$. Notice that the leading coefficient of $Q_{k}$ is thus $\alpha^{(k)}_{k}$ and for $Q_{k-1}$ it is $\alpha^{(k-1)}_{k-1}$. Using the recursion and performing a comparison of coefficients we see that

We have to understand the terms involving $m^{k-1}$. Thus we need to expand $(m-1)^{k}$ which we can do with the help of Lemma C.3:

We also need to expand the next polynomial as follows:

Collecting all the coefficients, we end up with the following recursion for the second coefficient:

Thus $\alpha_{0}^{(1)}=1$, we conclude that

Finally, we need a result on exponential functions and their limit definition:

This can be found in standard analysis books such as Rudin, (1976). ∎

C.5 Proof of Theorem 5.1

We can now prove the convergence of the moments as follows.

Using the NTK parametrization for ReLU, i.e. $\sigma_{1}^{2}=1$ and $\sigma_{2}^{2}=\dots=\sigma_{l}^{2}=\frac{2}{m}$, this amounts to

We thus essentially need to understand the term

We first use Lemma C.4 to expand the ratio $\frac{\Gamma\left(k+\frac{r}{2}\right)}{\Gamma\left(\frac{r}{2}\right)}$ as a polynomial. Denote the coefficients by $\beta_{i}$ for $i=1,\dots,k$ ($i\not=0$ because the polynomial has no intercept). We then swap the two sums:

Now we can apply Lemma C.5 to expand the inner sum for each $i$, denoting the corresponding polynomials again by $Q_{i}$:

Notice that $mQ_{i-1}(m)$ is a polynomial of order $i$. For large $m$, the factor $\frac{1}{m^{k}}$ dominates all such polynomials except for the one with $i=k$. Thus in the large-width limit it holds

The only coefficients contributing to $\frac{1}{m}$ are the second highest coefficient of $Q_{k-1}$ and the highest coefficient of $Q_{k-2}$. Using Lemma C.5 and Lemma C.4, we hence find that

Finally, taking $X\sim\mathcal{N}(0,1)$, $Y\sim\mathcal{LN}(-\frac{5}{4}\gamma,\frac{5}{4}\gamma)$ and defining $Z=XY$, we can easily see that

Appendix D Additonal results and lemmas

Here we list some of the moments arising from a Binomial distribution of the form $U\sim\operatorname{Bin}(n,\frac{1}{2})$. We invite the reader to sanity-check our results in Lemma C.5 regarding the coefficients of $Q_{k}$.

Consider the random variable $U\sim\operatorname{Bin}(m,\frac{1}{2})$. We can calculate its first $4$ moments as