Uncertainty-Aware (UNA) Bases for Deep Bayesian Regression Using Multi-Headed Auxiliary Networks

TL;DR

This paper introduces Uncertainty-Aware (UNA) bases, a new training framework for Neural Linear Models that improves the estimation of predictive uncertainties, especially for out-of-distribution inputs, making them more reliable for risk-sensitive applications.

Contribution

The paper proposes a novel training method for Neural Linear Models that better captures predictive uncertainties, addressing limitations of traditional training procedures.

Findings

Traditional NLM training underestimates uncertainty on OOD inputs.

The proposed framework improves uncertainty estimation for downstream tasks.

Enhanced reliability of NLMs in risk-sensitive applications.

Abstract

Neural Linear Models (NLM) are deep Bayesian models that produce predictive uncertainties by learning features from the data and then performing Bayesian linear regression over these features. Despite their popularity, few works have focused on methodically evaluating the predictive uncertainties of these models. In this work, we demonstrate that traditional training procedures for NLMs drastically underestimate uncertainty on out-of-distribution inputs, and that they therefore cannot be naively deployed in risk-sensitive applications. We identify the underlying reasons for this behavior and propose a novel training framework that captures useful predictive uncertainties for downstream tasks.