Year: 2019

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.

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

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

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

heading: 2.1 BAYESIAN NEURAL NETWORKS AND VARIATIONAL INFERENCE
text: Consider a training set comprising of N input-output pairs, D = {xn, yn}Nn=1, and a neural network parameterized by weights and biases, w, that describes the distribution over an output yn given an input xn, p(yn|w,xn). 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 LMFVI(q(w)) ≈ −KL[q(w)||p(w)] + 1K ∑K k=1 ∑N n=1 log p(yn|wk,xn), where wk = µ + σ k and k ∼ N (0, I), i.e. wk 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.

heading: 2.2 META-REPRESENTING WEIGHTS AND NETWORKS
text: 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 zl,i, possibly of high dimensions Dz , 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, wl,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 Cw(z) consisting of the set of weight encodings for each individual weight in the network cwl,i,j = [ zl+1,i, zl,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:
p(w|z) = p(w|Cw(z)) = L−1∏ l=1 Hl∏ i=1 Hl+1∏ j=1 p(wl,i,j |zl+1,i, zl,j), (1)
where l denotes a visible or hidden layer and Hl 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 Hl ×Hl+1 weights compared to Hl +Hl+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(wl,i,j |zl+1,i, zl,j , θ) = N (wl,i,j ;µl,i,j , σ2l,i,j), where (µl,i,j , log σl,i,j) = NNθ([zl+1,i, zl,j ]), or models dependence between weights implicitly as p(w|Cw(z), θ) ∝ ∫ p( )
∏ l,i,j NNθ([zl+1,i, zl,j , ])d , where is a ran-
dom variate, wl,i,j = NNθ([zl+1,i, zl,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(zl) ] [p(w|Cw(z), θ)] [ N∏ n=1 p(yn|w,xn) ] . (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.

heading: 3 META-REPRESENTATIONS WITH GAUSSIAN PROCESS MAPPINGS AND CONTEXTUAL LATENT VARIABLES
text: 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.

heading: 3.1 META-REPRESENTING WEIGHTS WITH GAUSSIAN PROCESSES
text: 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,
p(wl,i,j |f, cwl,i,j ) = N (wl,i,j ; f([zl+1,i, zl,j ]), σ2w); p(f |γ) = GP(f ;0, kcw(·, ·|γ)), (3)
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,
P (w|z) = P (w|Cw(z)) = ∫ p(f) [ L−1∏ l=1 Hl+1∏ i=1 Hl∏ j=1 p(wl,i,j |f, zl+1,i, zl,j) ] df. (4)
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 2Dz , where Dz is the dimensionality of each latent variable z. 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 GPMetaPrior (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 factor-
ized kernel at the example of two weights in the network,
kcw(cw1 , cw2) = k([zl1+1,i1 , zl1,j1 ], [zl2+1,i2 , zl2,j2 ])
= kout(zl1+1,i1 , zl2+1,i2) · kin(zl1,j1 , zl2,j2). (5)
In this section and what follows, we will use the popular exponentiated quadratic kernel with ARD lengthscales, k(x1,x2) = σ2k exp (∑2Dz d=1 −(x1,d−x2,d)2 2l2d ) , where {ld}2Dzd=1 are the kernel lengthscales and σ2k is the kernel variance. For the model considered here, this admits the following separation in the kernel computation:
kwl1,i1,j1 ,wl2,i2,j2 = σ 2 k exp ( Dz∑ d=1 −(zl1,i1,d − zl2,i2,d)2 2l2d ) exp ( Dz∑ d=1 −(zl1+1,j1,d − zl2+1,j2,d)2 2l2Dz+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.

heading: 3.2 CONTEXTUAL VARIABLES FOR MODULATING FUNCTION PRIORS
text: 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 cwl,i,j = [ zl+1,i, zl,j , x ] . In turn, this yields input-conditional weight models p(wn,l,i,j |f, zl+1,i, zl,j ,xn). We again turn to compositional kernels and introduce a new input kernel Kx which we use as follows,
kcw(cw1 , cw2) = k([zl1+1,i1 , zl1,j1 , x1], [zl2+1,i2 , zl2,j2 , x2])
= kout(zl1+1,i1 , zl2+1,i2) · kin(zl1,j1 , zl2,j2) · kx(x1, x2). (6)
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 Kx. This, in turn, would modulate the entire kernel Kcw 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}Nn=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(Vxn), where V ∈ RDe×Dx , and De and Dx are the dimensionality of n and xn, respectively, and g(·) is an arbitrary transformation. We may also layer these transformations in general. We typically set De Dx so this transformation could be thought of as a dimensionality reduction operation. For low dimensional inputs, we can directly set n = xn. The auxiliary variable is then augmented to the input of the meta mapping as follows,
p(wn,l,i,j |f, zl+1,i, zl,j ,xn,V) = N (wn,l,i,j ; f([zl+1,i, zl,j , n]), σ2w), (7)
that is, the input dimension of the meta mapping is now 2Dz +De. The covariance matrices can be decomposed and computed efficiently as Kronecker products Kcw = Kout⊗Kin⊗Kx, as Kx 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 (WN ×WN ).1 Additionally, we also place a prior over the linear transformation: p(V) = N (V;0, I).

heading: 4 INFERENCE AND LEARNING USING STOCHASTIC STRUCTURED
text: 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(wn|f,Cw(z,xn,V))p(yn|w,xn) ] .
We first partition the space Z of inputs to the function f into a finite set of M variables called inducing inputs zu and the remaining inputs, Z = {zu,Z 6=zu}. The function f is partitioned identically, f = {u, f6=u}, where u = f(zu). We can then rewrite the GP prior as follows, p(f) = p(f6=u |u, zu)p(u|zu).2 The inducing inputs and outputs, {zu,u}, will be used to parameterize the approximation. In particular, a variational approximation is judiciously chosen to mirror the form of the joint density:
q(w, z, f,V|x) = q(z)q(V)p(f 6=u |u, zu)q(u)p(w|f, z,V,x), (8)
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(σ2z)) and q(V) = N (V;µV,diag(σ2V)), respectively, and q(u) is chosen to be a correlated multivariate Gaussian, q(u) = N (u;µu,Σu). This approximation allows convenient
1We can solve this in complexity dominated by the decomposition of the largest kernel in the product. 2The conditioning on Z6=zu in p(f 6=u |u, zu) is made implicit here and in the rest of this paper.
cancellations yielding a tractable variational lower bound as follows, LmetaGP(·) = ∫ w,z,f,V q(w, z, f,V|x) log p(z)p(V)(( ((((p(f6=u |u, zu)p(u|zu)((((( (( p(w|f, z,V,x)p(y|w,x) q(z)q(V)(((( ((p(f 6=u |u, zu)q(u)((((( (( p(w|f, z,V,x)
≈ −KL[q(z)||p(z)]−KL[q(V)||p(V)]−KL[q(u)||p(u|zu)]
+ 1
K K∑ k=1 ∫ w,f q(w, f |zk,Vk) log p(y|w,x)
where the last expectation has been partly approximated using simple Monte Carlo with the reparameterization trick, i.e. zk ∼ q(z) and Vk ∼ q(V). We will next discuss how to approximate the expectation Fk = ∫ w,f
q(w, f |zk,Vk,x) log p(y|w,x). Note that we split f into f6=u and u, and that we can integrate f 6=u out exactly to give, q(w|zk,u,Vk,x) = N (w;A(k)u,B(k)),
A(k) = K (k) f6=uu K−1uu; B (k) = K (k) l,f6=uf6=u −K(k)f6=uuK −1 uuK (k) uf6=u + σ2wI. (9)
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
q(w|zk,Vk,x) = N (w;A(k)µu,B(k) + A(k)ΣuAᵀ,(k)). (10)
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,
q̂(w|zk,Vk,x) = N (w;A(k)µu,diag(B(k) + A(k)ΣuAᵀ,(k))). (11)
Whilst the diagonal approximation above might look poor from the first glance, it is conditioned on a sample of the latent variables zk 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. Fk is approximated by Fk ≈ ∫ w q̂(w|zk,Vk,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.

heading: 5 RELATED WORK
text: 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.

heading: 6 EXPERIMENTS
text: 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.

heading: 6.1 TOY REGRESSION AND CLASSIFICATION EXAMPLES
text: 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.

heading: 6.2 INPUT DEPENDENT NEURAL NETWORKS FOR UNCERTAINTY QUANTIFICATION
text: 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.

heading: 6.3 MODELING INDUCTIVE BIASES FOR NEURAL NETWORKS WITH INPUT-DEPENDENT KERNELS
text: 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.

heading: 6.4 ACTIVE LEARNING
text: 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.

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

heading: A EXTRA EXPERIMENTAL RESULTS
text: A.1 AN EMPIRICAL EVALUATION OF VARIOUS APPROXIMATIONS FOR q(w|zk,Vk,x)
In this experiment, we analyze the impact of different approximations to the covariance matrix of q(w|zk,Vk,x):
q(w|zk,Vk,x) = N (w;A(k)µu,B(k) + A(k)ΣuAᵀ,(k)). If we use the exact, fully correlated Gaussian distribution above, it is necessary to sample from this distribution to evaluate the lower bound. This step costs O(W 3) where W is the number of parameters in the network.
The complexity can be greatly improved by making a diagonal approximation to B(k) as follows,
q̂FITC(w|zk,Vk,x) = N (w;A(k)µu, diag(B(k)) + A(k)ΣuAᵀ,(k)).
Sampling from this distribution can be done in O(WM2) where M is the number of pseudo-points.
This can be further approximated by assuming a diagonal covariance matrix,
q̂diag(w|zk,Vk,x) = N (w;A(k)µu, diag(B(k) + A(k)ΣuAᵀ,(k))). The variational bound can then be evaluated by drawing samples from q̂diag as in the above approximation, or by drawing activity samples by employing the local reparameterization trick (Kingma et al., 2015).
We evaluate the performance of using the exact and approximate conditional distributions above in a range of toy regression and classification, and show representative results in figs. 7 and 8. We note that the diagonal approximation is fast and gives qualitatively similar performance compared to more structured approximation or the exact case, in both cases where there is a singla GP for all weights in the network and there is multiple GPs, one for each weight layer in the network.
A.2 ROBUSTNESS IN VARIOUS DATA REGIMES FOR A TOY REGRESSION PROBLEM
In this experiment, we evaluate the qualitative performance of various methods, including MFVI, NN-Metarep, GP-Metarep and GP-Metarep with local, input-dependent kernel, on a toy regression problem, in different data regimes. In particular, we considers 10, 20, and 50 training points respectively, and plot the predictions in fig. 9. GP-metarep demonstrates consistent performance across all data regimes, comparable to that of NN-metarep. The input-dependent kernel helps the performance further in the out-of-distribution area.
A.3 EFFECT OF INPUT-DEPENDENT KERNEL AND APPLICATIONS TO FORWARD TRANSFER/FAST ADAPTATION/MULTI-TASK LEARNING
In this section, we investigate the effect of the input-dependent kernel by probing the learnt kernel hyperparameters. In particular, we first train a model with an periodic input-dependent kernel on a sinusoidal 1D data set, as show in the main text. We then vary the period hyper-parameter in the kernel whilst keeping other hyperparameters and variational parameters fixed. The predictions for a few hyperparameters are shown in fig. 10. We note the variation/period in the data is captured by weight modulation, governed by the input-dependent kernel. Changing the period hyperparameter affects how fast or slow the weights are changing wrt the input.
We further investigate using the proposed model for multi-task learning. In particular, the latent variable z and corresponding hyper-parameters and variational parameters can be shared across different tasks whilst the meta mapping and the input-dependent kernel are private to each individual task. We first train the model on four regression tasks, each corresponds to a sinusoid of a particular frequency. At test time, a novel test set is shown to the model. The hyper-parameters of the input kernel and variational parameters corresponding to this new test set are optimised while other hyper-parameters and the latent variables are kept fixed. We evaluate the performance of the model on the novel test sets to see how the latent variables can be reused and shared across tasks to facilitate fast adaptation to new settings. The performance of the model on the tasks used for training and new tasks at test time is shown in fig. 11. This result demonstrates the ability of the model trained with multiple similarly related tasks to faithfully and quickly adapt to new settings.
A.4 MNIST EXPERIMENT: A COMPARISON TO DEEP KERNEL LEARNING AND FULL FIGURES
In this section, we include the full figures of the MNIST out-of-distribution uncertainty experiment, as well as additional results using deep kernel learning (Wilson et al., 2016). In particular, we employ the same network architecture with the last layer being replaced by multiple independent GPs, one for each class (output dimension). As exact inference is intractable, variational inference based on inducing points is employed – we used 50 inducing points for each output. The full results of all models/methods considered are shown in fig. 12. For clarify, the results of deep kernel learning and GP-metarep are shown in fig. 13. GP-metarep with the input-dependent kernel shows good performance, outperforming deep kernel learning in all cases. In addition, we include the full figures for the predictive distributions on representative test examples in figs. 14 and 15.
A.5 ROBUSTNESS OF GP-METAREP WITH NETWORK ARCHITECTURES
In this experiment, we compare the performance of GP-metarep for various numbers of hidden units (20, 50 and 100) and two activation functions (tanh and ReLU) on a toy regression problem. The observation noise is fixed in this experiment. We observe that the performance of the models is in general consistent across different activation functions and numbers of hidden units. We show the results in fig. 16.
A.6 A TOY ACTIVE LEARNING SETTING
Please see fig. 17 and the associated caption.
A.7 PROBABILITY DENSITIES OF WEIGHTS IN THE TOY 1D REGRESSION EXPERIMENT
Please see figs. 18 and 19 and the associated captions.
1 0
d a ta
p o in
ts 2
0 d
a ta
p o in
ts 5
0 d
a ta
p o in
ts
M FV
I
N N
-M e ta
re p
G P -M
e ta
re p
G P -M
e ta
re p
+ l o ca
l
Fi gu
re 9:
Pe rf
or m
an ce
of m
ea n-
fie ld
va ri
at io
na li
nf er
en ce
,N N
-M et
ar ep
w ith
va ri
at io
na li
nf er
en ce
an d
G P-
M et
ar ep
w ith
va ri
at io
na li
nf er
en ce
on a
to y
re gr
es si on pr ob le m w ith va ri ou s nu m be ro ft ra in in g po in ts .B es tv ie w ed in co lo ur .
T R
A IN
IN G
Ta sk
1 Ta
sk 2
Ta sk
3 Ta
sk 4
T E S
T IN
G
20 data points
Ta sk
5 Ta
sk 6
Ta sk
7 Ta
sk 8
50 data points 200 data points
Fi gu
re 11
:T ra
in in
g on
m ul
tip le
re la
te d
ta sk
s an
d ad
ap ta
tio n
to no
ve lt
as ks
at te
st tim
e. In
th is
ca se
,t he
la te
nt va
ri ab
le s
(a s
w el
la s
w ei
gh tc
od e
hy pe
rp ar
am et
er s)
ar e
sh ar
ed ac
ro ss
ta sk
s w
hi le
ea ch
in di
vi du
al ha
s its
ow n
in pu
tde
pe nd
en tk
er ne
l. A
tt es
tt im
e, on
ly th
e pr
iv at
e pa
ra m
et er
s fo
rt he
ne w
ta sk
ar e
re -i
ni tia
lis ed
an d
op tim
is ed
.B es
tv ie
w ed
in co
lo ur
.
Fi gu
re 18
:L ef
t: Tr
ai ni
ng da
ta an
d pr
ed ic
tio ns
us in
g a
B ay
es ia
n ne
ur al
ne tw
or k
w ith
a m
ea n-
fie ld
va ri
at io
na lG
au ss
ia n
ap pr
ox im
at io
n. T
he ne
tw or
k ha
s on
e hi
dd en
la ye
r of
50 re
ct ifi
ed lin
ea r
un its
,h en
ce th
er e
ar e
50 in
pu tw
ei gh
ts an
d 50
ou tp
ut w
ei gh
ts .
R ig
ht :
th e
m ar
gi na
lv ar
ia tio
na la
pp ro
xi m
at io
ns ov
er th
e w
ei gh
ts af
te r
tr ai
ni ng
. M
os th
id de
n un
its ar
e pr
un ed
ou ta
ft er
tr ai
ni ng
,a s
ob se
rv ed
in Tr
ip pe
& Tu
rn er
(2 01
8) .
In th
is ca
se ,t
he re
ar e
12 ac
tiv e
hi dd
en un
its .
W e
in cl
ud e
th e
hi st
og ra
m of
th e
w ei
gh ts
in ea
ch la
ye ri
n fig
.3 .
Fi gu
re 19
:M id
dl e:
Tr ai
ni ng
da ta
an d
pr ed
ic tio
ns us
in g
a B
ay es
ia n
ne ur
al ne
tw or
k w
ith an
in pu
tde
pe nd
en tm
et a
G P
pr io
r. Si
m ila
rt o
fig .1
8, th
e ne
tw or
k ha
s on e hi dd en la ye ro f5 0 un its .T he di st rib ut io ns of th e w ei gh ts in th e ne tw or k, gi ve n an in pu t, ar e in cl ud ed .N ot e th at ,f or ea ch w ei gh t, w e fir st sa m pl e th e la te nt va ria bl e z , an d th en pr op ag at e th is sa m pl e th ro ug h th e m et a G P m ap pi ng to ob ta in a di st ri bu tio n of ea ch w ei gh t. In ea ch su bp lo t, w e pl ot 10 di st ri bu tio ns co rr es po nd in g to 10 di ff er en tz sa m pl es . O n th e le ft fig ur es ,a s pa rt of th e in pu tt o th e G P is fa r aw ay fr om th e tr ai ni ng in pu ts ,t he w ei gh ts be co m e un ce rt ai n w hi ch le ad s to hi gh un ce rta in ty in th e pr ed ic tio n. O n th e co nt ra ry ,a ro un d th e tra in in g po in ts ,t he G P ou tp ut s ar e m or e ce rta in an d m or e di ve rs e. N ot e th at ,t he m ar gi na lo fe ac h w ei gh ti s an in fin ite m ix tu re of G au ss ia n de ns iti es .B es tv ie w ed in co lo ur .

Review summaries:
This paper proposes a hierarchical Bayesian approach to hyper-networks by placing a Gaussian process prior over the latent representation for each weight. A stochastic variational inference scheme is then proposed to infer the posterior over both the Gaussian process and the weights themselves. Experiments are performed on toy regression, classification, (edit: post rebuttal) and transfer learning tasks, as well as an uncertainty quantification experiment on MNIST.