Latent Derivative Bayesian Last Layer Networks

Joe Watson*, Jihao Andreas Lin, Pascal Klink, Joni Pajarinen, Jan Peters

*Corresponding author for this work

Research output: Chapter in Book/Report/Conference proceedingConference article in proceedingsScientificpeer-review

58 Downloads (Pure)

Abstract

Bayesian neural networks (BNN) are powerful parametric models for nonlinear regression with uncertainty quantification. However, the approximate inference techniques for weight space priors suffer from several drawbacks. The 'Bayesian last layer' (BLL) is an alternative BNN approach that learns the feature space for an exact Bayesian linear model with explicit predictive distributions. However, its predictions outside of the data distribution (OOD) are typically overconfident, as the marginal likelihood objective results in a learned feature space that overfits to the data. We overcome this weakness by introducing a functional prior on the model's derivatives w.r.t. the inputs. Treating these Jacobians as latent variables, we incorporate the prior into the objective to influence the smoothness and diversity of the features, which enables greater predictive uncertainty. For the BLL, the Jacobians can be computed directly using forward mode automatic differentiation, and the distribution over Jacobians may be obtained in closed-form. We demonstrate this method enhances the BLL to Gaussian process-like performance on tasks where calibrated uncertainty is critical: OOD regression, Bayesian optimization and active learning, which include high-dimensional real-world datasets.

Original languageEnglish
Title of host publication24TH INTERNATIONAL CONFERENCE ON ARTIFICIAL INTELLIGENCE AND STATISTICS (AISTATS)
EditorsA Banerjee, K Fukumizu
PublisherJMLR
Pages1198-1206
Number of pages13
Publication statusPublished - 2021
MoE publication typeA4 Conference publication
EventInternational Conference on Artificial Intelligence and Statistics - Virtual, Online
Duration: 13 Apr 202115 Apr 2021
Conference number: 24

Publication series

NameProceedings of Machine Learning Research
Volume130
ISSN (Print)2640-3498

Conference

ConferenceInternational Conference on Artificial Intelligence and Statistics
Abbreviated titleAISTATS
CityVirtual, Online
Period13/04/202115/04/2021

Fingerprint

Dive into the research topics of 'Latent Derivative Bayesian Last Layer Networks'. Together they form a unique fingerprint.

Cite this