Robust image representations with counterfactual contrastive learning

TL;DR

This paper introduces counterfactual contrastive learning, a novel approach that uses causal image synthesis to generate more realistic positive pairs, improving robustness and fairness in medical image representations.

Contribution

It presents a new framework leveraging causal image synthesis for contrastive learning, enhancing robustness to domain shifts and reducing subgroup disparities in medical imaging.

Findings

Outperforms standard contrastive learning on five medical datasets.

Achieves better robustness to acquisition shifts and external dataset performance.

Reduces subgroup disparities related to biological sex.

Abstract

Contrastive pretraining can substantially increase model generalisation and downstream performance. However, the quality of the learned representations is highly dependent on the data augmentation strategy applied to generate positive pairs. Positive contrastive pairs should preserve semantic meaning while discarding unwanted variations related to the data acquisition domain. Traditional contrastive pipelines attempt to simulate domain shifts through pre-defined generic image transformations. However, these do not always mimic realistic and relevant domain variations for medical imaging, such as scanner differences. To tackle this issue, we herein introduce counterfactual contrastive learning, a novel framework leveraging recent advances in causal image synthesis to create contrastive positive pairs that faithfully capture relevant domain variations. Our method, evaluated across five datasets encompassing both chest radiography and mammography data, for two established contrastive objectives (SimCLR and DINO-v2), outperforms standard contrastive learning in terms of robustness to acquisition shift. Notably, counterfactual contrastive learning achieves superior downstream performance on both in-distribution and external datasets, especially for images acquired with scanners under-represented in the training set. Further experiments show that the proposed framework extends beyond acquisition shifts, with models trained with counterfactual contrastive learning reducing subgroup disparities across biological sex.