Scale Mixtures of Neural Network Gaussian Processes

Recent works have revealed that infinitely-wide feed-forward or recurrent neural networks of any architecture correspond to Gaussian processes referred to as Neural Network Gaussian Processes ( NNGPs ) . While these works have extended the class of neural networks converging to Gaussian processes significantly , however , there has been little focus on broadening the class of stochastic processes that such neural networks converge to . In this work , inspired by the scale mixture of Gaussian random variables , we propose the scale mixture of NNGPs for which we introduce a prior distribution on the scale of the last-layer parameters . We show that simply introducing a scale prior on the last-layer parameters can turn infinitely-wide neural networks of any architecture into a richer class of stochastic processes . With certain scale priors , we obtain heavy-tailed stochastic processes , and in the case of inverse gamma priors , we recover Student ’ s t processes . We further analyze the distributions of the neural networks initialized with our prior setting and trained with gradient descents and obtain similar results as for NNGPs . We present a practical posterior-inference algorithm for the scale mixture of NNGPs and empirically demonstrate its usefulness on regression and classification tasks . In particular , we show that in both tasks , the heavy-tailed stochastic processes obtained from our framework are robust to out-of-distribution data . 1 INTRODUCTION . There has been growing interest in the literature on the equivalence between wide deep neural networks and Gaussian Processes ( GPs ) . Neal ( 1996 ) showed that a shallow but infinitely-wide Bayesian Neural Network ( BNN ) with random weights and biases corresponds to a GP . This result was extended to fully-connected deep neural networks of any depth ( Lee et al. , 2018 ; Matthews et al. , 2018 ) , which are shown to converge to GPs as the width grows . Similar results were later obtained for deep Convolutional Neural Networks ( CNNs ) ( Novak et al. , 2018 ; Garriga-Alonso et al. , 2019 ) and attention networks ( Hron et al. , 2020 ) . In fact , Yang ( 2019 ) showed that wide feed-forward or recurrent neural networks of any architecture converge to GPs and presented a generic method for computing kernels for such GPs . Under this correspondence , the posterior inference of an infinitelywide BNN boils down to the posterior inference of the corresponding GP called NNGP for which a closed-form posterior can be computed exactly . Our goal is to advance this line of research by going beyond GPs . We present a simple yet flexible recipe for constructing infinitely-wide BNNs that correspond to a wide range of stochastic processes . Our construction includes heavy-tailed stochastic processes such as Student ’ s t processes which have been demonstrated to be more robust than GPs under certain scenarios ( Shah et al. , 2014 ) . Our construction is inspired by a popular class of distributions called scale mixtures of Gaussians ( Andrews & Mallows , 1974 ) ; such a distribution is obtained by putting a prior on the scale or variance parameter of a Gaussian distribution . We extend this scale mixing to NNGPs , where we introduce a prior distribution on the scale of the parameters for the last layer ( which is often referred to as readout layer ) in a wide neural network . We show that simply introducing a scale prior on the last layer can turn infinitely-wide BNNs of any architecture into a richer class of stochastic processes , which we name as scale mixtures of NNGPs . These scale mixtures include a broad class of stochastic processes , such as heavy-tailed processes which are shown to be robust to outliers . In particular , when the prior on the scale is inverse gamma , the scale mixture of NNGPs becomes Stduent ’ s t process ( Shah et al. , 2014 ) . We demonstrate that , despite increasing flexibility , mixing NNGPs with a prior on the scale parameter does not increase the difficulty of posterior inference much or at all in some cases . When we mix the scale parameter with inverse gamma ( so that the mixture becomes a t process ) , we can compute the kernel efficiently and infer the exact posterior , as in the case of NNGPs . For generic scale priors and regression tasks with them , we present an efficient approximate posterior-inference algorithm based on importance sampling , which saves computation time by reusing the shared covariance kernels . For classification tasks with categorical likelihood ( for which an exact posterior is not available even for the original NNGP ) , we present an efficient stochastic variational inference algorithm . We further analyze the distributions of the neural networks initialized with our prior setting and trained with gradient descents and obtain results similar to the ones for NNGPs . For NNGP , it has been shown ( Matthews et al. , 2017 ; Lee et al. , 2019 ) that when a wide neural network is initialized with the NNGP specification and then trained only for the last layer ( with all the other layers fixed ) , its fully trained version becomes a random function drawn from the NNGP posterior . Similarly , we analyze the distribution of a wide neural network of any architecture initialized with our prior specification and trained only for the last layer , and show that it becomes a sample from a scale mixture of NNGPs with some scaling distribution . Interestingly , the scaling distribution is a prior , not a posterior . For the fully-connected wide neural networks , we extend the analysis to the case where all the layers are trained with gradient descent , and show that the limiting distribution is again a scale mixture of GPs with a prior scaling distribution and each GPs using a kernel called Neural Tangent Kernel ( NTK ) . In the case of an inverse gamma prior , the limiting distribution becomes Student ’ s t process , which can be computed analytically . We empirically show the usefulness of our construction on various real-world regression and classification tasks . We demonstrate that , despite the increased flexibility , the scale mixture of NNGPs is readily applicable to most of the problems where NNGPs are used , without increasing the difficulty of inference . Moreover , the heavy-tailed processes derived from our construction are shown to be more robust than NNGPs for out-of-distribution or corrupted data while maintaining similar performance for the normal data . Our empirical analysis suggests that our construction is not merely a theoretical extension of the existing framework , but also provides a practical alternative to NNGPs . 1.1 RELATED WORKS . NNGP and NTK Our construction heavily depends on the tensor-program framework ( Yang , 2019 ) which showed that a wide BNN of any feed-forward or recurrent architecture converges in distribution to an NNGP , as its width increases . Especially , we make use of the so-called master theorem and its consequence ( reviewed in Section 1.1 ) to derive our results . Building on the seminal work on NTK ( Jacot et al. , 2018 ) , Lee et al . ( 2019 ) analyzed the dynamics of fully-connected neural networks trained with gradient descent and showed that fully-trained infinitely-wide networks are GPs . In particular , they confirmed the result in Matthews et al . ( 2017 ) that when a fully-connected neural network is initialized from a specific prior ( often referred to as the NTK parameterization ) and trained only for the last layer under gradient descent and squared loss , its fully-trained version becomes a posterior sample of the corresponding NNGP . Lee et al . ( 2019 ) further analyzed the case where all the layers are trained with gradient descent and showed that the network also converges to a GP with specific parameters computed with the so called NTK kernel . We extend these results to our scale mixture of NNGPs in Section 3 . Heavy-tailed stochastic processes from infinitely-wide BNNs The attempts to extend the results on NNGPs to heavy-tailed processes have been made in the past , although not common . The representative case is the work by Favaro et al . ( 2020 ) , which showed that under an alternative prior specification , a wide fully-connected neural network converges to stable processes , as the widths increase . This result was later extended to deep CNNs in ( Bracale et al. , 2021 ) . What distinguishes our approach from these works is the simplicity of our construction ; it simply puts a prior distribution on the scale of the last-layer parameters of a network and makes the scale a random variable , while the constructions in those works replace priors for entire network parameters from Gaussian to other distributions . This simplicity has multiple benefits . First , most of the nice properties of NNGPs , such as easy-to-compute posteriors and the correspondence to gradient descent training , continue to hold in our approach as a version for mixture distributions , while it is at least difficult to find similar adjustments of these results in those works . Second , our approach is applicable to arbitrary architectures , while those works considered only fully-connected networks and CNNs . 2 PRELIMINARIES . We start with a quick review on a few key concepts used in our results . For M ∈ N , let [ M ] be the set { 1 , . . . , M } . Also , write R+ for the set { x ∈ R | x > 0 } . 2.1 TENSOR PROGRAMS AND NEURAL NETWORK GAUSSIAN PROCESSES . Tensor programs ( Yang , 2019 ) are a particular kind of straight-line programs that express computations of neural networks on fixed inputs . In a sense , they are similar to computation graphs used by autodiff packages , such as TensorFlow and PyTorch , but differ from computation graphs in two important aspects . First , a single tensor program can express the computations of infinitely-many neural networks , which share the same architecture but have different widths . This capability comes from the parameter n of the tensor program , which determines the dimensions of vectors and matrices used in the program . The parameter corresponds to the width of the hidden layers of a neural network , so that changing n makes the same tensor program model multiple neural networks of different widths . Second , tensor programs describe stochastic computations , where stochasticity comes from the random initialization of network weights and biases . These two points mean that we can use a single tensor program to model the sequence of BNNs of the same architecture but with increasing widths , and understand the limit of the sequence by analyzing the program . The syntax of and other details about tensor programs are given in Appendix A . We use tensor programs because of the so called master theorem , which provides a method for computing the infinite-width limits of BNN sequences expressed by tensor programs . For the precise formulation of the theorem and the method , again see Appendix A . To explain the master theorem more precisely , assume that we are given a family of BNNs { fn } n∈N that are indexed by their widths n and have the following form : fn ( − ; vn , Ψn ) : RI → R , fn ( x ; vn , Ψn ) = 1√ n ∑ α∈ [ n ] vn , α · φ ( gn ( x ; Ψn ) ) α , where vn ∈ Rn and Ψn ∈ Rn×P ( for some fixed P ) are the parameters of the n-th network , the former being the ones of the readout layer and the latter all the other parameters , I is the dimension of the inputs , φ is a non-linear activation function of the network , and gn ( x ; Ψn ) is the output of a penultimate linear layer , which feeds directly to the readout layer after being transformed by the activation function . In order for the computations of these BNNs on some inputs to be represented by a tensor program , the components of the networks have to satisfy at least the following two conditions . First , the entries of vn and Ψn are initialized with samples drawn independently from ( possibly different ) zero-mean Gaussian distributions . For the entries of vn , the distributions are the same and have the variance σ2v for some σv > 0 . For the entries of Ψn , the distributions may be different , but if we restrict them to those in each column j of Ψn , they become the same and have the variance σj/n for some σj > 0 . Second , φ ( z ) is controlled , i.e. , it is bounded by a function exp ( C|z|2− + c ) for some C , , c > 0 . This property ensures that φ ( z ) for any Gaussian variable z has a finite expectation . Let x1 , . . . , xM ∈ RI be M inputs . When the computations of the BNNs on these inputs are represented by a single tensor program , the master theorem holds . It says that there is a general method for computing µ ∈ RM and Σ ∈ RM×M inductively on the syntax of the program such that µ and Σ characterize the limit of ( gn ( x1 ; Ψn ) , . . . , gn ( xM ; Ψn ) ) n∈N in the following sense . Theorem 2.1 ( Master Theorem ) . Let h : RM → R be a controlled function . Then , as n tends to∞ , 1 n n∑ α=1 h ( gn ( x1 ; Ψn ) α , . . . , gn ( xM ; Ψn ) α ) a.s.−−→ EZ∼N ( µ , Σ ) [ h ( Z ) ] , ( 1 ) where a.s.−−→ refers to almost sure convergence . The next corollary describes an important consequence of the theorem . Corollary 2.2 . As n tends to∞ , the joint distribution of ( fn ( xm ; vn , Ψn ) ) m∈ [ M ] converges weakly to the multivariate Gaussian distribution with zero mean and the following covariance matrix K ∈ RM×M : K ( i , j ) = σ2v · EZ∼N ( µ , Σ ) [ φ ( Zi ) φ ( Zj ) ] ( 2 ) where Zi and Zj are the i-th and j-th components of Z . The above corollary and the fact that K depend only on the shared architecture of the BNNs , not on the inputs x1 : M , imply that the BNNs converge to a GP whose mean is the constant zero function and whose kernel is κ ( x , x′ ) = σ2v · EZ∼N ( µ , Σ ) [ φ ( Z1 ) φ ( Z2 ) ] ( 3 ) where ( µ , Σ ) is constructed as described above for the case that M = 2 and ( x1 , x2 ) = ( x , x′ ) . This GP is called Neural Network Gaussian Process ( NNGP ) in the literature . 2.2 STUDENT ’ S t PROCESSES Student ’ s t processes feature prominently in our results to be presented shortly , because they are marginals of the scale mixtures of NNGPs when the mixing is done by inverse gamma . These processes are defined in terms of multivariate Student ’ s t distributions defined as follows . Definition 2.1 . A distribution on Rd is a multivariate ( Student ’ s ) t distribution with degree of freedom ν ∈ R+ , location µ ∈ Rd and positive-definite scale matrix Σ ∈ Rd×d if it has the following density : p ( x ) = G ( ( ν + d ) /2 ) ( νπ ) d 2 ·G ( ν/2 ) · |Σ| 12 ( 1 + 1 ν ( x− µ ) > Σ−1 ( x− µ ) ) − ν+d2 where G ( · ) is the gamma function . To express a random variable drawn from this distribution , we write x ∼ MVTd ( ν , µ , Σ ) . When ν > 1 , the mean of the distribution MVTd ( ν , µ , Σ ) exists , and it is µ . When ν > 2 , the covariance of distribution also exists , and it is νν−2Σ . As for the Gaussian distribution , the multivariate t distribution is also closed under marginalization and conditioning . More importantly , it can be obatined as a scale mixture of Gaussians ; if σ2 ∼ InvGamma ( a , b ) and x|σ2 ∼ N ( µ , σ2Σ ) , then the marginal distribution of x is MVT ( 2a , µ , baΣ ) . For the definition of the inverse gamma distribution , see Definition B.1 . Student ’ s t process is a real-valued random function such that the outputs of the function on a finite number of inputs are distributed by a multivariate t distribution ( Shah et al. , 2014 ) . Consider a set X ( for inputs ) with a symmetric positive-definite function κ : X × X → R. Definition 2.2 . A random function f : X → R is Student ’ s t process with degree of freedom ν ∈ R+ , location functionM : X → R , and scale function κ , if for all d and inputs x1 , . . . , xd ∈ X , the random vector ( f ( x1 ) , . . . , f ( xd ) ) has the multivariate t distribution MVTd ( ν , µ , Σ ) , where µ = ( M ( x1 ) , . . . , M ( xd ) ) and Σ is the d× d matrix defined by Σ ( i , j ) = k ( xi , xj ) .

Summary: Although the connection between Bayesian neural networks and Gaussian processes has been a topic of interest for several years, the work on tensor programs facilitating the formulation of the Neural Network-Gaussian Process (NNGP) model has been influential in further emphasising the correspondence between the two classes of models. In this work, the authors propose an extension to the NNGP model that allows for the formulation of more general stochastic processes that may follow alternative distributions such as Student-$t$. This is achieved by way of introducing a prior on the scale of the parameters in the last layer of the neural network. This formulation allows for greater flexibility without requiring the more involved changes proposed in related works such as Favran et al. (2020) and Bracale et al. (2021). Following the analysis carried out in the initial formulation of NNGPs, the authors also investigate the correspondence between GPs and infinitely-wide BNNs configured with the proposed set-up and trained using gradient descent. The experimental evaluation features a mixture of regression and classification tasks, with a particular emphasis on how the use of models with heavier-tailed distributions can be better suited to datasets with challenging properties such as corrupted data and label imbalance.