Function Space Diversity for Uncertainty Prediction via Repulsive Last-Layer Ensembles

TL;DR

This paper introduces a novel function space ensemble method that improves uncertainty estimation in neural networks by using a repulsive last-layer ensemble, compatible with large and pretrained models, and effective for out-of-distribution detection and active learning.

Contribution

It proposes a minimal-parameter, scalable ensemble approach in function space that enhances uncertainty estimation and integrates seamlessly with pretrained networks.

Findings

Improved uncertainty estimates for out-of-distribution detection.

Effective disentanglement of aleatoric and epistemic uncertainty.

Compatible with large and pretrained neural networks.

Abstract

Bayesian inference in function space has gained attention due to its robustness against overparameterization in neural networks. However, approximating the infinite-dimensional function space introduces several challenges. In this work, we discuss function space inference via particle optimization and present practical modifications that improve uncertainty estimation and, most importantly, make it applicable for large and pretrained networks. First, we demonstrate that the input samples, where particle predictions are enforced to be diverse, are detrimental to the model performance. While diversity on training data itself can lead to underfitting, the use of label-destroying data augmentation, or unlabeled out-of-distribution data can improve prediction diversity and uncertainty estimates. Furthermore, we take advantage of the function space formulation, which imposes no restrictions on network parameterization other than sufficient flexibility. Instead of using full deep ensembles to represent particles, we propose a single multi-headed network that introduces a minimal increase in parameters and computation. This allows seamless integration to pretrained networks, where this repulsive last-layer ensemble can be used for uncertainty aware fine-tuning at minimal additional cost. We achieve competitive results in disentangling aleatoric and epistemic uncertainty for active learning, detecting out-of-domain data, and providing calibrated uncertainty estimates under distribution shifts with minimal compute and memory.