Bayesian Autoencoder for Medical Anomaly Detection: Uncertainty-Aware Approach for Brain 2 MRI Analysis

TL;DR

The paper tackles unsupervised anomaly detection in brain MRI by explicitly modeling uncertainty. It introduces a Bayesian Variational Autoencoder with multi-head attention to capture both epistemic and aleatoric uncertainty, optimized via an ELBO with beta warm-up. An uncertainty-aware anomaly score combines reconstruction error and predictive uncertainty, and the approach achieves ROC AUC around $0.834$ and PR AUC around $0.833$ on BraTS2020, with ablations confirming the value of both attention and uncertainty components. The method yields interpretable uncertainty maps to assist clinicians, though it faces computational and 2D-processing limitations that offer clear directions for future work in 3D and multi-modal extensions.

Abstract

In medical imaging, anomaly detection is a vital element of healthcare diagnostics, especially for neurological conditions which can be life-threatening. Conventional deterministic methods often fall short when it comes to capturing the inherent uncertainty of anomaly detection tasks. This paper introduces a Bayesian Variational Autoencoder (VAE) equipped with multi-head attention mechanisms for detecting anomalies in brain magnetic resonance imaging (MRI). For the purpose of improving anomaly detection performance, we incorporate both epistemic and aleatoric uncertainty estimation through Bayesian inference. The model was tested on the BraTS2020 dataset, and the findings were a 0.83 ROC AUC and a 0.83 PR AUC. The data in our paper suggests that modeling uncertainty is an essential component of anomaly detection, enhancing both performance and interpretability and providing confidence estimates, as well as anomaly predictions, for clinicians to leverage in making medical decisions.

Figures (5)

• Figure 1: ROC Curve (Optimal threshold = 0.097) and Precision-Recall Curve (F1 = 0.818). The ROC curve shows an AUC of 0.834, while the PR curve shows an AUC of 0.805.
• Figure 2: Uncertainty distributions for normal and abnormal samples: (a) Aleatoric Uncertainty Distribution showing higher values for abnormal samples, (b) Epistemic Uncertainty Distribution showing a similar pattern with clear separation between normal and abnormal cases.
• Figure 3: ROC Curves for Different Uncertainty Types. Total Uncertainty (AUC = 0.846), Aleatoric (AUC = 0.846), Epistemic (AUC = 0.644), Combined Score (AUC = 0.834), and Random baseline.
• Figure 4: Training dynamics showing total uncertainty across epochs. The uncertainty decreases over training time, with some fluctuations, eventually stabilizing in the later epochs.
• Figure 5: Visualization of model outputs for normal and abnormal brain MRI slices. From left to right: Original image, reconstructed image, aleatoric uncertainty, epistemic uncertainty and total uncertainty and error map.