# ContriMix: Scalable stain color augmentation for domain generalization without domain labels in digital pathology

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 26 references that are in the map.
- First posted on arXiv 7 Jun 2023.

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

- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Attention-based Deep Multiple Instance Learning](https://mlsanity.com/paper/1802.04712/) (2018)
- [DRIT++: Diverse Image-to-Image Translation via Disentangled Representations](https://mlsanity.com/paper/1905.01270/) (2019)
- [Improving Out-of-Distribution Robustness via Selective Augmentation](https://mlsanity.com/paper/2201.00299/) (2022)
- [Quantifying the effects of data augmentation and stain color normalization in convolutional neural networks for computational pathology](https://mlsanity.com/paper/1902.06543/) (2019)
- [Extending the WILDS Benchmark for Unsupervised Adaptation](https://mlsanity.com/paper/2112.05090/) (2021)
- [StainGAN: Stain Style Transfer for Digital Histological Images](https://mlsanity.com/paper/1804.01601/) (2018)
- [Pareto Invariant Risk Minimization: Towards Mitigating the Optimization Dilemma in Out-of-Distribution Generalization](https://mlsanity.com/paper/2206.07766/) (2022)
- [Structure-Preserving Multi-Domain Stain Color Augmentation using Style-Transfer with Disentangled Representations](https://mlsanity.com/paper/2107.12357/) (2021)
- [StainNet: a fast and robust stain normalization network](https://mlsanity.com/paper/2012.12535/) (2020)
- [Learning domain-agnostic visual representation for computational pathology using medically-irrelevant style transfer augmentation](https://mlsanity.com/paper/2102.01678/) (2021)
- [Pix2Pix-based Stain-to-Stain Translation: A Solution for Robust Stain Normalization in Histopathology Images Analysis](https://mlsanity.com/paper/2002.00647/) (2020)
- [Contextual Vision Transformers for Robust Representation Learning](https://mlsanity.com/paper/2305.19402/) (2023)
- [SC-MIL: Supervised Contrastive Multiple Instance Learning for Imbalanced Classification in Pathology](https://mlsanity.com/paper/2303.13405/) (2023)

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

- [PLUTO: Pathology-Universal Transformer](https://mlsanity.com/paper/2405.07905/) (2024)
- [Supervised Contrastive Block Disentanglement](https://mlsanity.com/paper/2502.07281/) (2025)
- [Stylizing ViT: Anatomy-Preserving Instance Style Transfer for Domain Generalization](https://mlsanity.com/paper/2601.17586/) (2026)
- [Uncertainty-Guided Selective Adaptation Enables Cross-Platform Predictive Fluorescence Microscopy](https://mlsanity.com/paper/2511.12006/) (2025)
- [When normalization hallucinates: unseen risks in AI-powered whole slide image processing](https://mlsanity.com/paper/2512.07426/) (2025)
- [Robustifying pathology foundation models via fine-tuning](https://mlsanity.com/paper/2607.22861/) (2026)
- [Simple, Safe, and Overlooked: Reclaiming Sustainable Domain Generalization with Statistical Color Matching](https://mlsanity.com/paper/2608.18915/) (2026)

## Abstract (arXiv)

Differences in staining and imaging procedures can cause significant color variations in histopathology images, leading to poor generalization when deploying deep-learning models trained from a different data source. Various color augmentation methods have been proposed to generate synthetic images during training to make models more robust, eliminating the need for stain normalization during test time. Many color augmentation methods leverage domain labels to generate synthetic images. This approach causes three significant challenges to scaling such a model. Firstly, incorporating data from a new domain into deep-learning models trained on existing domain labels is not straightforward. Secondly, dependency on domain labels prevents the use of pathology images without domain labels to improve model performance. Finally, implementation of these methods becomes complicated when multiple domain labels (e.g., patient identification, medical center, etc) are associated with a single image. We introduce ContriMix, a novel domain label free stain color augmentation method based on DRIT++, a style-transfer method. Contrimix leverages sample stain color variation within a training minibatch and random mixing to extract content and attribute information from pathology images. This information can be used by a trained ContriMix model to create synthetic images to improve the performance of existing classifiers. ContriMix outperforms competing methods on the Camelyon17-WILDS dataset. Its performance is consistent across different slides in the test set while being robust to the color variation from rare substances in pathology images. We make our code and trained ContriMix models available for research use. The code for ContriMix can be found at https://gitlab.com/huutan86/contrimix