A Kernel Framework to Quantify a Model's Local Predictive Uncertainty under Data Distributional Shifts

TL;DR

This paper introduces a kernel-based framework for quantifying a neural network's local predictive uncertainty, especially under data distributional shifts, by estimating the model's prediction PDF in a Gaussian RKHS, leading to improved detection of prediction errors.

Contribution

The paper proposes a novel kernel framework that explicitly estimates the model's prediction PDF in a Gaussian RKHS, enhancing uncertainty quantification and shift detection over Bayesian methods.

Findings

Kernel framework provides more precise uncertainty estimates.

Better detection of distributional shifts in test data.

Outperforms existing methods on benchmark datasets.

Abstract

Traditional Bayesian approaches for model uncertainty quantification rely on notoriously difficult processes of marginalization over each network parameter to estimate its probability density function (PDF). Our hypothesis is that internal layer outputs of a trained neural network contain all of the information related to both its mapping function (quantified by its weights) as well as the input data distribution. We therefore propose a framework for predictive uncertainty quantification of a trained neural network that explicitly estimates the PDF of its raw prediction space (before activation), p(y'|x,w), which we refer to as the model PDF, in a Gaussian reproducing kernel Hilbert space (RKHS). The Gaussian RKHS provides a localized density estimate of p(y'|x,w), which further enables us to utilize gradient based formulations of quantum physics to decompose the model PDF in terms of multiple local uncertainty moments that provide much greater resolution of the PDF than the central moments characterized by Bayesian methods. This provides the framework with a better ability to detect distributional shifts in test data away from the training data PDF learned by the model. We evaluate the framework against existing uncertainty quantification methods on benchmark datasets that have been corrupted using common perturbation techniques. The kernel framework is observed to provide model uncertainty estimates with much greater precision based on the ability to detect model prediction errors.