Causality-Inspired Single-Source Domain Generalization for Medical Image Segmentation

Deep learning models usually suffer from domain shift issues, where models trained on one source domain do not generalize well to other unseen domains. In this work, we investigate the single-source domain generalization problem: training a deep network that is robust to unseen domains, under the condition that training data is only available from one source domain, which is common in medical imaging applications. We tackle this problem in the context of cross-domain medical image segmentation. Under this scenario, domain shifts are mainly caused by different acquisition processes. We propose a simple causality-inspired data augmentation approach to expose a segmentation model to synthesized domain-shifted training examples. Specifically, 1) to make the deep model robust to discrepancies in image intensities and textures, we employ a family of randomly-weighted shallow networks. They augment training images using diverse appearance transformations. 2) Further we show that spurious correlations among objects in an image are detrimental to domain robustness. These correlations might be taken by the network as domain-specific clues for making predictions, and they may break on unseen domains. We remove these spurious correlations via causal intervention. This is achieved by stratifying the appearances of potentially correlated objects. The proposed approach is validated on three crossdomain segmentation tasks: cross-modality (CT-MRI) abdominal image segmentation, cross-sequence (bSSFP-LGE) cardiac MRI segmentation, and cross-center prostate MRI segmentation. The proposed approach yields consistent performance gains compared with competitive methods when tested on unseen domains.

[1]  Jinhui Tang,et al.  Causal Intervention for Weakly-Supervised Semantic Segmentation , 2020, NeurIPS.

[2]  Loïc Le Folgoc,et al.  Attention U-Net: Learning Where to Look for the Pancreas , 2018, ArXiv.

[3]  Ce Liu,et al.  Supervised Contrastive Learning , 2020, NeurIPS.

[4]  M. Bethge,et al.  Shortcut learning in deep neural networks , 2020, Nature Machine Intelligence.

[5]  Sotirios A. Tsaftaris,et al.  Semi-supervised Meta-learning with Disentanglement for Domain-generalised Medical Image Segmentation , 2021, MICCAI.

[6]  Nima Tajbakhsh,et al.  UNet++: A Nested U-Net Architecture for Medical Image Segmentation , 2018, DLMIA/ML-CDS@MICCAI.

[7]  Simon Andermatt,et al.  AirLab: Autograd Image Registration Laboratory , 2018, ArXiv.

[8]  Daniel C. Castro,et al.  Machine Learning with Multi-Site Imaging Data: An Empirical Study on the Impact of Scanner Effects , 2019, ArXiv.

[9]  Juan Cao,et al.  Progressive Domain Expansion Network for Single Domain Generalization , 2021, 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[10]  Eric P. Xing,et al.  Self-Challenging Improves Cross-Domain Generalization , 2020, ECCV.

[11]  Rui Li,et al.  Generalize Ultrasound Image Segmentation via Instant and Plug & Play Style Transfer , 2021, ArXiv.

[12]  Guillaume Lemaitre,et al.  Computer-Aided Detection and diagnosis for prostate cancer based on mono and multi-parametric MRI: A review , 2015, Comput. Biol. Medicine.

[13]  Jimmy Ba,et al.  Adam: A Method for Stochastic Optimization , 2014, ICLR.

[14]  Daniel C. Castro,et al.  Domain Generalization via Model-Agnostic Learning of Semantic Features , 2019, NeurIPS.

[15]  Yang Zou,et al.  Domain Adaptation for Semantic Segmentation via Class-Balanced Self-Training , 2018, ArXiv.

[16]  Adrian V. Dalca,et al.  A Learning Strategy for Contrast-agnostic MRI Segmentation , 2020, MIDL.

[17]  C. R. Deboor,et al.  A practical guide to splines , 1978 .

[18]  Balaji Lakshminarayanan,et al.  AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty , 2020, ICLR.

[19]  Tolga Tasdizen,et al.  Regularization With Stochastic Transformations and Perturbations for Deep Semi-Supervised Learning , 2016, NIPS.

[20]  Florian Jung,et al.  Evaluation of prostate segmentation algorithms for MRI: The PROMISE12 challenge , 2014, Medical Image Anal..

[21]  Sergey Levine,et al.  Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks , 2017, ICML.

[22]  Jakub M. Tomczak,et al.  Selecting Data Augmentation for Simulating Interventions , 2021, ICML.

[23]  Charles Blundell,et al.  Representation Learning via Invariant Causal Mechanisms , 2020, ICLR.

[24]  Jacob Steinhardt,et al.  Limitations of Post-Hoc Feature Alignment for Robustness , 2021, 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[25]  Fabio Maria Carlucci,et al.  Domain Generalization by Solving Jigsaw Puzzles , 2019, 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[26]  Anna Khoreva,et al.  Grid Saliency for Context Explanations of Semantic Segmentation , 2019, NeurIPS.

[27]  Shin Ishii,et al.  Virtual Adversarial Training: A Regularization Method for Supervised and Semi-Supervised Learning , 2017, IEEE Transactions on Pattern Analysis and Machine Intelligence.

[28]  Bernhard Schölkopf,et al.  Domain Generalization via Invariant Feature Representation , 2013, ICML.

[29]  Trevor Darrell,et al.  Deep Domain Confusion: Maximizing for Domain Invariance , 2014, CVPR 2014.

[30]  Konstantinos Kamnitsas,et al.  Unsupervised domain adaptation in brain lesion segmentation with adversarial networks , 2016, IPMI.

[31]  Serge J. Belongie,et al.  Arbitrary Style Transfer in Real-Time with Adaptive Instance Normalization , 2017, 2017 IEEE International Conference on Computer Vision (ICCV).

[32]  Daguang Xu,et al.  Generalizing Deep Learning for Medical Image Segmentation to Unseen Domains via Deep Stacked Transformation , 2020, IEEE Transactions on Medical Imaging.

[33]  智一 吉田,et al.  Efficient Graph-Based Image Segmentationを用いた圃場図自動作成手法の検討 , 2014 .

[34]  Quoc V. Le,et al.  EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks , 2019, ICML.

[35]  Yoshua Bengio,et al.  Learning deep representations by mutual information estimation and maximization , 2018, ICLR.

[36]  Koby Crammer,et al.  A theory of learning from different domains , 2010, Machine Learning.

[37]  Sotirios A. Tsaftaris,et al.  Multimodal MR Synthesis via Modality-Invariant Latent Representation , 2018, IEEE Transactions on Medical Imaging.

[38]  Gal Chechik,et al.  A causal view of compositional zero-shot recognition , 2020, NeurIPS.

[39]  Zijian Wang,et al.  Learning to Diversify for Single Domain Generalization , 2021, 2021 IEEE/CVF International Conference on Computer Vision (ICCV).

[40]  Marcus A. Badgeley,et al.  Variable generalization performance of a deep learning model to detect pneumonia in chest radiographs: A cross-sectional study , 2018, PLoS medicine.

[41]  Thomas Brox,et al.  U-Net: Convolutional Networks for Biomedical Image Segmentation , 2015, MICCAI.

[42]  Ioannis Mitliagkas,et al.  Manifold Mixup: Better Representations by Interpolating Hidden States , 2018, ICML.

[43]  Pheng-Ann Heng,et al.  Shape-aware Meta-learning for Generalizing Prostate MRI Segmentation to Unseen Domains , 2020, MICCAI.

[44]  Víctor M. Campello,et al.  Cardiac Segmentation on Late Gadolinium Enhancement MRI: A Benchmark Study from Multi-Sequence Cardiac MR Segmentation Challenge , 2020, Medical Image Anal..

[45]  Daniel Rueckert,et al.  Realistic Adversarial Data Augmentation for MR Image Segmentation , 2020, MICCAI.

[46]  Geoffrey E. Hinton,et al.  ImageNet classification with deep convolutional neural networks , 2012, Commun. ACM.

[47]  Sen Wu,et al.  On the Generalization Effects of Linear Transformations in Data Augmentation , 2020, ICML.

[48]  Hao Chen,et al.  Unsupervised Cross-Modality Domain Adaptation of ConvNets for Biomedical Image Segmentations with Adversarial Loss , 2018, IJCAI.

[49]  David Lopez-Paz,et al.  Invariant Risk Minimization , 2019, ArXiv.

[50]  Christoph Meinel,et al.  Deep Learning for Medical Image Analysis , 2018, Journal of Pathology Informatics.

[51]  Guillermo Sapiro,et al.  Deep Neural Networks with Random Gaussian Weights: A Universal Classification Strategy? , 2015, IEEE Transactions on Signal Processing.

[52]  Andreas Nürnberger,et al.  CHAOS Challenge - Combined (CT-MR) Healthy Abdominal Organ Segmentation , 2020, Medical Image Anal..

[53]  Xi Peng,et al.  Learning to Learn Single Domain Generalization , 2020, 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[54]  Marc Niethammer,et al.  Robust and Generalizable Visual Representation Learning via Random Convolutions , 2020, ICLR.

[55]  Victor S. Lempitsky,et al.  Unsupervised Domain Adaptation by Backpropagation , 2014, ICML.

[56]  J. Pearl Causal diagrams for empirical research , 1995 .

[57]  Silvio Savarese,et al.  Generalizing to Unseen Domains via Adversarial Data Augmentation , 2018, NeurIPS.

[58]  Judy Hoffman,et al.  Learning to Balance Specificity and Invariance for In and Out of Domain Generalization , 2020, ECCV.

[59]  D. Rueckert,et al.  Self-Supervision with Superpixels: Training Few-shot Medical Image Segmentation without Annotation , 2020, ECCV.

[60]  Hongyi Zhang,et al.  mixup: Beyond Empirical Risk Minimization , 2017, ICLR.

[61]  Shruti Tople,et al.  Domain Generalization using Causal Matching , 2020, ICML.

[62]  Graham W. Taylor,et al.  Improved Regularization of Convolutional Neural Networks with Cutout , 2017, ArXiv.

[63]  Tao Xiang,et al.  Domain Generalization with MixStyle , 2021, ICLR.

[64]  Yongxin Yang,et al.  Episodic Training for Domain Generalization , 2019, 2019 IEEE/CVF International Conference on Computer Vision (ICCV).

[65]  Daniel Rueckert,et al.  Cooperative Training and Latent Space Data Augmentation for Robust Medical Image Segmentation , 2021, MICCAI.

[66]  Yan Wang,et al.  The Medical Segmentation Decathlon , 2021, ArXiv.

[67]  Daniel C. Castro,et al.  Causality matters in medical imaging , 2019, Nature Communications.