Amortised Inference in Neural Networks for Small-Scale Probabilistic
  Meta-Learning
The global inducing point variational approximation for BNNs is based on
using a set of inducing inputs to construct a series of conditional
distributions that accurately approximate the conditionals of the true
posterior distribution. Our key insight is that these inducing inputs can be
replaced by the actual data, such that the variational distribution consists of
a set of approximate likelihoods for each datapoint. This structure lends
itself to amortised inference, in which the parameters of each approximate
likelihood are obtained by passing each datapoint through a meta-model known as
the inference network. By training this inference network across related
datasets, we can meta-learn Bayesian inference over task-specific BNNs.
Matthew Ashman, Tommy Rochussen, Adrian Weller
[

###### Abstract

In many machine learning applications, well-calibrated posterior predictive distributions are required for a number of closely-related datasets. Given similarity between datasets, it is natural to wish to develop meta-learning algorithms that utilise other datasets to reduce the computational complexity and / or improve predictive performance when deploying models on newly-seen datasets at test time. There have been a number of significant recent developments in meta-learning for predictive distributions, most notably that of the neural process (NP) family (Garnelo et al., 2018, 2018; Foong et al., 2020; Gordon et al., 2018, 2019). Despite the utility of these methods on large-scale meta-datasets, they perform poorly in settings where the number of datasets and the total number of datapoints is small. We argue that this is a result of the large number of shared model parameters overfitting to the meta-dataset. A natural solution is to remove these shared model parameters, and instead train a meta-model to learn to approximate fully Bayesian inference over task-specific model parameters.

Recently, Ober and Aitchison (2021) developed a variational approximation for Bayesian neural networks (BNNs) based on using a set of inducing inputs to construct a series of conditional distributions that accurately approximate the conditionals of the true posterior distribution. Notably, the variational distribution consists of the prior multiplied by a set of approximate likelihoods for each inducing input. Our key insight is that these inducing inputs can be replaced by the actual data, such that the variational distribution consists of a set of approximate likelihoods for each datapoint. This structure lends itself to amortised inference, in which the parameters of each approximate likelihood are obtained by passing each datapoint through a meta-model known as the inference network. By training this inference network across related datasets, we can meta-learn Bayesian inference over task-specific BNNs, addressing the challenge above.

A]Amortised Inference in Neural Networks for Small-Scale Probabilistic Meta-Learning Matthew Ashman]Matthew Ashman\({}^{*}\) mca39@cam.ac.uk

Tommy Rochussen\({}^{*}\) tnr22@cam.ac.uk

A]Marian Weller aw665@cam.ac.uk

## 1 Introduction

In many machine learning applications, well-calibrated posterior predictive distributions are required for a number of closely-related datasets. Given similarity between datasets, it is natural to wish to develop meta-learning algorithms that utilise other datasets to reduce the computational complexity and / or improve predictive performance when deploying models on newly-seen datasets at test time. There have been a number of significant recent developments in meta-learning for predictive distributions, most notably that of the neural process (NP) family (Garnelo et al., 2018, 2018; Foong et al., 2020; Gordon et al., 2018, 2019). Despite the utility of these methods on large-scale meta-datasets, they perform poorly in settings where the number of datasets and the total number of datapoints is small. We argue that this is a result of the large number of shared model parameters overfitting to the meta-dataset. A natural solution is to remove these shared model parameters, and instead train a meta-model to learn to approximate fully Bayesian inference over task-specific model parameters.

Recently, Ober and Aitchison (2021) developed a variational approximation for Bayesian neural networks (BNNs) based on using a set of inducing inputs to construct a series of conditional distributions that accurately approximate the conditionals of the true posterior distribution. Notably, the variational distribution consists of the prior multiplied by a set of approximate likelihoods for each inducing input. Our key insight is that these inducing inputs can be replaced by the actual data, such that the variational distribution consists of a set of approximate likelihoods for each datapoint. This structure lends itself to amortised inference, in which the parameters of each approximate likelihood are obtained by passing each datapoint through a meta-model known as the inference network. By training this inference network across related datasets, we can meta-learn Bayesian inference over task-specific BNNs, addressing the challenge above.

## 2 Related Work

Neural processesOur work is most similar to the NP family (Garnelo et al., 2018, 2018), which seeks to meta-learn predictive distributions either through maximisation of the posterior predictive likelihood or variational inference (VI) (Foong et al., 2020). Similar to our method, NPs utilise an encoder to create embeddings for each datapoint. These embeddings
are then aggregated to form a distribution over a latent variable which is then sampled and passed, together with a test datapoint, through a decoder. Volpp et al. (2020) propose the use of Bayesian aggregation, in which embeddings of individual datapoints take the form of approximate likelihoods which are multiplied together with the prior to form an approximate posterior distribution over the latent variable. Whilst these methods differ to ours in their use of shared model parameters, our method can be reinterpreted as a member of the NP family in which the latent variables are the parameters of the decoder. Through this perspective we can train our model in an identical way to NPs.

Meta-learning neural networksMeta-learning for neural networks had received a significant amount of attention from the research community. Notable examples include MAML Finn et al. (2017) and its extensions (Yoon et al., 2018; Antoniou et al., 2018), which seek good parameter initialisation, and those which explicitly condition on the dataset to obtain task specific parameters (Requeira et al., 2019; Gordon et al., 2018). Whilst conceptually similar to our approach, these methods differ in their use of shared model parameters--the task specific parameters amount to a small subset of the overall model parameters. In addition to requiring a large meta-dataset, this limits these methods to meta-datasets in which the individual datasets are very similar. By contrast, our approach does not use any shared model parameters but rather meta-learns inference in a BNN. We discuss the relationship between our method and NPs in Appendix A, demonstrating that they conceptually differ only in a change of objective functions.

## 3 Background

In this section we review the GI-BNN of Ober and Aitchison (2021) and NP family (Garnelo et al., 2018). Throughout, let \(\mathbf{\Xi}=\{\mathcal{D}\}\) denote a meta-dataset of \(|\Xi|\) datasets, and \(\mathcal{D}=\{\mathbf{X},\mathbf{y}\}\) denote a dataset consisting of inputs \(\mathbf{X}\in\mathbb{R}^{N\times D_{0}}\) and outputs \(\mathbf{y}\in\mathbb{R}^{N\times P}\).

### Global Inducing Point Variational Posteriors for BNNs

Let \(\mathbf{W}=\{\mathbf{W}^{\ell}\}_{l=1}^{L}\) denote the weights of an \(L\)-layer neural network, such that \(\mathbf{W}^{\ell}\in\mathbb{R}^{D^{\ell-1}\times D^{\ell}}\) where \(D^{\ell}\) denotes the dimensions of the \(\ell\)-th hidden layer, and \(\psi(\cdot)\) denote the element-wise activation function acting between layers. Ober and Aitchison (2021) introduce the global inducing point variational approximation for BNNs, in which the variational approximation to the posterior \(p(\mathbf{W}|\mathcal{D})\) is defined recursively as \(q_{\phi}(\mathbf{W})=\prod_{\ell=1}^{L}q_{\phi}(\mathbf{W}^{\ell}|\{\mathbf{W }^{\ell}\}_{\ell^{\prime}=1}^{\ell-1},\mathbf{U}^{0})\), where

\[q_{\phi}\left(\mathbf{W}^{\ell}|\{\mathbf{W}^{\ell^{\prime}}\}_{\ell^{\prime }=1}^{\ell-1},\mathbf{U}_{0}\right)\propto\prod_{d=1}^{D^{\ell}}p\left( \mathbf{w}_{d}^{\ell}\right)\underbrace{\mathcal{N}\left(\mathbf{v}_{d}^{\ell };\psi(\mathbf{U}^{\ell})\mathbf{w}_{d}^{\ell},\left[\mathbf{\Lambda}_{d}^{ \ell}\right]^{-1}\right)}_{t_{d}^{\ell}(\mathbf{w}_{d}^{\ell})}. \tag{1}\]

This mirrors the structure of the true posterior in the sense that it is equivalent to the product of the prior and an _approximate likelihood_, \(t_{d}^{\ell}(\mathbf{w}_{d}^{\ell})\). The variational parameters \(\phi\) of this approximation are the parameters of each approximate likelihood, \(\mathbf{v}_{d}^{\ell}\in\mathbb{R}^{M}\) and \(\mathbf{\Lambda}_{d}^{\ell}\in\mathbb{R}^{M\times M}\)--which themselves can be interpreted as _pseudo observations_--and the _global
inducing locations_, \(\mathbf{U}_{0}\in\mathbb{R}^{M\times D_{0}}\), which are used to define \(\{\mathbf{U}^{\ell}\}_{\ell=1}^{L}\) according to

\[\mathbf{U}^{1}=\mathbf{U}^{0}\mathbf{W}_{1},\quad\mathbf{U}^{\ell}=\psi( \mathbf{U}^{\ell-1})\mathbf{W}^{\ell}\quad\ell=2,\ldots,L. \tag{2}\]

Optimisation of \(\phi\) is achieved through maximisation of the evidence lower bound (ELBO):

\[\mathcal{L}_{\mathrm{ELBO}}(\phi;\mathcal{D})=\mathbb{E}_{q_{\phi}(\mathbf{W}) }\left[\log p(\mathbf{y}|\mathbf{W},\mathbf{X})\right]-\mathrm{KL}\left[q_{ \phi}(\mathbf{W})||p(\mathbf{W})\right]. \tag{3}\]

We refer to this variational approximation as pseudo-observation variational inference for BNNs (POVI-BNN).

Ober and Aitchison (2021) demonstrate the efficacy of POVI-BNNs relative to mean-field Gaussian variational approximations for BNNs, achieving state-of-the-art performance on a number of regression and classification experiments. Their effectiveness is further demonstrated by Bui (2022), who shows that the estimate of the marginal likelihood provided by the POVI-BNN approximation is close to the true value, indicating the approximation is close to the true posterior.

## 4 Amortising Inference in Bayesian Neural Networks

In this section, we build upon the work of Ober and Aitchison (2021) to develop an effective method of performing amortised inference in BNNs.

Consider the same variational approximation described in Section 3.1, except with diagonal precision matrices \(\boldsymbol{\Lambda}_{d}^{\ell}\) and inducing locations \(\mathbf{U}\) replaced by training locations \(\mathbf{X}\), such that

\[q\left(\mathbf{W}^{\ell}|\{\mathbf{W}^{\ell^{\prime}}\}_{\ell^{\prime}=1}^{ \ell-1},\mathcal{D}\right)\propto\prod_{d=1}^{D_{\ell}}p\left(\mathbf{w}_{d}^ {\ell}\right)\prod_{n=1}^{N}\underbrace{\mathcal{N}\left(v_{d,n}^{\ell};x_{d,n }^{\ell},\sigma_{d,n}^{\ell}\right)}_{t_{d,n}^{\ell}(\mathbf{w}_{d}^{\ell})} \tag{4}\]

where

\[\mathbf{x}_{n}^{1}=\mathbf{W}^{1}\mathbf{x}_{n},\quad\mathbf{x}_{n}^{\ell}= \mathbf{W}^{\ell}\psi(\mathbf{x}_{n}^{\ell-1})\quad\forall\ell=2,\ldots,L. \tag{5}\]

This form of variational approximation enables the use of per-datapoint amortised inference. Specifically, rather than treating the variational parameters \(\{\{\{\{\{\{\{\{\{\{\{\{\{\{{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{ }}}}}}}}}}}}}\}\}\}\} \mathbf{\}}\mathbf{\{\{\{{{\{{\{{{\{{\mathbf{\mathbf{{\mathbf{{\mathbf{{\mathbf{ \mathbf{{\mathbf{\mathbf{{}}}}}}}}}}}}}}}\}\}\}\}}}}\mathbf{\{\{{{{{{ \{{\{{{\mathbf{\mathbf{{\mathbf{{\mathbf{{\mathbf{\mathbf{{\mathbf{{\mathbf{\mathbf{{ }}}}}}}}}}}}}}\mathbf{{}}}}\}}}}}}}}}\mathbf{\{ \{{{{{\{{{\{{{\mathbf{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{ \mathbf{{\mathbf{{\mathbf{{\mathbf{{}}}}}}}}}}}}}}}}}\mathbf{}}}\}\}}}}}}} }\mathbf{\{\{{{{{\{{{\{{{\{{\mathbf{\mathbf{{\mathbf{{\mathbf{{\mathbf{{ \mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{ }}}}}}}}}}}}}}}}}}\mathbf{{}}}\mathbf{{{{{{ {\mathbf{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{ }}}}}}}}}}}}}\mathbf{{{{{{{{ \mathbf{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}} \mathbf{\{\{{{{{{\mathbf{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{ \mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{{\mathbf{{ }}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}\mathbf{\{ \{{{{{\{{{\{{{\{{{\{{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{ \mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{\mathbf{} }}                 }}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}\}\\\\\\}\}\\\\\\\\\\\}\\\\\\\\\\\}}\\\\\\\\\\\\\            \}\ \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\\\\\\}\\\\\\\\}\\\\\\\\\\\\\\}\\\\\\\\\\\\\\\\}\\\\\\\\\\\\\\}\\\\\\\\\}\\\\\\\\\\}\\\\\\\\\}\\\\\\\\\\\\\}\\\\\\\\\\\}\\\\\\\\\\\\}\\\\\\\\\\\}\\\\\\\\\\\}\\\\\\\\\\\}\\\\\\\\\\\}\\\\\\\\\\}\\\\\\\\\}\\\\\\\\\\}\\\\\\\\\\}\\\\\\\\\}\\\\\\\\\\}\\\\\\\\\}\\\\\\\\\\}\\\\\\\\\}\\\\\\\\\\\\}\\\\\\\\}\\\\\\\\\\\\}\\\\\\\\}\\\\\\\\\\\}\\\\\\\\
An important limitation of this approach is that performing stochastic optimisation through mini-batching datapoints is not possible, as to compute the \(q(\mathbf{W}|\mathcal{D})\) we require passing the entire dataset \(\mathcal{D}\) through the inference network. Nonetheless, this limitation is only significant for large datasets--provided the entire dataset can be passed through the network at once (i.e. in the small to medium-sized dataset regime which we consider here) this is not an issue. At test time, we can obtain an approximate posterior \(q(\mathbf{W}|\mathcal{D}_{*})\) with a single pass of the dataset through the inference networks.

## 5 Results and Discussion

In this section, we evaluate the performance of our model in the meta-learning setting. We consider a synthetic meta-dataset consisting of samples from a GP with a squared-exponential (SE) kernel. Each dataset consists of between 10 and 20 training datapoints. We compare A-POVI-BNNs to amortised mean-field variational inference for BNNs (A-MFVI-BNN), which we detail in Appendix B, and a ConvCNP (Gordon et al., 2019). All

Figure 1: Posterior predictive distributions for A-POVI-BNN (top), A-MFVI-BNN (middle), and ConvCNP (bottom) after training on a meta-dataset of size \(|\Xi|=1\) (left) and \(|\Xi|=100\) (right). Data points are shown as black dots, the predictive mean is shown as a dark blue line and the 95% confidence interval is shown as shaded blue.
NNs used in the amortised BNN architectures (both the model and inference network) consist of two layers of 50 hidden units and ReLU activation functions. The ConvCNP implementation and architecture is identical to that provided in [https://github.com/cambridge-mlg/convcnp/blob/master/convcnp_regression.ipynb](https://github.com/cambridge-mlg/convcnp/blob/master/convcnp_regression.ipynb).

Figures 1(_a_) to 1(_f_) compare the posterior predictive distributions of each method on an unseen dataset drawn from a GP with the same hyperparameters as those in the meta-dataset. We evaluate the predictions of each method trained on a meta-dataset of size \(|\Xi|=1\) and \(|\Xi|=100\). We see that only A-POVI-BNN obtains a sensible predictive posterior on the unseen dataset in both cases. The ConvCNP performs poorly for \(|\Xi|=1\), which is unsurprising given the large number of model parameters increasing it susceptibility to overfitting. The A-MFVI-BNN performs significantly better for \(|\Xi|=100\), yet, in both cases the quality of its predictive posterior is poor relative to both the A-POVI-BNN and ConvCNP.

Despite these results being very preliminary, they are encouraging and suggest that A-POVI-BNNs may provide a more effective alternative to NPs when the size of meta-datasets are small. We intend to explore the effectiveness of A-POVI-BNNs in more diverse settings, such as image completion, in future work.
