# Cooperative Training and Latent Space Data Augmentation for Robust Medical Image Segmentation

arXiv: https://arxiv.org/abs/2107.01079 · HTML: https://mlsanity.com/paper/2107.01079/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 38 references that are in the map.
- First posted on arXiv 2 Jul 2021.

## Builds on (most cited first, 26 of 26)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [A Survey on Deep Learning in Medical Image Analysis](https://mlsanity.com/paper/1702.05747/) (2017)
- [Virtual Adversarial Training: A Regularization Method for Supervised and Semi-Supervised Learning](https://mlsanity.com/paper/1704.03976/) (2017)
- [A Survey on Multi-Task Learning](https://mlsanity.com/paper/1707.08114/) (2017)
- [Self-Challenging Improves Cross-Domain Generalization](https://mlsanity.com/paper/2007.02454/) (2020)
- [Domain Generalization via Model-Agnostic Learning of Semantic Features](https://mlsanity.com/paper/1910.13580/) (2019)
- [Deep Learning for Medical Image Analysis](https://mlsanity.com/paper/1708.08987/) (2017)
- [Generalizing Across Domains via Cross-Gradient Training](https://mlsanity.com/paper/1804.10745/) (2018)
- [Efficient Object Localization Using Convolutional Networks](https://mlsanity.com/paper/1411.4280/) (2014)
- [TorchIO: A Python library for efficient loading, preprocessing, augmentation and patch-based sampling of medical images in deep learning](https://mlsanity.com/paper/2003.04696/) (2020)
- [Models Genesis: Generic Autodidactic Models for 3D Medical Image Analysis](https://mlsanity.com/paper/1908.06912/) (2019)
- [Robust and Generalizable Visual Representation Learning via Random Convolutions](https://mlsanity.com/paper/2007.13003/) (2020)
- [Learning to Balance Specificity and Invariance for In and Out of Domain Generalization](https://mlsanity.com/paper/2008.12839/) (2020)
- [Maximum-Entropy Adversarial Data Augmentation for Improved Generalization and Robustness](https://mlsanity.com/paper/2010.08001/) (2020)
- [Learning from Extrinsic and Intrinsic Supervisions for Domain Generalization](https://mlsanity.com/paper/2007.09316/) (2020)
- [Improving Robustness Without Sacrificing Accuracy with Patch Gaussian Augmentation](https://mlsanity.com/paper/1906.02611/) (2019)
- [Hardness-Aware Deep Metric Learning](https://mlsanity.com/paper/1903.05503/) (2019)
- [Realistic Adversarial Data Augmentation for MR Image Segmentation](https://mlsanity.com/paper/2006.13322/) (2020)
- [Automatic Liver and Lesion Segmentation in CT Using Cascaded Fully Convolutional Neural Networks and 3D Conditional Random Fields](https://mlsanity.com/paper/1610.02177/) (2016)
- [Improving out-of-distribution generalization via multi-task self-supervised pretraining](https://mlsanity.com/paper/2003.13525/) (2020)
- [Improving the generalizability of convolutional neural network-based segmentation on CMR images](https://mlsanity.com/paper/1907.01268/) (2019)
- [Anatomical Priors for Image Segmentation via Post-Processing with Denoising Autoencoders](https://mlsanity.com/paper/1906.02343/) (2019)
- [Cardiac MRI Segmentation with Strong Anatomical Guarantees](https://mlsanity.com/paper/1907.02865/) (2019)

## Built on by (most cited first, 9 of 9)

- [Causality-inspired Single-source Domain Generalization for Medical Image Segmentation](https://mlsanity.com/paper/2111.12525/) (2021)
- [Domain and Content Adaptive Convolution based Multi-Source Domain Generalization for Medical Image Segmentation](https://mlsanity.com/paper/2109.05676/) (2021)
- [vMFNet: Compositionality Meets Domain-generalised Segmentation](https://mlsanity.com/paper/2206.14538/) (2022)
- [Improved post-hoc probability calibration for out-of-domain MRI segmentation](https://mlsanity.com/paper/2208.02870/) (2022)
- [Fighting the scanner effect in brain MRI segmentation with a progressive level-of-detail network trained on multi-site data](https://mlsanity.com/paper/2211.02400/) (2022)
- [Navigating Distribution Shifts in Medical Image Analysis: A Survey](https://mlsanity.com/paper/2411.05824/) (2024)
- [Shape-intensity knowledge distillation for robust medical image segmentation](https://mlsanity.com/paper/2409.17503/) (2024)
- [SUPRA: Superpixel Guided Loss for Improved Multi-modal Segmentation in Endoscopy](https://mlsanity.com/paper/2211.04658/) (2022)
- [Compositionally Equivariant Representation Learning](https://mlsanity.com/paper/2306.07783/) (2023)

## Abstract (arXiv)

Deep learning-based segmentation methods are vulnerable to unforeseen data distribution shifts during deployment, e.g. change of image appearances or contrasts caused by different scanners, unexpected imaging artifacts etc. In this paper, we present a cooperative framework for training image segmentation models and a latent space augmentation method for generating hard examples. Both contributions improve model generalization and robustness with limited data. The cooperative training framework consists of a fast-thinking network (FTN) and a slow-thinking network (STN). The FTN learns decoupled image features and shape features for image reconstruction and segmentation tasks. The STN learns shape priors for segmentation correction and refinement. The two networks are trained in a cooperative manner. The latent space augmentation generates challenging examples for training by masking the decoupled latent space in both channel-wise and spatial-wise manners. We performed extensive experiments on public cardiac imaging datasets. Using only 10 subjects from a single site for training, we demonstrated improved cross-site segmentation performance and increased robustness against various unforeseen imaging artifacts compared to strong baseline methods. Particularly, cooperative training with latent space data augmentation yields 15% improvement in terms of average Dice score when compared to a standard training method.