SASWISE-UE: Segmentation and Synthesis with Interpretable Scalable Ensembles for Uncertainty Estimation

TL;DR

This paper presents SASWISE-UE, a scalable ensemble framework that improves interpretability and uncertainty estimation in medical imaging models, demonstrating robust performance on segmentation and synthesis tasks with enhanced reliability assessments.

Contribution

The paper introduces a novel ensemble method that generates diverse models from a single checkpoint, enabling uncertainty estimation and improved interpretability in medical deep learning applications.

Findings

Achieved a mean Dice coefficient of 0.814 in segmentation.

Attained a Mean Absolute Error of 88.17 HU in synthesis.

Maintained correlation between uncertainty and error under data corruption.

Abstract

This paper introduces an efficient sub-model ensemble framework aimed at enhancing the interpretability of medical deep learning models, thus increasing their clinical applicability. By generating uncertainty maps, this framework enables end-users to evaluate the reliability of model outputs. We developed a strategy to develop diverse models from a single well-trained checkpoint, facilitating the training of a model family. This involves producing multiple outputs from a single input, fusing them into a final output, and estimating uncertainty based on output disagreements. Implemented using U-Net and UNETR models for segmentation and synthesis tasks, this approach was tested on CT body segmentation and MR-CT synthesis datasets. It achieved a mean Dice coefficient of 0.814 in segmentation and a Mean Absolute Error of 88.17 HU in synthesis, improved from 89.43 HU by pruning. Additionally, the framework was evaluated under corruption and undersampling, maintaining correlation between uncertainty and error, which highlights its robustness. These results suggest that the proposed approach not only maintains the performance of well-trained models but also enhances interpretability through effective uncertainty estimation, applicable to both convolutional and transformer models in a range of imaging tasks.