ConDiSR: Contrastive Disentanglement and Style Regularization for Single Domain Generalization

TL;DR

This work introduces a novel SDG method for medical image classification, utilizing channel-wise contrastive disentanglement, and demonstrates that this method consistently outperforms the SOTA independently on the choice of the source domain while exhibiting greater performance stability.

Abstract

Medical data often exhibits distribution shifts, which cause test-time performance degradation for deep learning models trained using standard supervised learning pipelines. This challenge is addressed in the field of Domain Generalization (DG) with the sub-field of Single Domain Generalization (SDG) being specifically interesting due to the privacy- or logistics-related issues often associated with medical data. Existing disentanglement-based SDG methods heavily rely on structural information embedded in segmentation masks, however classification labels do not provide such dense information. This work introduces a novel SDG method aimed at medical image classification that leverages channel-wise contrastive disentanglement. It is further enhanced with reconstruction-based style regularization to ensure extraction of distinct style and structure feature representations. We evaluate our method on the complex task of multicenter histopathology image classification, comparing it against state-of-the-art (SOTA) SDG baselines. Results demonstrate that our method surpasses the SOTA by a margin of 1% in average accuracy while also showing more stable performance. This study highlights the importance and challenges of exploring SDG frameworks in the context of the classification task. The code is publicly available at https://github.com/BioMedIA-MBZUAI/ConDiSR

Figures (5)

• Figure 1: Overview of the proposed method. The feature disentanglement module ($\mathcal{F}_{DM}$) processes the original and augmented images to produce structure-related features ($f{str}$) and style-related features ($f_{sty}$). The structure-related features ($f_{str}$) are passed further to the classification network ($\mathcal{F}_C$) for computing the classification loss, while the style-related ones ($f_{sty}$) go through the reconstruction network ($\mathcal{F}_R$) for further computation of the reconstruction loss. Additionally a contrastive loss is applied to minimize the distance between similar structure-related components and maximize the distance between different style-related components.
• Figure 2: Overview of the feature disentanglement module of the network. Here $\sigma$ represents the softmax operator, applied across the channel-related dimension of the parameter $\theta_{d}$
• Figure 3: Sample images from the domains of the five centers of Camelyon17-WILDS.
• Figure 4: Sample images from the DR datasets.
• Figure 5: Qualitative performance comparison between ConDiSR and C$^2$SDG hu2023devil via Grad-CAM Selvaraju_2019 with the tumor presense masks taken from the original Camelyon17 challenge dataset bandi2018detection.