Bayesian Deep Ensembles via the Neural Tangent Kernel

TL;DR

This paper introduces a modification to deep ensembles that allows them to be interpreted as Bayesian posterior distributions in the infinite width limit, improving uncertainty estimation and out-of-distribution performance.

Contribution

A simple, computationally-efficient modification to deep ensembles that enables a Bayesian posterior interpretation via the Neural Tangent Kernel, bridging deep learning and Gaussian processes.

Findings

Bayesian deep ensembles approximate posterior predictive distributions.

They produce more conservative predictions than standard ensembles.

They outperform standard ensembles in out-of-distribution tasks.

Abstract

We explore the link between deep ensembles and Gaussian processes (GPs) through the lens of the Neural Tangent Kernel (NTK): a recent development in understanding the training dynamics of wide neural networks (NNs). Previous work has shown that even in the infinite width limit, when NNs become GPs, there is no GP posterior interpretation to a deep ensemble trained with squared error loss. We introduce a simple modification to standard deep ensembles training, through addition of a computationally-tractable, randomised and untrainable function to each ensemble member, that enables a posterior interpretation in the infinite width limit. When ensembled together, our trained NNs give an approximation to a posterior predictive distribution, and we prove that our Bayesian deep ensembles make more conservative predictions than standard deep ensembles in the infinite width limit. Finally, using finite width NNs we demonstrate that our Bayesian deep ensembles faithfully emulate the analytic posterior predictive when available, and can outperform standard deep ensembles in various out-of-distribution settings, for both regression and classification tasks.