Global inducing point variational posteriors for Bayesian neural networks and deep Gaussian processes

1 INTRODUCTION . Deep models , formed by stacking together many simple layers , give rise to extremely powerful machine learning algorithms , from deep neural networks ( DNNs ) to deep Gaussian processes ( DGPs ) ( Damianou & Lawrence , 2013 ) . One approach to reason about uncertainty in these models is to use variational inference ( VI ) ( Jordan et al. , 1999 ) . VI in Bayesian neural networks ( BNNs ) requires the user to specify a family of approximate posteriors over the weights , with the classical approach being independent Gaussian distributions over each individual weight ( Hinton & Van Camp , 1993 ; Graves , 2011 ; Blundell et al. , 2015 ) . Later work has considered more complex approximate posteriors , for instance using a Matrix-Normal distribution as the approximate posterior for a full weight-matrix ( Louizos & Welling , 2016 ; Ritter et al. , 2018 ) . By contrast , DGPs use an approximate posterior defined over functions — the standard approach is to specify the inputs and outputs at a finite number of “ inducing ” points ( Damianou & Lawrence , 2013 ; Salimbeni & Deisenroth , 2017 ) . Critically , these classical BNN and DGP approaches define approximate posteriors over functions that are independent across layers . An approximate posterior that factorises across layers is problematic , because what matters for a deep model is the overall input-output transformation for the full model , not the input-output transformation for individual layers . This raises the question of what family of approximate posteriors should be used to capture correlations across layers . One approach for BNNs would be to introduce a flexible “ hypernetwork ” , used to generate the weights ( Krueger et al. , 2017 ; Pawlowski et al. , 2017 ) . However , this approach is likely to be suboptimal as it does not sufficiently exploit the rich structure in the underlying neural network . For guidance , we consider the optimal approximate posterior over the top-layer units in a deep network for regression . Remarkably , the optimal approximate posterior for the last-layer weights given the earlier weights can be obtained in closed form without choosing a restrictive family of distributions . In particular , the optimal approximate posterior is given by propagating the training inputs through lower layers to compute the top-layer representation , then using Bayesian linear regression to map from the top-layer representation to the outputs . Inspired by this result , we use Bayesian linear regression to define a generic family of approximate posteriors for BNNs . In particular , we introduce learned “ pseudo-data ” at every layer , and compute the posterior over the weights by performing linear regression from the inputs ( propagated from lower layers ) onto the pseudo-data . We reduce the burden of working with many training inputs by summarising the posterior using a small number of “ inducing ” points . We find that these approximate posteriors give excellent performance in the non-tempered , no-data-augmentation regime , with performance on datasets such as CIFAR-10 reaching 86.7 % , comparable to SGMCMC ( Wenzel et al. , 2020 ) . Our approach can be extended to DGPs , and we explore connections to the inducing point GP literature , showing that inference in the two classes of models can be unified . 2 METHODS . We consider neural networks with lower-layer weights { W ` } L ` =1 , W ` ∈ RN ` −1×N ` , and top-layer weights , WL+1 ∈ RNL×NL+1 , where the activity , F ` , at layer ` is given by , F1 = XW1 , F ` = φ ( F ` −1 ) W ` for ` ∈ { 2 , . . . , L } , ( 1 ) where φ ( · ) is an elementwise nonlinearity . The outputs , Y ∈ RP×NL+1 , depend on the top-level activity , FL , and the output weights , WL+1 , according to a likelihood , P ( Y|WL+1 , FL ) . In the following derivations , we will focus on ` > 1 ; corresponding expressions for the input layer can be obtained by replacing φ ( F0 ) with the inputs , X ∈ RP×N0 . The prior over weights is independent across layers and output units ( see Sec . 2.3 for the form of S ` ) , P ( W ` ) = ∏N ` λ=1 P ( w ` λ ) , P ( w ` λ ) = N ( w ` λ ∣∣∣0 , 1N ` −1 S ` ) , ( 2 ) where w ` λ is a column of W ` , i.e . all the input weights to unit λ in layer ` . To fit the parameters of the approximate posterior , Q ( { W } L+1 ` =1 ) , we maximise the evidence lower bound ( ELBO ) , L = EQ ( { W } L+1 ` =1 ) [ log P ( Y , { W } L+1 ` =1 |X ) − log Q ( { W } L+1 ` =1 ) ] . ( 3 ) To build intuition about how to parameterise Q ( { W } L+1 ` =1 ) , we consider the optimal Q ( WL+1| { W ` } L ` =1 ) for any given Q ( { W ` } L ` =1 ) . We begin by simplifying the ELBO by incorporating terms that do not depend on WL+1 into a constant , c , L = EQ ( { W ` } L+1 ` =1 ) [ log P ( Y , WL+1|X , { W ` } L ` =1 ) − log Q ( WL+1| { W ` } L ` =1 ) + c ] . ( 4 ) Rearranging these terms , we find that all WL+1 dependence can be written in terms of the KL divergence between the approximate posterior of interest and the true posterior , L = EQ ( { W ` } L ` =1 ) [ log P ( Y|X , { W ` } L ` =1 ) − DKL ( Q ( WL+1| { W ` } L ` =1 ) ||P ( WL+1|Y , X , { W ` } L ` =1 ) ) + c ] . ( 5 ) Thus , the optimal approximate posterior is , Q ( WL+1| { W ` } L ` =1 ) = P ( WL+1|Y , X , { W ` } L ` =1 ) ∝ P ( Y|WL+1 , FL ) P ( WL+1 ) , ( 6 ) and where the final proportionality comes by applying Bayes theorem and exploiting the model ’ s conditional independencies . For regression , the likelihood is Gaussian , P ( Y|WL+1 , FL ) = ∏NL+1 λ=1 N ( yλ ; φ ( FL ) w L+1 λ , Λ −1 L+1 ) , ( 7 ) where yλ is the value of a single output channel for all training inputs , and ΛL+1 is a precision matrix . Thus , the posterior is given in closed form by Bayesian linear regression ( Rasmussen & Williams , 2006 ) . 2.1 DEFINING THE FULL APPROXIMATE POSTERIOR WITH GLOBAL INDUCING POINTS AND PSEUDO-DATA . We adapt the optimal scheme above to give a scalable approximate posterior over the weights at all layers . To avoid propagating all training inputs through the network , which is intractable for large datasets , we instead propagate M global inducing locations , U0 , U1 = U0W1 , U ` = φ ( U ` −1 ) W ` for ` = 2 , . . . , L+ 1 . ( 8 ) Next , the optimal posterior requires outputs , Y . However , no outputs are available at inducing locations for the output layer , let alone for intermediate layers . We thus introduce learned variational parameters to mimic the form of the optimal posterior . In particular , we use the product of the prior over weights and a “ pseudo-likelihood ” , N ( v ` λ ; u ` λ , Λ −1 ` ) , representing noisy “ pseudo-observations ” of the outputs of the linear layer at the inducing locations , u ` λ = φ ( U ` −1 ) w ` λ . Substituting u ` λ into the pseudo-likelihood the approximate posterior becomes , Q ( W ` | { W ` ′ } ` −1 ` ′=1 ) ∝ ∏N ` λ=1N ( v ` λ ; φ ( U ` −1 ) w ` λ , Λ −1 ` ) P ( w ` λ ) , Q ( W ` ∣∣∣ { W ` ′ } ` −1 ` ′=1 ) = ∏N ` λ=1N ( w ` λ∣∣∣Σw ` φ ( U ` −1 ) T Λ ` v ` λ , Σw ` ) , Σw ` = ( N ` −1S −1 ` + φ ( U ` −1 ) T Λ ` φ ( U ` −1 ) ) −1 . ( 9 ) where v ` λ and Λ ` are variational parameters . Therefore , our full approximate posterior factorises as Q ( { W ` } L+1 ` =1 ) = L+1∏ ` =1 Q ( W ` ∣∣∣ { W ` ′ } ` −1 ` ′=1 ) . ( 10 ) Thus , the full ELBO can be written , L = EQ ( { W } L+1 ` =1 ) [ log P ( Y|X , { W } L+1 ` =1 ) + log P ( { W ` } L+1 ` =1 ) − log Q ( { W ` } L+1 ` =1 ) ] ( 11 ) = EQ ( { W } L+1 ` =1 ) log P ( Y , |X , { W } L+1 ` =1 ) + L+1∑ ` =1 log P ( W ` ) Q ( W ` | { W ` } ` −1 ` =1 ) . ( 12 ) The forms of the ELBO and approximate posterior suggest a sequential procedure to evaluate and subsequently optimise it : we alternate between sampling the weights using Eq . ( 9 ) and propagating the data and inducing locations using Eq . ( 8 ) ( see Alg . 1 ) . In summary , the parameters of the approximate posterior are the global inducing inputs , U0 , and the pseudo-data and precisions at all layers , { V ` , Λ ` } L+1 ` =1 . As each factor Q ( W ` ∣∣∣ { W ` ′ } ` −1 ` ′=1 ) is a Gaussian , these parameters can be optimised using standard reparameterised variational inference ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) in combination with the Adam optimiser ( Kingma & Ba , 2014 ) ( Appendix A ) . Importantly , by placing inducing inputs on the training data ( i.e . U0 = X ) , and setting v ` λ = yλ this approximate posterior matches the optimal top-layer posterior ( Eq . 6 ) . Finally , we note that while this posterior is conditionally Gaussian , the full posterior over all { W ` } L+1 ` =1 is non-Gaussian , and is thus potentially more flexible than a full-covariance Gaussian over all weights . Algorithm 1 : Global inducing points for neural networks Parameters : inducing inputs , U0 , inducing outputs and precisions , { V ` , Λ ` } L ` =1 , at all layers . Neural network inputs : ( e.g . MNIST digits ) F0 Neural network outputs : ( e.g . classification logits ) FL+1 L+∞← 0 for ` ∈ { 1 , . . . , L+ 1 } do Compute the mean and covariance over the weights at this layer Σw ` = ( N ` −1S −1 ` + φ ( U ` −1 ) T Λ ` φ ( U ` −1 ) ) −1 M ` = Σ w ` φ ( U ` −1 ) T Λ ` V ` Sample the weights and compute the ELBO W ` ∼ N ( M ` , Σw ` ) = Q ( W ` ∣∣∣ { W ` ′ } ` −1 ` ′=1 ) L ← L+ log P ( W ` ) − logN ( W ` |M ` , Σw ` ) Propagate the inputs and inducing points using the sampled weights , U ` = φ ( U ` −1 ) W ` F ` = φ ( F ` −1 ) W ` L ← L+ log P ( Y|FL+1 )

Summary: This paper proposes a posterior approximation for BNN that models correlations between the layers weights. The paper begins by pointing out that, for any posterior distribution approximation, the optimal conditional posterior distribution over the top-layer weights given the weights of the previous layers has a closed-form in the case of Gaussian likelihood, which, due to Bayes rule, turns out to be product of the likelihood and the top-layer prior. Based on this insight the paper proposes to model each conditional posterior distribution over intermediate layers weights given previous layers weights following the same structure, that is, as a product of a 'pseudo-likelihood' over unobserved noisy activations and the prior for that layer. In order to make inference tractable, the paper proposes the use of global inducing points as well as noisy pseudo-observations of the activations of intermediate layers which are treated as variational parameters.  The paper also describes how such procedure applies to convolutional neural networks and how it can be applied for DGPs as well.