Last Layer Empirical Bayes

TL;DR

Uncertainty quantification for neural networks is challenging; Bayesian neural networks and deep ensembles provide predictive uncertainty via a weight distribution $q^*$, but ensembles often outperform BNNs at higher cost. The paper introduces Last Layer Empirical Bayes (LLEB), which learns a data dependent prior as a normalizing flow applied to the last-layer weights and optimizes the ELBO while restricting Bayesian treatment to the last layer. Empirically, LLEB yields competitive results with baselines of similar cost and provides an empirical Bayes perspective that may guide future uncertainty quantification work. Overall, the work highlights empirical Bayes as a promising direction to improve uncertainty quantification without the training burden of large ensembles.

Abstract

The task of quantifying the inherent uncertainty associated with neural network predictions is a key challenge in artificial intelligence. Bayesian neural networks (BNNs) and deep ensembles are among the most prominent approaches to tackle this task. Both approaches produce predictions by computing an expectation of neural network outputs over some distribution on the corresponding weights; this distribution is given by the posterior in the case of BNNs, and by a mixture of point masses for ensembles. Inspired by recent work showing that the distribution used by ensembles can be understood as a posterior corresponding to a learned data-dependent prior, we propose last layer empirical Bayes (LLEB). LLEB instantiates a learnable prior as a normalizing flow, which is then trained to maximize the evidence lower bound; to retain tractability we use the flow only on the last layer. We show why LLEB is well motivated, and how it interpolates between standard BNNs and ensembles in terms of the strength of the prior that they use. LLEB performs on par with existing approaches, highlighting that empirical Bayes is a promising direction for future research in uncertainty quantification.