SUPER-Net: Trustworthy Image Segmentation via Uncertainty Propagation in Encoder-Decoder Networks

TL;DR

SUPER-Net advances trustworthy medical image segmentation by analytically propagating the first two moments of the variational posterior through an encoder–decoder network, yielding both a segmentation and a pixel-level uncertainty map without Monte Carlo sampling. The framework relies on a variational inference objective ($ELBO$) and a first-order Taylor expansion to propagate mean and covariance across convolutions, activations, and decoder operations. It demonstrates robust performance across MRI/CT datasets under Gaussian noise and adversarial perturbations, and provides uncertainty maps that correlate with errors and atypical image regions, supporting clinical trust and human–AI collaboration. While incurring modestly higher inference time and memory, SUPER-Net offers intrinsic uncertainty estimation, improved robustness, and broad applicability to medical imaging tasks, with clear avenues for future improvement in priors and scalability.

Abstract

Deep Learning (DL) holds great promise in reshaping the industry owing to its precision, efficiency, and objectivity. However, the brittleness of DL models to noisy and out-of-distribution inputs is ailing their deployment in sensitive fields. Current models often lack uncertainty quantification, providing only point estimates. We propose SUPER-Net, a Bayesian framework for trustworthy image segmentation via uncertainty propagation. Using Taylor series approximations, SUPER-Net propagates the mean and covariance of the model's posterior distribution across nonlinear layers. It generates two outputs simultaneously: the segmented image and a pixel-wise uncertainty map, eliminating the need for expensive Monte Carlo sampling. SUPER-Net's performance is extensively evaluated on MRI and CT scans under various noisy and adversarial conditions. Results show that SUPER-Net outperforms state-of-the-art models in robustness and accuracy. The uncertainty map identifies low-confidence areas affected by noise or attacks, allowing the model to self-assess segmentation reliability, particularly when errors arise from noise or adversarial examples.

Figures (11)

• Figure 1: An illustration of the SUPER-Net model, where all mathematical operations are performed on random variables. At each step, both the mean and covariance are propagated (indicated by the two arrows). The output of SUPER-Net consists of the predicted segmented image and a covariance matrix, which is used to generate the associated uncertainty map.
• Figure 2: Performance of the four networks, i.e., U-Net (blue), Bayes U-Net (red), Ensemble U-Net (gray), and SUPER U-Net (black), under various levels of Gaussian noise added to the (a) entire image, (b) Anterior pixels only, and (c) Posterior pixels only of the Hippocampus test data. We plot Dice Similarity Coefficient (DSC) versus Signal to Noise Ratios (SNRs) for the Anterior and Posterior hippocampus.
• Figure 3: Performance of the four networks, i.e., U-Net (blue), Bayes U-Net (red), Ensemble U-Net (gray) and SUPER U-Net (black), under various levels of Gaussian noise added to (a) the tumor pixels only and (b) all pixels of the BraTS test data. The three sub-plots show the Dice Similarity Coefficient (DSC) values for a range of Signal to Noise Ratios (SNRs) for three different tumor regions: whole tumor, core, and enhancing.
• Figure 4: Performance of four networks, i.e., U-Net (blue), Bayes U-Net (red), Ensemble U-Net (gray) and SUPER U-Net (black), under various levels of untargeted attacks to the Lungs test data. We display Dice Similarity Coefficient (DSC) values for a range of Signal to Noise Ratio (SNR).
• Figure 5: Performance of the four networks, i.e., U-Net (blue), Bayes U-Net (red), Ensemble U-Net (gray), and SUPER U-Net (black), under various levels of adversarial attacks applied to the Hippocampus test data. We show targeted adversarial attacks with (a) source: label 1, target: label 2, (b) viceversa. The two subplots show the Dice Similarity Coefficient (DSC) values for the Anterior and Posterior hippocampus measured using Signal to Noise Ratios (SNRs).