JUCAL: Jointly Calibrating Aleatoric and Epistemic Uncertainty in Classification Tasks

TL;DR

JUCAL is a new calibration method that jointly adjusts aleatoric and epistemic uncertainties in ensemble classifiers, significantly improving predictive uncertainty estimates and reducing inference costs.

Contribution

We introduce JUCAL, a simple calibration algorithm that jointly calibrates aleatoric and epistemic uncertainties for any trained ensemble, outperforming state-of-the-art methods.

Findings

JUCAL reduces negative log-likelihood by up to 15%.

JUCAL decreases predictive set size by up to 20%.

JUCAL enables smaller ensembles to outperform larger temperature-scaled ensembles.

Abstract

We study post-calibration uncertainty for trained ensembles of classifiers. Specifically, we consider both aleatoric (label noise) and epistemic (model) uncertainty. Among the most popular and widely used calibration methods in classification are temperature scaling (i.e., pool-then-calibrate) and conformal methods. However, the main shortcoming of these calibration methods is that they do not balance the proportion of aleatoric and epistemic uncertainty. Not balancing these uncertainties can severely misrepresent predictive uncertainty, leading to overconfident predictions in some input regions while being underconfident in others. To address this shortcoming, we present a simple but powerful calibration algorithm Joint Uncertainty Calibration (JUCAL) that jointly calibrates aleatoric and epistemic uncertainty. JUCAL jointly calibrates two constants to weight and scale epistemic and aleatoric uncertainties by optimizing the negative log-likelihood (NLL) on the validation/calibration dataset. JUCAL can be applied to any trained ensemble of classifiers (e.g., transformers, CNNs, or tree-based methods), with minimal computational overhead, without requiring access to the models' internal parameters. We experimentally evaluate JUCAL on various text classification tasks, for ensembles of varying sizes and with different ensembling strategies. Our experiments show that JUCAL significantly outperforms SOTA calibration methods across all considered classification tasks, reducing NLL and predictive set size by up to 15% and 20%, respectively. Interestingly, even applying JUCAL to an ensemble of size 5 can outperform temperature-scaled ensembles of size up to 50 in terms of NLL and predictive set size, resulting in up to 10 times smaller inference costs. Thus, we propose JUCAL as a new go-to method for calibrating ensembles in classification.