Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 1.1 RELATED WORKS
text: 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.

heading: 2 PRELIMINARIES
text: 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}.

heading: 2.1 TENSOR PROGRAMS AND NEURAL NETWORK GAUSSIAN PROCESSES
text: 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).

heading: 3 RESULTS
text: We will present our results for the regression setting. Throughout this section, we consider only those BNNs whose computations over fixed inputs are representable by tensor programs. Assume that we are given a training set Dtr = {(xk, yk) ∈ RI × R | 1 ≤ k ≤ K} and a test set Dte = {(xK+`, yK+`) | 1 ≤ ` ≤ L}. Our idea is to treat the variance σ2v 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 σ2v and the random initialization of the network parameters). To state this more precisely, consider the computations of the BNNs {fn(−; vn,Ψn)}n∈N on the inputs x1:K+L in the training and test sets such that n is the width of the hidden layers of fn and these
computations are representable by a tensor program. Our idea is to change the distribution for the initial value of vn from a Gaussian to a scale mixture of Gaussians, so that at initialization, the BNNs on x1:K+L become the following random variables: for each n ∈ N, σ2v ∼ H, vn,α|σ2v ∼ N (0, σ2v) for α ∈ [n], Ψn,(α,j) ∼ N (0, σ2j /n) for α ∈ [n], j ∈ [P ],
fn(xi; vn,Ψn) = 1√ n ∑ α∈[n] vn,α · φ(gn(xi; Ψn))α for i ∈ [K + L],
whereH is a distribution on positive real numbers, such as the inverse-gamma distribution, P is the number of columns of Ψn, and σ2j and gn are, respectively, the variance specific to the j-th column of Ψn and the output of the penultimate linear layer, as explained in our review on tensor programs.
Our first result says that as the width of the BNN grows to∞, the BNN converges in distribution to a mixture of NNGPs, and, in particular, to Student’s t process if the distributionH is inverse gamma. Theorem 3.1 (Convergence). As n tends to∞, the random variable (fn(xi; vn,Ψn))i∈[K+L] converges in distribution to the following random variable (f∞(xi))i∈[K+L]:
σ2v ∼ H, (f∞(xi))i∈[K+L]|σ2v ∼ N (0, σ2v · K),
where K(i,j) = EZ∼N (µ,Σ)[φ(Zi)φ(Zj)] for i, j ∈ [K + L]. Furthermore, if H is the inversegamma distribution with shape a and scale b, the marginal distribution of (f∞(xi))i∈[K+L] is the following multivariate t distribution:
(f∞(xi))i∈[K+L] ∼ MVTK+L(2a, 0, (b/a) · K), where MVTd(ν′, µ′,K′) denotes the d-dimensional Student’s t distribution with degree of freedom ν′, location µ′, and scale matrix K′.
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 σ2v 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 finite dimensional case with the inputs x1:K+L, it implies the convergence of the BNNs with random σ2v to a scale mixture of GPs or to Student’s t process. Note that σ2v · K in the theorem is precisely the covariance matrix K of the NNGP kernel on x1:K+L in Corollary 2.2. We thus call the limiting stochastic process as scale mixture of NNGPs.
Our next results characterize a fully-trained BNN at the infinite-width limit under two different training schemes. They are the versions of the standard results in the literature, generalized from GPs to scale mixtures of GPs. Assume the BNNs under the inputs x1: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 BNNs using the training set Dtr under mean squared loss and infinitesimal step size. Formally, this means the network parameters are evolved under the following differential equations:
(v(0)n ,Ψ (0) n ) = (vn,Ψn),
dv (t) n
dt = K∑ i=1 ( yi − fn(xi; v(t)n ,Ψ (t) n ) )2φ(gn(xi; Ψ(t)n )) K √ n , dΨ(t)n dt = 0.
Then, the time and width limit of the random variables (fn(xK+i; v (t) n ,Ψ (t) n ))i∈[L], denoted (f∞∞ (xK+i))i∈[L], is distributed by the following mixture of Gaussians:
σ2v ∼ H, (f∞∞ (xK+i))i∈[L]|σ2v ∼ N (( Kte,trK −1 tr,trYtr ) , σ2v ( Kte,te −Kte,trK −1 tr,trKtr,te )) ,
where Ytr consists of the y values in Dtr (i.e., Ytr = (y1, . . . , yK)) 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, ifH is the inverse-gamma distribution with shape a and scale b, the marginal distribution of (f∞∞ (xK+i))i∈[L] is the following multivariate t distribution:
(f∞∞ (xK+i))i∈[L] ∼ MVTL ( 2a, ( Kte,trK −1 tr,trYtr ) , b
a
( Kte,te −Kte,trK −1 tr,trKtr,te )) .
One important point about this theorem is that the distribution of (f∞∞ (xK+i))i∈[L] is not the posterior of (f∞(xK+i))i∈[L] in Theorem 3.1 under the condition on Ytr. The former uses the prior on σ2v for mixing the distribution of (f∞(xK+i))i∈[L] | (σ2v , Ytr), while the latter uses the posterior on σ2v 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. (2019) about the training of infinitely-wide neural networks. Concretely, we use the so called NTK parametrization for Ψn. This means that the BNNs on x1:K+L are now intialized to be the following random variables: for each n ∈ N,
σ2v ∼ H, vn,α|σ2v ∼ N (0, σ2v), Ψn,(α,j) ∼ N (0, σ2j ) for α ∈ [n], j ∈ [P ],
fn(xi; vn, 1√ n Ψn) = 1√ n ∑ α∈[n] vn,α · φ(gn(xi; 1√ n Ψn))α for i ∈ [K + L].
When we adopt this NTK parameterization and the BNNs 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 BNNs: as n tends to infinity, for all i, j ∈ [K + L],〈
∇(vn,Ψn)fn(xi; vn, 1√ n Ψn), ∇(vn,Ψn)fn(xj ; vn, 1√ n
Ψn) 〉 a.s.−−→ Θ(i,j) (4)
for some Θ(i,j) determined by the architecture of the BNNs. Let Θ be the matrix (Θ(i,j))i,j∈[K+L]. The next theorem describes the outcome of our analysis for the fully-trained infinite-width BNN. Theorem 3.3 (Convergence under General Training). Assume that the BNNs are multi-layer perceptrons and we train all of their parameters using the training set Dtr under mean squared loss and infinitesimal step size. Formally, this means the network parameters are evolved under the following differential equations:
v(0)n = vn, dv
(t) n
dt =
2
K √ n K∑ i=1 ( yi − fn(xi; v(t)n , 1√ n Ψ(t)n ) ) φ(gn(xi; 1√ n Ψ(t)n )),
Ψ(0)n = Ψn, dΨ(t)n dt = 2 K K∑ i=1 ( yi − fn(xi; v(t)n , 1√ n Ψ(t)n ) ) ∇Ψfn(xi; v(t)n , 1√ n Ψ(t)n ).
Let Θ̄ be Θ with σ2v in it set to 1 (so that Θ = σ 2 vΘ̄). Then, the time and width limit of the random variables (fn(xK+i; v (t) n , 1√ n Ψ(t)n ))i∈[L] has the following distribution:
σ2v ∼ H, (f∞∞ (xK+i))i∈[L]|σ2v ∼ N (µ′, σ2vΘ′) where
µ′ = Θ̄te,trΘ̄ −1 tr,trYtr, Θ′ = Kte,te + ( Θ̄te,trΘ̄ −1 tr,trKtr,trΘ̄−1tr,trΘ̄tr,te ) − ( Θ̄te,trΘ̄ −1 tr,trKtr,te +Kte,trΘ̄−1tr,trΘ̄tr,te ) .
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∞∞ (xK+i))i∈[L] is the following multivariate t distribution:
(f∞∞ (xK+i))i∈[L] ∼ MVTL(2a, µ′, (b/a) ·Θ′). (5)
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 σ2v of the readout layer’s parameters, the networks in our setup fit into the standard setting for NNGPs and NTKs, 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 σ2v 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.

heading: 3.1 POSTERIOR INFERENCE
text: Gaussian likelihood As for the NNGPs, when the scale is mixed by the inverse gamma prior, we can exactly compute the posterior predictives of the scale mixture of NNGPs 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 specifically for scale mixtures of NNGPs.
Consider a scale mixture of NNGPs with a prior H on the variance σ2v . 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 ∈ RI under the condition onDtr. The importance sampling with prior as proposal computes the estimator ( ∑N i=1 wih(yi))/ ∑N j=1 wj where βi and yi for i = 1, . . . , N are independent samples from the prior H and the posterior of the Gaussian distribution, respectively, and the wi’s are the importance weights of the βi’s:
βi ∼ H, wi = N (Ytr; 0, βiKtr,tr), yi ∼ N (Kx,trK −1 tr,trYtr, βi(Kx,x −Kx,trK −1 tr,trKtr,x)).
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 inefficient calculation of the likelihoods N (Ytr; 0, βiKtr,tr) of the sampled βi’s. To see this, note that the log likelihood of βi can be decomposed as follows:
logN (Ytr; 0, βiKtr,tr) = − K 2 log(2π)− 1 2 log det
( Ktr,tr ) − K
2 log(βi)−
1
2βi Y >tr K −1 tr,trYtr.
The terms K2 log(2π), 1 2 log det(Ktr,tr), and Y > tr K −1 tr,trYtr 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 NNGPs 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.

heading: 4 EXPERIMENTS
text: We empirically evaluated the scale mixtures of NNGPs on various synthetic and real-world tasks. We tested the scale mixture of NNGPs 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).

heading: 4.1 EMPIRICAL VALIDATION OF THE THEORIES
text: 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.052/n). For the last layer, we sampled σ2v ∼ InvGamma(a, b) with varying (a, b) values and sampled weights from N (0, σ2v). 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 first 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.

heading: 4.2 REGRESSION WITH GAUSSIAN LIKELIHOODS
text: We tested the scale mixtures of NNGPs 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 NNGPs to be better calibrated than NNGP due to their heavy-tail behaviors. To see this, we fitted the scale mixture of NNGPs 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. (2020). In summary, ours generally performed similarly to NNGP except for some datasets for which ours significantly outperformed NNGP. Considering the fact that Student’s t processes include GPs as limiting cases, this result suggests that one can use the scale mixtures as an alternative to NNGPs even for the datasets not necessarily including heavy-tailed noises. See Appendix G for detailed settings for the training.

heading: 4.3 CLASSIFICATION WITH GAUSSIAN LIKELIHOOD
text: Following the literature, we apply the NNGPs and the scale mixtures of NNGPs to classification 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.

heading: 4.4 CLASSIFICATION WITH CATEGORICAL LIKELIHOOD
text: We compared the NNGP and the scale mixture of NNGPs with inverse-gamma priors on the image classification task. Following the standard setting in image classification with deep neural networks, we computed the posterior predictive distributions of the BNNs under categorical likelihood. Since both NNGPs and the scale mixtures of NNGPs 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 NNGPs (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 classification 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. Table 2 summarizes the results, which show that SVTP and SVGP perform similarly for in-distribution data,
1https://archive.ics.uci.edu/ml/datasets.php
but SVTP significantly outperforms SVGP for OOD data in terms of uncertainty calibration (measured by negative log-likelihoods).

heading: 4.5 CLASSIFICATION BY FINITE MODEL ENSEMBLE
text: 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 finite 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.

heading: 5 CONCLUSION
text: In this paper, we proposed a simple extension of NNGPs by introducing a scale prior on the lastlayer weight parameters. The resulting method, entitled as the scale mixture of NNGPs, defines 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 infinitely-wide BNN of any architecture constructed in a specific way corresponds to the scale mixture of NNGPs. Also, we have extended the existing convergence results of infinitely-wide BNNs trained with gradient descent to GPs so that the results hold for our construction. Our empirical evaluation validates our theory and shows promising results for multiple real-world regression and classification tasks. Especially, it shows that the heavy-tailed processes from our construction are robust to the out-of-distribution data.

heading: A DETAILS ON TENSOR PROGRAMS
text: 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 =∑l i=1 wix
i where y, x1, . . . , xl are G-variables, and w1, . . . , wl ∈ R are constants. The final Nonlin assignments apply (usually nonlinear) scala functions on the values of G-variables coordinatewise. That is, they have the form y = φ(x1, . . . , xl), where y is an H-variable, x1, . . . , xl are G-variables, φ : Rl → R is a possibly nonlinear function, and φ(x1, . . . , xl) means the lifted application of φ to vector arguments.
Every tensor program ends with a return statement of the form:
return(v1>y1/ √ n, . . . , vl>yl/ √ n), (6)
where v1, . . . , vl are G-variables not appearing anywhere in the program, and y1, . . . , yl 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 σ2W /n for some σW > 0 (i.e., Wij ∼ N (0, σ2W /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 x1, . . . , xm (which store vectors in Rn) 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 vi 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, σ2v). Definition 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 infinite-width limit (Theorem A.2). The computation is defined inductively as follows. Let g1, . . . , gN be all the G-variables in the program, ordered as they appear in the program. For any G-variable gk, µ(gk) is defined by:
µ(gk) =  µin(gk) if gk is an input G-variable∑m i=1 wiµ(x i) if gk = ∑m i=1 aix i
0 otherwise .
Here x1, . . . , xm are some G-variables. For any pair of G-variables gk, gl, Σ(gk, gl) is defined by:
Σ(gk, gl) =  Σin(gk, gl) if gk, gl are input G-variables∑m i=1 wiΣ(x i, gl) if gk = ∑m i=1 wix i∑m i=1 wiΣ(g k, xi) if gl = ∑m i=1 wix i
σ2WEZ [φ(Z)φ′(Z)] if gk = Wy, gl = Wy′ for the same W , with Z,φ,φ′ below 0 otherwise
.
Here x1, . . . , xm are some G-variables, and H-variables y and y′ in the second last case are introduced by the Nonlin assignments y = φ(gi1 , . . . , gip) and y′ = φ′(gi1 , . . . , gip), respectively, for
some i1, . . . , ip ∈ [max(k − 1, l − 1)]. We have adjusted the inputs of φ and φ′ 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 (µ′,Σ′) where µ′ and Σ′ are the mean vector and covariance matrix for gi1 , . . . , gip by the previous step of this inductive construction. Definition A.2. Given h : Rn → R and for any x ∈ Rn if |h(x)| is upper bounded by a function exp(C‖x‖2− + c) with C, c, > 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 g1, . . . , gN be all of the G-variables that appear in this program, and h : RN → R be a controlled function. Then, as n tends to∞,
1
n n∑ α=1 h(g1α, . . . , g N α ) a.s.−−→ EZ∼N (µ,Σ)[h(Z1, . . . , ZN )].
Here a.s.−−→ refers to almost sure convergence, and Z = (Z1, . . . , ZN ) ∈ RN . Also, µ and Σ are from Definition A.1.
Corollary A.3. Consider a tensor program satisfying Assumption A.1 and using only controlled φ’s. Let g1, . . . , gN be all of the G-variables that appear in this program. Assume that the program returns a vector (v>y1/ √ n, . . . , v>yk/ √ n) for some H-variables yi’s such that each yi is updated by yi = φi(g1, . . . , gN ) 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:
K(i,j) = σ2v · EZ∼N (µ,Σ)[φi(Z)φj(Z)] for all i, j ∈ [k],
where µ and Σ are from Definition A.1.

heading: B PROOFS
text: In this section, we provide the proofs of our three main results. We first 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
x = [ x1 x2 ] ∈ Rd1+d2 ∼ MVTd1+d2 ( ν, [ µ1 µ2 ] , [ Σ11 Σ12 Σ21 Σ22 ]) , (7)
we have that x1 ∼ MVTd1(ν, µ1,Σ11) and x2|x1 ∼ MVTd2 ( ν + d1, µ ′ 2, ν + c
ν + d1 Σ′22
) (8)
where
µ′2 = µ2 + Σ21Σ −1 11 (x1 − µ1), c = (x1 − µ1)>Σ −1 11 (x1 − µ1), Σ′22 = Σ22 − Σ21Σ −1 11 Σ12.
Definition 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:
p(x) = ba
G(a)
( 1
x
)a+1 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 Σ ∈ Rd×d be a postive definite matrix. Consider µ ∈ Rn and a, b ∈ R+. If
σ2 ∼ InvGamma(a, b), x|σ2 ∼ N (µ, σ2Σ),
then the marginal distribution of x is MVTd(2a, µ, baΣ).
Now we proceed to the proofs of main theorems. Theorem 3.1 (Convergence). As n tends to∞, the random variable (fn(xi; vn,Ψn))i∈[K+L] converges in distribution to the following random variable (f∞(xi))i∈[K+L]:
σ2v ∼ H, (f∞(xi))i∈[K+L]|σ2v ∼ N (0, σ2v · K),
where K(i,j) = EZ∼N (µ,Σ)[φ(Zi)φ(Zj)] for i, j ∈ [K + L]. Furthermore, if H is the inversegamma distribution with shape a and scale b, the marginal distribution of (f∞(xi))i∈[K+L] is the following multivariate t distribution:
(f∞(xi))i∈[K+L] ∼ MVTK+L(2a, 0, (b/a) · K), where MVTd(ν′, µ′,K′) denotes the d-dimensional Student’s t distribution with degree of freedom ν′, location µ′, and scale matrix K′.
Proof. Let oi[n] = fn(xi; vn,Ψn) for i ∈ [K + L] and n. Pick an arbitrary bounded continuous function h : RK+L → R. For each n, define Gn to be the σ-field generated by gn(x1; Ψn), . . . , gn(xK+L; Ψn).
We claim that as n tends to∞, E [ h ( o1[n], . . . , oK+L[n] )] → Eσ2v∼H [ EY 1:K+L∼N(0,σ2vK) [ h ( Y 1, . . . , Y K+L )]] .
The theorem follows immediately from this claim. The rest of the proof is about showing the claim.
Note that
E [ h ( o1[n], . . . , oK+L[n] )] = E [ E [ h ( 1√ n v>n φ 1:K+L )∣∣∣∣G, σ2v]] , where φ1:K+L = [φ (gn (x1; Ψn)) , . . . , φ (gn (xK+L; Ψn))] ∈ Rn×(K+L). Conditioned on G and σ2v , the random variable 1√ n v>n φ
1:K+L in the nested expectation from above is distributed by a multivariate Gaussian. The mean of this Gaussian distribution is:
E [
1√ n v>n φ 1:K+L ∣∣∣∣G, σ2v] = E [v>n ∣∣G, σ2v]× 1√nφ1:K+L = 0 ∈ RK+L.
The covariance is: K[n](i,j) = E [(
1√ n v>n φ (gn (xi; Ψn)) )( 1√ n v>n φ (gn (xj ; Ψn)) )∣∣∣∣G, σ2v] = 1
n n∑ α=1 n∑ β=1 E [ vn,αvn,β | G, σ2v ] × (φ (gn (xi; Ψn)))α (φ (gn (xj ; Ψn)))β
= 1
n σ2v n∑ α=1 (φ (gn (xi; Ψn)))α (φ (gn (xj ; Ψn)))α .
Using what we have observed so far, we compute the limit in our claimed convergence:
lim n→∞
E [ h ( o1[n], . . . , oK+L[n] )] = lim n→∞ E [ Eu1:K+L∼N(0,K[n]) [ h ( u1:K+L )]] = E [ lim n→∞ Eu1:K+L∼N(0,K[n]) [ h(u1:K+L) ]] .
The last equality follows from the dominated convergence theorem and the fact that h is bounded. To go further in our computation, we observe thatK[n] converges almost surely to σ2vK by Theorem 2.1 and so
Eu1:K+L∼N (0,K[n])[h(u 1:K+L)] =
∫ h(u1:K+L)N (u1:K+L; o,K[n])du1:K+L
a.s.−−→ ∫ h(u1:K+L)N (u1:K+L; 0, σ2vK)du1:K+L
= EY 1:K+L∼N (0,σ2vK) [ h ( Y 1:K+L )] .
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:
E [
lim n→∞
Eu1:K+L∼N (0,K[n]) [ h(u1:K+L) ]] = E [ EY 1:K+L∼N (0,σ2vK) [ h ( Y 1:K+L )]] = Eσ2v∼H [ EY 1:K+L∼N (0,σ2vK) [ h ( Y 1:K+L )]] .
Theorem 3.2 (Convergence under Readout-Layer Training). Assume that we train only the readout layers of the BNNs using the training set Dtr under mean squared loss and infinitesimal step size. Formally, this means the network parameters are evolved under the following differential equations:
(v(0)n ,Ψ (0) n ) = (vn,Ψn),
dv (t) n
dt = K∑ i=1 ( yi − fn(xi; v(t)n ,Ψ (t) n ) )2φ(gn(xi; Ψ(t)n )) K √ n , dΨ(t)n dt = 0.
Then, the time and width limit of the random variables (fn(xK+i; v (t) n ,Ψ (t) n ))i∈[L], denoted (f∞∞ (xK+i))i∈[L], is distributed by the following mixture of Gaussians:
σ2v ∼ H, (f∞∞ (xK+i))i∈[L]|σ2v ∼ N (( Kte,trK −1 tr,trYtr ) , σ2v ( Kte,te −Kte,trK −1 tr,trKtr,te )) ,
where Ytr consists of the y values in Dtr (i.e., Ytr = (y1, . . . , yK)) 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, ifH is the inverse-gamma distribution with shape a and scale b, the marginal distribution of (f∞∞ (xK+i))i∈[L] is the following multivariate t distribution:
(f∞∞ (xK+i))i∈[L] ∼ MVTL ( 2a, ( Kte,trK −1 tr,trYtr ) , b
a
( Kte,te −Kte,trK −1 tr,trKtr,te )) .
Our proof of this theorem reuses the general structure of the proof of the corresponding theorem for NNGPs 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 σ2v , 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 φ ( gn ( xi; Ψ (t) n )) = 1√ n φ ( gn ( xi; Ψ (0) n )) for i ∈ [K + L],
Φ̄tr[n] = [ φ̄1[n], . . . , φ̄K [n] ]> ∈ RK×n, Φ̄te[n] = [φ̄K+1[n], . . . , φ̄K+L[n]]> ∈ RL×n. Using this notation, we rewrite the differential equation for v(t)n as follows:
dv̄ (t) n
dt = − 2 K Φ̄tr[n]
> ( Φ̄tr[n] v̄ (t) n − Ytr ) .
This ordinary differential equation has a closed-form solution, which we write below:
v̄(t)n = v̄ (0) n + Φ̄tr[n] > (Φ̄tr[n]Φ̄tr[n]>)−1(exp(− 2 K tΦ̄tr[n]Φ̄tr[n] > ) − I )( Φ̄tr[n] v̄ (0) n − Ytr ) .
We multiply both sides of the equation with the Φ̄te[n] matrix, and get
Φ̄te[n] v̄ (t) n = Φ̄te[n] v̄ (0) n +K[n]te,trK[n]−1tr,tr
( exp ( − 2 K tK[n]tr,tr ) − I )( Φ̄tr[n] v̄ (0) n − Ytr ) where K[n]te,tr = Φ̄te[n]Φ̄tr[n]> and K[n]tr,tr = Φ̄tr[n]Φ̄tr[n]>. Then, we send t to∞, and obtain
Φ̄te[n] v̄ (∞) n = Φ̄te[n] v̄ (0) n −K[n]te,trK[n]−1tr,tr ( Φ̄tr[n] v̄ (0) n − Ytr ) = K[n]te,trK[n]−1tr,trYtr + Φ̄te[n] v̄(0)n −K[n]te,trK[n]−1tr,tr ( Φ̄tr[n] v̄ (0) n ) .
But
(fn(xK+i; v (∞) n ,Ψ (∞) n ))i∈[L] = Φ̄te[n] v̄ (∞) n .
Thus, we can complete the proof of the theorem if we show that for every bounded continuous function h : RL → R,
lim n→∞
E [ h ( K[n]te,trK[n]−1tr,trYtr + Φ̄te[n] v̄(0)n −K[n]te,trK[n]−1tr,tr ( Φ̄tr[n] v̄ (0) n ))] = E [ h (
(f∞∞ (xK+i))i∈[L]
)] .
Pick a bounded continuous function h : RL → R. For each n, let Gn be the σ-field generated by gn(x1; Ψn), . . . , gn(xK+L; Ψn). Also, define K[n]tr,te = Φ̄tr[n]Φ̄te[n]>. We show the sufficient condition mentioned above:
lim n→∞
E [ h ( K[n]te,trK[n]−1tr,trYtr + Φ̄te[n] v̄(0)n −K[n]te,trK[n]−1tr,tr ( Φ̄tr[n] v̄ (0) n ))]
= lim n→∞
E [ E [ h ( K[n]te,trK[n]−1tr,trYtr + Φ̄te[n] v̄(0)n −K[n]te,trK[n]−1tr,tr ( Φ̄tr[n] v̄ (0) n ))∣∣∣∣∣Gn, σ2v ]]
= lim n→∞ E [ EZ∼N (K[n]te,trK[n]−1tr,trYtr, σ2v(K[n]te,te−K[n]te,trK[n]−1tr,trK[n]tr,te)) [h(Z)] ] = E [ lim n→∞ EZ∼N (K[n]te,trK[n]−1tr,trYtr, σ2v(K[n]te,te−K[n]te,trK[n]−1tr,trK[n]tr,te)) [h(Z)] ]
= E [ E Z∼N (Kte,trK −1 tr,trYtr, σ 2 v(Kte,te−Kte,trK −1 tr,trKtr,te)) [h(Z)] ]
= E [ h (
(f∞∞ (xK+i))i∈[L]
)] .
The second equality uses the fact that v(0)n consists of iid samples from N (0, σ2v), 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 definition of (f∞∞ (xK+i))i∈[L].
If σ2v ∼ InvGamma(a, b), we get the following closed-form marginal distribution by Lemma B.2: (f∞∞ (xK+i))i∈[L] ∼ MVTL ( 2a, Kte,trK −1 tr,trYtr, b
a (Kte,te −Kte,trK
−1 tr,trKtr,te) ) .
Our proof of Theorem 3.3 uses the corresponding theorem for the standard NTK setup in Lee et al. (2019). We restate this existing theorem in our setup using our notation:
Theorem B.3. Assume that the BNNs are multi-layer perceptrons, and we train these models using the training set Dtr under the mean squared loss and infinitesimal step size. Regard σ2v as a deterministic value. Then as n → ∞ and t → ∞, (fn(xK+i; v(t)n , 1√nΨ (t) n ))i∈[L] converges in distribution to the following random variable:
(f∞∞ (xK+i))i∈[L] ∼ N (µ∞te ,Σ∞te )
where
µ∞te = Θte,trΘ −1 tr,trYtr, Σ∞te = Kte,te + Θte,trΘ−1tr,trKtr,trΘ−1tr,trΘtr,te − (Θte,trΘ−1tr,trKtr,te +Kte,trΘ−1tr,trΘtr,te).
Theorem 3.3 (Convergence under General Training). Assume that the BNNs are multi-layer perceptrons and we train all of their parameters using the training set Dtr under mean squared loss and
infinitesimal step size. Formally, this means the network parameters are evolved under the following differential equations:
v(0)n = vn, dv
(t) n
dt =
2
K √ n K∑ i=1 ( yi − fn(xi; v(t)n , 1√ n Ψ(t)n ) ) φ(gn(xi; 1√ n Ψ(t)n )),
Ψ(0)n = Ψn, dΨ(t)n dt = 2 K K∑ i=1 ( yi − fn(xi; v(t)n , 1√ n Ψ(t)n ) ) ∇Ψfn(xi; v(t)n , 1√ n Ψ(t)n ).
Let Θ̄ be Θ with σ2v in it set to 1 (so that Θ = σ 2 vΘ̄). Then, the time and width limit of the random variables (fn(xK+i; v (t) n , 1√ n Ψ(t)n ))i∈[L] has the following distribution:
σ2v ∼ H, (f∞∞ (xK+i))i∈[L]|σ2v ∼ N (µ′, σ2vΘ′)
where
µ′ = Θ̄te,trΘ̄ −1 tr,trYtr, Θ′ = Kte,te + ( Θ̄te,trΘ̄ −1 tr,trKtr,trΘ̄−1tr,trΘ̄tr,te ) − ( Θ̄te,trΘ̄ −1 tr,trKtr,te +Kte,trΘ̄−1tr,trΘ̄tr,te ) .
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∞∞ (xK+i))i∈[L] is the following multivariate t distribution:
(f∞∞ (xK+i))i∈[L] ∼ MVTL(2a, µ′, (b/a) ·Θ′). (5)
Proof. Assume that we condition on σ2v . Then, by Theorem B.3, as we send n to∞ and then t to ∞, the conditioned random variable (fn(xK+i; v(t)n , 1√nΨ (t) n ))i∈[L] | σ2v converges in distribution to the following random variable:
(f∞∞ (xK+i))i∈[L] | σ2v ∼ N (µ∞te ,Σ∞te ),
where
µ∞te = Θ̄te,trΘ̄ −1 tr,trYte,
Σ∞te = σ 2 v ( Kte,te + Θ̄te,trΘ̄−1tr,trKtr,trΘ̄−1tr,trΘ̄tr,te − (Θ̄te,trΘ̄−1tr,trKtr,te +Kte,trΘ̄−1tr,trΘ̄tr,te) ) .
Thus, by Lemma B.4, if we remove the conditioning on σ2v (i.e., we marginalize over σ 2 v), we get the convergence of the unconditioned (fn(xK+i; v (t) n , 1√ n Ψ(t)n ))i∈[L] to the following random variable:
σ2v ∼ H, (f∞∞ (xK+i))i∈[L] ∼ N (µ′, σ2vΘ′)
where
µ′ = Θ̄te,trΘ̄ −1 tr,trYtr, Θ′ = Kte,te + ( Θ̄te,trΘ̄ −1 tr,trKtr,trΘ̄−1tr,trΘ̄tr,te ) − ( Θ̄te,trΘ̄ −1 tr,trKtr,te +Kte,trΘ̄−1tr,trΘ̄tr,te ) .
In particular, whenH = InvGamma(a, b), the exact marginal distribution of (f∞∞ (xK+i))i∈[L] has the following form by Lemma B.2:
MVT ( 2α, µ′, b a Θ′ ) .
Lemma B.4. Let (Xn)n∈N be a sequence of p-dimensional random vectors, X be a p-dimensional random variable, and Y be a random variable. If Xn | Y converges in distribution to X | Y almost surely (with respect to the randomness of Y ) as n tends to∞, then Xn converges in distribution to X .
Proof. Let h : Rp → R be a bounded continuous function. We have to show that lim n→∞ E[h(Xn)] = E[h(X)].
The following calculation shows this equality:
lim n→∞ E[h(Xn)] = lim n→∞
E[E[h(Xn) | Y ]] = E [
lim n→∞
E[h(Xn) | Y ] ] = E[E[h(X) | Y ]]
= E[h(X)]. The first 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 Xn | Y converges in distribution to X|Y almost surely.

heading: C STOCHASTIC VARIATIONAL GAUSSIAN PROCESS
text: Let (X,y) be training points and (X∗, y∗) be test points. And let Z be inducing points. We want to make q(f, fZ) which well approximates p(f, fZ |y) by maximizing ELBO. Then we will construct q(f, fZ) = p(f |fZ)q(fZ) where p(f |fZ) = N (f |KXZK−1ZZfZ ,KXX − KXZK −1 ZZKZX) and q(fZ) = N (fZ |µ,Σ). Then log p(y) ≥ Ef,fZ∼q(f,fZ)[log p(y|f, fZ)]−KL(q(fZ)||p(fZ)). Here we can calculate likelihood term as∫ log p(y|f)N (f ;Aµ,AΣA> +D)df
where A = KXZK−1ZZ , B = KXX −KXZK −1 ZZKZX . Also we can calculate KL term as 1
2
[ log(
detKZZ det Σ )− nZ + Tr{K−1ZZΣ}+ µ >K−1ZZµ
] .
We find optimal µ,Σ and inducing points Z by using stochastic gradient descent. And with these optimal values, we can calculate predictive distribution p(f∗|y).
p(f∗|y) = ∫ ∫ p(f∗, f, fZ |y)dfdfZ
= ∫ ∫ p(f∗|f, fZ)q(f, fZ)dfdfZ
= ∫ ∫ p(f∗|f, fZ)p(f |fZ)q(fZ)dfdfZ
= ∫ p(f∗|fZ)q(fZ)dfZ .
This equation can be write as p(f∗|y) = N (f∗;K∗ZK−1ZZµ,K∗ZK −1 ZZΣ(K∗ZK −1 ZZ) > + K∗∗ − K∗ZK −1 ZZK > ∗Z). Now if we use reparametrization trick at fZ , we can write fZ = Lu, where L is the lower triangular matrix from the cholesky decomposition of the matrix KZZ . Then
p(fZ) = N (fZ ; 0, LL>) = N (fZ ; 0,KZZ) q(fZ) = N (fZ ;Lµu, LΣuL>).
In this case ELBO changes into∫ log(p(y|f))N (f ;ALµu, ALΣuL>A> +B)df −KL(N (u;µu,Σu)||N (u; 0, I)).
With softmax likelihood, we can use MC approximation which is
ELBO ' 1 T N B T∑ t=1 B∑ i=1 log p(yi|f (t)i )−KL(N (u;µu,Σu)||N (u; 0, I))
where B is minibatch and T is sample number of f which sampled from
N (f ;ABLµu, ABLΣuL>A>B +BB).
Here AB = KXBZK −1 ZZ and BB = KXBXB −KXBZK −1 ZZKZXB . Finally we can calculate predictive distribution with this reparametrization trick.
p(f∗|y) = ∫ p(f∗|fZ)q(fZ)dfZ
= ∫ N (f∗;K∗ZK−1ZZfZ ,K∗∗ −K∗ZK −1 ∗ZKZ∗)N (fZ ;Lµu, LΣuL >)dfZ
=N (f∗;K∗ZK−1ZZLµu,K∗ZK −1 ZZLΣuL >(K∗ZK −1 ZZ) > +K∗∗ −K∗ZK−1ZZK > ∗Z)
Thus the final predictive distribution p(y∗|x∗) is
log p(y∗|x∗) = log M∏ j p(yj∗|xj∗)
= M∑ j log p(yj∗|xj∗)
= M∑ j log ∫ p(yj∗|f∗)p(f∗|y)df∗ ' M∑ j log 1 N N∑ i p(yj∗|f i∗)
where f i∗s are sampled from N (f∗;K∗ZK−1ZZLµu,K∗ZK −1 ZZLΣuL >(K∗ZK −1 ZZ) > + K∗∗ − K∗ZK −1 ZZK > ∗Z).
D STOCHASTIC VARIATIONAL STUDENT t PROCESS
We want to make q(f, fZ , σ2) which well approximates p(f, fZ , σ2|y) by maximizing ELBO. First p(f, fZ , σ
2) = p(f |fZ , σ2)p(fZ |σ2)p(σ2) where p(f |fZ , σ2) = N (f |KXZK−1ZZfZ , σ2(KXX − KXZK −1 ZZKZX)), p(fZ |σ2) = N (fZ |0, σ2KZZ) and p(σ2) = Γ−1(σ2|α, β). And q(f, fZ , σ2) = p(f |fZ , σ2)q(fZ |σ2)q(σ2) where q(fZ |σ2) = N (fZ |µ, σ2Σ) and q(σ2) = Γ−1(σ2|a, b). Then ELBO can be computed as
log p(y) ≥ Ef,fZ ,σ2∼q(f,fZ ,σ2)[log p(y|f, fZ , σ 2)]−KL(q(f, fZ , σ2)||p(f, fZ , σ2)).
Here KL divergence can be reformulated as
KL(q(f, fZ , σ 2)||p(f, fZ , σ2)) = ∫ ∫ q(fZ , σ 2) log q(fZ , σ 2)
p(fZ , σ2) dfZdσ
2
= KL(q(fZ , σ 2)||p(fZ , σ2)).
Now let’s compute likelihood part first.
Ef,fZ ,σ2∼q(f,fZ ,σ2)[log p(y|f, fZ , σ 2)] =
∫ q(f) log p(y|f)df
where q(f) = ∫ ∫ p(f |fZ , σ2)q(fZ |σ2)q(σ2)dfZdσ2. And this is
q(f) = MVT(f |2a,KXZK−1ZZµ, b
a (KXZK
−1 ZZΣK −1 ZZKZX +KXX −KXZK −1 ZZKZX)).
By slightly modifying the KL divergence between normal-gamma distributions in Soch & Allefeld (2016), we get
KL(q(fZ , σ 2)||p(fZ , σ2)) = ∫ ∫ q(fZ , σ 2) log q(fZ , σ 2)
p(fZ , σ2) dfZdσ
2
= 1
2
a b µ>K−1ZZµ+ 1 2 Tr(K−1ZZΣ) + 1 2 log |KZZ | |Σ| − nZ 2 + α log b β
− log Γ(a) Γ(α) + (a− α)ψ(a) + (β − b)a b
where ψ(·) is digamma function. Now if we use reparametrization trick at fZ , we can write fZ = Lu, where L is the lower triangular matrix from the cholesky decomposition of the matrix KZZ . Then
p(fZ |σ2) = N (fZ ; 0, σ2LL>) = N (fZ ; 0, σ2KZZ) q(fZ |σ2) = N (fZ ;Lµu, σ2LΣuL>).
In this case ELBO changes into∫ log(p(y|f))MVT(f ; 2a,ALµu, b
a (ALΣuL
>A> +B))df
−KL(N (u;µu, σ2Σu)Γ−1(σ2; a, b)||N (u; 0, σ2I)Γ−1(σ2, α, β)). With softmax likelihood, we can use MC approximation which is
ELBO ' 1 T N B T∑ t=1 B∑ i=1 log p(yi|f (t)i )
−KL(N (u;µu, σ2Σu)Γ−1(σ2; a, b)||N (u; 0, σ2I)Γ−1(σ2, α, β)) where B is minibatch and T is sample number of f which sampled from
MVT(f ; 2a,ABLµu, b
a (ABLΣuL
>A>B +BB)).
Here AB = KXBZK −1 ZZ and BB = KXBXB −KXBZK −1 ZZKZXB . Finally we can calculate predictive distribution with this reparametrization trick.
p(f∗|y) = ∫ ∫ p(f∗|fZ , σ2)q(fZ |σ2)q(σ2)dfZdσ2
= ∫ ∫ N (f∗;K∗ZK−1ZZfZ , σ 2(K∗∗ −K∗ZK−1∗ZKZ∗))N (fZ ;Lµu, σ 2(LΣuL >))Γ−1(a, b)dfZdσ 2
=MVT(f∗; 2a,K∗ZK −1 ZZLµu,
b a (K∗ZK −1 ZZLΣuL >(K∗ZK −1 ZZ) > +K∗∗ −K∗ZK−1ZZK > ∗Z)).
Thus the final predictive distribution p(y∗|x∗) is
log p(y∗|x∗) = log M∏ j p(yj∗|xj∗)
= M∑ j log p(yj∗|xj∗)
= M∑ j log ∫ p(yj∗|f∗)p(f∗|y)df∗ ' M∑ j log 1 N N∑ i p(yj∗|f i∗)
where f i∗s are sampled from MVT(f∗; 2a,K∗ZK −1 ZZLµu, b a (K∗ZK −1 ZZLΣuL >(K∗ZK −1 ZZ) > + K∗∗ −K∗ZK−1ZZK>∗Z)).
E INFERENCE ALGORITHMS
E.1 INFERENCE FOR REGRESSION
In order to calculate the predictive posterior distribution for an intractable marginal distribution, we can use self-normalized importance sampling. Consider a scale mixture of NNGPs with a priorH on the variance σ2v . 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 ∈ RI under the condition on Dtr. Then, the expectation of h(y) is
E[h(y)] = ∫ h(y)p(y|Dtr)dy
= ∫ ∫ h(y)p(y|Dtr, σ2v)p(σ2v |Dtr)dσ2vdy
= ∫ ∫ h(y)p(y|Dtr, σ2v)
p(σ2v , Ytr|Xtr) p(Ytr|Xtr) dσ2vdy
= 1
Z
∫ ∫ h(y)p(y|Dtr, σ2v)
p(σ2v , Ytr|Xtr) p(σ2v) p(σ2v)dσ 2 vdy
where Z = ∫ p(σ2v , Ytr|Xtr)dσ2v . We can approximate Z as follows:
Z =
∫ p(σ2v , Ytr|Xtr)
p(σ2v) p(σ2v)dσ 2 v
' 1 N N∑ i=1 p(βi, Ytr|Xtr) p(βi)
= 1
N N∑ i=1 p(Ytr|Xtr, βi)
where βis are sampled independently from H. Using this approximation, we can also approximate expectation of h(y) as follows:
E[h(y)] = 1
Z
∫ ∫ h(y)p(y|Dtr, σ2v)
p(σ2v , Ytr|Xtr) p(σ2v) p(σ2v)dσ 2 vdy
' N∑ i=1 wi∑N j=1 wj h(yi)
where the yi’s are sampled from the posterior of the Gaussian distribution and the wi’s are the corresponding importance weights for all i ∈ [N ]:
βi ∼ H, wi = N (Ytr; 0, βiKtr,tr), yi ∼ N (Kx,trK −1 tr,trYtr, βi(Kx,x −Kx,trK −1 tr,trKtr,x)).
TheK 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 yi’s. First, for importance weights, we observe that the log likelihood of βi has the form:
logN (Ytr; 0, βiKtr,tr) = − K 2 log(2π)− 1 2 log det(Ktr,tr)− K 2 log(βi)− 1 2βi Y >tr K −1 tr,trYtr,
and the three terms K2 log(2π), 1 2 log det(Ktr,tr), and Y > tr K −1 tr,trYtr 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 yi’s, we first draw N samples ȳis from N (0, Kx,x − Kx,trK −1 tr,trKtr,x), then multiply each of these samples with √ βi, and finally add Kx,trK −1 tr,trYtr to each of the results. Since Kx,trK −1 tr,trYtr + √ βiȳi ∼ N (Kx,trK −1 tr,trYtr, βi(Kx,x − Kx,trK −1 tr,trKtr,x)) for all i ∈ [N ], we can use these final outcomes of these computations as the yi’s.
E.2 TIME COMPLEXITY ANALYSIS
For time complexity, as we mentioned in Appendix E, our posterior-predictive algorithm based on importance sampling does not induce significant overhead thanks to the reuse of the shared terms for calculation. More specifically, our algorithm with K sample variances spends O(K + N3) time, instead of O(KN3), for computing a posterior predictive, where N is the number of training points. Compare this with the usual time complexity O(N3) of the standard algorithm for Gaussian processes. When it comes to SVGP and SVTP, one update step of both SVGP and SVTP takes O(BM2 + M3) time, where B is the number of the batch size of the input dataset and M is the number of the inducing points.

heading: F CLASSIFICATION WITH GAUSSIAN LIKELIHOOD
text: We compare the NNGPs and the scale mixture of NNGPs with Inverse Gamma prior and Burr Type XII prior for the classification tasks. Following the convention (Lee et al., 2018; Novak et al., 2018; Garriga-Alonso et al., 2019), we treat the classification 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. (2020) and we choose both c and k from [0.5, 1., 2., 3., 4.]. We do not particularly expect the scale-mixtures of NNGPs to outperform NNGP in terms of predictive accuracy, but we expect them to excel in terms of calibration due to its flexibility 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 NNGPs outperform NNGP in terms of NLL.

heading: G EXPERIMENTAL DETAILS
text: 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.8/0.1/0.1. 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 NNGPs. We choose the best hyperparameters based on validation NLL values and measured NLL values on the test set with permuted train sets.
For the classification experiments, we divide each dataset into train/test sets as provided by TensorFlow Datasets2, and further divide train sets into train/validation sets with the ratio 0.9/0.1. 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
2https://www.tensorflow.org/datasets
of the model, we used the corrupted variants of the dataset and made three new dataset from the base datasets. For the corrupted variants, we trained the model on the original version of the datasets, and tested on the provided corrupted verions. For the out-of-distribution variants, 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 classfication 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 classification by finite model ensemble experiments, each CNN layer has 128 channels.

heading: H ADDITIONAL EXPERIMENTS
text: The experiment code is available at GitHub3. We used a server with Intel Xeon Silver 4214R CPU and 128GB of RAM to evaluate the classification 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. 3. 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. (2020) and we use same settings as described in Section 4.2.
3https://github.com/Anonymous0109/Scale-Mixtures-of-NNGP
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.

heading: I ADDITIONAL INFORMATION
text: 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 identifiable information or offensive content.

Review summaries:
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.

The paper extends upon the corpus of works that consider the infinite-width limit of a deep net. The foundation of the work is the Neural Network Gaussian Process (NNGP) model, which is taken as the infinite-width limit of a deep net under certain (mild) assumptions; that is the so-called Master Theorem.   The here proposed approach treats the variances, σ^2, of the penultimate layer weights as random variables. Then, it imposes an appropriate prior distribution on these and proceeds with a reiteration of the Master Theorem (and how this varies in this setting).   They show that, under this setting, the infinite-width limit of the network is a scale-mixture of NNGPs. In the special case of an imposed Gamma-prior, the so-obtained scale mixture reduces to a Student's-t process; this is a well-known result.  

Building on the recent results on the correspondence between infinite-width neural networks and Gaussian processes (NNGPs), this paper proposes and studies a simple extension, looking at a scale mixture of such processes. The key idea is to introduce a scale prior on the parameters of the last layer, allowing construction of a richer classes of stochastic processes, particularly those with heavy tails. Some convergence results for these general processes are obtained by applying the tensor program and the Master Theorem, under various settings. Empirical results are also provided, showing the promise of the approach and robustness to out-of-distribution data.

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.