Bayesian uncertainty-weighted loss for improved generalisability on polyp segmentation task

TL;DR

This paper introduces a Bayesian uncertainty-weighted loss function to enhance the generalisability of polyp segmentation models across diverse datasets, addressing variability and out-of-distribution challenges in clinical settings.

Contribution

It proposes a novel Bayesian uncertainty-based training method that improves model robustness and fairness across multi-center datasets for polyp segmentation.

Findings

Enhanced generalisability on multi-center datasets.

Maintained state-of-the-art performance.

Better focus on underrepresented samples.

Abstract

While several previous studies have devised methods for segmentation of polyps, most of these methods are not rigorously assessed on multi-center datasets. Variability due to appearance of polyps from one center to another, difference in endoscopic instrument grades, and acquisition quality result in methods with good performance on in-distribution test data, and poor performance on out-of-distribution or underrepresented samples. Unfair models have serious implications and pose a critical challenge to clinical applications. We adapt an implicit bias mitigation method which leverages Bayesian predictive uncertainties during training to encourage the model to focus on underrepresented sample regions. We demonstrate the potential of this approach to improve generalisability without sacrificing state-of-the-art performance on a challenging multi-center polyp segmentation dataset (PolypGen) with different centers and image modalities.