Semantic Data Augmentation Enhanced Invariant Risk Minimization for Medical Image Domain Generalization

TL;DR

The paper tackles domain shifts in medical image classification by enhancing IRM through a domain-oriented semantic data augmentation (MedSDA) framework. It replaces VIRM's random augmentation with a director that uses inter-domain covariance to steer augmentations toward target domains, aided by a shared SDA estimator. Empirical results on GRDBench show consistent improvements over state-of-the-art methods, supported by ablations, OTDD analyses, and feature visualizations that confirm better cross-domain alignment. The approach offers a practical path to robust medical image generalization under limited data and diverse imaging conditions.

Abstract

Deep learning has achieved remarkable success in medical image classification. However, its clinical application is often hindered by data heterogeneity caused by variations in scanner vendors, imaging protocols, and operators. Approaches such as invariant risk minimization (IRM) aim to address this challenge of out-of-distribution generalization. For instance, VIRM improves upon IRM by tackling the issue of insufficient feature support overlap, demonstrating promising potential. Nonetheless, these methods face limitations in medical imaging due to the scarcity of annotated data and the inefficiency of augmentation strategies. To address these issues, we propose a novel domain-oriented direction selector to replace the random augmentation strategy used in VIRM. Our method leverages inter-domain covariance as a guider for augmentation direction, guiding data augmentation towards the target domain. This approach effectively reduces domain discrepancies and enhances generalization performance. Experiments on a multi-center diabetic retinopathy dataset demonstrate that our method outperforms state-of-the-art approaches, particularly under limited data conditions and significant domain heterogeneity.

Figures (7)

• Figure 1: Data augmentation methods expand feature overlap. However, the relatively small scale of medical image datasets and the clustered nature of features often result in many meaningless directions when random data augmentation is applied.
• Figure 2: Overview of the proposed method framework. Our method includes an inter-domain covariance-guided direction selector and a distribution estimator that augmentable the range of distribution. $\mathbf{x}^i$ denotes the input image from domain $i$, $\mathbf{z}^i$ denotes the feature representation of $\mathbf{x}^i$, $\mathbf{d}^i_j$ denotes the augment direction for domain $i$ to domian $j$, and $\tilde{z}^i_j$ denotes the augmented feature representation of $\mathbf{z}^i$ in domain $j$.
• Figure 3: Samples of different center dataset of GRDBench dgdr2023.
• Figure 4: Heatmap of the Optimal Transport Dataset Distance (OTDD)alvarez2020geometric for different domain pairs in the GRDBench dataset dgdr2023.
• Figure 5: Visualizing deep features on GRDBench dataset with two domains using UMAPmcinnes2018umap. For easier demonstration of the effectiveness of the proposed method, we visualized data from only two domains, although the original setting included four domains. As shown in the green box in \ref{['fig:vis_aug_dir_ours']}, the proposed method consistently enhances features towards the target domain. In contrast, the VIRM method (as shown in \ref{['fig:vis_aug_dir_virm']}) produces many random and meaningless directions, especially in the small-batch data scenario, leading to less stable feature representations.

Theorems & Definitions (1)

• Definition 1: VIRM