Robust image representations with counterfactual contrastive learning

2024-09-16

arXiv.org

AI Summary

Domain shift arising from scanner heterogeneity in medical imaging undermines conventional contrastive learning, which relies on predefined image augmentations incapable of modeling realistic clinical domain variations—thereby limiting representation robustness and downstream generalizability. To address this, we propose a counterfactual contrastive learning framework—the first to integrate causal inference into contrastive representation learning. Our method employs causal image synthesis to generate semantically consistent positive pairs exhibiting authentic domain variations (e.g., vendor-specific scanner characteristics), precisely emulating clinically relevant shifts. It jointly optimizes SimCLR and DINO-v2 contrastive objectives. Evaluated across five multi-center chest X-ray and mammography datasets, our approach significantly improves robustness to domain shift. Downstream task performance surpasses standard contrastive baselines, particularly enhancing generalization on underrepresented scanner types and reducing performance disparities across sex-stratified subgroups.

Technology Category

• Computer Vision: Adversarial Attacks & Robustness
• Machine Learning: Causal Learning
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Improves robustness to medical image acquisition shifts
• Enhances generalisation with counterfactual contrastive learning
• Reduces subgroup disparities in medical imaging models

Innovation

• Counterfactual contrastive learning for robust representations
• Leverages causal image synthesis for domain variations
• Outperforms standard contrastive learning in robustness