BayesNetCNN: incorporating uncertainty in neural networks for image-based classification tasks

TL;DR

This paper introduces BayesNetCNN, a method to convert standard neural networks into Bayesian models that estimate prediction uncertainty, improving classification accuracy and enabling selective manual review in medical imaging.

Contribution

The paper presents a novel approach to incorporate uncertainty into neural networks and combines it with a rejection strategy for better decision-making in image classification.

Findings

Classification accuracy improved from 0.86 to 0.95.

Retained 75% of the test set with high-confidence predictions.

Enabled case selection for manual review based on uncertainty.

Abstract

The willingness to trust predictions formulated by automatic algorithms is key in a vast number of domains. However, a vast number of deep architectures are only able to formulate predictions without an associated uncertainty. In this paper, we propose a method to convert a standard neural network into a Bayesian neural network and estimate the variability of predictions by sampling different networks similar to the original one at each forward pass. We couple our methods with a tunable rejection-based approach that employs only the fraction of the dataset that the model is able to classify with an uncertainty below a user-set threshold. We test our model in a large cohort of brain images from Alzheimer's Disease patients, where we tackle discrimination of patients from healthy controls based on morphometric images only. We demonstrate how combining the estimated uncertainty with a rejection-based approach increases classification accuracy from 0.86 to 0.95 while retaining 75% of the test set. In addition, the model can select cases to be recommended for manual evaluation based on excessive uncertainty. We believe that being able to estimate the uncertainty of a prediction, along with tools that can modulate the behavior of the network to a degree of confidence that the user is informed about (and comfortable with) can represent a crucial step in the direction of user compliance and easier integration of deep learning tools into everyday tasks currently performed by human operators.