Scale Mixtures of Neural Network Gaussian Processes
Authors: Hyungi Lee
Eunggu Yun
Hongseok Yang
Juho Lee
Venue: ICLR 2022 Poster (2022)
Keywords: Neural Network Gaussian Processes
Infinitely-wide Neural Networks
Scale Mixtures of Gaussians
Heavy-tailed Stochastic Processes

Abstract: Recent works have revealed that infinitely-wide feed-forward or recurrent neural networks of any architecture correspond to Gaussian processes referred to as NNGP. 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 NNGP 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 NNGP. We present a practical posterior-inference algorithm for the scale mixture of NNGP 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.

Published as a conference paper at ICLR 2022
SCALE MIXTURES OF NEURAL NETWORK GAUSSIAN PROCESSES
Hyungi Lee', Eunggu Yun', Hongseok Yang!.23, Juho Lee'
Kim Jaechul Graduate School of AI, KAIST, South Korea
School of Computing, KAIST, South Korea
Discrete Mathematics Group, Institute for Basic Science (IBS), Daejeon, South Korea AITRICS, Seoul, South Korea
(1hk2708, eunggu, yun, hongseok, yang, juholee) @ ac.kr
ABSTRACT
Recent works have revealed that infinitely-wide feed-forward or recurrent neu- ral 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, how- ever, 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 introduce a prior distribution 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 richer class of stochas- tic processes. With certain scale priors, we obtain heavy-tailed stochastic pro- cesses, and in the case of inverse gamma priors, we recover Student's 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 practical posterior-inference algorithm for the scale mixture NNGPs and empirically demonstrate its usefulness on regression and classifi- cation tasks. In particular, show that in both tasks, the heavy-tailed stochastic processes obtained from our framework are robust to -iitt--disttibttionn data.
INTRODUCTION
There has been growing interest in the literature 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 GP. This result was extended to fully-connected deep neural networks of any depth (Lee et al., 2018; Matthews et al., 2018), which shown converge GPs the width grows. Similar results for deep Convolutional Neural Networks (CNNS) (Novak et al., 2018; Garriga-Alonso et al., 2019) and attention networks (Hron al., 2020). 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 infinitely- wide BNN boils down to the posterior inference of the corresponding GP called NNGP for which closed-form posterior can be computed exactly.
later obtained
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 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 al., 2014).
Our construction is inspired by popular class of distributions called scale mixtures of Gaus- sians (Andrews & Mallows, 1974); such distribution is obtained by putting a prior on the seale variance parameter of Gaussian distribution. We extend this scale mixing to NNGPs, where introduce prior distribution on the scale of the parameters for the last layer (which is often re- ferred to readout layer) in wide neural network. We show that simply introducing scale prior

---

Published as a conference paper at ICLR 2022
on the last layer can turn infinitely-wide BNNs of any architecture into 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 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 at all in some cases. When we mix the scale parameter with inverse gamma (so that the mixture becomes a l 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-inferene 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 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 scale mixture of NNGPS with some scaling distribution. Interestingly, the scaling distribution is a prior, not posterior. For the fully-connected wide neural networks, we extend the analysis 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 /. process, which can be computed analytically.
We empirically show the usefulness of our construction on various real-world regression and classi- fication 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 mance 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.
corrupted data while maintaining similar perfor-
1.1 RELATED WORKS
NNGP and NTK Our construction heavily depends on the tensor-progran framework (Yang, 2019) which showed that a wide BNN of any feed-forward or recurrent architecture converges in distribu- tion 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 the seminal work on NTK (Jacot et al., 2018), Lee et al. (2019) analyzed the dynamics of fully fully-connected neural net- works 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- collecconeeeted neu- ral network is initialized from specific prior (often referred as the NTK parameterization) and trained only for the last layer under gradient descent and squared loss, its fully-trained version be- comes 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 rep- resentative case is the work by Favaro et al. (2020), which showed that under an alternative prior specification, a wide fully- fully-conneeted neural network converges to stable processes, as the widths in- crease. 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- parameters of network and makes the scale random variable, while
2

---

Published as a conference paper at ICLR 2022
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 casy-to-compute posteriors and the correspondence to gradient descent training, continue to hold in our approach as a version for mixture distributions, while it is at least difficult to find similar adjustments of these results in those works. Second, our approach is applicable to arbitrary architectures, while those works considered only fully-connected networks and CNNS.
2
PRELIMINARIES
We start with a quick review on a few key concepts used in our results. For M E N, let [M] be the set M.....,M.. Also, write R2 for the set {ir E IR x > 0}.
2.1 TENSOR PROGRAMS AND NEURAL NETWORK GAUSSIAN PROCESSES
Tensor programs (Yang. 2019) are a particular kind of straight-line programs that express compu- tations 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 of the tensor program, which determines the dimensions of vectors and ma- trices used in the program. The parameter corresponds to the width of the hidden layers of a neural network, 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-wide 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 {f JunNN that are indexed by their widths n and have the following form:
[n(=:Va, Vn))
): IR! + TR,
fn(x; Va, V,
Σ Vn, (gg va))a aE/n)
where Vn E RR and , E |RTX (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 9n (x; , is the output of 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 tensor program, the components of the networks have to satisfy at least the following two conditions. First, the entries of vn and V, are initialized with samples drawn independently from (possibly different) zero-mean Gaussian distributions. For the entries of v,. the distributions are the same and have the variance of for some a, > O. For the entries of I,, the distributions may be different, but if we restrict them to those in each column j of I,,, they become the same and have the variance oj/n for some σj 0. Second, d(z) is controlled, i.e., it is bounded by a function exp(CC2|[2 + c) for some C, e, >> 0. This property ensures that ((2) for any Gaussian variable has finite expectation.
Let XI...., XM E R be AM inputs. When the computations of the BNNS on these inputs are represented by single tensor program, the master theorem holds. It says that there is a general method for computing /a E IR and Σ E RM M inductively on the syntax of the program such that I and Σ characterize the limit of (gn(xi;
..... go (xM; V Nn))n in the following sense. Theorem 2.1 (Master Theorem). Let h RM + I. be controlled function. functionn Then, as n tends to ox,
h [g,(xx; V Wnya.... In (xM: L Wn)aa
a.e. EZ-N(()) [M(2)).
(1)
3

---

Published as a conference paper at ICLR 2022
where - refers almost convergence.
The next corollary describes an important consequence of the theorem.
Corollary 2.2. As tends to bo, the joint distribution of (Jn (xmi Vny ) converges weakly to the multivariate Gaussian distribution with zero mean and the following covariance matrix K. R MM,
K(()) = o RZ~N(a 10 (Z))))Z)))
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 X.M,, imply that the BNNs converge to GP whose mean is the constant zero function and whose kernel
<(x,'' x') o. (Z-N((H>))) (o( (Z))))))
(3)
where (n, Σ) is constructed as described above for the case that M 2 and (x1, = (x,x). This GP called Neural Network Gaussian Process (NNGP) in the literature.
2.2 STUDENT'S S PROCESSES
Student's & processes feature prominently in our results to be presented shortly, because they marginals of the scale mixtures of NNGPS when the mixing is done by inverse gamma. These pro- cesses are defined terms of multivariate Student's distributions defined follows.
Definition 2.1. A distribution on I is multivariate (Student's 's) t distribution with degree of free- dom V E R4, location E Rd and positive-defintte scale matrix E E xdd it has the following density:
4
G((v + d)/2) p(x) (Ux) / G(v/2) [2|
(1+ (x+)
m1)
where G(-) is the gamma function. To express a random variable drawn from this distribution, we write x MVTa(P, ))
When v > 1, the mean of the distribution MVTa(K, 2) exists, and it is When /' > 2, the covariance of distribution also exists, and it -ーكX. As for the Gaussian distribution, the multi- variate distribution is also closed under marginalization and conditioning. More importantly, it can be obatined as a scale mixture of Gaussians; if o² ~ Inv InyGamma(a, b) and x|a² ~ (Z), then the marginal distribution of x is MVT(2a, Σ). For the definition of the inverse gamma distribution, see Definition B.1.
Student's process is a real-valued random function such that the outputs of the function on a finite number of inputs distributed by a multivariate distribution (Shah al., 2014). Consider X' (for inputs) with a symmetric positive-definit function A: X' x X R.
Definition 2.2. A random function f X + R Student's 's process with degree of freedom E R2, location function M X + R, and scale function if for all d and inputs .1...., E X, the random vector (f(z))..... (xx2) has the multivariate t distribution MVT a μ,, E), where (M M (xx)) and Σ is the d x matrix defined by E(i)) k(x,,
RESULTS
We will present results for the regression setting. Throughout this section, consider only those BNNs whose computations over fixed inputs are representable by tensor programs. Assume that we are given training Drr {(x,, We) E X k < K) and test Dic
((XK+& vK++) < E
Our idea is treat the variance of for the parameters of the readout layer in BNN random variable, not a deterministic value, so that the BNN represents mixture of random functions (where the randomness comes from that of of and the random initialization of the network parameters). To
4

---

Published as a conference paper at ICLR 2022
state this more precisely, consider the computations of the BNNS (fn(-: Wn)hne 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
representable by tensor program. Our idea is to change the distribution for the
initial value of v from a Gaussian to a scale mixture of Gaussians, so that at initialization, the BNNS on XL become the following random variables: for each N, o H,
Vn ~ (N(0,, for a E [n],
V (n,(a.)) ~ *N(0, o3/n) for E [m]. j e [P],
(x, Vn. vn) Vn aE(N)
Σ Vn,a (xgn(x;; vn))a for i E [K + L).
where H is distribution on positive real numbers, such as the inverse-gamnn distribution, P is the number of columns of n, and of and In are, respectively, the variance specific to the j-th column of L, and the output of the penultimate linear layer, explained in review tensor programs.
Our first result says that as the width of the BNN grows to oc, the BNN converges in distribution to a mixture of NNGPS, and, in particular, to Student's t process if the distribution H is inverse gamma. Theorem 3.1 (Convergence). As n tends to oc, the random variable (fn(x,; Vn. (n)),e(K+1) con- verges in distribution to the following random variable (xx))))))
0 ~ H,
(=((xi))e ~ N KC),
where ((.,)) - EZ-N( ( (Z)))((z) for i.j E [K + L}. Furthermore, if H. is the inverse- gamma distribution with shape and scale b, the marginal distribution of (z is the following multivariate distribution:
(xi))x)))((Ka)) MVTK+L(24) O, (b/a). -K),
where MVTa(U/ ,,',K') denotes the d-dimensional Student's 's distribution with degree of freedom J, location pl, 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 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 of 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 of to a scale mixture GPs or to Student's 's t process. Note that o? T in the theorem is precisely the covariance matrix of the NNGP kernel X1:K-1 in Corollary 2.2. We thus call the limiting stochastic process scale mixture of NNGPs.
Our next results characterize fully- srrined BNN the infinite- limit under two different train- ing 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 x that we have been using far.
Theorem 3.2 (Convergence under Readout-Layer Training). Assume that we train only the readout layers of the BNNs using the training set Drr under mean squared loss and infinitesimal step size. Formally, this means the network parameters are evolved under the following differential equations:
dv (1)
K
20 (2))
d
(v) (0)) (Vns In);
)
Kvn
dt
Then, the time and width limit of the random variables (f (xK+5 vn (1) (f% (xx+;)))ie()) is distributed by the following mixture of Gaussians:
DiELL). denoted (t)
H. (x(K+))))) ~ (RRRRE of (Krese R -n
where Y, consists of the y values in Dn (i.e., Y, (V.....,9K) and the K's with different sub- scripts are the restrictions of K in Theorem 3.1 with training or test inputs as directed by those sub- seripts. Furthermore, H is the inverse- garrma distribution with shape and scale b, the marginal distribution of (xx+x)) is the following multivariate t distribution:
( (xx++)))) ~ MVT, (2a, (Kre b (Kiese kre
5

---

Published as a conference paper at ICLR 2022
One important point about this theorem is that the distribution of (xX++)) is not the pos- terior of (J>a (xx) in Theorem 3.1 under the condition on Yu. The former uses the prior on og for mixing the distribution of (Joo (xK+1))e[21) (o Yu), while the latter uses the posterior on og 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 al. (2019) about the training of infinitely-wide neural networks. Concretely, we use the so called NTK parametrization for L ,.. This means that the BNNS on X1:K+L L are now intialized be the following random variables: for each N.
ohe ~ H. fu(x; Vni
on ~ -V(0,
L Vn,(as) ~ ~N(0, a E [n], j E [P],
Σ aE(n]
Vn, 0(gn(x M vn)), for [K + L}.
vn
When we adopt this NTK parameterization and the BNNS are multi-layer perceptrons (more gener- ally, they share architecture for which the limit theory of neural tangent kernlls has been devel- oped), we can analyze their distributions after the training of all parameters. Our analysis uses the following famous result for such BNNS: tends infinity, for all i, [K+ L.,
1 (Viv..Wa).fa Vn. Vn L nn), V(r.. (......).ee. (n).. (xj; Vn Vn P,,)) a.*. (
1
(4)
for some (H.j) determined by the architecture of the BNNS. Let θ be the matrix (O JEEKK))))))) The next theorem describes the outcome of our analysis for the fully-trained trainedid infinite- wiith BNN
Theorem 3.3 (Convergence under General Training). Assume that the BNNS are multi-layer percep- trons and we train all of their parameters using the training set Dur under mean squared loss and infinitesimal step size. Formally, this means the network parameters are evolved under the following differential equations:
Vay
dvn
dt
(t)
KVn
K
2
Σ
K
(xx:
fn
1 J (1!"))
Vn
(ec)))((n((() (x;:
p(gn
Vn
1
(())),
i=1
(1)))
L
d
(t)
2
U/ fn
(1) Vn
1 T (
dt
K
Vn
Let ② be θ with o in it set to (so that θ = O. Then, the time and width limit of the random
variables (fn (xx+ii (()) 1 )<e(r) has the following distribution:
vm
d ~ H.
(F2 ~N(((0700)
where
n1 Oml
O' K + (r Ko O 10 + O
Here the K's with subseripts are the ones Theorem 3.1 and the O's with subscripts are similar restrictions of θ with training test inputs. Furthermore, if H is the inverse gamma with shape a and scale b, the exact marginal distribution of (JS (x is the following multivariate t distribution:
(XX+))) MVTr, n¹, (b/a)
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 o of the readout layer's parameters, the networks in our setup fit into the standard setting for NNGPS and NTKs, that they satisfy the existing results in the setting. This conditioning trick means that most of the results in the standard setting be carried over to setup in scale-mixture forn, as explained in this section.
Second, if the prior on of is inverse gamma, the marginalization of of in those results can be calculated analytically and have form of Student's distribution or process.

---

Published as a conference paper at ICLR 2022
3.1 POSTERIOR INFERENCE
Gaussian likelihood As for the NNGPs, when the scale is mixed by the inverse gania prior, can exactly compute the posterior predictives of the scale mixture NNGPs for Gaussian likeli- hood since it corresponds the Student's & process having closed-form characterizations for the posteriors. For generic prior settings for the scales, we no longer have such closed-form formu- las for predictive posteriors, and have 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 prior H. on the variance o. Assume that we want to estimate the expectation of h(y) for some h IR R where the random variable is drawn from the predictive posterior of the mixture at some input E RR' under the condition on Du. The importance sampling with prior proposal computes the estimator (E w,h(y:))/ Σ₁= Wy where Bg and Y/ for ....., N are independent samples from the prior H and the posterior of the Gaussian distribution, respectively, and the wy are the importance weights of the 3,'s:
3 H,
w; (Yrr Bir,,)) Ye ~ N( B. Eu
The K is the covariance matrix computed as in Theorem 3.1 except that x is now a test input.
naive implementation of this importance sampler is slow unnecessarily due to the inefficient cal- culation of the likelihoods (Yx: 0, (.Ku..) of the sampled B.'s. To see this, note that the log likelihood of B; can be decomposed as follows:
K 2
K 1 log det (Ku.w)
log( 1 Y Ku Yu. 23g
log N
log(2x)
2
2
The terms 2 so need not be computed everytime we draw Br. To avoid these duplicated calculations, we compute these shared terms heforehand, so that the log-likleihood for each B3 can be computed in O(1) time.
K log( 2 det(Kaa) and Y,, K Yu are shared across all the samples {3
Generic likelihoods For generic likelihoods other than Gaussian, even the vanilla NNGPS do not admit closed-form posterior predictives. In such cases, can employ the Sparse Variational Gaus- sian Process (SVGP) (Titsias, 2009) for approximate inference. We present similar algorithm for the inverse-gamma mixing case which leads to Student's process. See Appendices C and D for the detailed description.
4 EXPERIMENTS
We empirically evaluated the scale mixtures of NNGPS on various synthetic and real-world tasks. We tested the scale mixture of NNGPS with inverse-gammnn -arsma prior corresponding to Student's 's pro- cesses, 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.1 EMPIRICAL 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 (000, 8/n) and the biases were initialized with (0, 0.052/n) For the last layer, we sampled o² ~ Inv InvGammçaa, b) with varying (a. b) values and sampled weights from (0, To check Theorem 3.1, we initialize 1.000 models and computed the distribution of out- puts 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 a, values, and averaged the results to get function value zero. Similarly, we trained all the layers to check Theorem 3.3. As shown in Fig. 1, the theoretical limit mell-matched the empirically obtained distributions for all settings we tested. See Appendix G for the details and more results.
7

---

Published as a conference paper at ICLR 2022
y- tP
1-0
8:00
Proandie
-25
810
02
-10
50
Figure 1: (1st-2nd column) Posterior predictives obtained from NNGP (left) and scale mixture of NNGPs with inverse gamma prior (right). (3rd-5th column) correspondence between wide finite model vs theoretical limit.
Table 1:
NLL values on UCI dataset. (m, d) denotes number of data points and features, respec- tively. We take results from Adlam al. (2020) except our model.
PBP-MV
Ensembles
REF
NNGP
Ours
Boston Housing
(G06, 133)
2.54 + 0.08
2.40 + 0.04
2:41 + 0.25
2.63
0.08
2.65 = 0.13
2.72 = 0.06
Concrete Streagth
(3020, 85
3.04 + 0.03
193 0.02
3.06 + 0.18
3.52
0.11
3.19 = 0.05
3.13
0.04
Lificiency
(768.
0.01
1.21 0.001
1,38 1:
0.78 0.0.6
1.01
0.04
Kin8nm
-1.20
0.01
Propulsion
(11934,
-4.83 0.06
-4.45
0.01
0.01
0.04
Power Plant
(96888 4)
2.78 +: 0.01
2.50 0.01
2.79 4: 0.00
2.94 0.01
2.77 = 0.02
2.66
0.01
Wine Quality Red
(1588, 11)
0.97
1. 0.01
0.93 0.01
0.94 1 0.12
-0.78 0.07
4,9 0.06
-0,77
0.07
Yacht Hydrodymamics
(308.
1.64
0.002
1,25 0.001
1.811.2.1
0.09.1.0.08
1,07
0.27
0.17
0.25
4.2 REGRESSION WITH GAUSSIAN LIKELIHOODS
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 and measured the negative log- likelihood values on the test set. Table 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 fact that Student's 's t processes include GPs limiting cases, this result suggests that one 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.
4.3 CLASSIFICATION WITH GAUSSIAN LHKEL
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 brief summary, ours with heavy-tailed priors (inverse gamma, Burr Type XII) outperformed NNGP terms of uncertainty calibration for various datasets including corrupted datasets.
4.4 CLASSIFICATION WITH CATEGORICAL LIKELIHOOD
We compared the NNGP and the scale mixture of NNGPs with inverse-gamma priors the image classification task. Following the standard setting in image classification with deep neural networks, 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 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 approximates the posteriors whose limit corresponds to Student's process. We com- pared SVGP and SVTP on multiple image classification benchmarks, including MNIST, CIFAR 10, and SVHN. We also evaluated the predictive performance on Out-Of-Diisrrbuttion (OOD) data for which we intentionally removed three training classes to save OOD classes to be tested. We used four-layer CNN as a base model to compute NNGP kernels. See Appendix G for details. Table summarizes the results, which show that SVTP and SVGP perform similarly for in-distribution data,
hhttps://archis iss usi edd/mm datasets..ssssse
8

---

Published as a conference paper at ICLR 2022
Table 2: Classification accuracy and NLL SVGP and SVTP for image datasets and their variants. NLL values are multiplied by 10?.
SVOP
SVTP (Das)
SVGP
SVIP (Ous)
Dataet
NLL (X100))
Accurses (4)
NLL (X110"
Accuayy (5)
Tateest
[x
Accency (5)
Accency (5)
MNJST
5.00 0.32
07.73 0.03
5.90 2: 0.00
97.78 2:0.5
CIFARI
131.96 = -0.5
$4.00 2.0.2
182.33 2 0.24
54.10-2:00 0.115
+She
24.23
0:00
28.28 2.0.5
94.63 1:000
+She
143:11 2: 0.29
08:43
142.9
1.18
48:13
Impplac
56.33
0.05
07.91
0:558
0.11
0.108
160.7
0.75
-Spaie
Clave Hhrr
103.65
2255
62.63 0.65
07.19 2.25
63.52 1.77
213.50
25.47
4-1.6
0000
215.58 6 1.83
67.70 0.03
206.66 +: 1.85
67.07 0.08
SSan 5
105 44 0.51
41447
2: 0.227
166.1 = 0.25
41.47 2 0.11
w. 0000
341.09
1.75
41,33 0.11
533.70
3 1.53
01.19 0.17
V. OOO
XMNIST
205.41 2.40
53M3 =: 0.13
63.92
68.76 4: 0.06
0.09
257.86 d: 1.85
53. 85
0.30
60.45 0.05
$3.96 2:
v. 000
EMNIST
183.26 = 1.40
1/-19 1:21
B1.25 1 0.22
72.58 0.25
177 2 0.43
5492 1 0.91
72.35 0.1.5
$4.55 1. 0:32
Faslion MNISTT
94.30 = 0.10
87.84 0.15
34:15 # 0.15
$7.90 +: 0.00
SYHN 000
100.8- = 0.17
71.67 0:1.11
101.04 = 0.13
71.71 +-0.0
W. DOR
252.61
3 3.51
01.29 1 0.01
241.69 1. 2:45
02.21 1.006
379.76
9:24
0.19
40.43
Table 3: Classification accuracy and NLL of ensemble models for image datasets and their variants. We used 8 models of 4-layer CNN for our base ensemble model. NLL values are multiplied by 10a.
Gressies
Insene Garana Pios (Ous))
Gazoin
Insens Cartro Plie (Oun))
30")
1x10
NLL(X10")
Accerscy
Accussy
MNIST
0:01
0.01
-Sher
1.42 10114
98.21 0.44
1.42
0:17
96.21
0.46
OOD
78.64 2:2.44
65.64-1 0.07
7686 2 8.14
05.01 0.15
Imguloo
3.91 1 0.000
KL0001 1.35
6.37
0:78
62:90
1:51
W: Iattabenn
8.17 1 0:27
77.01 0.95
612 1 0.57
71.23
1:67
Sycte:
0.76 +- 0.02
97.65
= 0.22
0.74
0.02
91.33
0.00
W. Nocy Label
10.78 0.00
$4.43 =: 0.10
1867 2 0.05
84.90 0.12
Glae Wir
60.71 0:31
6a.50
0.20
90.22
w. OOD
Fesioon MNIST
2.3.2
0.11
2.57 0.00
01.52
K: OLI
02.17 2 0.50
65,47
0.00
58.37 1 0.50
$4.45
0.00
Nokey Label
7.70 3: 0.00
99.02
0.000 7,63
07.36
K. Intalance
5.54 0.05
$3,89
0.12
5,83 2: 0.0s
65.100
0.11
R. Noccy Label
0.50 0.00
87 00
0.07
9.29 2: 0.00
$7.27 0.10
CIRRRO
R,6 0.00
71,77 0.31
5.75
0.00
70-22
1.30
Star.
16,00 4 0:27
51:48 2 0.34
15.24
0.38
53.24
0.57
EMNIST OOD
0.76 0.00
91.
0.13
91.30
52.29 2: 1:11
53.51
1:24
29.20
1.40
38
1.36
5.57 2:
70.50 0.09
5.25 2.117
76-71
0.05
Sputier
12:85
0:221
50.00
12.36
0.37
39.36
W. Vorcy Tabel
3.23 1.23..00.00
56.64 0.17
318 0 0.01
66.25 0.20
Tog."
18.18 0:26
0.08
4.70
61.57 0.01
SUM
8862
0.50
OLID
105.17 1: 1.75
56,92 0.05
101.7 1: 1:00
56,11 0.05
W. Inbulaace
10.22 4: 0.12
58.73 1.35
19.05
0.18
21.34
0.57
Inbalance
14.20 2: 0.17
68.64 0.82
14.58 + 0.119
61,97 0.92
M. Nokiy Label
15.88
61 0.00
$4.0 = 0.34
15.00 8 0.00
64.51 = 0.55
W. Nocry Labe
11.22 2 0.11
$0.24 1. 0.20
12:35 1 0.07
79.84 0.221
but SVTP significantly outperforms SVGP for OOD data in terms of uncertainty calibration (measured by negative log- likeliiiiooo....
4.5 CLASSIFICATION BY FINITE MODEL ENSEMBLE
Although we proved Theorem 3.3 only for fully- fully-connected networks trained with squared losses, Lee et al. (2019) empirically observed that there still seem to be 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 ini- tialized 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, constructed ensem- bles of 8 models, and compared the performance on MNIST, MNIST-like variants, CIFAR10 10, 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.
5 CONCLUSION
In this paper. we proposed simple extension of NNGPS by introducing a scale prior on the last- layer 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 f processes. Based on the result in (Yang, 2019), we have shown that an infinitely-wide BNN of any architecture constructed in specific way corresponds to the scale mixture of NNGPs. Also, we have exlended 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 robust to the out-o--distribution data.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGEMENT AND DISCLOSURE OF FUNDING
We would like to thank Jachoon 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 (IIT)) grant funded by the Korea government (MSIT) (No.2019.0-00075, Artificial Intelligence Graduate School Program(KAIS))
No. 2021-0-02068, Artificial Intelligence Inno- vation Hub), National Research Foundation of Korea (NRF) funded by the Ministry of Education HY was supported by the Engineering Research Center Program through the National Research Foundation of Korea (NRF) funded by the Korean Government MSIT 1TE-20088RAAAA000000 and also by the Institute for Basic Science (IBS-R029-C1).
INRF202IRIFIA 1061655,
REFERENCES
Ben Adlam, Jaehoon Lee, Lechao Xiao, Jeffrey Pennington, and Jasper Snoek. Exploring the un- certainty properties of neural networks' implicit priors in the infinite-width limit, arXiv preprint arXiv:2010.07355, 2020.
D. F. Andrews and C. L. Mallows. Scale mixtures of normal distributions. Journal of the Royal Statistical Society: Series B (Methodologiial), 36, 1974.
D. Bracale, S. Favaro, S. Fortini, and S. Peluchetti. Infinite-channel deep stable convolutional neural networks. arXiv:2021.03739 2021.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Min Milne,, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http: //github. com/ coaom//oolle//eoo
S. Favaro, S. Fortini, and S. Peluchetti. Stable behaviour of infinitely wide deep neural networks. In Proceedings of The 23rd International Conference on Artificial Intelligence and Statistics (AIS- TATS 2020), 2020.
Adrià Garriga-AAonso, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional net- works as shallow Gaussian processes. In International Conference on Learning Representations (ICLR), 2019.
Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: NNGP and NTK for deep attention networks. Proceedings of The 37th International Conference on Machine Learning (ICML 2020), pp. 4376-4386. PMLR, 2020.
A. Jacot, E. Gabriel, and C. Hongler. Neural tangent kernel: convergence and generalization in neural networks. In Advances Neural Information Processing Systems 31 (NeurIPS 2018), 2018.
J. Lee, L. Xiao, S. S. Schoenholz, Y. Bahri, R. Novak. J. Sohl-Dickstein, and J. Pennington. Wide neural networks of any depth evolve linear models under gradient descent. In Advances in Neural Information Processing Systems 32 (NeurIPS 2019), 2019.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as Gaussian processes. In International Conference on Learning Representations (ICLR), 2018.
A. G. de G. Matthews, J. Hron, R. E. Turner, and Z. Ghahramani. Sample-then-optimize posterior sampling for Bayesian linear models. NeurIPS Workshop on Advances in Approximate Bayesian Inference, 2017.
Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard E Turner, and Zoubin Ghahra- mani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations (ICLR), 2018.
R. M. Neal. Priors for infinite networks. Bayesian Learning for Neural Networks, pp. 29-53.
Springer, 1996.
10

---

Published as a conference paper at ICLR 2022
Roman Novak, Lechao Xiao, Jachoon Lee, Yasaman Bahri, Greg Yang. Jiri Hron, Daniel A Abo- lafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are Gaussian processes. In International Conference on Learning Representations (ICLR), 2018.
Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference on Learning Representations (ICLR), 2020. URL https:// gi thub. wtt.ccom//oooooo e/ menuralllaaggnnee
A. Shah, A. Wilson, and Z. Ghahramani. Student-t processes as alternatives to Gaussian processes. In Proceedings of The 17th International Conference Artificial Intelligence and Statistics (AIS- TATS 2014), 2014.
Joram Soch and Carsten Allefeld. Kullback-leibler divergence for the normal-gamma distribution. arXiv preprini orXiv: 1.01437. 2016.
M. Titsias. Variational learning of inducing variables in sparse Gaussian processes. In Proceedings of The 12th International Conference on Artificial Intelligence and Statistics (AISTATS 2009),
2009.
Greg Yang. Tensor programs I: Wide feedforward or recurrent neural networks of Any architecture Gaussian processes. arXiv preprint arXiv: 12478, 2019.

---

Published as a conference paper at ICLR 2022
A DETAILS ON TENSOR PROGRAMS
In this section, we will review tensor programs from (Yang, 2019), which are sequences of assign- ment 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. G- and H-variables are vector-type variables, and A-variables are matrix-type variables. Each tensor program is parameterized by n E 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, respectively. Due to the random initialization of some of these variables, the contents of these variables are random in general. G-variables denote daussian-distriutter 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- Garaiabbl and all H- H-variables may not be assigned 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 = Wr where y is G-variable, W is an A-variable, and is a H-variable. The LinComb assignments compute linear combinations, and they have the form y
E- was' where y ..... x' are G-variables, and wi,..., Wi E R are constants. The final Nonlin assignments apply (usually nonlinear) scala functions on the values of G-variables coordinatewise. ,..,=,,,.......... 6 R' + R is possibly nonlinear function, and 6(x' ,...,"") x²) means the lifted application of ø to vector arguments.
That is, they have the form y d(x*, ..... x²), where y is H-variable, x² ..... x² are G-variables,
Every tensor program ends with a return statement of the form:
return(v" y² /VT..... of y VV)
(6)
where u',....,u ot are G-variables not appearing anywhere in the program, and y², y²,...,y variables.
of are H-
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 pw/n ' for some ow O (i.e., Wag (0 for j j [ri), and these samples are independent with those used for other A-variables
and impur G-variables. For every a E [n)) the components of all imput G-variables x......3 ... (which store vectors in R") are initialized with independent samples from N(r" Enn) for some mean 1d and (possibly singular) covariance E over all input G-variables. Note that the mean and the covariance do not depend on the component index a E [n]. Finally, the input G-variables e' 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 (O0
Definition A.1. Given 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 the vector is the number N of all G-variables. The computed and Σ specify the multivariate Gaussian distribution over the a-th components of G-varablles that arises at the infinite-widtt limit (Theo- rem A.2). The computation is defined inductively as follows. Let "t..g....... ."...""" be all the G-variables in the program, ordered as they appear in the program. For any G-variable g², j((g)) is defined by:
(0*)
if g is an input G-variable
p(g*)=
Em wap(x")
if g = Em a otherwise
Here x²1 sm are some G-variables. For any pair of G-variables g E.(g" g²) is defined by:
(g ,g')
if g gf are input G-variables
De w; Exx . Wgg ( )
if g² Em waz
Eig
if g = _i war
ow E (Z)
if g = W y. gf = Wy for the same W, with Z below otherwise
Here x!, ....... th are some G-variables, and H-variables y and y' in the second last case are intro- duced by the Nonlin assignments y = Ø(g"1 ...""""''''" grr) and y' =
g²r). respectively, for
12

---

Published as a conference paper at ICLR 2022
some i1, :1....:2p [max(kk - 1, 1)]. We have adjusted the inputs of and ø' ensure that they take the same set of G-variables. The random variable in the second last case is a random Gaussian vector drawn from MV'',, Σ') where μ' and E' are the mean vector and covariance matrix for "'',....''' by the previous step of this inductive construction.
Definition A.2. Given h R - IR and for any E R if [h(x)] is upper bounded by a function exp(C"|=<||>- + e) with C, O, then this function is controlled.
Theorem A.2 (Master Theorem). Consider tensor progrum satisfying Assumption A. and using only controlled o Let g¹ "'....""" be all of the G-variables that appear in this program, and h: T - I. be controlled function. Then, as n. tends to c0,
Σ higa......g. ........))))
) a
SE.NWu.Dill/2I....... kN)].
Here 9.21, refers almost sure convergence, and Z (ZI...., ΣN) RN Also, μ and E are from Definition A.
Corollary A.3. Consider a tensor program satisfying Assumption A.I and using only controlled be all of the G-variables that appear in this program. Assume that the program (1) for some H-variables y' 's such that each yi is updated for a controlled o' Then, as n approaches to oc, the joint distribution of this vector converges weakty to the multivariate Gaussian distribution (0, with the following covariance matrix:
é's. Let g'...., .....g" returns a vector (v Vt........ v'y' n²t
,,
by y' = "'..'.....g''
K(.)) o BZ=N for all i, [N).
where u and E are from Definition A.
B PROOFS
In this section, we provide the proofs of our three main results. We first describe lemma related to a basic property of Student's distribution:
Lemma B.1. The class of multivariate distributions is closed under marginalization and condi- fioning. That is, when
Xi
E RRdi+d2 ~ MVTdi+da
Σ1
Σ12
x =
X2
14,
U2
Σg1
Σ22
we have that
MVT,, E1, and x2/x1 MVTa, (u+ d1, pby
2222
(8)
+ di
where
_ = 122 + E21 E2 (x1- -1), c= c= (x1 -1) ETI (1)= E' = Egg - Σ21 E E12
Definition B.1. A distribution on R4, is a inverse gamma distribution with shape parameter a E R. and scale parameter E R4 if it has the following density;
6a
1
a+1
p(x)
e
b/z
G(a)
x
where G(-) is the gamma function. To express a random variable drawn from this distribution, we write & ~ Inv LnyGamma(a,b b).
Lemma B.2. Let Σ E Raxd be postive definite matrix. Consider μ E R and a. b E IR., If
o² InyGamma(a,, b),
x/o² ~ N(H,
then the marginal distribution of x is MVT (2a, n. 5.).
Now proceed to the proofs of main theorems.
13

---

Published as a conference paper at ICLR 2022
Theorem 3.1 (Convergence). As tends to the random variable (fn(x; Un))ie[K++)) con- verges in distribution to the following random variable (x=(()))
0g H,
(x)) ~NO,0, R),
where K(t.j) Ez 2) ())((2))((Z)) for i, E [K + LJ. Furthermore, if H is the inverse-
gamma distribution with shape and scale b, the marginal distribution of (x)(()))) is the following multivariate distribution:
(foc (x<c(x)))
MVT (2o, 0, (b/a) -K),
where MVTa(V/ (V. ,p', K') denotes the d-dimensional Student's 's distribution with degree of freedom v, location pl, and scale matrix K'.
Proof. Let [mn fn(x;Vn, Vn) for E [K + I.] and n. Pick an arbitrary bounded con- tinuous function
gK+L + R. For each n, define Gn to be the o-field generated by In (X1; . In (XK+Li n.
We claim that as n tends to oo,
E h (o [n]....,.. o [ml]] + Eag-H ((.0)) [h (Y ...., r..,,,++
The theorem follows immediately from this claim. The rest of the proof is about showing the claim. Note that
hh (o' K+L [nl)] = ( plT
h
vn -
g, g
@ (gn, (XK+L; Dn [=*( Conditioned on g and d the random variable Vn dl:kL in the nested expectation from above is distributed by multivariate Gaussian. The mean of this Gaussian distribution is:
where @ gl: = T6 (gn (x): 1
E
LL o
v G
1 g1:KLL
Vn
=0ERK+E
The covariance is:
K(n((s)) = E
vn (gn (x: (x::))))
1 ( (on (x): (x:: M.i))
g 9.
Vn
n
E Vo.o Vn,al x (ob (gn (x: v.))), (o (gn (xzi vn))
n
9=1
0
( (gn (x;; V (,))),, (o (gn (xy; V I.))
Using what we have observed so far, compute the limit in our claimed convergence:
lim E (o'(n, K [n]]] = lim E 20000
E B [ [1]]
n-
E lim E (. [1(12:*44]]]
The last equality follows from the dominated convergence theorem and the fact that bounded. To go further in our computation, observe that K[n]] converges almost surely to of by Theorem 2.1 and so
E B (L.K: 2x+t)
;o, 2|n)dduu
a.a. /
hlutkt ou
=BnnXte (x (y1:KL]]].
14

---

Published as a conference paper at ICLR 2022
Here use A 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:
lim N2L [h(u [[]]] -E - [ (yI +00+00
= RoghH [YLKI [ (YIK+h)]]
Theorem 3.2 (Convergence under Readout-Layer Training). Assume that we train only the readout layers of the BNNs using the training set DDr under mean squared loss and infinitesimal step size. Formally, this means the network parameters are evolved under the following differential equations:
dv
(2)
K
(x::
d = 0.
(v!") T (0)) (Vn,, V n)
V (n
dt
KVn
dt
Then, the time and width limit of the random variables (f (fn(xx+i (f% (xx+)))e(())- is distributed by the following mixture of Gaussians:
DEEL)) denoted vn (
~ (2(xx+)))t N (N(KKesSKT of (Kusso Ke Kun.KooaKKan))
where Y, consists of the y values in Dur (i.e., Yo (HI...., 9BK) and the K's with different sub- scripts are the restrictions of K. in Theorem 3.1 with training test inputs as directed by those sub- scripts. Furthermore, if is the inverse- peeerse-gammaa distribution with shape a and scale b, the marginal distribution of (xz(x))))) is the following multivariate t distribution:
(f=((x<)))E)) ~ MVT (2a, (Kre.tEiiili.). b
(Ku.. (e
Our proof of this theorem reuses the general structure of the proof of the corresponding theorem for NNGPS in (Lee al., 2019), but 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 og so that the existing results such as the master theorem be used in our case.
Proof. Note that y (2) = (0) for all because d' (t) /dt = Let
5. [n]
Vn ①
x;; (o)
Vn & (gn x ::0) i |K + LJ.
[n]
[01(m], [m],..., xx [n]]
E
TR. K
( [0k+1[m),.... 5K+L E RRax
Using this notation, we rewrite the differential equation for vi follows:
dv
2
2nin]
[n v/( Yun)
dt K
This ordinary differential equation has closed-fore formm-form solution, which we write below:
(0) = (0)) n + [n] [n] (.__n]__,[]]'''' exp (-
K t [n] -1)
( (Φ.,[n) v! Yu
We multiply both sides of the equation with the Prc[n]] [n] matrix, and get
[n((n] v(0- = [r((r) (()) + Mmpleen
[K(m)) xK)ninn)--)
( ( v(
exp
K
where Klnlie,r [n (e][n]20rl [n[]]] and K[n]ir = [m[[m] [r][ T. Then, we send to co, and obtain
Pie|n]] ((0) Luln ((()) K/rlan ( [n) v! ---y Yir)
- + (n( (0)) Imniit [1
( v!!0)
15

---

Published as a conference paper at ICLR 2022
But
(xx 10(0))))111) Pio (n] v(00),
Thus, can complete the proof of the theorem if show that for every bounded continuous function h. : IR. lk,
(Fr(r) v )
lim KI + Pie[n] v (e 1 10000
=E (xx+))))
Pick bounded continuous function h R + R. For each n, let In be the a- o-field generated by In (x1; in)..... In (x*+1; nn). Also, define K[nllma =
We show the sufficient
condition mentioned above:
([ [e[[n] [n] M ( v w
lim n-000
[M)
lim n-+00
E
E Klulaeaannlaatt + Pw(ni v Anll @ [m] v
In. o
lim
Eg Ez~N Ep og( -E ((ml ( [{(2)]
n-+00
E
lim 11-000
E Ez W Tp o(T (m.l [h [b(2)]
E
Yon of k [h())]]
E ((
The second equality uses the fact that 1 consists of iid samples from (O, of),, and it is conse- quence of straightforward calculation. The third equality uses the dominated convergence theorem and the fact that à 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 (JS (xx++))
o3 InvGamma(a, b), we get the following closed-form marginal distribution by Lemma B.2:
b (x((x))))ee[4) MVT, (2a, Ree - (Ku,,o Kuc.aFu
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 setup using notation:
Theorem B.3. Assume that the BNNs are multi-layer perceptrons, and we train these models us- ing the training set Du under the mean squared loss and infinitesimal step size. Regard o? as deterministic value. Then as
and I + (Jn(Xx+1 vio). 1 (()))) converges in distribution the following random variable:
Vn.
(xx4))))) N )
where
1 Ir,, O
X Kre,s + Ore O 00 ( + K,e,
Theorem 3.3 (Convergence under General Training). Assume that the BNNS are multi-layer percep- trons and train all of their parameters using the training set DDr under mean squared loss and
16

---

Published as a conference paper at ICLR 2022
infinitesimal step size. Formally, this means the network parameters are evolved under the following differential equations:
(0)
Va.
dv
dt
K
Kvn
Σ
K
-
In (x(x:V
/ vn w(a)
((
vn
1
i=1
(0 V n
dvp (t)
2
K
fn(x;v
Vn !
1
v (
dt
K
1/i
Vn
Let ④ be O with of in it set (so that @ dg Then, the time and width limit the random
variables (fn (XK+ii vn (t)
J (=E111 has the following distribution:
dg H.
(f2 ~N(H',000)
where
n¹ O
O' K + (Gre On (f + Ru
Here the K's with subscripts are the ones in Theorem 3.1 and the O's with subscripts are similar restrictions of θ with training or test inputs. Furthermore, if is the inverse gamma with shape and scale b, the exact marginal distribution of (f (xe+)))ee[n] is the following multivariate t distribution:
(xK+)))))EE MVT pL, (b/a) 07.
(5)
Proof. Assume that we condition on of Then, by Theorem B.3, as we send n to oo and then to oo, the conditioned random variable (Jn (xx(x w vn the following random variable:
(((1) )e<{)) af converges in distribution
(x((x)))) of ~ N(4e spe)
where
Mhe = 11 Yu,
E of (Kiete +
(Orenn - + Kr.
Thus, by Lemma B.4, if remove the conditioning convergence of the unconditioned (Xx4siv
of (i.e., marginalize over of), we get the I (t) vn )aa(r to the following random variable:
2 ~ H,
(x ~ ~i(y'o030)
where
n1 O O..V
O' = Ke + u.u Ku.
(Orc.n + Ki., Ord
In particular, when H. InvGamma(a, b). the exact marginal distribution of (f (xx+))) has the following form by Lemma B.2:
MVT 20, n², b 2
Lemma B.4. Let (Xn)neN be a sequence of p-dimensional random vectors. X be a p-dimensional random variable, and Y be random variable. If Xn, Y converges in distribution X Y almost surely (with respect to the randomness of Y as n tends to oc, then X,, converges in distribution to X.
17

---

Published as a conference paper at ICLR 2022
Proof. Let h R" - IR be bounded continuous function. We have to show that
lim [[,( [X)(
n--00
The following calculation shows this equality:
lim [[h(X..]] lim B[ Y] n2 22000
lim B()(X,,)) |Y] = =[ Y]] +00+00
n-000
= =]](
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 X,, Y converges in distribution to X|Y almost surely.
C STOCHASTIC VARIATIONAL GAUSSIAN PROCESS
Let (X, be training points and (Xx, yz) be test points. And let Z be inducing points. We want to make f() which well approximates p(f. fzly) by maximizing ELBO. Then we will construct q(f, fz) p(f(fzz((fz)) where p(flfz) 2(f Kxx KxzKzzKzn and ((fz) = (fzlpp, Σ). Then
logp(v) 2
KL(g((2)ll((2)).
Here we can calculate likelihood term as
Ap: AEAT D)df
where A = Kx2k 22. B Kxx - Kx2K K-1 22 Kzx. Also we can calculate KL term
det Kzz det. Σ
log(
11% Tr Tr[ + p K
We find optimal μu, Σ and inducing points Z by using stochastic gradient descent. And with these optimal values, calculate predictive distribution p(f.ly)-
/
p(f+ly)
p pls.,.,f2/w)aflzz
/ / PlfAls.falalff2)aafa
/ pFAAE.22ao(1122222222244 I
This equation can be write as p(f-ly = W(fu:K.Zk 724, K. K / + K K, Kr Now if we reparametrization trick at Sz. we write fz Lu where is the lower triangular matrix from the cholesky decomposition of the matrix K zzz Then
p(fz) = (ff;; LL N[f2;; K2z)
(fz)) = NNzz; )
In this case ELBO changes into
I log(p(ulf)) ALI,, + B)df E.,
With softmax likelihood, we can use MC approximation which
N log logp(uil.!"") KL(N(u; (2,
ELBO~ ~ T 1=1
18

---

Published as a conference paper at ICLR 2022
where B is minibatch and T is sample number of f which sampled from
NU: Ag AgLE, L AB Bg).
Here Ag Kx
22 22 and B3 x Kxazk 2x Zzx Finally we can calculate predic- tive distribution with this reparametrization trick.
p(f-ly) J D(F-)f2)(((2)11z
= | NU.KK2KZzd2. K.. K,
LE jdfz
22 K. 2Z) T + K.- K.zKz2K.
Thus the final predictive distribution xx.) is
logp(y (x.) log II p(v!l=!)
log lgg((vl
Σ log / p(s
M
M
N
Σ
log
Σ p(v111t)
where fis are sampled from N(f.: K-zk Litues K, K 2K (KazKzz)) + K
K.zK 22K22,,),
D STOCHASTIC VARIATIONAL STUDENT PROCESS
We want to make o(f,Iz,o") 0²) which well approximates f(f, obly) by maximizing ELBO. First where p(flʃz,0²) = oz Kx (ixK2) p(f2(=²)= =N 1f.2200,o,220 ozKzz) and p(o²) = I B). And q(f. =
plJ,Iz,0²) = o')l
plʃlíz, (g((fl)( (o)) (o²²) where {(2200²) = N(Iz(1,,²E)) and q(0²) = (- a, b). Then
ELBO can be computed
logp(y) 2 1f http::/pw....aaa.k-k.k...aaa1aa.
Here KL divergence can be reformulated as
/ o( o(f2,03) p(fz,00) df2do²
KT(g(S2, (ll((t,,))))
Now let's compute likelihood part first.
Bst20r--aff22eeppogblll......... 1 (())ogrlll))
where o And this
(()) = MVT(/12,, KWTT//2,,Kx2K311/1 ATT. + Kxx KxzKz
19

---

Published as a conference paper at ICLR 2022
By slightly modifying the KL divergence between normal- -amaaaa distributions in Soch & Allefeld (2016), we get
/ glf2, o log ((2,,0²)
KL(a(fz, o1)p((z,
plfz, o²) df2do²
b 14
a T K
K
+ Tr
Tr(K 27 2) +
1
log
K
RZ
log
b 3
2
2
[
2
a()) I(a)
log (a- -
+
-
where (-) is digamma function. Now if use reparametrizatioon trick at fz, can write fz Lu, where L is the lower triangular matrix from the cholesky decomposition of the matrix Kzz Then
= -
q(fz Ni Lpa. Lpa L.LELL)
In this case ELBO changes into
log(o(wl/))MVV((j..a) ALp (ALE,, A B))df KL(N(u; Mu, o² 2.)) (0 3))
With softmax likelihood, we can use MC approximation which
1 Σ Σ log p(uilf))
B
ELBO
TB
a, where is minibatch and T is sample number of f which sampled from
KL(N(u; pa, oFE
a,)
MVT(J; 20, A - AA T. AB + -BB)).
Here AB KX z K Z2 and BB = Kx8X8 Kx K k 2xg. Finally we can calculate predic- tive distribution with this reparametrization trick.
p(f=lu)= / o'
Idf2do2
=f N..:KK 02
(( (1 b)dfzdo²
-VVT(-;; 2a, K. K K b -K.z LEu (K.z (K.
K.. - K.zk 27
Thus the final predictive distribution p(y, xx,) is
logp(y, (x.) log II p(y!
Σ log p(v!
Σ log / p(
M
M
N
Σ log Σ p(v(151)
where fts are sampled from MVT(f.;20 K. z K Li (K.2 LEN (K.z (K Tz +
K.. K.z KKIKziK] kTz))
20

---

Published as a conference paper at ICLR 2022
E INFERENCE ALGORITHMS
E. INFERENCE FOR REGRESSION
In order to calculate the predictive posterior distribution for an intractable marginal distribution, we can use self-normalized importance sampling. Consider scale mixture of NNGPs with prior H the variance o. Assume that we would like to estimate the expectation of lv(y) for some function h : I R where the random variable y is drawn from the predictive posterior of the mixture some input E R' under the condition on Du- Then, the expectation of h(y) is
B(h(y)l = / h(y)p(ul Du)dy
/ h(y)pullee ol
/ J Mu)p(y plo. Yirl (ir) do
p (u))
2 J MM)ppyyDDm ot)
o Yu Xu))
plo plo ((o?)
where 2 = f(o Xn)
We can approximate Z as follows:
p(o
Yirl plo
Z
p(o2))
N
Yul,)
N
p(B.)
N
PYYll Bi)
N
where Bis are sampled independently from H. Using this approximation, can also approximate expectation of (y) as follows:
//1 h(w)p(ulD,, of) plo Yu Xw)
//
E(h(y)] 7
plo?) plo
h(3)) w
where the y 's are sampled from the posterior of the Gaussian distribution and the w. 's are corresponding importance weights for all i [N]:
3, H. w; (Yrr; 0, B.K.)))) y NTKx Bi(Kx,x Kx
The KC 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 's. First, for importance weights, we observe that the log likelihood of (3 has the form:
log N 0,
K log(2xr) log det (c.,.)
K log (3))
1 23,
Y, L
2 log(2 log det and Y.T TrTKi here are independent with B,. Thus, using this observation, we compute these three terms beforehand only once, and calculate the log likelihood of each sample 33 in O(1) time. Next, for sampling the y. 's, we first draw N samples jas from (O, Æxx - Inxx then multiply each of these samples with vBI, and finally add Khy to each of the results. Since Ex /Bar (Kx B, (r for all i [N],
and the three terms K
these final outcomes of these computations as the y, 's.
21

---

Published as a conference paper at ICLR 2022
Table 4: Experimental results of Classification with Gaussian Likelihood
Dataset
Accuracy
NNGP
Inverse Gamma
Burr
Type XII
MNIST
96.8 +: 0.2
9.29 +: 0.002
4,89 + 0.001
4.96 +: 0.001
Shot Noise
94.9 1: 0.4
9.32 1: 0.001
4.88 0.001
4.96 0.001
Impulse Noise
84.4 2.3
9.63 1: 0.000
5.17 0.001
5.25
0.001
Spatter
95.4
9.27 + 0.001
4,44 0.002
0.002
4.49 + 0.007
Glass Blur
90.5 0.7
9.12 + 0.001
4.20 0.005
4,00 + 0.017
Zigzag
9.51 1 0.001
4.62 1: 0.006
4,49 + 0.021
EMNIST
0.9
9.40 0.001
4.15 + 0.001
4.15 + 0.013
Fashion MNIST
61.3
9.34 + 0.002
4,96 0.002
4.97 0.002
KMNIST
81.1 1 1.3
9.48 1 0.001
4.60 0.002
4.34 1 0.013
SVHN
42.5 1.5
6.01 + 0.062
4.16:
4.15 4: 0.000
E.2 TIME COMPLEXITY ANALYSIS
For time complexity, as we mentioned in Appendix E. our posterior-predictee 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 + Nº) time, instead of N*), for computing a posterior predictive, where N is the number of training points. Compare this with the usual time complexity ((N)) 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² + S time, where B is the number of the batch size of the input dataset and M the number of the inducing points.
F CLASSIFICATION WITH GAUSSIAN LIKELIHOOD
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 al. (2020) and we choose both e and k from [0.5, 1., 2., 4.]. We do not particularly expect the scale-mixtures of NNGPs to outperform NNGP 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 terms of NL..
of NNGPS outperform NNGP in
G EXPERIMENTAL DETAILS
In Fig. 1, we used the inverse gamma prior with hyperparameter setting 2 and 9 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/tess sets with the ratio 0.8/0.1/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 NNGPS. We choose the best hyperparameters based on validation NLL values and measured NL.L values on the test set with permuted train sets.
For the classification experiments, divide each dataset into train/test sets provided by Ten- sorFlow Datasets², and further divide train sets into train/validation sets with the ratio 0.9/0.1. We choose the best hyperparameters based the validation NLLL values, and measure NLL and aceu- racy values on the test set with different initialization seeds. To measure the uncertainty calibration
Phttps://wer http::/wwwweensooooooooooog/////
22

---

Published as a conference paper at ICLR 2022
Table 5: RMSE values on UCI dataset. (m, denotes number of data points and features, respec- tively. We take results from Adlam et al. (2020) except our model.
Daraset
PBP-MV
Dropout
Ensembles
RBF
NNGP
Ours
Boston Housing
(506, 13)
3.11 +: 0.15
2990 +: 0.18
3.28
1.00
3.24 0.21
3.07
0.24
3.30 :: 0.03
Strength
8)
0.16
6.03
0.58
5.63 0.24
Energy Efficiency
(768, 8)
0.45
+ 0.01
0.54
-: 0.06
2.09
+ 0.29
0.50 2: 0.01
0.57 1: 0.02
0.44 :: 0.03
Kin8tm
(8192.8)
0.07 0.00
0.08 0.000
0.09
0.00
0.07 0.00
0,07 0.00
0:00
Naval Propulsion
(11934, 16)
0.00 1: 0.00
0.00 1: 0.00
0.00
0.00
0.00 0.00
0.00 1: 0.00
0.00 1: 0.00
Power Plant
(9568. 4)
3.91 + 0.04
4.01 + 0.04
4.11 0.17
3.82 0.04
3.61
+ 0.04
3,53 : 0.04
Wine Quality
(1588,
0.01
0.04
0.64
Yacht Hydrodymamics
(308, 6)
0.61 +: 0.06
0.67 + 0.05
1.58 0.48
0.60 0.07
0.41 0.04
0.35 :: 0.04
Table 6: Additional regression results for Burr Type XII prior distribution.
(y,, d)
NNGP
Gamma prior
Burr Type
prior
Boston Housing
(506. 13)
2.65 +: 0.13
2.72 = 0.05
2.77
0.02
Concrete Strength
(1030. 8)
3.19 1 0.05
3.13 0.04
3.29
0.09
Energy Efficiency
(768, 8)
1.01 + 0.04
0.67 = 0.04
0.70
0.04
Kin8am
(8192. 8)
-1.15
0.01
-1.18
-1.23
Naval Propulsion
(11934, 16)
-10.01
+ 0.01
-8.04 = 0.04
-4.38 0.05
Plant
(9568, 4)
2.77 0.02
0.01
Wine Quality Red
(1 5888, 11)
-0.98 + 0.06
-0.77 = 0.07
-0.16 0.05
Yacht Hydrodynamics
(308, 6)
1.07 = 0.27
= 0.25
0.00
0.15
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 the provided corrupted verions. For the out-of-distribution variants, removed 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 compute the kernel. Due 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.
H ADDITIONAL EXPERIMENTS
The experiment code available at GitHub'. 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 theory, if a is smaller, the output distribution is more heavy-tailed.
empirically see that consistently with the
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 behavior predicted by theorem. The empirical validation with for ResNet shown Fig. 3. We set a = and 3 = 4 for the inverse gamma prior in this experiment.
Additional results of regression with Gaussian likelihood In addition to Table 1, we also mea- sured 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.
https:: thw.. hww..ttoo...oooo/oooooaaaaaaaaaa
Socaleeeee Mi-Mi www.ee- -
23

---

Published as a conference paper at ICLR 2022
Initill
Last bayr training
Full byyer traiinng
o- 2-32
2-2,1-2
4.
0.1
64
Probabilty
0.4
03
0.
20.0
-25
20.0
-10
260
Figure 2: Impact of the prior hyperparameters for fully connected neural network initial, last layer
training and full layer training.
Full layer training
Full layer training
Prodi ctod
a=4, 0-1/2, S-1/2
Probability
11-
Output range
Output range
Figure 3: (left) Correspondence between wide finite ResNet model vs theoretical limit. (right) impact of the prior hyperparameters.
Additional regression experiment As an additional regression experiments, we test the models which use Burr Type XII distribution as last layer variance's prior. We use Appendix order to calculate the predictive posterior distribution. Our results are in Table 6.
I ADDITIONAL INFORMATION
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 't use any data which contains personally identifiable informa- tion or offensive content.
24


Reviews:
[{"review_id": "upqCykHKGeN", "reviewer": "Reviewer_Z24X", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\n\n-\tTensor programs for expressing GPs using the NNGP formulation are a fairly recent development, and the work presented here is consequently a timely contribution to a topic that has accrued great interest within the community.\n-\tAlthough the core contribution feels fairly incremental in nature, both the execution and the corresponding analysis are non-trivial. Although I was not familiar enough with tensor programs to properly go through and understand all the derivations featured in the paper, the exposition appears to be quite thorough all throughout the paper, with sufficient detail also being provided in the supplementary material.\n-\tThe experiments are broad and varied, covering a wide selection of possible practical challenges (noisy data/unreliable labels, etc), where the proposed technique might be particularly effective.\n\n### Weaknesses\n\n-\tWhile well-written overall, I often found the paper to be quite difficult to follow properly. The paper contains a lot of set-up, whereas the paper’s own contributions are only presented from the very end of page 4 onwards. Even then, the majority of the paper’s contributions are all grouped together in a single ‘Results’ section that are only slightly disambiguated under three different theorems. Although the paper already contains a very extensive and thorough appendix, I do believe that the paper would benefit from some additional restructuring and refactoring.\n-\tIt would also be nice to expand on some of the related work (in particular the competing approaches highlighted for CNNs) in order to better emphasise how the approach being proposed in this paper is ‘simpler’ while being just as effective.\n-\tI have a few questions on the Experiments section:\n--\tIs there a particular reason for not showing RMSE results for the regression datasets? While the benefits of using a heavier-tailed distribution may indeed be more pronounced in the MNLL metrics, it would still be interesting to see whether this comes at the expense of possibly inferior mean predictions. \n--\tFor the results on the regression datasets shown in Table 1, the authors comment that the results for competing methods are copied from the work of Adlam et al (2020). However, given that there are no predefined train/test splits for these UCI datasets, how are you ensuring that the same splits are used as for this earlier work?\n--\tWould it be possible to compare against the other heavy-tailed stochastic processes described in the Related Work section? This should at least be possible for the image classification experiments right? There’s also a small error in the caption of Figure 1 which contains references to ‘top’ and ‘bottom’ figures.\n-\tThere are a few typos remaining in the paper, but these should be fairly easy to with a proper read-through of the paper. A few of the titles in the references also require capitalization for words such as Gaussian, etc.\n", "summary_of_the_paper": "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.", "summary_of_the_review": "The contributions featured in this paper provide a valuable and non-trivial extension to the literature on NNGPs. Given that this is a fairly new area of research, such work is especially timely. Even so, I still believe that the paper’s writing could be heavily improved in order to make the exposition more clear. I also have some remaining concerns on the experimental evaluation that I would like to see addressed in either the rebuttal or a future revision of the paper before tending more heavily towards acceptance.\n\n** Rebuttal Update **\n\nRaised score to 6 following clarifications and additional experiment results provided by authors.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "SDYWw4d82y", "reviewer": "Reviewer_axPe", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths: The approach is interesting, somewhat novel, and easily applicable to a broad range of deep nets. The provided theorems are sufficient for supporting the method and correct. The derivations pertaining to posterior inference are correct and result in efficient algorithms.  The results are broad enough, with many comparisons and many datasets considered; these vouch for the usefulness of the method.\nCons: Some discussion on complexity is missing.", "summary_of_the_paper": "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. \n\nThe 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). \n\nThey 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.\n\n", "summary_of_the_review": "A strong paper lacking some discussion on complexity. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "EIvhnJ6_N4g", "reviewer": "Reviewer_C7Q4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Font in tables is too small to be readable. Same for Fig. 1.\n\nThe paper in general clear and well written. The experiments are also extensive. The theorems seem correct although I did not check them thoroughly.\n\nMy main point of criticism of this paper is that there is little novelty in it. Although it provides interesting theoretical results, the main application, i.e., the use of Student's T process is already known in the machine learning community. So there are little new practical applications of this work.\n\nSumming up, I believe that although this paper is interesting it is not expected to have a big impact in the machine learning community.\n\n\n\n", "summary_of_the_paper": "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.\n", "summary_of_the_review": "Nice theoretical results that lead to already known methods in the machine learning community which questions the novelty of the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "C4JqHn8NugU", "reviewer": "Reviewer_QZea", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:  \n- Interesting extension of NNGPs to consider limiting processes which can have heavy tails, backed by rigorous convergence results and empirical evaluations\n- The paper is mostly well presented and easy to follow\n\nWeaknesses:\n- There are empirical findings showing that the scale mixture of NNGPs leads to robustness to the out-of-distribution data, but no theoretical results are provided to justify these (would be nice to at least give some intuition)\n\nComments/questions:\n- It is not clear to me how and why the heavy tail nature of the NNGPs leads to the robustness (also, why not robustness to noisy labels, imbalance, etc.). What is the key intuition there?\n- As for the empirical results, the scale mixture of NNGPs seems to give mixed performance (both accuracy and robustness to various perturbations) for different datasets. How should one decide which model to use for a given task to achieve desirable accuracy and robustness? Are there some trade-offs?\n\nMinor comments:\n- It may be a good idea to define the inverse Gamma distribution introduced on page 4\n- There is a missing comma after the first equation in Theorem 3.2", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, this is a well written paper that studies an interesting extension of NNGPs to a scale mixture version, which allows realization of a richer class of stochastic processes such as those with heavy tails. Rigorous results are provided and they seem to be technically correct (although I have not checked all the details). These are the key strengths of the paper and the main factors behind my rating (an acceptance). However, no intuition or theory is provided to justify the empirical findings for the robustness, which seems to be an important practical utility of the proposed approach. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a new formulation for the infinitely wide limiting case of deep networks as Gaussian processes, i.e. NNGPs.  The authors extend the existing case to incorporate a scale term at the penultimate layer of the network, which results in a scale mixture of NNGPs or a Student-t process in a specific case.  This formulation allows for a more heavy tailed output distribution which e.g. can be more robust to outliers.  The four reviews averaged just above borderline, with a 5, 8, 6, 6.  The reviewers found the approach to be sensible, technically correct and timely given the recent literature.  They found the experiments to be compelling for the most part, demonstrating the added robustness of this approach over the baseline NNGP.  The main concern raised by the reviewers is that the work is incremental, given that both NNGPs and Student-t processes are already established.
Decision: Accept (Poster)