[Paper Review] Predictive Uncertainty Estimation via Prior Networks

Andrey Malinin, Mark Gales|arXiv (Cornell University)|Feb 28, 2018

Adversarial Robustness in Machine Learning28 references357 citations

TL;DR

This paper introduces Prior Networks (PNs) to explicitly model distributional uncertainty separately from data and model uncertainty, enabling better out-of-distribution detection and misclassification detection, with Dirichlet Prior Networks (DPNs) applied to MNIST and CIFAR-10.

ABSTRACT

Estimating how uncertain an AI system is in its predictions is important to improve the safety of such systems. Uncertainty in predictive can result from uncertainty in model parameters, irreducible data uncertainty and uncertainty due to distributional mismatch between the test and training data distributions. Different actions might be taken depending on the source of the uncertainty so it is important to be able to distinguish between them. Recently, baseline tasks and metrics have been defined and several practical methods to estimate uncertainty developed. These methods, however, attempt to model uncertainty due to distributional mismatch either implicitly through model uncertainty or as data uncertainty. This work proposes a new framework for modeling predictive uncertainty called Prior Networks (PNs) which explicitly models distributional uncertainty. PNs do this by parameterizing a prior distribution over predictive distributions. This work focuses on uncertainty for classification and evaluates PNs on the tasks of identifying out-of-distribution (OOD) samples and detecting misclassification on the MNIST dataset, where they are found to outperform previous methods. Experiments on synthetic and MNIST and CIFAR-10 data show that unlike previous non-Bayesian methods PNs are able to distinguish between data and distributional uncertainty.

Motivation & Objective

• Motivate the need to separate three sources of predictive uncertainty: model (epistemic), data (aleatoric), and distributional (dataset shift).
• Propose Prior Networks that parameterize a distribution over predictive distributions to isolate distributional uncertainty.
• Develop and evaluate Dirichlet Prior Networks (DPNs) for classification tasks, focusing on OOD detection and misclassification detection.
• Offer uncertainty measures derived from the PN framework and compare against Bayesian/multi-model baselines.

Proposed method

• Introduce Prior Networks (PNs) that explicitly model a distribution over predictive distributions p(mu|x, theta).
• Use a Dirichlet distribution to parameterize p(mu|x; theta), with alpha = f(x; theta), enabling sharp corners for confident in-domain predictions and flat priors for out-of-distribution inputs.
• Train Dirichlet Prior Networks (DPNs) via a multi-task objective that minimizes KL divergence to sharp Dirichlet targets for in-domain data and flat Dirichlet targets for out-of-domain data (eq. 12).
• Regularize and smooth in-distribution targets (eq. 15) to avoid delta-function targets and optionally use teacher-student smoothing.
• Discuss different marginalizations of the PN hierarchy (data, distributional, model uncertainty) and derive uncertainty measures (entropy, mutual information) from these marginals.
• Evaluate PN/Dirichlet PN on synthetic data, MNIST, and CIFAR-10, comparing against standard DNNs and MC-Dropout ensembles.

Experimental results

Research questions

• RQ1 Can Prior Networks separately model data uncertainty, distributional uncertainty, and model uncertainty in classification tasks?
• RQ2 Do Dirichlet Prior Networks improve out-of-distribution detection and misclassification detection compared to baselines like DNNs and MC-Dropout ensembles?
• RQ3 Which uncertainty measures (entropy, mutual information, differential entropy) best reflect the different sources of uncertainty under PN framework?
• RQ4 How do PN-based methods perform on MNIST and CIFAR-10, including noisy/augmented scenarios and various real-world OOD datasets?

Key findings

• Dirichlet Prior Networks yield more accurate distributional uncertainty estimates than MC-Dropout and standard DNNs for OOD detection on MNIST/CIFAR-10.
• PNs outperform baselines on misclassification detection across MNIST and CIFAR-10.
• Differential entropy of the Dirichlet prior is particularly effective for OOD detection when class distinctions are weak or noisy.
• Uncertainty measures derived from PN framework can be computed analytically at test time, offering lower computational cost than ensembles.
• In synthetic data, PN’s ability to distinguish in-distribution vs. out-of-distribution improves when class overlap is high, unlike standard entropy measures.
• Entropy and maximum posterior probability remain strong simple indicators, with differential entropy providing advantages in certain OOD scenarios (especially with less distinct classes).