[Paper Review] An Uncertainty-aware Loss Function for Training Neural Networks with Calibrated Predictions

Afshar Shamsi, Hamzeh Asgharnezhad+3 more|arXiv (Cornell University)|Oct 7, 2021

Adversarial Robustness in Machine Learning4 citations

TL;DR

This paper proposes two novel uncertainty-aware loss functions that combine cross entropy with Expected Calibration Error (ECE) or Predictive Entropy (PE) to train Monte Carlo Dropout (MC-Dropout) networks for calibrated uncertainty predictions. The method significantly reduces overlap between uncertainty distributions for correct and incorrect predictions, improving reliability in safety-critical applications without sacrificing accuracy.

ABSTRACT

Uncertainty quantification of machine learning and deep learning methods plays an important role in enhancing trust to the obtained result. In recent years, a numerous number of uncertainty quantification methods have been introduced. Monte Carlo dropout (MC-Dropout) is one of the most well-known techniques to quantify uncertainty in deep learning methods. In this study, we propose two new loss functions by combining cross entropy with Expected Calibration Error (ECE) and Predictive Entropy (PE). The obtained results clearly show that the new proposed loss functions lead to having a calibrated MC-Dropout method. Our results confirmed the great impact of the new hybrid loss functions for minimising the overlap between the distributions of uncertainty estimates for correct and incorrect predictions without sacrificing the model's overall performance.

Motivation & Objective

• To address the lack of calibrated uncertainty predictions in deep learning models, especially in safety-critical domains like medical diagnosis and autonomous vehicles.
• To reduce the overlap between uncertainty distributions for correct and incorrect predictions, enabling better detection of risky outputs.
• To improve calibration of Monte Carlo Dropout (MC-Dropout) without degrading model accuracy.
• To develop a training framework that enables neural networks to auto-generate higher uncertainty estimates for erroneous predictions.

Proposed method

• Proposes a hybrid loss function combining cross entropy with Expected Calibration Error (ECE) to minimize miscalibration in predictions.
• Introduces a second hybrid loss function combining cross entropy with Predictive Entropy (PE) to enhance uncertainty separation.
• Uses MC-Dropout as the uncertainty quantification method, applying stochastic forward passes during inference to estimate predictive uncertainty.
• Employs a multi-objective training approach where both accuracy (via cross entropy) and calibration (via ECE or PE) are optimized simultaneously.
• Applies the new loss functions during training to encourage low uncertainty for correct predictions and high uncertainty for incorrect ones.
• Validates the method on two synthetic datasets (two-moon and blob) using repeated random splits to ensure robustness.

Experimental results

Research questions

• RQ1 Can combining cross entropy with ECE improve the calibration of Monte Carlo Dropout predictions?
• RQ2 Can combining cross entropy with Predictive Entropy enhance the separation between uncertainty distributions for correct and incorrect predictions?
• RQ3 Does the proposed loss function maintain or improve model accuracy while improving uncertainty calibration?
• RQ4 How does the proposed method compare to standard MC-Dropout and ensemble methods in terms of uncertainty separation and calibration?

Key findings

• The MC-Dropout model trained with the Predictive Entropy-based loss function achieved the largest distance (0.401) between mean uncertainty of correct and incorrect predictions, indicating superior separation.
• The ECE-based loss function reduced calibration error more effectively than standard MC-Dropout, with ECE decreasing across both datasets.
• The PE-based loss function achieved the lowest uncertainty for correct predictions (μ₁ = 0.069%) and highest for incorrect ones (μ₂ = 0.470%), minimizing overlap.
• The proposed method maintained high accuracy (98.333% for PE-based loss) comparable to standard MC-Dropout (98.333%) and slightly improved it with ECE-based loss (98.667%).
• Uncertainty-aware (UA) scores showed that the PE-based loss function best captured uncertainty, especially under increasing aleatoric noise levels.
• The ECE plots confirmed that both proposed loss functions led to consistently lower calibration error across both datasets and noise levels.