SASWISE-UE: Segmentation and synthesis with interpretable scalable ensembles for uncertainty estimation

Area of Science:

• Medical Imaging
• Artificial Intelligence
• Deep Learning

Background:

• Deep learning models in medicine lack interpretability, limiting clinical adoption.
• Assessing the reliability of medical AI predictions is crucial for patient safety.
• Existing methods often require extensive retraining or separate uncertainty models.

Purpose of the Study:

• To introduce an efficient sub-model ensemble framework for enhanced interpretability of medical deep learning models.
• To enable end-users to evaluate the reliability of model outputs through uncertainty maps.
• To demonstrate the framework's applicability to both convolutional and transformer-based models.

Main Methods:

• Developed a strategy to generate diverse models from a single checkpoint, creating a model family.
• Implemented a method to produce multiple outputs from one input, fuse them, and estimate uncertainty via output disagreement.
• Utilized U-Net and UNETR architectures for segmentation and synthesis tasks on CT and MR-CT datasets.

Main Results:

• Achieved a mean Dice coefficient of 0.814 for CT body segmentation.
• Reduced Mean Absolute Error in MR-CT synthesis from 89.43 HU to 88.17 HU through pruning.
• Demonstrated robustness under image corruption and data undersampling, maintaining uncertainty-error correlation.

Conclusions:

• The proposed framework maintains model performance while significantly enhancing interpretability via uncertainty estimation.
• This approach improves the clinical applicability of deep learning in medical imaging.
• The method is versatile, applicable to various deep learning architectures and imaging tasks.