A Rate-Distortion View of Uncertainty Quantification

TL;DR

This paper introduces Distance Aware Bottleneck (DAB), a simple method for neural networks to estimate uncertainty by measuring input distance from a learned codebook, improving out-of-distribution detection.

Contribution

The paper proposes DAB, a novel approach that enhances neural networks with uncertainty estimates using a learned codebook, outperforming existing methods in OOD detection.

Findings

DAB provides deterministic uncertainty estimates with a single forward pass.

DAB outperforms ensemble methods, Gaussian Processes, and IB-based approaches in OOD detection.

The method is simple to train and integrates seamlessly with existing neural networks.

Abstract

In supervised learning, understanding an input's proximity to the training data can help a model decide whether it has sufficient evidence for reaching a reliable prediction. While powerful probabilistic models such as Gaussian Processes naturally have this property, deep neural networks often lack it. In this paper, we introduce Distance Aware Bottleneck (DAB), i.e., a new method for enriching deep neural networks with this property. Building on prior information bottleneck approaches, our method learns a codebook that stores a compressed representation of all inputs seen during training. The distance of a new example from this codebook can serve as an uncertainty estimate for the example. The resulting model is simple to train and provides deterministic uncertainty estimates by a single forward pass. Finally, our method achieves better out-of-distribution (OOD) detection and misclassification prediction than prior methods, including expensive ensemble methods, deep kernel Gaussian Processes, and approaches based on the standard information bottleneck.