S CALE M IXTURES OF N EURAL N ETWORK G AUSSIAN P ROCESSES
A BSTRACT
Recent works have revealed that inﬁnitely-wide feed-forward or recurrent neural networks of any architecture correspond to Gaussian processes referred to as
Neural Network Gaussian Processes ( NNGP s).
While these works have extended the class of neural networks converging to Gaussian processes signiﬁcantly, 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 NNGP s 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 inﬁnitely-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
NNGP s. We present a practical posterior-inference algorithm for the scale mixture of NNGP s and empirically demonstrate its usefulness on regression and classiﬁcation 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.
I NTRODUCTION
There has been growing interest in the literature on the equivalence between wide deep neural networks and Gaussian Processes ( GP s).
Neal (1996) showed that a shallow but inﬁnitely-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 GP s as the width grows.
Similar results were later obtained for deep Convolutional Neural Networks ( CNN s) (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 GP s and presented a generic method for computing kernels for such GP s. Under this correspondence, the posterior inference of an inﬁnitelywide 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 GP s. We present a simple yet ﬂexible recipe for constructing inﬁnitely-wide BNN s 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 GP s 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 NNGP s, 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 inﬁnitely-wide BNN s of any architecture into a richer class of stochastic processes, which we name as scale mixtures of NNGP s. 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 NNGP s becomes
Stduent’s t process (Shah et al., 2014).
We demonstrate that, despite increasing ﬂexibility, mixing NNGP s with a prior on the scale parameter does not increase the difﬁculty 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 efﬁciently and infer the exact posterior, as in the case of NNGP s. For generic scale priors and regression tasks with them, we present an efﬁcient approximate posterior-inference algorithm based on importance sampling, which saves computation time by reusing the shared covariance kernels.
For classiﬁcation tasks with categorical likelihood (for which an exact posterior is not available even for the original NNGP ), we present an efﬁcient 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 NNGP s. 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 speciﬁcation and then trained only for the last layer (with all the other layers ﬁxed), 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 speciﬁcation and trained only for the last layer, and show that it becomes a sample from a scale mixture of NNGP s 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 GP s with a prior scaling distribution and each GP s 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 classiﬁcation tasks.
We demonstrate that, despite the increased ﬂexibility, the scale mixture of NNGP s is readily applicable to most of the problems where NNGP s are used, without increasing the difﬁculty of inference.
Moreover, the heavy-tailed processes derived from our construction are shown to be more robust than NNGP s 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 NNGP s.
1. R ELATED W ORKS 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 inﬁnitely-wide networks are GP s.
In particular, they conﬁrmed the result in Matthews et al.
(2017) that when a fully-connected neural network is initialized from a speciﬁc 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 speciﬁc parameters computed with the so called NTK kernel.
We extend these results to our scale mixture of NNGP s in Section 3.
Heavy-tailed stochastic processes from inﬁnitely-wide BNN s The attempts to extend the results on NNGP s 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 speciﬁcation, a wide fully-connected neural network converges to stable processes, as the widths increase.
This result was later extended to deep CNN s 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 beneﬁts.
First, most of the nice properties of NNGP s, 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 difﬁcult to ﬁnd 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 CNN s.
P RELIMINARIES
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. T ENSOR P ROGRAMS AND N EURAL N ETWORK G AUSSIAN P ROCESSES
Tensor programs (Yang, 2019) are a particular kind of straight-line programs that express computations of neural networks on ﬁxed 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 inﬁnitely-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 BNN s 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 inﬁnite-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 BNN s { f n } n ∈ N that are indexed by their widths n and have the following form:
[Equation]
where v n ∈ R n and Ψ n ∈ R n × P (for some ﬁxed 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 g n ( 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 BNN s 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 v n and Ψ n are initialized with samples drawn independently from (possibly different) zero-mean Gaussian distributions.
For the entries of v n , the distributions are the same and have the variance σ 2 v 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 − (cid:15) + c ) for some C, (cid:15), c > 0 .
This property ensures that φ ( z ) for any Gaussian variable z has a ﬁnite expectation.
Let x 1 , . . .
, x M ∈ R I be M inputs.
When the computations of the BNN s on these inputs are represented by a single tensor program, the master theorem holds.
It says that there is a general method for computing µ ∈ R M and Σ ∈ R M × M inductively on the syntax of the program such that µ and Σ characterize the limit of ( g n ( x 1 ; Ψ n ) , . . .
, g n ( x M ; Ψ n ))
n ∈ N in the following sense.
Theorem 2.1 (Master Theorem) .
Let h : R M → R be a controlled function.
Then, as n tends to ∞ ,
[Equation]
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 ( f n ( x m ; v n , Ψ n )) m ∈
[ M ] converges weakly to the multivariate Gaussian distribution with zero mean and the following covariance matrix K ∈ R M × M :
[Equation]
where Z i and Z j 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 BNN s, not on the inputs x 1: M , imply that the BNN s converge to a GP whose mean is the constant zero function and whose kernel is
[Equation]
where ( µ, Σ) is constructed as described above for the case that M = 2 and ( x 1 , x 2 ) =
( x , x (cid:48) ) .
This GP is called Neural Network Gaussian Process ( NNGP ) in the literature.
2. S TUDENT ’ S t P ROCESSES
Student’s t processes feature prominently in our results to be presented shortly, because they are marginals of the scale mixtures of NNGP s when the mixing is done by inverse gamma.
These processes are deﬁned in terms of multivariate Student’s t distributions deﬁned as follows.
Deﬁnition 2.1.
A distribution on R d is a multivariate (Student’s) t distribution with degree of freedom ν ∈ R + , location µ ∈ R d and positive-deﬁnite scale matrix
Σ ∈ R d × d if it has the following density:
[Equation]
where G ( · ) is the gamma function.
To express a random variable drawn from this distribution, we write x ∼ MVT d ( ν, µ, Σ) .
When ν > 1 , the mean of the distribution MVT d ( ν, µ, Σ) 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(2 a, µ, ba Σ) .
For the deﬁnition of the inverse gamma distribution, see Deﬁnition B.1.
Student’s t process is a real-valued random function such that the outputs of the function on a ﬁnite number of inputs are distributed by a multivariate t distribution (Shah et al., 2014).
Consider a set
X (for inputs) with a symmetric positive-deﬁnite function κ : X × X → R .
Deﬁnition 2.2.
A random function f : X → R is Student’s t process with degree of freedom ν ∈ R + , location function M : X → R , and scale function κ , if for all d and inputs x 1 , . . .
, x d ∈ X , the random vector ( f ( x 1 ) , . . .
, f ( x d )) has the multivariate t distribution MVT d ( ν, µ, Σ) , where µ = ( M ( x 1 ) , . . .
, M ( x d )) and Σ is the d × d matrix deﬁned by Σ ( i,j ) =
k ( x i , x j ) .
R ESULTS
We will present our results for the regression setting.
Throughout this section, we consider only those BNN s whose computations over ﬁxed inputs are representable by tensor programs.
Assume that we are given a training set D tr = { ( x k ,
y k ) ∈
R I × R | 1 ≤ k ≤ K } and a test set D te = { ( x K + (cid:96) , y K + (cid:96) ) | 1 ≤ (cid:96) ≤ L } .
Our idea is to treat the variance σ 2 v for the parameters of the readout layer in a BNN as a random variable, not a deterministic value, so that the BNN represents a mixture of random functions (where the randomness comes from that of σ 2 v and the random initialization of the network parameters).
To state this more precisely, consider the computations of the BNN s { f n ( − ; v n , Ψ n ) } n ∈ N on the inputs x 1: K + L in the training and test sets such that n is the width of the hidden layers of f n and these computations are representable by a tensor program.
Our idea is to change the distribution for the initial value of v n from a Gaussian to a scale mixture of Gaussians, so that at initialization, the BNN s on x 1: K + L become the following random variables: for each n ∈ N , σ
[Equation]
where H is a distribution on positive real numbers, such as the inverse-gamma distribution, P is the number of columns of Ψ n , and σ 2 j and g n are, respectively, the variance speciﬁc to the j -th column of Ψ n and the output of the penultimate linear layer, as explained in our review on tensor programs.
Our ﬁrst result says that as the width of the BNN grows to ∞ , the BNN converges in distribution to a mixture of NNGP s, and, in particular, to Student’s t process if the distribution H is inverse gamma.
Theorem 3.1 (Convergence) .
As n tends to ∞ , the random variable ( f n ( x i ; v n , Ψ n ))
i ∈
[ K + L ] converges in distribution to the following random variable ( f ∞ ( x i ))
i ∈
[ K + L ] :
[Equation]
where K ( i,j ) =
E Z ∼N ( µ, Σ) [ φ ( Z i ) φ ( Z j )] for i, j ∈ [ K + L ] .
Furthermore, if H is the inversegamma distribution with shape a and scale b , the marginal distribution of ( f ∞ ( x i )) i ∈
[ K + L ] is the following multivariate t distribution:
[Equation]
where MVT d ( ν (cid:48) , µ (cid:48) , K (cid:48) ) denotes the d -dimensional Student’s t distribution with degree of freedom ν (cid:48) , location µ (cid:48) , and scale matrix K (cid:48) .
This theorem holds mainly because of the master theorem for tensor programs (Theorem 2.1).
In fact, the way that its proof uses the master theorem is essentially the same as the one of the NNGP convergence proof (i.e., the proof of Corollary 2.2), except for one thing: before applying the master theorem, the proof of Theorem 3.1 conditions on σ 2 v and removes its randomness.
The detailed proof can be found in the appendix.
As in the case of Corollary 2.2, although the theorem is stated for the ﬁnite dimensional case with the inputs x 1: K + L , it implies the convergence of the BNN s with random σ 2 v to a scale mixture of GP s or to Student’s t process.
Note that σ 2 v · K in the theorem is precisely the covariance matrix K of the NNGP kernel on x 1: K + L in Corollary 2.2.
We thus call the limiting stochastic process as scale mixture of NNGP s.
Our next results characterize a fully-trained BNN at the inﬁnite-width limit under two different training schemes.
They are the versions of the standard results in the literature, generalized from GP s to scale mixtures of GP s. Assume the BNN s under the inputs x 1: K + L that we have been using so far.
Theorem 3.2 (Convergence under Readout-Layer Training) .
Assume that we train only the readout layers of the BNN s using the training set D tr under mean squared loss and inﬁnitesimal step size.
Formally, this means the network parameters are evolved under the following differential equations: ( v [Equation] ) n , Ψ [Equation]
n
[Equation]
n
[Equation] Ψ n [Equation] Then, the time and width limit of the random variables ( f n ( x K + i ; v ( t ) n , Ψ ( t ) n ))
i ∈
[ L ] , denoted ( f ∞∞ ( x K + i ))
i ∈
[ L ] , is distributed by the following mixture of Gaussians:
[Equation]
where Y tr consists of the y values in D tr (i.e., Y tr = ( y 1 , .
. . , y K ) )
and the K ’s with different subscripts are the restrictions of K in Theorem 3.1 with training or test inputs as directed by those subscripts.
Furthermore, if H is the inverse-gamma distribution with shape a and scale b , the marginal distribution of ( f ∞∞ ( x K + i ))
i ∈
[ L ] is the following multivariate t distribution:
[Equation]
, tr K tr , te (cid:17)(cid:19) .
One important point about this theorem is that the distribution of ( f ∞∞ ( x K + i ))
i ∈
[ L ] is not the posterior of ( f ∞ ( x K + i ))
i ∈
[ L ] in Theorem 3.1 under the condition on Y tr .
The former uses the prior on σ 2 v for mixing the distribution of ( f ∞ ( x K + i ))
i ∈
[ L ] | ( σ 2 v , Y tr ) , while the latter uses the posterior on σ 2 v for the same purpose.
Our last result assumes a change in the initialization rule for our models.
This is to enable the use of the results in Lee et al.
about the training of inﬁnitely-wide neural networks.
Concretely, we use the so called NTK parametrization for Ψ n .
This means that the BNN s on x 1: K + L are now intialized to be the following random variables: for each n ∈ N ,
[Equation]
When we adopt this NTK parameterization and the BNN s are multi-layer perceptrons (more generally, they share an architecture for which the limit theory of neural tangent kernels has been developed), we can analyze their distributions after the training of all parameters.
Our analysis uses the following famous result for such BNN s: as n tends to inﬁnity, for all i, j ∈ [ K + L ] ,
[Equation]
for some Θ ( i,j ) determined by the architecture of the BNN s. Let Θ be the matrix (Θ ( i,j ) )
i,j ∈
[ K + L ] .
The next theorem describes the outcome of our analysis for the fully-trained inﬁnite-width BNN .
Theorem 3.3 (Convergence under General Training) .
Assume that the BNN s are multi-layer perceptrons and we train all of their parameters using the training set D tr under mean squared loss and inﬁnitesimal step size.
Formally, this means the network parameters are evolved under the following differential equations:
[Equation]
Let ¯Θ be Θ with σ 2 v in it set to 1 (so that Θ =
σ 2 v ¯Θ ).
Then, the time and width limit of the random variables ( f n ( x K + i ; v ( t ) n , 1 √ n Ψ ( t ) n )) i ∈
[ L ] has the following distribution:
[Equation]
where
[Equation]
Here the K ’s with subscripts are the ones in Theorem 3.1 and the ¯Θ ’s with subscripts are similar restrictions of ¯Θ with training or test inputs.
Furthermore, if H is the inverse gamma with shape a and scale b , the exact marginal distribution of ( f ∞∞ ( x K + i ))
i ∈
[ L ] is the following multivariate t distribution:
[Equation]
Remark (Unifying High-level Principles) .
All of our results and their proofs are derived from two high-level principles.
First, once we condition on the variance σ 2 v of the readout layer’s parameters, the networks in our setup ﬁt into the standard setting for NNGP s and NTK s, so that they satisfy the existing results in the setting.
This conditioning trick means that most of the results in the standard setting can be carried over to our setup in a scale-mixture form, as we explained in this section.
Second, if the prior on σ 2 v is inverse gamma, the marginalization of σ 2 v in those results can be calculated analytically and have a form of Student’s t distribution or process.
3. P OSTERIOR INFERENCE
Gaussian likelihood As for the NNGP s, when the scale is mixed by the inverse gamma prior, we can exactly compute the posterior predictives of the scale mixture of NNGP s for Gaussian likelihood since it corresponds to the Student’s t process having closed-form characterizations for the posteriors.
For generic prior settings for the scales, we no longer have such closed-form formulas for predictive posteriors, and have to rely on approximate inference methods.
We describe one such method, namely, self-normalizing importance sampling with prior as proposal, together with a technique for optimizing the method speciﬁcally for scale mixtures of NNGP s.
Consider a scale mixture of NNGP s with a prior H on the variance σ 2 v .
Assume that we want to estimate the expectation of h ( y ) for some h : R → R where the random variable y is drawn from the predictive posterior of the mixture at some input x ∈ R
I under the condition on D tr .
The importance sampling with prior as proposal computes the estimator ( (cid:80) Ni =1 w i h
( y i )) /
(cid:80) Nj =1 w j
where β
i
and y i for i
= 1 , . .
, N are independent samples from the prior H and the posterior of the Gaussian distribution, respectively, and the w i ’s are the importance weights of the β
i ’s: [Equation] i ∼
[Equation] [Equation]
The K is the covariance matrix computed as in Theorem 3.1 except that x is now a test input.
A naive implementation of this importance sampler is slow unnecessarily due to the inefﬁcient calculation of the likelihoods N ( Y tr ; 0 , β i K tr , tr ) of the sampled β i ’s. To see this, note that the log likelihood of β i can be decomposed as follows:
[Equation]
The terms K 2 log(2 π ) , 12 log det( K tr , tr ) , and Y (cid:62) tr K − 1 tr , tr Y tr are shared across all the samples { β i } Ni =1 , so need not be computed everytime we draw β i .
To avoid these duplicated calculations, we compute these shared terms beforehand, so that the log-likleihood for each β i can be computed in O (1) time.
Generic likelihoods For generic likelihoods other than Gaussian, even the vanilla NNGP s do not admit closed-form posterior predictives.
In such cases, we can employ the Sparse Variational Gaussian Process ( SVGP ) (Titsias, 2009) for approximate inference.
We present a similar algorithm for the inverse-gamma mixing case which leads to Student’s t process.
See Appendices C and D for the detailed description.
E XPERIMENTS
We empirically evaluated the scale mixtures of NNGP s on various synthetic and real-world tasks.
We tested the scale mixture of NNGP s with inverse-gamma prior corresponding to Student’s t processes, and with another heavy-tailed prior called Burr Type XII distribution (Appendix H).
Our implementation used Neural Tangents library (Novak et al., 2020) and JAX (Bradbury et al., 2018).
4. E MPIRICAL VALIDATION OF THE THEORIES
To empirically validate our theories, we set up a fully connected network having three layers with 512 hidden units and erf activation function.
Except for the last layer, the weights of the network were initialized with N (0 , 8 /n ) and the biases were initialized with N (0 , 0 .
05 2 /n ) .
For the last layer, we sampled σ 2 v ∼ InvGamma( a, b ) with varying ( a, b ) values and sampled weights from
N (0 , σ 2 v ) .
To check Theorem 3.1, we initialize 1,000 models and computed the distribution of outputs evaluated at zero.
For Theorem 3.2, using y = sin( x ) as the target function, we ﬁrst initialized the parameters, trained only the last layer for multiple
σ v values, and averaged the results to get function value at zero.
Similarly, we trained all the layers to check Theorem 3.3.
As shown in
Fig.
1, the theoretical limit well-matched the empirically obtained distributions for all settings we tested.
See Appendix G for the details and more results.
[Figure]
[-Student's t 5.0 2.5 0.0 2.5 5.
-]
[Figure]
[Table]
[Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
4. R EGRESSION WITH G AUSSIAN LIKELIHOODS
We tested the scale mixtures of NNGP s and other models on various real-world regression tasks.
All the models are trained with the squared loss function.
We expect the scale mixtures of NNGP s to be better calibrated than NNGP due to their heavy-tail behaviors.
To see this, we ﬁtted the scale mixture of NNGP s with inverse gamma prior on eight datasets collected from UCI repositories 1 and measured the negative log-likelihood values on the test set.
Table 1 summarizes the result.
The results other than ours are borrowed from Adlam et al.
In summary, ours generally performed similarly to NNGP except for some datasets for which ours signiﬁcantly outperformed NNGP .
Considering the fact that Student’s t processes include GP s as limiting cases, this result suggests that one can use the scale mixtures as an alternative to NNGP s even for the datasets not necessarily including heavy-tailed noises.
See Appendix G for detailed settings for the training.
4. C LASSIFICATION WITH G AUSSIAN LIKELIHOOD
Following the literature, we apply the NNGP s and the scale mixtures of NNGP s to classiﬁcation problems with Gaussian likelihoods (squared loss).
We summarize the results in Appendix F.
As a brief summary, ours with heavy-tailed priors (inverse gamma, Burr Type XII) outperformed NNGP in terms of uncertainty calibration for various datasets including corrupted datasets.
4. C LASSIFICATION WITH CATEGORICAL LIKELIHOOD
We compared the NNGP and the scale mixture of NNGP s with inverse-gamma priors on the image classiﬁcation task.
Following the standard setting in image classiﬁcation with deep neural networks, we computed the posterior predictive distributions of the BNN s under categorical likelihood.
Since both NNGP s and the scale mixtures of NNGP s do not admit closed-form posteriors, we employed
SVGP as an approximate inference method for NNGP (Appendix C), and extended the SVGP for the scale mixture of NNGP s (Appendix D).
We call the latter as Sparse Variational Student’s t Process ( SVTP ) since it approximates the posteriors whose limit corresponds to Student’s t process.
We compared SVGP and SVTP on multiple image classiﬁcation benchmarks, including MNIST, CIFAR10, and SVHN.
We also evaluated the predictive performance on Out-Of-Distribution ( OOD ) data for which we intentionally removed three training classes to save as OOD classes to be tested.
We used four-layer CNN as a base model to compute NNGP kernels.
See Appendix G for details.
Tablesummarizes the results, which show that SVTP and SVGP perform similarly for in
[Table] [Table]
but SVTP signiﬁcantly outperforms SVGP for OOD data in terms of uncertainty calibration (measured by negative log-likelihoods).
4. C LASSIFICATION BY F INITE M ODEL E NSEMBLE
Although we proved Theorem 3.3 only for fully-connected networks trained with squared losses,
Lee et al.
(2019) empirically observed that there still seem to be a similar correspondence for neural networks trained with cross-entropy loss.
To this end, we tested the ensembles of wide ﬁnite CNNs with 4 layers trained by gradient descent over the cross-entropy loss.
Each model in our test is initialized by the NTK parameterization with random scales drawn from inverse gamma prior on the last layer (ours) or without such random scales ( NNGP ).
For each prior setting, we constructed ensembles of 8 models, and compared the performance on MNIST, MNIST-like variants, CIFAR10, and
SVHN.
We also compared the models under the presence of corruption, OOD data, label imbalance, and label noises.
See Appendix G for detailed description for the experimental setting.
As for the previous experiments, Table 3 demonstrates that ours with inverse gamma prior largely outperforms the baseline, especially for the OOD or corrupted data.
C ONCLUSION
In this paper, we proposed a simple extension of NNGP s by introducing a scale prior on the lastlayer weight parameters.
The resulting method, entitled as the scale mixture of NNGP s, deﬁnes a broad class of stochastic processes, especially the heavy-tailed ones such as Student’s t processes.
Based on the result in (Yang, 2019), we have shown that an inﬁnitely-wide BNN of any architecture constructed in a speciﬁc way corresponds to the scale mixture of NNGP s. Also, we have extended the existing convergence results of inﬁnitely-wide BNN s trained with gradient descent to GP s so that the results hold for our construction.
Our empirical evaluation validates our theory and shows promising results for multiple real-world regression and classiﬁcation tasks.
Especially, it shows that the heavy-tailed processes from our construction are robust to the out-of-distribution data.
A CKNOWLEDGEMENTS AND D ISCLOSURE OF F UNDING
We would like to thank Jaehoon Lee for helping us to understand his results on NNGP and NTK.
This work was partly supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2019-0-00075, Artiﬁcial
Intelligence Graduate School Program(KAIST),
No. 2021-0-02068, Artiﬁcial Intelligence Innovation Hub), National Research Foundation of Korea (NRF) funded by the Ministry of Education (NRF2021R1F1A1061655, NRF-2021M3E5D9025030).
HY was supported by the Engineering Research Center Program through the National Research Foundation of Korea (NRF) funded by the Korean Government MSIT (NRF-2018R1A5A1059921) and also by the Institute for Basic Science (IBS-R029-C1) A D ETAILS ON T ENSOR P ROGRAMS
In this section, we will review tensor programs from (Yang, 2019), which are sequences of assignment statements followed by a return statement where each variable is updated at most once.
Tensor programs use the three types of variables: G-variables , H-variables , and A-variables .
Gand H-variables are vector-type variables, and A-variables are matrix-type variables.
Each tensor program is parameterized by n ∈ N , which determines the dimension of all vectors (i.e., G- and
H-variables) and matrices (i.e., A-variables) in the program to n and n × n , respectively.
Due to the random initialization of some of these variables, the contents of these variables are random in general.
G-variables denote Gaussian-distributed random variables, and H-variables store the results of applying (usually nonlinear) functions on G-variables.
A-variables store random matrices whose entries are set with iid samples from a Gaussian distribution.
All H-variables in a tensor program are set their values by some assignments, but some G-variables and all H-variables may not be assigned to in the program.
Such G-variables are called input G-variables.
Assignments in tensor programs have one of the three kinds: MatMul , LinComb , and Nonlin .
The
MatMul assignments have the form y
=
Wx where y is a G -variable, W is an A-variable, and x is a H-variable.
The LinComb assignments compute linear combinations, and they have the form y = (cid:80) li
=1 w i x i where y, x 1 , . . .
, x l are G-variables, and w 1 , . . .
, w l ∈ R are constants.
The ﬁnal Nonlin assignments apply (usually nonlinear) scala functions on the values of G-variables coordinatewise.
That is, they have the form y
= φ ( x 1 , . . .
, x l ) , where y is an H-variable, x 1 , . .
, x l are G-variables, φ : R l → R is a possibly nonlinear function, and φ ( x 1 , . . .
, x l ) means the lifted application of φ to vector arguments.
Every tensor program ends with a return statement of the form:
[Equation]
where v 1 , .
, v l are G-variables not appearing anywhere in the program, and y 1 , . . .
, y l are Hvariables.
Assumption A.1.
Every A-variable W in a tensor program is initialized with iid samples from the Gaussian distribution with mean zero and covariance σ 2 W /n for some σ W > 0 (i.e., W ij ∼
N (0 , σ 2 W /n ) for i, j ∈ [ n ] ), and these samples are independent with those used for other A-variables and input G-variables.
For every α ∈ [ n ] , the components of all input G-variables x 1 , . . .
, x m (which store vectors in R n ) are initialized with independent samples from N ( µ in , Σ in ) for some mean µ in and (possibly singular) covariance Σ in over all input G-variables.
Note that the mean and the covariance do not depend on the component index α ∈ [ n ] .
Finally, the input G-variables v i used in the return statement of the program are independent with all the other random variables in the program, and their components are set by iid samples from N (0 , σ 2 v ) .
Deﬁnition A.1.
Given a tensor program satisfying Assumption A.1, we can compute µ and Σ for not only input G-variables but all the G-variables in the program.
Here the dimension of the vector µ is the number N of all G-variables.
The computed µ and Σ specify the multivariate Gaussian distribution over the α -th components of G-variables that arises at the inﬁnite-width limit (Theorem A.2).
The computation is deﬁned inductively as follows.
Let g 1 , . .
, g N be all the G-variables in the program, ordered as they appear in the program.
For any G-variable g k , µ ( g k ) is deﬁned by: [Equation] if
[Equation]
is an
[Equation]
otherwise .
Here x 1 , . . .
, x m are some G-variables.
For any pair of G-variables g k , g l , Σ( g k , g l ) is deﬁned by: Σ( g k , g l ) =
[Equation]
l are input
[Equation]
variables
[Equation]
l
[Equation]
for the same W , with Z , φ , φ (cid:48) below otherwise .
Here x 1 , . . .
, x m are some G-variables, and H-variables y and y (cid:48) in the second last case are introduced by the Nonlin assignments
y
=
φ ( g i 1 , . .
, g i p ) and y (cid:48) =
φ (cid:48)
( g i 1 , . . .
, g i p ) , respectively, for some i 1 , . . .
, i p ∈
[max( k − 1 , l − 1)] .
We have adjusted the inputs of φ and φ (cid:48) to ensure that they take the same set of G-variables.
The random variable Z in the second last case is a random
Gaussian vector drawn from N ( µ (cid:48) , Σ (cid:48) )
where µ (cid:48) and Σ (cid:48) are the mean vector and covariance matrix for g i 1 , . .
, g i p by the previous step of this inductive construction.
Deﬁnition A.2.
Given h : R n → R and for any x ∈ R n if | h ( x ) | is upper bounded by a function exp( C (cid:107) x (cid:107) 2 − (cid:15) + c ) with C, c, (cid:15) > 0 , then this function is controlled .
Theorem A.2 (Master Theorem) .
Consider a tensor program satisfying Assumption A.1 and using only controlled φ ’s. Let g 1 , . .
, g N be all of the G-variables that appear in this program, and h : R N → R be a controlled function.
Then, as n tends to ∞ ,
[Equation]
Here a.s. −−→ refers to almost sure convergence, and Z =
( Z 1 , . . .
, Z N ) ∈ R N .
Also, µ and Σ are from Deﬁnition A.1.
Corollary A.3.
Consider a tensor program satisfying Assumption A.1 and using only controlled φ ’s. Let g 1 , . .
, g N be all of the G-variables that appear in this program.
Assume that the program returns a vector ( v (cid:62) y 1 / √ n, . .
. , v (cid:62) y k / √ n ) for some H-variables y i ’s such that each y i is updated by y i =
φ i ( g 1 , . . . , g N ) for a controlled φ i .
Then, as n approaches to ∞ , the joint distribution of this vector converges weakly to the multivariate Gaussian distribution N (0 , K ) with the following covariance matrix:
[Equation]
where µ and Σ are from Deﬁnition A.1.
B P ROOFS
In this section, we provide the proofs of our three main results.
We ﬁrst describe a lemma related to a basic property of Student’s t distribution:
Lemma B.1.
The class of multivariate t distributions is closed under marginalization and conditioning.
That is, when
[Equation]
we have that
[Equation]
where
[Equation]
Deﬁnition B.1.
A distribution on R + is a inverse gamma distribution with shape parameter a ∈ R + and scale parameter b ∈ R + if it has the following density:
[Equation]
a + e − b/x where G ( · ) is the gamma function.
To express a random variable drawn from this distribution, we write x ∼ InvGamma( a, b ) .
Lemma B.2.
Let Σ ∈ R
d
×
d be a postive deﬁnite matrix.
Consider µ ∈ R n and a, b ∈ R + .
If
[Equation]
then the marginal distribution of x is MVT d (2 a, µ, ba Σ) .
Now we proceed to the proofs of main theorems.
Theorem 3.1 (Convergence) .
As n tends to ∞ , the random variable ( f n ( x i ; v n , Ψ n ))
i ∈
[ K + L ] converges in distribution to the following random variable ( f ∞ ( x i ))
i ∈
[ K + L ] :
[Equation]
where K ( i,j ) =
E Z ∼N ( µ, Σ) [ φ ( Z i ) φ ( Z j )] for i, j ∈ [ K + L ] .
Furthermore, if H is the inversegamma distribution with shape a and scale b , the marginal distribution of ( f ∞ ( x i )) i ∈
[ K + L ] is the following multivariate t distribution:
[Equation]
where MVT d ( ν (cid:48) , µ (cid:48) , K (cid:48) ) denotes the d -dimensional Student’s t distribution with degree of freedom ν (cid:48) , location µ (cid:48) , and scale matrix K (cid:48) .
Proof.
Let o i [ n ] = f n ( x i ; v n , Ψ n ) for i ∈
[ K + L ] and n .
Pick an arbitrary bounded continuous function h : R K + L → R .
For each n , deﬁne G n to be the σ -ﬁeld generated by g n ( x 1 ; Ψ n ) , . .
g n ( x K + L ; Ψ n ) .
We claim that as n tends to ∞ ,
[Equation]
The theorem follows immediately from this claim.
The rest of the proof is about showing the claim.
Note that
[Equation]
where φ 1: K + L =
[ φ ( g n ( x 1 ; Ψ n )) , . .
, φ ( g n ( x K + L ; Ψ n ))] ∈ R n × ( K + L ) .
Conditioned on G and σ 2 v , the random variable 1 √ n v (cid:62) n φ 1: K + L in the nested expectation from above is distributed by a multivariate Gaussian.
The mean of this Gaussian distribution is:
[Equation]
The covariance is:
[Equation]
Using what we have observed so far
[Equation]
we compute the limit in our claimed convergence:
[Equation]
The last equality follows from the dominated convergence theorem and the fact that h is bounded.
To go further in our computation, we observe that K [ n ] converges almost surely to σ 2 v K by Theorem
[Equation]
and so
[Equation]
Here we use N not just for the multivariate Gaussian distribution but also for its density, and derive the almost sure convergence from the dominated convergence theorem.
Using this observation, we complete our computation:
[Equation]
Theorem 3.2 (Convergence under Readout-Layer Training) .
Assume that we train only the readout layers of the BNN s using the training set D tr under mean squared loss and inﬁnitesimal step size.
Formally, this means the network parameters are evolved under the following differential equations:
[Equation]
Then, the time and width limit of the random variables ( f n ( x K + i ; v ( t ) n , Ψ ( t ) n ))
i ∈
[ L ] , denoted ( f ∞∞ ( x K + i ))
i ∈
[ L ] , is distributed by the following mixture of Gaussians:
[Equation]
, where Y tr consists of the y values in D tr (i.e., Y tr = ( y 1 , . .
, y K ) ) and the K ’s with different subscripts are the restrictions of K in Theorem 3.1 with training or test inputs as directed by those subscripts.
Furthermore, if H is the inverse-gamma distribution with shape a and scale b , the marginal distribution of ( f ∞∞ ( x K + i ))
i ∈
[ L ] is the following multivariate t distribution:
[Equation]
Our proof of this theorem reuses the general structure of the proof of the corresponding theorem for NNGP s in (Lee et al., 2019), but it adjusts the structure slightly so as to use the master theorem for tensor programs and cover general architectures expressed by those programs.
Another new important part of the proof is to condition on the random scale σ 2 v , so that the existing results such as the master theorem can be used in our case.
Proof.
Note that Ψ ( t ) n = Ψ (0) n for all t because d Ψ ( t ) n /dt = 0 .
Let ¯ φ i [ n ] = 1 √ nφ (cid:16) g n (cid:16) x [Equation] Ψ ( t ) n (cid:17)(cid:17) [Equation]
[Equation]
n (cid:17)(cid:17) for i ∈
[ K + L ] , ¯ Φ tr [ n ] = (cid:2) ¯ φ 1 [ n ] , . .
, ¯ φ K [ n ] (cid:3)
[Equation]
Using this notation, we rewrite the differential equation for v ( t ) n as follows:
[Equation]
This ordinary differential equation has a closed-form solution, which we write below:
[Equation]
We multiply both sides of the equation with the ¯ Φ te [ n ] matrix, and get
[Equation]
where K [ n ] te , tr = ¯ Φ te [ n ] ¯ Φ tr [ n ] (cid:62) and K [ n ] tr , tr = ¯ Φ tr [ n ] ¯ Φ tr [ n ] (cid:62) .
Then, we send t to ∞ , and obtain ¯ Φ te [ n ] ¯ v ( ∞ ) n
[Equation]
But
[Equation]
Thus
[Equation]
we can complete the proof of the theorem if we show that for every bounded continuous function h : R L → R ,
[Equation]
Pick a bounded continuous function h : R L → R .
For each n , let G n be the σ -ﬁeld generated by g n ( x 1 ; Ψ n ) , . .
. , g n ( x K + L ; Ψ n ) .
Also, deﬁne K [ n ] tr , te = ¯ Φ tr [ n ] ¯ Φ te [ n ] (cid:62) .
We show the sufﬁcient condition mentioned above:
[Equation]
The second equality uses the fact that v (0) n consists of iid samples from N (0 , σ 2 v ) , and it is a consequence of straightforward calculation.
The third equality uses the dominated convergence theorem and the fact that h is bounded.
The fourth equality uses the master theorem for tensor programs (Theorem 2.1), the boundedness of h , and the dominated convergence theorem.
The last equality follows from the deﬁnition of ( f ∞∞ ( x K + i ))
i ∈
[ L ] .
If σ 2 v ∼ InvGamma( a, b ) , we get the following closed-form marginal distribution by Lemma B.2:
[Equation]
Our proof of Theorem 3.3 uses the corresponding theorem for the standard NTK setup in Lee et al.
We restate this existing theorem in our setup using our notation:
Theorem B.3.
Assume that the BNN s are multi-layer perceptrons, and we train these models using the training set D tr under the mean squared loss and inﬁnitesimal step size.
Regard σ 2 v as a deterministic value.
Then as n → ∞ and t → ∞ , ( f n ( x K + i ; v ( t ) n , 1 √ n Ψ ( t ) n ))
i ∈
[ L ] converges in distribution to the following random variable:
[Equation]
where
[Equation]
Theorem 3.3 (Convergence under General Training) .
Assume that the BNN s are multi-layer perceptrons and we train all of their parameters using the training set D tr under mean squared loss and inﬁnitesimal step size.
Formally, this means the network parameters are evolved under the following differential equations:
[Equation]
Let ¯Θ be Θ with σ 2 v in it set to 1 (so that Θ =
σ 2 v ¯Θ )
[Equation]
Then, the time and width limit of the random variables ( f n ( x K + i ; v ( t ) n , 1 √ n Ψ ( t ) n )) i ∈
[ L ] has the following distribution:
[Equation]
where
[Equation]
Here the K ’s with subscripts are the ones in Theorem 3.1 and the ¯Θ ’s with subscripts are similar restrictions of ¯Θ with training or test inputs.
Furthermore, if H is the inverse gamma with shape a and scale b , the exact marginal distribution of ( f ∞∞ ( x K + i ))
i ∈
[ L ] is the following multivariate t distribution:
[Equation]
Proof.
Assume that we condition on σ 2 v .
Then, by Theorem B.3, as we send n to ∞ and then t to ∞ , the conditioned random variable ( f n ( x K + i ; v ( t ) n , 1 √ n Ψ ( t ) n )) i ∈
[ L ] | σ 2 v converges in distribution to the following random variable:
[Equation]
where
[Equation]
Thus, by Lemma B.4, if we remove the conditioning on σ 2 v (i.e., we marginalize over σ 2 v ), we get the convergence of the unconditioned ( f n ( x K + i ; v ( t ) n , 1 √ n Ψ ( t ) n ))
i ∈
[ L ] to the following random variable:
[Equation]
where
[Equation]
In particular, when H =
InvGamma( a, b ) , the exact marginal distribution of ( f ∞∞ ( x K + i ))
i ∈
[ L ] has the following form by Lemma B.2:
[Equation]
Lemma B.4.
Let ( X n ) n ∈
N be a sequence of p -dimensional random vectors, X be a p -dimensional random variable, and Y be a random variable.
If X n | Y converges in distribution to X | Y almost surely (with respect to the randomness of Y ) as n tends to ∞ , then X n converges in distribution to
X .
Proof.
Let h : R p
→ R be a bounded continuous function.
We have to show that
[Equation]
The following calculation shows this equality:
[Equation]
The ﬁrst and last equalities follow from the standard tower law for conditional expectation, the second equality uses the dominated convergence theorem and the boundness of h , and the third equality uses the fact that X n | Y converges in distribution to X | Y almost surely.
C S TOCHASTIC V ARIATIONAL G AUSSIAN P ROCESS
Let ( X, y ) be training points and ( X ∗ , y ∗ ) be test points.
And let Z be inducing points.
We want to make q ( f, f Z ) which well approximates p ( f, f Z | y ) by maximizing ELBO.
Then we will construct q ( f, f Z ) =
p ( f | f Z ) q ( f Z ) where p ( f | f Z ) =
N ( f | K XZ K − 1 ZZ f Z , K XX − K XZ K − 1 ZZ K ZX ) and q ( f Z ) =
N ( f Z | µ, Σ) .
Then
[Equation]
Here we can calculate likelihood term as
[Equation]
where A = K XZ K − 1 ZZ ,
[Equation]
=
[Equation]
XZ K −
[Equation]
ZZ K ZX .
Also we can calculate KL term as
[Equation]
We ﬁnd optimal µ, Σ and inducing points
Z by using stochastic gradient descent.
And with these optimal values, we can calculate predictive distribution p ( f ∗ | y ) .
[Equation]
This equation can be write as p
[Equation]
K ∗ Z K − 1 ZZ K (cid:62)∗ Z ) .
Now if we use reparametrization trick at f Z , we can write f Z =
Lu , where L is the lower triangular matrix from the cholesky decomposition of the matrix K ZZ .
Then
[Equation]
In this case ELBO changes into (cid:90) log( p ( y | f
[Equation]
N
[Equation]
, I )) .
With softmax likelihood, we can use MC approximation which is
[Equation]
where B is minibatch and T is sample number of f which sampled from
[Equation]
Here A B =
K X
B Z K − 1 ZZ and B B = K X B X B − K X B Z K − 1 ZZ K ZX B .
Finally we can calculate predictive distribution with this reparametrization trick.
[Equation]
Thus the ﬁnal predictive distribution p ( y ∗ | x ∗ ) is
[Equation]
where f i ∗ s are sampled from N ( f ∗ ; K ∗ Z K −
[Equation]
ZZ Lµ
[Equation]
Z
[Equation]
−
K
[Equation]
Z ) .
D S TOCHASTIC V ARIATIONAL STUDENT t P ROCESS
We want to make q ( f, f Z , σ 2 ) which well approximates p ( f, f Z , σ 2 | y ) by maximizing ELBO.
First p ( f, f Z , σ 2 ) = p ( f | f Z , σ 2 ) p ( f Z | σ 2 ) p ( σ 2 ) where p ( f | f Z , σ 2 ) =
N ( f | K XZ K − 1 ZZ f Z , σ 2 ( K XX −
K XZ K − 1 ZZ K ZX )) , p ( f Z | σ 2 ) =
N ( f Z | 0 , σ 2 K ZZ ) and p ( σ 2 ) =
Γ − 1 ( σ 2 | α, β ) .
And q ( f, f Z , σ 2 ) =
p ( f | f Z , σ 2 ) q
( f Z | σ 2 ) q ( σ 2 ) where q ( f Z | σ 2 ) =
N ( f Z | µ, σ 2 Σ) and q ( σ 2 ) =
Γ − 1 ( σ 2 | a, b ) .
Then
ELBO can be computed as
[Equation]
Here KL divergence can be reformulated as
[Equation]
Now let’s compute likelihood part ﬁrst.
[Equation]
where q ( f ) = (cid:82) (cid:82) p ( f | f Z , σ 2 ) q
( f Z | σ 2 ) q ( σ 2 )d f Z d σ 2 .
And this is q ( f )
=
MVT( f | 2 a, K XZ K − 1 ZZ µ, b a ( K XZ K − 1 ZZ Σ K − 1 ZZ K ZX + K XX − K XZ K − 1 ZZ K ZX )) .
By slightly modifying the KL divergence between normal-gamma distributions in Soch & Allefeld (2016), we get
[Equation]
where ψ ( · ) is digamma function.
Now if we use reparametrization trick at f Z
[Equation]
we can write f Z =
Lu , where L is the lower triangular matrix from the cholesky decomposition of the matrix K ZZ .
Then
[Equation]
In this case ELBO changes into
[Equation]
With softmax likelihood
[Equation]
we can use MC approximation which is
[Equation]
[Equation]
[Equation]
[Equation]
where B is minibatch and T is sample number of f which sampled from
MVT
[Equation]
[Equation]
Here A B =
K X
B Z K − 1 ZZ and B B = K X B X B − K X B Z K − 1 ZZ K ZX B .
Finally we can calculate predictive distribution with this reparametrization trick.
[Equation]
Thus the ﬁnal predictive distribution p ( y ∗ | x ∗ ) is
[Equation]
i ∗ ) where f i ∗ s are sampled from MVT( f ∗ ; 2 a, K ∗ Z K − 1 ZZ Lµ u , ba ( K ∗ Z K − 1 ZZ L Σ u L (cid:62) ( K ∗ Z K − 1 ZZ ) (cid:62) +
K ∗∗ − K ∗ Z K − 1 ZZ K (cid:62) ∗ Z )) .
E I NFERENCE A LGORITHMS E. I NFERENCE FOR R EGRESSION
In order to calculate the predictive posterior distribution for an intractable marginal distribution
[Equation]
we can use self-normalized importance sampling.
Consider a scale mixture of NNGP s with a prior H on the variance σ 2 v .
Assume that we would like to estimate the expectation of h ( y ) for some function h : R → R where the random variable y is drawn from the predictive posterior of the mixture at some input x ∈ R
I under the condition on D tr .
Then, the expectation of h ( y ) is
[Equation]
where Z = (cid:82) p ( σ 2 v
[Equation]
Y tr | X tr )d σ 2 v .
We can approximate Z as follows:
[Equation]
where β i s are sampled independently from H .
Using this approximation, we can also approximate expectation of h ( y ) as follows:
[Equation]
where the y i ’s are sampled from the posterior of the Gaussian distribution and the w i ’s are the corresponding importance weights for all i ∈
[ N ] :
[Equation]
i
[Equation]
The K is the covariance matrix computed with test input as in Theorem 3.1.
To speed up calculation, we removed duplicated calculations in our importance weights and also during the sampling of the y i ’s. First, for importance weights, we observe that the log likelihood of β i has the form:
[Equation]
, and the three terms K 2 log(2 π ) , 1 2 log det( K tr , tr ) , and Y (cid:62) tr K − 1 tr , tr Y tr here are independent with β i .
Thus, using this observation, we compute these three terms beforehand only once, and calculate the log likelihood of each sample β i in O (1) time.
Next, for sampling the y i ’s, we ﬁrst draw N samples
¯ y
i s from N (0 , K x , x − K x , tr K − 1 tr , tr K tr , x ) , then multiply each of these samples with √ β i , and ﬁnally add K x , tr K − 1 tr , tr Y tr to each of the results.
Since K x , tr K − 1 tr , tr Y tr + √ β
[Table]
E.
T IME C OMPLEXITY A NALYSIS
For time complexity, as we mentioned in Appendix E
[Table]
our posterior-predictive algorithm based on importance sampling does not induce signiﬁcant overhead thanks to the reuse of the shared terms for calculation.
More speciﬁcally, our algorithm with K sample variances spends O ( K + N 3 ) time, instead of O ( KN 3 ) , for computing a posterior predictive, where N is the number of training points.
Compare this with the usual time complexity O ( N 3 ) of the standard algorithm for Gaussian processes.
When it comes to SVGP and SVTP, one update step of both SVGP and SVTP takes
O ( BM 2 + M 3 ) time, where B is the number of the batch size of the input dataset and M is the number of the inducing points.
F C LASSIFICATION WITH G AUSSIAN L IKELIHOOD
We compare the NNGP s and the scale mixture of NNGP s with Inverse Gamma prior and Burr Type
XII prior for the classiﬁcation tasks.
Following the convention (Lee et al., 2018;
Novak et al., 2018;
Garriga-Alonso et al., 2019), we treat the classiﬁcation problem as a regression problem where the regression targets are one-hot vectors of class labels and computed the posteriors induced from squared-loss.
We searched hyperparameters refer to Adlam et al.
and we choose both c and k from [0 . 5 , 1 ., 2 ., 3 ., 4 . ] .
We do not particularly expect the scale-mixtures of NNGP s to outperform
NNGP in terms of predictive accuracy, but we expect them to excel in terms of calibration due to its ﬂexibility in describing the variances of class probabilities.
To better highlight this aspect, for
MNIST data, we trained the models on clean training data but tested on corrupted data to see how the models would react under such distributional shifts.
The results are summarized in Table 4.
Due to the limitation of the resources for computing full data kernel matrix, we use 5000 samples as train set, and 1000 samples as test set.
For all datasets, the scale-mixture of NNGP s outperform NNGP in terms of NLL.
G E XPERIMENTAL D ETAILS
In Fig.
1, we used the inverse gamma prior with hyperparameter setting α = 2 and β = 2 for the experiments validating convergence of initial distribution (Theorem 3.1) and last layer training (Theorem 3.2).
For the full layer training (Theorem 3.3), we used α = 1 and β = 1 .
For the regression experiments, we divide each dataset into train/validation/test sets with the ratio 0 .
We performed gradient descent in order to update parameters of our models except number of layers which is discrete value.
We referred Adlam et al.
(2020) for initializing parameters of NNGP s.
We choose the best hyperparameters based on validation NLL values and measured NLL values on the test set with permuted train sets.
For the classiﬁcation experiments, we divide each dataset into train/test sets as provided by TensorFlow Datasets 2 , and further divide train sets into train/validation sets with the ratio 0 .
We choose the best hyperparameters based on the validation NLL values, and measure NLL and accuracy values on the test set with different initialization seeds.
To measure the uncertainty calibration
[Table] [Table]
of the model
[Table]
we used the corrupted variants of the dataset and made three new dataset from the base datasets.
For the corrupted variants
[Table]
we trained the model on the original version of the datasets
[Table]
and tested on the provided corrupted verions.
For the out-of-distribution variants
[Table]
we removed 3 classses on the train set.
For imbalance variants, we limited the samples per class with exponential ratio on the train set.
For noisy label variants, we selected 50% of labels and assigned uniformly selected new labels from the train set.
Except the corrupted variants, we tested the model on the original test set.
For the classﬁcation by categorical likelihood experiments, we use the CNN model with two layers to compute the kernel.
Due to the limitations of computing resources, we can only use 400 inducing points for the variational inference which leads slight degradation on the performance.
For the classiﬁcation by ﬁnite model ensemble experiments, each CNN layer has 128 channels.
H A DDITIONAL E XPERIMENTS
The experiment code is available at GitHub 3 .
We used a server with Intel Xeon Silver 4214R CPU and 128GB of RAM to evaluate the classiﬁcation with Gaussian likelihood experiment, and used
NVIDIA GeForce RTX 2080Ti GPU to conduct other experiments.
Impact of the prior hyperparameters With same experimental setting with Section 4.1, in order to inspect the impact of the prior hyperparameters, we sampled 1000 models for initial, last layer training and full layer training.
In Fig.
2, here we can empirically see that consistently with the theory, if α is smaller, the output distribution is more heavy-tailed.
Full-layer training correspondence for ResNet Even though we only proved the general training result (Theorem 3.3) only for the fully-connected neural networks, we experimentally found that
ResNet also shows a behavior predicted by our theorem.
The empirical validation with for ResNet is shown in Fig.
We set α = 4 and β = 4 for the inverse gamma prior in this experiment.
Additional results of regression with Gaussian likelihood In addition to Table 1, we also measured the root mean square error values on the test set.
Table 5 summarizes the results.
The results other than ours are borrowed from Adlam et al.
and we use same settings as described in
Section 4.2.
[Figure] [Figure]
Additional regression experiment As an additional regression experiments, we test the models which use Burr Type XII distribution as its last layer variance’s prior.
We use Appendix E in order to calculate the predictive posterior distribution.
Our results are in Table 6.
I A DDITIONAL I NFORMATION
We refer to the source of each dataset through footnotes with URL links.
We don’t use any data which obtained from people and don’t use any data which contains personally identiﬁable information or offensive content.