MaskMedPaint: Masked Medical Image Inpainting with Diffusion Models for Mitigation of Spurious Correlations

TL;DR

This paper introduces MaskMedPaint, a diffusion model-based method for medical image inpainting that reduces spurious correlations, improving model generalization across different domains with limited unlabeled data.

Contribution

It proposes MaskMedPaint, a novel inpainting approach using diffusion models to augment training data and mitigate spurious correlations in medical imaging.

Findings

Improves domain generalization in medical image classification.

Effective with limited unlabeled target domain images.

Enhances robustness across natural and medical datasets.

Abstract

Spurious features associated with class labels can lead image classifiers to rely on shortcuts that don't generalize well to new domains. This is especially problematic in medical settings, where biased models fail when applied to different hospitals or systems. In such cases, data-driven methods to reduce spurious correlations are preferred, as clinicians can directly validate the modified images. While Denoising Diffusion Probabilistic Models (Diffusion Models) show promise for natural images, they are impractical for medical use due to the difficulty of describing spurious medical features. To address this, we propose Masked Medical Image Inpainting (MaskMedPaint), which uses text-to-image diffusion models to augment training images by inpainting areas outside key classification regions to match the target domain. We demonstrate that MaskMedPaint enhances generalization to target domains across both natural (Waterbirds, iWildCam) and medical (ISIC 2018, Chest X-ray) datasets, given limited unlabeled target images.