Amortised Inference in Bayesian Neural Networks

TL;DR

This paper introduces APOVI-BNN, a data-efficient Bayesian neural network approach that uses amortised inference for improved probabilistic meta-learning, especially with limited data, outperforming traditional methods.

Contribution

The paper presents APOVI-BNN, a novel amortised inference method for Bayesian neural networks that enhances data efficiency in probabilistic meta-learning tasks.

Findings

APOVI-BNN achieves comparable or better posterior quality than traditional variational inference.

The model performs best in limited data scenarios across regression and image completion tasks.

APOVI-BNN can be viewed as a new neural process family member, with potential for improved predictive performance.

Abstract

Meta-learning is a framework in which machine learning models train over a set of datasets in order to produce predictions on new datasets at test time. Probabilistic meta-learning has received an abundance of attention from the research community in recent years, but a problem shared by many existing probabilistic meta-models is that they require a very large number of datasets in order to produce high-quality predictions with well-calibrated uncertainty estimates. In many applications, however, such quantities of data are simply not available. In this dissertation we present a significantly more data-efficient approach to probabilistic meta-learning through per-datapoint amortisation of inference in Bayesian neural networks, introducing the Amortised Pseudo-Observation Variational Inference Bayesian Neural Network (APOVI-BNN). First, we show that the approximate posteriors obtained under our amortised scheme are of similar or better quality to those obtained through traditional variational inference, despite the fact that the amortised inference is performed in a single forward pass. We then discuss how the APOVI-BNN may be viewed as a new member of the neural process family, motivating the use of neural process training objectives for potentially better predictive performance on complex problems as a result. Finally, we assess the predictive performance of the APOVI-BNN against other probabilistic meta-models in both a one-dimensional regression problem and in a significantly more complex image completion setting. In both cases, when the amount of training data is limited, our model is the best in its class.