Gaussian Process Meta-Representations Of Neural Networks

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 = { 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 . 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 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 . 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 .

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