CROCODILE: Causality aids RObustness via COntrastive DIsentangled LEarning

TL;DR

This paper tackles domain shift in medical image analysis by proposing CROCODILE, a framework that fuses causal feature disentanglement, dual disease- and domain-prediction branches, contrastive learning, and a prior-knowledge injection to improve robustness across unseen domains. It introduces a Transformer-based architecture that yields disentangled causal ($Q^{ca}$) and spurious ($Q^{sp}$) embeddings, and employs latent causal intervention via backdoor adjustment alongside a Relational Scorer to align cross-branch representations. A task-prior mechanism leverages a causality map over chest X-ray findings to inject medical knowledge and further stabilize learning. Empirically, CROCODILE provides strong OOD generalization and fairness improvements on multi-dataset chest X-ray classification, albeit with some trade-off on in-domain performance, suggesting that causality-driven disentanglement can yield safer, more generalizable medical AI systems. The approach offers a general bias-m mitigation strategy that can extend to other CAD tasks and cross-domain medical imaging applications.

Abstract

Due to domain shift, deep learning image classifiers perform poorly when applied to a domain different from the training one. For instance, a classifier trained on chest X-ray (CXR) images from one hospital may not generalize to images from another hospital due to variations in scanner settings or patient characteristics. In this paper, we introduce our CROCODILE framework, showing how tools from causality can foster a model's robustness to domain shift via feature disentanglement, contrastive learning losses, and the injection of prior knowledge. This way, the model relies less on spurious correlations, learns the mechanism bringing from images to prediction better, and outperforms baselines on out-of-distribution (OOD) data. We apply our method to multi-label lung disease classification from CXRs, utilizing over 750000 images from four datasets. Our bias-mitigation method improves domain generalization and fairness, broadening the applicability and reliability of deep learning models for a safer medical image analysis. Find our code at: https://github.com/gianlucarloni/crocodile.

Figures (4)

• Figure 1: A causal view on classifying medical images $\mathcal{I}$ coming from different domains $\mathcal{D}$ for the presence of diseases $\mathcal{Y}$. By applying the latent causal intervention (hammer), the backdoor path through the spurious features is cut off.
• Figure 2: CROCODILE involves two branches to learn robust, invariant features for predicting the labels from medical images (e.g., multi-label findings from CXRs) while disregarding confounding features. We disentangle causal features determining the label from spurious features associated with the label due to domain shift. We exploit images from multiple domains in a contrastive learning scheme and propose a new way to inject prior knowledge. Best seen in color.
• Figure 3: Our Relational Scorer stratifies and concatenates every combination of causal and spurious features across both tasks. With a fully connected layer and a consecutive sigmoid($\cdot$), it maps each pair to a relational score between 0 and 1. We use an MSE loss regressing the relational scores to the ground truth. The model learns to compare the four sets of disentangled features. Best in color.
• Figure 4: Causal graphical model among the CXR findings of interest (blue) and the ground-truth causality map defined from that graph. Gray boxes represent additional findings or risk factors (not investigated in this study) that might be associated with the desired ones.