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

We adapt the optimal top-layer approximate posterior 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, $\mathbf{U}_{0}\in\mathbb{R}^{M\times N_{0}}$,

Next, the optimal posterior requires outputs, $\mathbf{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 an “inducing-likelihood”, $\mathcal{N}\left(\mathbf{v}_{\lambda}^{\ell};\mathbf{u}_{\lambda}^{\ell},\mathbf{\Lambda}_{\ell}^{-1}\right)$, representing noisy “pseudo-outputs” of the outputs of the linear layer at the inducing locations, $\mathbf{u}_{\lambda}^{\ell}=\phi\left(\mathbf{U}_{\ell-1}\right)\mathbf{w}_{\lambda}^{\ell}$. Substituting $\mathbf{u}_{\lambda}^{\ell}$ into the inducing-likelihood the approximate posterior becomes,

where $\mathbf{v}_{\lambda}^{\ell}$ and $\mathbf{\Lambda}_{\ell}$ are variational parameters. For clarity, we have: $\mathbf{u}_{\lambda}^{\ell},\mathbf{v}_{\lambda}^{\ell}\in\mathbb{R}^{M}$, so that $\mathbf{U}_{\ell-1}\in\mathbb{R}^{M\times N_{\ell-1}}$ and $\mathbf{V}_{\ell}\in\mathbb{R}^{M\times N_{\ell}}$ are formed by stacking these vectors, and $\mathbf{w}_{\lambda}^{\ell}\in\mathbb{R}^{N_{\ell-1}}$, with $\mathbf{S}_{\ell},\mathbf{\Sigma}^{\mathbf{w}}_{\ell}\in\mathbb{R}^{N_{\ell-1}\times N_{\ell-1}}$ and $\mathbf{\Lambda}_{\ell}\in\mathbb{R}^{M\times M}$. Therefore, our full approximate posterior factorises as

Substituting this approximate posterior and the factorised prior into the ELBO (Eq. 3), the full ELBO can be written,

where $\operatorname{P}\left(\mathbf{W}_{\ell}\right)$ is given by Eq. (2) and $\operatorname{Q}\left(\mathbf{W}_{\ell}|\left\{\mathbf{W}_{\ell}\right\}_{\ell=1}^{\ell-1}\right)$ is given by Eq. (11). 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. (11) and propagating the data and inducing points (Eq. 1 and Eq. 10; see Alg. 1). In summary, the parameters of the approximate posterior are the global inducing inputs, $\mathbf{U}_{0}$, and the pseudo-outputs and precisions at all layers, $\{\mathbf{V}_{\ell},\mathbf{\Lambda}_{\ell}\}_{\ell=1}^{L+1}$. As each factor $\operatorname{Q}\left(\mathbf{W}_{\ell}\middle|\left\{\mathbf{W}_{\ell^{\prime}}\right\}_{\ell^{\prime}=1}^{\ell-1}\right)$ is 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 B). Importantly, by placing inducing inputs on the training data (i.e. $\mathbf{U}_{0}=\mathbf{X}$), and setting $\mathbf{v}_{\lambda}^{\ell}=\mathbf{y}_{\lambda}$ 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 $\{\mathbf{W}_{\ell}\}_{\ell=1}^{L+1}$ is non-Gaussian, and is thus potentially more flexible than a full-covariance Gaussian defined jointly over all weights at all layers.

The previous sections were valid for a fully connected network. The extension to convolutional networks is straightforward in principle: we transform the convolution into a matrix multiplication by treating each patch as a separate input feature-vector, flattening the spatial and channel dimensions together into a single vector. Thus, the feature-vectors have length in_channels $\times$ kernel_width $\times$ kernel_height, and the matrix $\mathbf{U}_{\ell}$ contains patches_per_image $\times$ minibatch patches. Likewise, we now have inducing outputs, $\mathbf{v}_{\lambda}^{\ell}$, at each location in all the inducing images, so this again has length patches_per_image $\times$ minibatch. After explicitly extracting the patches, we can straightforwardly apply standard Bayesian linear regression.

However, explicitly extracting image patches is very memory intensive in a DNN. If we consider a standard convolution with a $3\times 3$ convolutional kernel, then there is a $3\times 3$ patch centred at each pixel in the input image, meaning a factor of $9$ increase in memory consumption. Instead, we note that computing the matrices required for linear regression, $\phi\left(\mathbf{U}_{\ell-1}\right)^{T}\mathbf{\Lambda}_{\ell}\phi\left(\mathbf{U}_{\ell-1}\right)$ and $\phi\left(\mathbf{U}_{\ell}\right)^{T}\mathbf{\Lambda}_{\ell}\mathbf{V}_{\ell}$, does not require explicit extraction of image-patches. Instead, these matrices can be computed by taking the autocorrelation of the image/feature map, i.e. a convolution operation where we treat the image/feature map, as both the inputs and the weights (Appendix C for details).

We consider four priors in this work, which we refer to using the class names in the bayesfunc library published alongside this paper. We are careful to ensure that all parameters in the model have a prior and approximate posterior, which is necessary to ensure that ELBOs are comparable across models.

First, we consider a Gaussian prior with fixed scale, NealPrior, so named because it is necessary to obtain meaningful results when considering infinite networks (Neal, 1996),

though it bears strong similarities to the “He” initialisation (He et al., 2015). NealPrior is defined so as to ensure that the activations retain a sensible scale as they propagate through the network. We compare this to the standard $\mathcal{N}(0,1)$ (StandardPrior), which causes the activations to increase exponentially as they propagate through network layers (see Eq. 2):

Next, we consider ScalePrior, which defines a prior and approximate posterior over the scale,

where here we parameterise the Gamma distribution with the shape and rate parameters, and $\alpha_{\ell}$ and $\beta_{\ell}$ are non-negative learned parameters of the approximate posterior over $s_{\ell}$. Finally, we consider SpatialIWPrior, which allows for spatial correlations in the weights (e.g. see Fortuin et al., 2021, for a more restrictive spatial prior over weights). In particular, we take the covariance to be the Kronecker product of an identity matrix over channel dimensions, and a Wishart-distributed matrix, $\mathbf{L}^{-1}_{\ell}$, over the spatial dimensions,

where $\mathcal{W}^{-1}$ is the inverse-Wishart distribution, and the non-negative real number, $\nu$, and the positive definite matrix, $\mathbf{\Psi}$, are learned parameters of the approximate posterior (see Appendix D).

It is a remarkable but underappreciated fact that BNNs are special cases of DGPs, with a particular choice of kernel (Louizos & Welling, 2016; Aitchison, 2019). Combining Eqs. (1) and (2),

Here, we generalise our approximate posterior to the DGP case and link to the DGP literature. In a DGP there are no weights; instead we work directly with inducing outputs $\{\mathbf{U}_{\ell}\}_{\ell=1}^{L+1}$,

Note that here, we take the “global” inducing approach of using the inducing outputs from the previous layer, $\mathbf{U}_{\ell-1}$ as the inducing inputs for the next layer. In this case, we need only learn the original inducing inputs, $\mathbf{U}_{0}$. This contrasts with the standard “local” inducing formulation, (as in Salimbeni & Deisenroth, 2017), which learns separate inducing inputs at every layer, $\mathbf{Z}_{\ell-1}$, giving $\operatorname{P}\left(\mathbf{U}_{\ell}\middle|\mathbf{Z}_{\ell-1}\right)={\textstyle\prod_{\lambda=1}^{N_{\ell}}}\mathcal{N}\left(\mathbf{u}_{\lambda}^{\ell}\middle|{\bf{0}},\mathbf{K}\left(\mathbf{Z}_{\ell-1}\right)\right)$.

As usual in DGPs (Salimbeni & Deisenroth, 2017), the approximate posterior over $\mathbf{U}_{\ell}$ induces an approximate posterior on $\mathbf{F}_{\ell}$ through the prior correlations. However, it is important to remember that underneath the tractable distributions in Eqs. (20) and (21), there is an infinite dimensional GP-distributed function, $\mathcal{F}_{\ell}$, such that $\mathbf{F}_{\ell}=\mathcal{F}_{\ell}(\mathbf{F}_{\ell-1})$. Standard local inducing point methods specify a factorised approximate posterior over $\mathcal{F}_{\ell}$ by specifying the function’s inducing outputs, $\mathbf{U}_{\ell}=\mathcal{F}_{\ell}\left(\mathbf{Z}_{\ell-1}\right)$, at a finite number of inducing input locations, $\mathbf{Z}_{\ell-1}$. Importantly, the approximate posterior over a function, $\mathcal{F}_{\ell}$, depends only on $\mathbf{Z}_{\ell-1}$, and $\mathbf{U}_{\ell}$. Thus, standard, local inducing, DGP approaches (e.g.  Salimbeni & Deisenroth, 2017), give a layerwise-independent approximate posterior over $\{\mathcal{F}_{\ell}\}_{\ell=1}^{L+1}$, as they treat the inducing inputs, $\{\mathbf{Z}_{\ell-1}\}_{\ell=1}^{L+1}$, as fixed, learned parameters and use a layerwise-independent approximate posterior over $\{\mathbf{U}_{\ell}\}_{\ell=1}^{L+1}$ (Appendix G).

Next, we need to choose the approximate posterior on $\{\mathbf{U}_{\ell}\}_{\ell=1}^{L+1}$. However, if our goal is to introduce dependence across layers, it seems inappropriate to use the standard layerwise-independent approximate posterior over $\{\mathbf{U}_{\ell}\}_{\ell=1}^{L+1}$. Indeed, in Appendix G, we show that such a posterior implies functions in non-adjacent layers (e.g. $\mathcal{F}_{\ell}$ and $\mathcal{F}_{\ell+2})$ are marginally independent, even with global inducing points.

To obtain more appropriate approximate posteriors, we consider the optimal top-layer posterior for DGPs, which involves GP regression from activations propagated from lower layers onto the output data (Appendix F). Inspired by the form of the optimal posterior we again define an approximate posterior by taking the product of the prior and a “inducing-likelihood”,

where $\mathbf{v}_{\lambda}^{\ell}$ and $\mathbf{\Lambda}_{\ell}^{-1}$ are learned parameters, and in our global inducing method, the inducing inputs, $\mathbf{U}_{\ell-1}$, are propagated from lower layers (Eq. 21). Importantly, setting the inducing inputs to the training data and $\mathbf{v}_{\lambda}^{L+1}=\mathbf{y}_{\lambda}$, the approximate posterior captures the optimal top-layer posterior for regression (Appendix F). Under this approximate posterior, dependencies in $\mathbf{U}_{\ell}$ naturally arise across all layers, and hence there are dependencies between functions $\mathcal{F}_{\ell}$ at all layers (Appendix G).

In summary, we propose an approximate posterior over inducing outputs that takes the form

As before, the parameters of this approximate posterior are the global inducing inputs, $\mathbf{U}_{0}$, and the pseudo-outputs and precisions at all layers, $\{\mathbf{V}_{\ell},\mathbf{\Lambda}_{\ell}\}_{\ell=1}^{L+1}$. The full ELBO, (see App.E for further details), takes the form

where $\operatorname{P}\left(\mathbf{U}_{\ell}\middle|\mathbf{U}_{\ell-1}\right)$ is given by Eq. (21) and $\operatorname{Q}\left(\mathbf{U}_{\ell}\middle|\mathbf{U}_{\ell-1}\right)$ is given by Eq. (22).

We provide a full description of our method as applied to DGPs in App. E.

In the deep GP case, the complexity for global inducing is exactly that of standard inducing point Gaussian processes, i.e. $\mathcal{O}(M^{3}+PM^{2})$ where $M$ is the number of inducing points, and $P$ can be taken to be the number of training inputs, or the size of a minibatch, as appropriate. The first term, $M^{3}$, comes from computing and sampling the posterior over $\mathbf{U}_{\ell}$ based on the inducing points (e.g. inverting the covariance). The second term, and $PM^{2}$ comes from computing the implied distribution over $\mathbf{F}_{\ell}$.

In the fully-connected BNN case, we have three terms, $\mathcal{O}(N^{3}+MN^{2}+PN^{2})$. The first term, $N^{3}$, where $N$ corresponds to the width of the network, arises from taking the inverse of the covariance matrix in Eq. (11), but is also the complexity e.g. for propagating the inducing points from layer to the next (Eq. 10). The second term, $MN^{2}$, comes from computing that covariance in Eq. (11), by taking the product of input features with themselves. The third term $PN^{2}$ comes from multiplying the training inputs/minibatch by the sampled inputs (Eq. 1).

We demonstrate the use of local and global inducing point methods in a toy 1-D regression problem, comparing it with fully factorised VI and Hamiltonian Monte Carlo (HMC; (Neal et al., 2011)). Following Hernández-Lobato & Adams (2015), we generate 40 input-output pairs $(x,y)$ with the inputs $x$ sampled i.i.d. from $\mathcal{U}([-4,-2]\cup[2,4])$ and the outputs generated by $y=x^{3}+\epsilon$, where $\epsilon\sim\mathcal{N}(0,3^{2})$. We then normalised the inputs and outputs. Note that we have introduced a ‘gap’ in the inputs, following recent work (Foong et al., 2019b; Yao et al., 2019; Foong et al., 2019a) that identifies the ability to express ‘in-between’ uncertainty as an important quality of approximate inference algorithms. We evaluated the inference algorithms using fully-connected BNNs with 2 hidden layers of 50 ReLU hidden units, using the NealPrior. For the inducing point methods, we used 100 inducing points per layer.

The predictive distributions for the toy experiment can be seen in Fig. 1. We observe that of the variational methods, the global inducing method produces predictive distributions closest to HMC, with good uncertainty in the gap. Meanwhile, factorised and local inducing fit the training data, but do not produce reasonable error bars, demonstrating an important limitation of methods lacking correlation structure between layers.

We provide additional experiments looking at the effect of the number of inducing points in Appendix I, and experiments looking at compositional uncertainty (Ustyuzhaninov et al., 2020) in both BNNs and DGPs for 1D regression in Appendix J.

The lack of correlations between layers might be expected to become more problematic in deeper networks. To isolate the effect of depth on different approximate posteriors, we considered deep linear networks trained on data generated from a toy linear model: 5 input features were mapped to 1 output feature, where the 1000 training and 100 test inputs are drawn IID from a standard Gaussian, and the true outputs are drawn using a weight-vector drawn IID from a Gaussian with variance $1/5$, and with noise variance of $0.1$. We could evaluate the model evidence under the true data generating process which forms an upper bound (in expectation) on the model evidence and ELBO for all models.

We found that the ELBO for methods that factorise across layers — factorised and local inducing — drops rapidly as networks get deeper and wider (Fig. 2). This is undesirable behaviour, as we know that wide, deep networks are necessary for good performance on difficult machine learning tasks. In contrast, we found that methods with global inducing points at the last layer decay much more slowly with depth, and perform better as networks get wider. Remarkably, global-inducing points gave good performance even with lower-layer weights drawn at random from the prior, which is not possible for any method that factorises across layers. We believe that fac $\rightarrow$ gi performed poorly at width $=250$ due to optimisation issues as rand $\rightarrow$ gi performs better yet is a special case of fac $\rightarrow$ gi.

We benchmark our methods on the UCI datasets in Hernández-Lobato & Adams (2015), popular benchmark regression datasets for BNNs and DGPs. Following the standard approach (Gal & Ghahramani, 2015), each dataset uses 20 train-test ‘splits’ (except for protein with 5 splits) and the inputs and outputs are normalised to have zero mean and unit standard deviation. We focus on the five smallest datasets, as we expect Bayesian methods to be most relevant in small-data settings (see App. K andM for all datasets). We consider two-layer fully-connected ReLU networks, using fully factorised and global inducing approximating families, as well as two- and five-layer DGPs with doubly-stochastic variational inference (DSVI) (Salimbeni & Deisenroth, 2017) and global inducing. For the BNNs, we consider the standard $\mathcal{N}(0,1)$ prior and ScalePrior.

We display ELBOs and average test log likelihoods for the un-normalised data in Fig. 3, where the dots and error bars represent the means and standard errors over the test splits, respectively. We observe that global inducing obtains better ELBOs than factorised and DSVI in almost every case, indicating that it does indeed approximate the true posterior better (since the ELBO is the marginal likelihood minus the KL to the posterior). While this is the case for the ELBOs, this does not always translate to a better test log likelihood due to model misspecification, as we see that occasionally DSVI outperforms global inducing by a very small margin. The very poor results for factorised on ScalePrior indicate that it has difficulty learning useful prior hyperparameters for prediction (see also Blundell et al. (2015)), which is due to the looseness of its bound to the marginal likelihood. We provide experimental details, as well as additional results with additional architectures, priors, datasets, and RMSEs, in Appendices K and M, for BNNs and DGPs, respectively.

For CIFAR-10, we considered a ResNet-inspired model consisting of conv2d-relu-block-avgpool2-block-avgpool2-block-avgpool-linear, where the ResNet blocks consisted of a shortcut connection in parallel with conv2d-relu-conv2d-relu, using 32 channels in all layers. In all our experiments, we used no data augmentation and 500 inducing points. Our training scheme (see App. O) ensured that our results did not reflect a ‘cold posterior’ (Wenzel et al., 2020). Our results are shown in Table 1. We achieved remarkable performance of $86.7\%$ predictive accuracy, with global inducing points used for all layers, and with a spatial inverse Wishart prior on the weights. These results compare very favourably with comparable Bayesian approaches, i.e. those without data augmentation or posterior sharpening: past work with deep GPs obtained $80.3\%$ (Shi et al., 2019), and work using infinite-width neural networks to define a GP obtained $81.4\%$ accuracy (Li et al., 2019). Remarkably, with only 500 inducing points we are approaching the accuracy of sampling-based methods (Wenzel et al., 2020), which are in principle able to more closely approximate the true posterior. Furthermore, we see that global inducing performs the best in terms of ELBO (per datapoint) by a wide margin, demonstrating that it gets far closer to the true posterior than the other methods. We provide additional results on uncertainty calibration and out-of-distribution detection in Appendix L.

Finally, while we have focused this work on achieving good results with a fully principled, Bayesian approach, we briefly consider training a full ResNet-18 (He et al., 2016) using more popular techniques such as data augmentation and tempering. While data augmentation and tempering are typically viewed as clouding the Bayesian perspective, there is work attempting to formalise both within the context of modified probabilistic generative models (Aitchison, 2020; Nabarro et al., 2021). Using a cold posterior with a factor of 20 reduction on the KL term, with horizontal flipping and random cropping, we obtained a test accuracy of $87.52\%$ and a test log likelihood of $-0.39$ nats with a standard fully factorised Gaussian posterior using SpatialIWPrior. Using fac $\rightarrow$ gi, we obtain a significant improvement of $92.41\%$ test accuracy and a test log likelihood of $-0.24$ nats. We attempted to train a full global inducing model; however, we encountered difficulties in scaling the method to the large widths of ResNet-18 layers. We believe this presents an avenue for fruitful future work in scaling global inducing point posteriors.

We would like to thank the anonymous reviewers for their valuable feedback which greatly improved the manuscript. We thank David R. Burt, Andrew Y. K. Foong, Siddharth Swaroop, Adrià Garriga-Alonso, Vidhi Lalchand, and Carl E. Rasmussen for valuable discussions and feedback. SWO acknowledges the support of the Gates Cambridge Trust in funding his doctoral studies. LA would like to thank the University of Bristol’s Advanced Computing Research Centre (ACRC) for computational resources.