Title: GAUSSIAN PROCESS META-REPRESENTATIONS FOR NEURAL NETWORKS

Abstract: Bayesian inference offers a theoretically grounded and general way to train neural networks and can potentially give calibrated uncertainty. It is, however, challenging to specify a meaningful and tractable prior over the network parameters. More crucially, many existing inference methods assume mean-field approximate posteriors, ignoring interactions between parameters in high-dimensional weight space. To this end, this paper introduces two innovations: (i) a Gaussian process-based hierarchical model for the network parameters based on recently introduced unit embeddings that can flexibly encode weight structures, and (ii) input-dependent contextual variables for the weight prior that can provide convenient ways to regularize the function space being modeled by the network through the use of kernels. Furthermore, we develop an efficient structured variational inference scheme that alleviates the need to perform inference in the weight space whilst retaining and learning non-trivial correlations between network parameters. We show these models provide desirable test-time uncertainty estimates, demonstrate cases of modeling inductive biases for neural networks with kernels and demonstrate competitive predictive performance of the proposed model and algorithm over alternative approaches on a range of classification and active learning tasks.

1. INTRODUCTION:

Bayesian neural networks (BNNs) (see e.g. MacKay, 1992; Neal, 1993; Ghahramani, 2016) are one of the research frontiers on combining Bayesian inference and deep learning, potentially offering flexible modelling power with calibrated predictive performance. In essence, applying probabilistic inference to neural networks allows all plausible network parameters, not just the most likely, to be used for predictions. Despite some earlier success in supervised learning tasks (Guyon et al., 2005; Quiñonero-Candela et al., 2006) , BNNs has enjoyed limited recent practical applications. This could be attributed to: (i) the need to choose an appropriate weight prior as commonly modelers care about functions priors rather than weight priors but require the power and expressivity of weight-based models, and (ii) recently developed approximate inference methods are not able to capture complicated posterior correlations in large networks. As such, addressing these challenges is an active research area. This paper attempts to remedy the aforementioned limitations for Bayesian neural networks by extending previous work on unit-based representations for neural networks with a unit-based Gaussian process hierarchical prior over network parameters, which flexibly models the correlations between weights in a layer and across layers while forgoing the need for a parametric hyper-network. Moreover, input-dependent contextual variables are used for modulating the GP-function and in turn specifying input-dependent weight priors which draw from the tools from the kernel literature to imbue neural networks with inductive biases. This may yield more desirable uncertainty estimates at test time, or restrict function spaces to generalize in particular ways. A structured variational inference approach is employed that side-steps the need to do inference in the weight space whilst retaining weight correlations in the approximate posterior. These theoretical benefits translate to competitive practical performance on a range of supervised and active learning tasks. The paper is organized as follows: in the next section, we will briefly review mean-field variational inference for BNNs and a recently proposed hierarchical model for BNNs. The proposed model and inference scheme are discussed in section 3 and section 4, respectively, followed by a discussion on related work in section 5. We then validate the performance of the proposed model and algorithm in section 6.

2. BACKGROUND:

In this section, we lay out the groundwork for the paper by providing a concise summary of variational approaches to Bayesian neural networks and a recent approach to learning meta-representations of neural networks.

2.1. BAYESIAN NEURAL NETWORKS AND VARIATIONAL INFERENCE:

Consider a training set comprising of N input-output pairs, D = {x n , y n } N n=1 , and a neural network parameterized by weights and biases, w, that describes the distribution over an output y n given an input x n , p(y n |w, x n ). We follow a Bayesian approach by placing a prior distribution over the network parameters, p(w), and obtaining the posterior distribution p(w|D), which involves calulation of the marginal likelihood p(D) = dwp(w)p(D|w). However, obtaining p(w|D) and p(D) exactly is intractable when N is large or when the network is large and as such, approximation methods are often required. In particular, mean-field Gaussian variational inference (MFVI) has recently become a method of choice for approximate inference for Bayesian neural networks due to its simplicity and the recently popularized reparameterization trick (Salimans & Knowles, 2013; Kingma & Welling, 2013; Titsias & Lázaro-Gredilla, 2014; Blundell et al., 2015) . MFVI sidesteps the intractability by positing a diagonal Gaussian approximation q(w) = N (w; µ µ µ, diag(σ σ σ 2 )) and optimising an approximate lower bound to the marginal likelihood

L MFVI (q(w)) ≈ -KL[q(w)||p(w)] + 1 K K k=1 N n=1 log p(y n |w k , x n ), where w k = µ µ µ + σ σ σ k and k ∼ N (0, I), i.e.

w k is a sample from q(w). Note that the mean-field variational Gaussian approximation with a standard normal prior, presented in is often outperformed by point estimation in certain settings (Trippe & Turner, 2018) . Despite being practical and able to give reasonable uncertainty estimates, improving MFVI is still an active research area, and the main focuses of which are (i) improving the reparameterization gradient estimator to enable faster convergence (Miller et al., 2017; Wu et al., 2018) , (ii) replacing the typical standard Normal prior, p(w) = N (w; 0, I) by a structured prior that better models the structures present in the weight a-priori (Ghosh et al., 2018; Neal, 2012; Blundell et al., 2015) , and (iii) using structured variational approximations that can potentially capture weight correlations in the posterior (Louizos & Welling, 2016; Zhang et al., 2017) . This paper builds on the two latter themes and proposes a hierarchical model for the prior and a structured variational scheme that explicitly model and infer weight structures.

2.2. META-REPRESENTING WEIGHTS AND NETWORKS:

Our work builds on and expands the class of hierarchical neural network models based on the concept of latent variables associated with units in a network as proposed in Karaletsos et al. (2018) . In these models, each unit (visible or hidden) of the l-th layer of the network has a corresponding latent hierarchical variable z l,i , possibly of high dimensions D z , where i denotes the index of the unit in a layer. Note that these latent variables do not describe the activation of units, but rather constitute latent features associated with a unit. These latent variables are then used to construct the weights in the network. Their construction is judiciously chosen such that a weight in the l-th weight layer, w l,i,j is linked to the latent variables z's of the i-th input unit and the j-th output unit of the weight layer. We can summarize this relationship by introducing the concept of weight codes C w (z) consisting of the set of weight encodings for each individual weight in the network c w l,i,j = z l+1,i , z l,j , which can be deterministically constructed from the collection of unit latent variables z by concatenating them correctly. Mathematically, the probabilistic description of the relationship between the weight codes (summarizing the structured latent variables) and the weights w is:

EQUATION

where l denotes a visible or hidden layer and H l is the number of units in that layer, and w denotes all the weights in this network.

The central motivations for this parameterization are three-fold. First, weights in a layer and across layers are explicitly correlated in the modelling stage, as they share latent variables if they are connected to the same unit, unlike the standard normal prior typically used as described in the last section. Second, the number of visible and hidden units in a neural network is typically much smaller than the number of weights. For example, for the l-th weight layer, there are H l × H l+1 weights compared to H l + H l+1 associated latent variables. This encourages the development of inference schemes that can work in the lower-dimensional hierarchical latent space directly without the need to infer the weights. Third, as hierarchical latent variables are associated with units in the network and have the same cardinalities, inferring these latent variables conditioning on observed data could be thought of as a form of network representation learning.

In Karaletsos et al. (2018) , a small parametric neural network regression model (conceptually a structured hyper-network) is chosen to map the latent variables to the weights. The network either assumes a Gaussian noise model which factorizes over weights, p(w l,i,j |z l+1,i , z l,j , θ) = N (w l,i,j ; µ l,i,j , σ 2 l,i,j ), where (µ l,i,j , log σ l,i,j ) = NN θ ([z l+1,i , z l,j ]), or models dependence between weights implicitly as p(w|C w (z), θ) ∝ p( ) l,i,j NN θ ([z l+1,i , z l,j , ])d , where is a random variate, w l,i,j = NN θ ([z l+1,i , z l,j , ]), and θ denotes the parameters of the neural network mapping. We will call this network a meta mapping. Note that given the collection of sampled unit variables z, the weights (or theirs mean and variance) can be obtained efficiently in parallel. A prior over the latent variables z, p(z) = N (z; 0, I), completes the model specification. We can thus write down the joint density of the resulting hierarchical model as follows,

p(y, w, z|x, θ) = L l=1 p(z l ) [p(w|C w (z), θ)] N n=1 p(y n |w, x n ) .

(2)

Variational inference was employed in prior work to infer z (and w implicitly), and to obtain a point estimate of θ, as a by-product of optimising the variational lower bound.

3. META-REPRESENTATIONS WITH GAUSSIAN PROCESS MAPPINGS AND CONTEXTUAL LATENT VARIABLES:

In this section, we present two novel extensions of the hierarchical model in section 2.2 that aim to increase the robustness in the small data settings and improve its out-of-sample uncertainty estimates.

Figure 1 : Graphical depiction of various models: vanilla BNNs, BNNs with hierarchical GP-MetaPriors, and BNNs with hierarchical GP-MetaPriors and auxiliary variables.

3.1. META-REPRESENTING WEIGHTS WITH GAUSSIAN PROCESSES:

Notice that in section 2.2, the meta mapping from the hierarchical latent variables to the weights is a parametric non-linear function, specified by a neural network. We replace the parametric neural network by a probabilistic functional mapping and place a nonparametric Gaussian process (GP) prior over this function. That is,

EQUATION

where we have assumed a zero-mean GP, k γ (•, •) is a covariance function and γ is a small set of hyper-parameters. The representational effect is that the latent function introduces correlations for the individual weight predictions, Here, we present a homoscedastic noise model for the weights, but the model is readily adaptable to a heteroscedastic noise model which we omit for clarity. For a more detailed discussion on GPs and their applications, please see Rasmussen & Williams (2005) . It is worth noting that whilst the number of latent variables and weights can be large, the input dimension to the GP mapping is only 2D z , where D z is the dimensionality of each latent variable z.

EQUATION

The GP mapping effectively performs one-dimensional regression from latent variables to individual weights while capturing their correlations. We will refer to this mapping as a GP-MetaPrior (metaGP).

In general, the form of the covariance function and the values of the hyperparameters encapsulate prior knowledge about the unknown function. We note that GPs benefit from the fact that kernels are compositional and can thus be used to represent correlation structures explicitly. We exploit this fact and define the following factorized kernel at the example of two weights in the network,

EQUATION

In this section and what follows, we will use the popular exponentiated quadratic kernel with ARD lengthscales, k(

x 1 , x 2 ) = σ 2 k exp 2Dz d=1 -(x 1,d -x 2,d ) 2 2l 2 d

, where {l d } 2Dz d=1 are the kernel lengthscales and σ 2 k is the kernel variance. For the model considered here, this admits the following separation in the kernel computation:

k w l 1 ,i 1 ,j 1 ,w l 2 ,i 2 ,j 2 = σ 2 k exp Dz d=1 -(z l1,i1,d -z l2,i2,d ) 2 2l 2 d exp Dz d=1 -(z l1+1,j1,d -z l2+1,j2,d ) 2 2l 2 Dz+d ,

that is, the computation of the covariance matrix between all the weights can be computed efficiently as a product of two sub-covariance matrices that are computed using the latent variables.

3.2. CONTEXTUAL VARIABLES FOR MODULATING FUNCTION PRIORS:

We first note that while the hierarchical latent variables and meta mappings introduce non-trivial coupling between the weights a priori, the weights and latent variables are inherently global. That is, a function drawn from the model, represented by a set of weights, does not take into account the inputs at which the function will be evaluated. Inspired by recent work on using contextual information or dataset summary to drive supervised learning predictions (Edwards & Storkey, 2016; Garnelo et al., 2018) , we introduce the input variable into the weight codes c w l,i,j = z l+1,i , z l,j , x . In turn, this yields input-conditional weight models p(w n,l,i,j |f, z l+1,i , z l,j , x n ). We again turn to compositional kernels and introduce a new input kernel K x which we use as follows,

EQUATION

As a result of having private contextual inputs to the meta mapping, the weights are now also local to each data point. In effect, each function drawn from this model explicitly needs the input locations we are interested in, in a similar vein to how functions are drawn from a Gaussian process.

What effects should we expect from such a modulation? As usual, the modeler has to choose the input kernel. Consider the use of an exponentiated quadratic kernel: we would expect data which lies far away from training data to receive small kernel values from K x . This, in turn, would modulate the entire kernel K cw for that data point to small values, leading to a weight model that reverts increasingly to the prior. We would hope such a model would help with modeling uncertainty by resetting weights to uninformative distributions away from training data. However, one may also want to use this mechanism to model inductive biases for the network, such as adding structure to the weight prior that can be captured with a kernel function. This is a potentially appealing avenue, as multiple useful kernels have been found in the GP literature that allow modelers to describe relationships between data, but was not a tool that was previously accessible to neural network modelers. We consider this a novel form of functional regularization, as the entire network can be given structure that will constrain its function space. Unfortunately, we cannot easily apply this method directly to high dimensional inputs as commonly used kernel functions fail in this setting. To overcome this limitation, we introduce { n } N n=1 , one for each training instance, which can be then used to as an extra input dimension for the meta mapping to generate weights instead of the actual inputs. In detail, each auxiliary input is obtained via a (potentially nonlinear) transformation applied to an input: n = g(Vx n ), where V ∈ R De×Dx , and D e and D x are the dimensionality of n and x n , respectively, and g(•) is an arbitrary transformation. We may also layer these transformations in general. We typically set D e D x so this transformation could be thought of as a dimensionality reduction operation. For low dimensional inputs, we can directly set n = x n . The auxiliary variable is then augmented to the input of the meta mapping as follows,

EQUATION

that is, the input dimension of the meta mapping is now 2D z + D e . The covariance matrices can be decomposed and computed efficiently as Kronecker products K cw = K out ⊗ K in ⊗ K x , as K x maps over data (N × N ) instead of over weights (W × W ) like the previous kernels and so their product requires efficient approximation to avoid calculating an object of size (W N × W N ).1 Additionally, we also place a prior over the linear transformation: p(V) = N (V; 0, I).

4. INFERENCE AND LEARNING USING STOCHASTIC STRUCTURED VARIATIONAL INFERENCE:

Performing inference is challenging due to the non-linearity of the neural network and the need to infer an entire latent function f . To address these problems, we derive a structured variational inference scheme that makes use of innovations from inducing point GP approximation literature (Titsias, 2009; Hensman et al., 2013; Quiñonero-Candela & Rasmussen, 2005; Matthews et al., 2016; Bui et al., 2017) and previous work on inferring meta-representations (Karaletsos et al., 2018) . As a reminder, we write down the joint density of all variables in the model specified in section 3:

p(y, w, z, f, V|x) = p(z)p(V)p(f )p(w|f, z, V, x)p(y|w, x) = p(z)p(V)p(f ) N n=1 p(w n |f, C w (z, x n , V))p(y n |w, x n ) .

We first partition the space Z of inputs to the function f into a finite set of M variables called inducing inputs z u and the remaining inputs, Z = {z u , Z =zu }. The function f is partitioned identically, f = {u, f =u }, where u = f (z u ). We can then rewrite the GP prior as follows, p(f ) = p(f =u |u, z u )p(u|z u ). 2 The inducing inputs and outputs, {z u , u}, will be used to parameterize the approximation. In particular, a variational approximation is judiciously chosen to mirror the form of the joint density:

EQUATION

where the variational distribution over w is made to explicitly depend on remaining variables through the conditional prior, q(z) and q(V) are chosen to be diagonal (mean-field) Gaussian densities, q(z) = N (z; µ µ µ z , diag(σ σ σ 2 z )) and q(V) = N (V; µ µ µ V , diag(σ σ σ 2 V )), respectively, and q(u) is chosen to be a correlated multivariate Gaussian, q(u) = N (u; µ µ µ u , Σ u ). This approximation allows convenient cancellations yielding a tractable variational lower bound as follows, ( ( ( ( ( p(w|f, z, V, x) p(y|w, x) q(z)q(V) ( ( ( ( ( ( p(f =u |u, z u )q(u) ( ( ( ( ( ( ( p(w|f, z, V, x)

L metaGP (•) = w,z,f,V q(w, z, f, V|x) log p(z)p(V) ( ( ( ( ( ( p(f =u |u, z u )p(u|z u ) ( (

≈ -KL[q(z)||p(z)] -KL[q(V)||p(V)] -KL[q(u)||p(u|z u )] + 1 K K k=1 w,f q(w, f |z k , V k ) log p(y|w, x)

where the last expectation has been partly approximated using simple Monte Carlo with the reparameterization trick, i.e. z k ∼ q(z) and V k ∼ q(V). We will next discuss how to approximate the expectation F k = w,f q(w, f |z k , V k , x) log p(y|w, x). Note that we split f into f =u and u, and that we can integrate f =u out exactly to give, q(w|z k , u, V k , x) = N (w; A (k) u, B (k) ),

EQUATION

At this point, we can either (i) sample u from q(u), or (ii) integrate u out analytically. We opt for the second approach, which gives

EQUATION

In contrast to GP regression and classification in which the likelihood term is factorized point-wise w.r.t. the parameters and thus their expectations only involve a low dimensional integral, we have to integrate out w in this case, which is of much higher dimensions. When necessary or practical, we resort to Kronecker factored models or make an additional diagonal approximation as follows,

EQUATION

Whilst the diagonal approximation above might look poor from the first glance, it is conditioned on a sample of the latent variables z k and thus the weights' correlations are retained after integrating out z. Such correlation is illustrated in fig. 2 where we show the marginal and conditional covariance structures for the weights of a small neural network, separated into diagonal and full covariance models. The diagonal approximation above has been observed to give pathological behaviours in the GP regression case (Bauer et al., 2016 ), but we did not observe these in practice. A comparison of this diagonal approximation to alternatives is provided in the appendix. F k is approximated by F k ≈ w q(w|z k , V k , x) log p(y|w, x) which can be subsequently efficiently estimated using the local reparameterization trick (Kingma et al., 2015) . The final lower bound is then optimized to obtain the variational parameters of q(u), q(z) and q(V), and estimates for the noise in the meta-GP model, the kernel hyper-parameters and the inducing inputs.

5. RELATED WORK:

Bayesian neural networks have recently garnered a resurgence of interests in the community. There is, however, a long history of research on developing approximate Bayesian inference methods for these networks. Notable work include extended Kalman filtering (Singhal & Wu, 1989 ), Laplace's approximation (MacKay, 1992), Hamiltonian Monte Carlo (Neal, 1993; 2012) , variational inference (Hinton & Van Camp, 1993; Barber & Bishop, 1998; Graves, 2011; Blundell et al., 2015; Gal & Ghahramani, 2016) , and approximate expectation propagation (Hernández-Lobato & Adams, 2015; Li et al., 2015; Hernández-Lobato et al., 2016) . The work in this paper is orthogonal to these as the model employs a hierarchical prior, and inference is done in a lower-dimensional latent space instead of the weight space. The variational approximation is chosen such that the marginal distribution over the weights is non-Gaussian and the correlations between weights are retained, in contrast to the popular mean-field Gaussian approximation. Additionally, a standard multivariate Normal distribution is typically used as the prior over the network weights and biases. Imposing additional structures over the weights a prior or carefully choosing the right prior variances have been observed to improve the predictive performance (Ghosh et al., 2018; Neal, 2012; Blundell et al., 2015) . In this vein, the proposed model in this paper imposes a class of flexible structures over the network parameters and learn a posterior over them conditioned on observed data. Another related theme is hyper-networks, the core idea of which is to generate network parameters using another network (see e.g. Ha et al., 2016; Stanley et al., 2009) . The proposed model could be thought of as a GP hyper-network, but utilizes latent variables for nodes as introduced in (Karaletsos et al., 2018) , thus avoiding over-parametrized, monolithic models over weight tensors. Instead, the model uses hierarchical modeling, conditional probabilities, and the unit nodes to perform one-dimensional regression to weights maintaining modeling correlations.

6. EXPERIMENTS:

In this section, we evaluate the proposed model and inference scheme on several toy and benchmark datasets. We compare the proposed approach primarily to maximum a posteriori estimation (MAP) and MFVI, both of which use a standard Normal prior over the weights. These were implemented using PyTorch (Paszke et al., 2017) and the code will be available upon acceptance. Additional results are included in the appendices. We use M = 50 inducing points for all experiments in this section. We first illustrate the performance of the proposed model on two toy examples, one regression and one classification. In both cases, a Bayesian neural network with one hidden layer of 50 hidden units is used. Additionally, the hierarchical latent variables have two dimensions. For the regression case, we generate 100 data points from a toy function suggested by Louizos et al. (2019) . For classification, a four-class dataset of 100 data points is used. Figure 3 shows the predictive performance of the proposed models, MFVI and a meta-NN on both tasks. Notably, the proposed model with contextual inputs (and thus input dependent network weights) gives uncertainty estimates that are reminiscent to that of a GP regression model, despite being a neural network under the hood. Following Bradshaw et al. (2017) , we also show the uncertainty for data further from the training instances. The model with auxiliary inputs remains uncertain as expected for these points whilst MFVI produces arguably overconfident predictions. We also attempt to summarize the distribution of the weights in the network in appendix A.7 in the appendices.

6.1. TOY REGRESSION AND CLASSIFICATION EXAMPLES:



6.2. INPUT DEPENDENT NEURAL NETWORKS FOR UNCERTAINTY QUANTIFICATION:

Motivated by the performance of the proposed metaGP model with auxiliary inputs in the toy examples in fig. 3 , this section tests the ability of this model class to produce calibrated predictive uncertainty to out-of-distribution samples. That is, for test samples that do not come from the same training distribution, a robust and well-calibrated model should produce uncertain predictive distribution and thus high predictive entropy. Such a model could find applications in safety-critical tasks where the cost of making confident but wrong predictions is high, or in area where detecting unfamiliar inputs is crucial such as active learning or reinforecement learning. In this work, we first train a neural network classifier with one hidden layer of 100 rectified linear units on the MNIST dataset, and apply the metaGP prior only to the last layer of the network. The dimensions of the latent variables and the auxiliary inputs are both 2. After training, we compute the entropy of the predictions on various test sets, including notMNIST, fashionMNIST, Kuzushiji-MNIST, and uniform and Gaussian noise inputs. Following (Lakshminarayanan et al., 2017; Louizos & Welling, 2017) , the CDFs of the predictive entropies for various methods are shown in fig. 4 . A calibrated classifier should give a CDF that bends towards the top-left corner of the plot for in-distribution examples and, vice versa, towards the bottom-right corner of the plot for out-of-distribution inputs. In most out-ofdistribution sets considered, except Gaussian random noise, metaGP and metaGP with local auxiliary variables demonstrate competitive performance to Gaussian MFVI. Notably, MAP estimation, often deployed in practice, tends to give wildly poor uncertainty estimates on out-of-distribution samples. We illustrate this behaviour and that of other models and methods on representative inputs of the MNIST and Kuzushiji-MNIST datasets in fig. 4 . 

6.3. MODELING INDUCTIVE BIASES FOR NEURAL NETWORKS WITH INPUT-DEPENDENT KERNELS:

We further explore the utility of the contextual variable towards modeling inductive biases for neural networks and evaluate on predictive performance on a regression example. In particular, we generate 100 training points from a synthetic sinusoidal function and create two test sets that contains in-sample inputs and out-of-sample inputs, respectively. We test an array of models and inference methods, including BNN with MFVI, metaGP and metaGP with contextual variables. In particular, as noted in section 3.1 we can choose the covariance function to be used for the auxiliary variables to encode our belief about how the weights should be modulated by the input. We pick exponentiated quadratic and periodic kernels (MacKay, 1998) in this example. The predictive performance are measured by the root mean squared error (RMSE) and the log-likelihood (LL) on the test examples. Figure 5 summarizes the results and illustrate the qualitative difference between models. It could be noted that the periodic kernel allows the model to discover and encode periodicity, allowing for more long-range confident predictions compared to that of the exponentiated quadratic kernel. In this example, MetaGP with contextual variables is superior to other methods, demonstrating good RMSE and LL on both in-distribution and out-of-distribution examples.

6.4. ACTIVE LEARNING:

We next stress-test the performance of the proposed model in a pool-based active learning setting for real-valued regression, where limited training data is provided initially and the target is to sequentially select points from a pool set to add to the training set. The criterion to select the next best point from the pool set is based on the entropy of the predictive distribution, i.e. we pick one with the highest entropy. In practice, we approximate the predictive density by a Gaussian density, which results in a tractable entropy computation. Note that this selection procedure can be interpreted as selecting points that maximally reduce the posterior entropy of the network parameters (Houlsby et al., 2011) . Four UCI regression datasets were considered, where each with 40 random train/test/pool splits. For each split, the initial train set has 20 data points, the test set has 100 data points, and the remaining points are used for the pool set, similar to the active learning set-up in Hernández-Lobato & Adams (2015). We compare the performance of the proposed model and inference scheme to that of Gaussian mean-field variational inference and show the average results in fig. 6 . Across all runs, we observe that active learning is superior to random selection and more crucially using the proposed model and inference scheme seems to yield comparable or better predictive errors with a similar number of queries. We also illustrate the procedure on a toy classification dataset and visualize the predictive uncertainty as more points are picked and used for training in fig. 17 in the appendix. This simple setting quantitatively reveals the inferior performance of MFVI, compared to MAP and metaGP.

7. SUMMARY:

We have developed a flexible hierarchical neural network model and a structured variational inference scheme that together (i) enable flexible modelling of the correlations between network parameters, (ii) allow weight structures to be retained in the approximate posterior without having to perform direct inference in the high-dimensional weight space, and (iii) allow input-dependent contextual variables to be used to generate private weights per data point which translates to more desirable uncertainty estimates in practice and the ability to generalize out of sample with inductive biases such as periodicity. We plan to evaluate the performance of the model on more challenging decision making tasks such as contextual bandits or reinforcement learning, and to extend the inference scheme to handle continual learning. Figure 18 : Left: Training data and predictions using a Bayesian neural network with a mean-field variational Gaussian approximation. The network has one hidden layer of 50 rectified linear units, hence there are 50 input weights and 50 output weights. Right: the marginal variational approximations over the weights after training. Most hidden units are pruned out after training, as observed in Trippe & Turner (2018) . In this case, there are 12 active hidden units. We include the histogram of the weights in each layer in fig. 3 .

Figure 19 : Middle: Training data and predictions using a Bayesian neural network with an input-dependent meta GP prior. Similar to fig. 18 , the network has one hidden layer of 50 units. The distributions of the weights in the network, given an input, are included. Note that, for each weight, we first sample the latent variable z, and then propagate this sample through the meta GP mapping to obtain a distribution of each weight. In each subplot, we plot 10 distributions corresponding to 10 different z samples. On the left figures, as part of the input to the GP is far away from the training inputs, the weights become uncertain which leads to high uncertainty in the prediction. On the contrary, around the training points, the GP outputs are more certain and more diverse. Note that, the marginal of each weight is an infinite mixture of Gaussian densities. Best viewed in colour.

:

We can solve this in complexity dominated by the decomposition of the largest kernel in the product.

The conditioning on Z =zu in p(f =u |u, zu) is made implicit here and in the rest of this paper.

Figures and Tables:

Figure 2: Marginal and conditional covariance structures over weights in a 1x50x1 neural network. Sampling from the posterior of the hierarchical model reveals that even a diagonal GP approximation can capture off-diagonal correlations induced through unit correlations. Also note the off-diagonal bands in the marginal plots above, which indicate the correlation structures induced by the latent variables of the hidden units connecting the layers.

Figure 3: Illustration of the performance of Bayesian neural networks with a standard Normal prior over the weights and a mean-field Gaussian variational approximation [MFVI], a meta-NN prior, a meta-GP prior and a meta-GP prior with auxiliary inputs. Top: 1-D regression on 100 noisy observations. The plots show the predictive mean and uncertainty as well as the training points. Middle and bottom: Multiclass classification on a toy four-class dataset of 100 examples. The background colour matches the entropy of the predictive probability distribution, e.g. darker is more certain. Each pair of plots show the behaviour around and away from the training points when zooming out. Best viewed in colour.

Figure 4: The CDFs of predictive entropies on in-distribution and out-of-distribution test sets for various methods [Left] and the predictive class probability by these methods for representative samples from an in-distribution test set [Top right] and an out-of-distribution test test [Bottom right]. Best viewed in colour.

Figure 5: Illustration of the effect of local variables and different kernels for the local variables.The underlying function is a sinusoid. Each subplot shows the training points (black pluses), the in-distribution test points (magenta crosses), the out-of-distribution test points (magenta filled circles), the predictive mean (connected line), two standard deviations (shaded area) and added observation noise (red dashed line). The title of the subplots include the predictive performance on the indistribution and out-of-distribution test points, demonstrating the ability of a model to interpolate and extrapolate, correspondingly. Best viewed in colour.

Figure 8: An evaluation of the covariance matrix approximations in a toy classification example. Top: objective function during training vs epoch/time. Bottom: Predictions after training using one of the approximations discussed in the text. Global: there is one GP for all weights in the network. Layer: there are multiple GPs, one for each weight layer in the network. Best viewed in colour.

Figure 17: Active learning with BNNs using maximum a posteriori estimation [BNN-MAP], meanfield Gaussian variational inference [BNN-MFVI] and a meta-GP hierarchical prior [BNN-MetaGP] on a toy multi-class classification task. For each plot, the filled circle markers are the current training points, with different colours illustrating different classes. The shaded crosses are the examples in the pool set, one of which we wish to pick and evaluate to be included in the training set. The unfilled circle markers are the examples from the pool set selected at a step. The objective function for selecting points from the pool set is the entropy of the predictive probability. Best viewed in colour.

Peer review (iclr2020):
This paper proposes an improvement over Probabilistic meta-representations of neural networks by replacing the NN parameterization of the network weights given latent variables by a probabilistic distribution whose mean is distributed by GP. Inference of the induced hierarchical model is achieved by variational inference and various approximations when needed. Comments: 1. The authors claim that the proposed method aims to increase the robustness in the small data settings and improve its out-of-sample uncertainty estimates. The second part is well justified. Could the author elaborate further on the high level intuition behind the first part? 2. What exactly is the gain obtained by replacing a NN parameterization with a GP parameterization? It seems like the proposed method gains the ability to model uncertainty, but potentially incurs performance trade-off from a series of approximation. 3. Following from comment #2, I am a bit surprised by the lack of comparison against the work of Karaletsos although this work was built on top of it. I am interested to understand the mentioned trade-off in practice. 4. Regarding the practical consideration in Eq.(11), I think it kind of defeats the purpose of setting up the latent variables so that *weights in a layer and across layers are explicitly correlated in the modelling stage* (Section 2.2). Do all experiments presented in the paper employ this practical consideration? 5. For the text before Eq. (8) should the inducing inputs be xu instead of zu? It seems like a systematic typo here because in the formulation for the lower bound it becomes p(u|xu) again instead of p(u|zu) 6. Is Kuu in Eq.(9) computed by taking Kronecker product of K_in, K_out and K(xu,xu) or just K(xu,xu). I am under the impression that it is the former, in which case taking the inversion is costly (cubic in the number of latent variables). This would not permit large NN anyway, which kind of defeats the purpose. It explains the choice of small architecture in your experiments as well. Overall comment: I think the paper presents an interesting idea but I have questions regarding its practical significance as highlighted in my specific comments above. I hope the authors would clarify these so I can converge on a final rating.