Precise characterization of the prior predictive distribution of deep ReLU networks

1 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 ; GarrigaAlonso 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 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 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 . Contributions 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 . 2 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 . 3 Background . 3.1 Fully Connected Neural Network . Given an input x ∈ Rd , we define a θ-parameterized L layer fully-connected neural network fθ ( x ) as the composition of layer-wise affine transformationsW ( l ) ∈ Rdl−1×dl and element-wise non-linearities σ : R→ R , fθ ( x ) = W ( L ) > σ ( W ( L−1 ) > . . . σ ( W ( 1 ) > x ) . . . ) , ( 1 ) where θ = ( W ( 1 ) , . . . , W ( L ) ) denotes the collection of all weights . Throughout this work , we assume standard initialization , i.e.W ( l ) ij ∼ N ( 0 , σ2l ) , where weights in each layer can have a different variance σ2l . Often , it will be more convenient to work with the corresponding unit-level formulation , expressed through the recursive equations : f ( l ) k ( x ) = dl−1∑ j=1 W ( l ) jk g ( l−1 ) j ( x ) , g ( l ) ( x ) = σ ( f ( l ) ( x ) ) , ( 2 ) When propagating an input x through the network we refer to f ( l ) ( x ) ∈ Rdl as the pre-activations and to g ( l ) ( x ) ∈ Rdl as the post-activations . When it is clear from the context , to enhance readability , we will occasionally abuse notation denoting f ( l ) : = f ( l ) ( x ) and g ( l ) : = g ( l ) ( x ) . We will also use the short-hand [ m ] = { 1 , . . . , m } . Imposing a probability distribution on the weights induces a distribution on the pre-activations f ( l ) ( x ) . Understanding the properties of this distribution is the main goal of this work . 3.2 Prior Predictive Distribution and Infinite Width . A precise characterization of the prior predictive distribution of a neural network has been established in the so-called infinite width setting ( Neal , 1995 ; Lee et al. , 2018 ; Matthews et al. , 2018 ) . By considering variances that scale inversely with the width , i.e . σ2l = 1 dl and fixed depth L ∈ N , it can be shown that the implied prior predictive distribution converges in law to a Gaussian process : fθ d−→ GP ( 0 , Σ ( L ) ) , ( 3 ) where Σ ( L ) : Rd × Rd , x , x′ 7→ Σ ( L ) ( x , x′ ) is the NNGP kernel ( Lee et al. , 2018 ) , available in closed-form through the recursion Σ ( 1 ) ( x , x′ ) = xTx′ , Σ ( l+1 ) ( x , x′ ) = Ez∼N ( 0 , Σ̃ ( l ) ) [ σ ( z1 ) σ ( z2 ) ] , ( 4 ) for l = 1 , . . . , L − 1 and where Σ̃ ( l ) = ( Σ ( l ) ( x , x ) Σ ( l ) ( x , x′ ) Σ ( l ) ( x′ , x ) Σ ( l ) ( x′ , x′ ) ) ∈ R2×2 . The proof relies on the multivariate central limit theorem in conjunction with an inductive technique , letting hidden layer widths go to infinity in sequence . Due to the Gaussian nature of the limit , exact Bayesian inference becomes tractable and techniques from the Gaussian process literature can be readily applied ( Rasmussen , 2003 ) . Although theoretically very appealing , from a practical point of view , infinite width networks are not as relevant due to their inferior performance ( Novak et al. , 2019 ) . Gaining a better understanding of this performance gap is hence of utmost importance .

Summary: The main result of the paper is a derivation of the exact characterization of priors in function space through Meijer G-functions. The results are obtained for Gaussian priors, ReLU and linear activation functions. The article also provides a solid analysis of the obtained characterization.   **Advantages**:  * It is the first accurate description of hidden units for standard feed-forward neural networks among the concurrent work mentioned in the paper.  * The results are in line with other works on the heavy-tailed nature of hidden units and the Gaussian process limit.  * There is a study on the role of width and depth demonstrating how wider networks lead to more Gaussian-like distributions.  * Authors also introduced generalized He priors that enable to tune the variance induced in the function space.    **Limitations**:  * Only linear and ReLU functions  * Only Gaussian priors  * Without the bias term