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: This paper proposes to extend the typical GP formulation as the limiting case of a NN with Gaussian weights when the number of hidden units tends to infinity. In particular, it considers a scale of mixture of Gaussians at the last layer for the prior. When this scale mixture of Gaussians is a Inv. Gamma distribution it is well known that the result is a Student-t process. The paper also gives some insights about the resulting process obtained by training the weights of the last layer and following a similar initialization. The resulting processes are heavy-tailed and more robust as shown in the experiments.