# Joint Optimization of Class-Specific Training- and Test-Time Data Augmentation in Segmentation

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 59 references that are in the map.
- First posted on arXiv 30 May 2023.

## Builds on (most cited first, 30 of 37)

- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation](https://mlsanity.com/paper/1606.06650/) (2016)
- [AutoAugment: Learning Augmentation Policies from Data](https://mlsanity.com/paper/1805.09501/) (2018)
- [Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization](https://mlsanity.com/paper/1603.06560/) (2016)
- [Bilevel Programming for Hyperparameter Optimization and Meta-Learning](https://mlsanity.com/paper/1806.04910/) (2018)
- [Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks](https://mlsanity.com/paper/1807.07356/) (2018)
- [Efficient Multi-Scale 3D CNN with Fully Connected CRF for Accurate Brain Lesion Segmentation](https://mlsanity.com/paper/1603.05959/) (2016)
- [Fast AutoAugment](https://mlsanity.com/paper/1905.00397/) (2019)
- [Hyperparameter optimization with approximate gradient](https://mlsanity.com/paper/1602.02355/) (2016)
- [Understanding and Robustifying Differentiable Architecture Search](https://mlsanity.com/paper/1909.09656/) (2019)
- [The KiTS19 Challenge Data: 300 Kidney Tumor Cases with Clinical Context, CT Semantic Segmentations, and Surgical Outcomes](https://mlsanity.com/paper/1904.00445/) (2019)
- [Data augmentation using learned transformations for one-shot medical image segmentation](https://mlsanity.com/paper/1902.09383/) (2019)
- [Shape-aware Meta-learning for Generalizing Prostate MRI Segmentation to Unseen Domains](https://mlsanity.com/paper/2007.02035/) (2020)
- [Better Aggregation in Test-Time Augmentation](https://mlsanity.com/paper/2011.11156/) (2020)
- [Causality-inspired Single-source Domain Generalization for Medical Image Segmentation](https://mlsanity.com/paper/2111.12525/) (2021)
- [Test-Time Adaptable Neural Networks for Robust Medical Image Segmentation](https://mlsanity.com/paper/2004.04668/) (2020)
- [DADA: Differentiable Automatic Data Augmentation](https://mlsanity.com/paper/2003.03780/) (2020)
- [Learning Loss for Test-Time Augmentation](https://mlsanity.com/paper/2010.11422/) (2020)
- [Realistic Adversarial Data Augmentation for MR Image Segmentation](https://mlsanity.com/paper/2006.13322/) (2020)
- [Cardiac Segmentation on Late Gadolinium Enhancement MRI: A Benchmark Study from Multi-Sequence Cardiac MR Segmentation Challenge](https://mlsanity.com/paper/2006.12434/) (2020)
- [Analyzing Overfitting under Class Imbalance in Neural Networks for Image Segmentation](https://mlsanity.com/paper/2102.10365/) (2021)
- [Studying Robustness of Semantic Segmentation under Domain Shift in cardiac MRI](https://mlsanity.com/paper/2011.07592/) (2020)

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

- [Test-time generative augmentation for medical image segmentation](https://mlsanity.com/paper/2406.17608/) (2024)
- [Deep Nets with Subsampling Layers Unwittingly Discard Useful Activations at Test-Time](https://mlsanity.com/paper/2410.01083/) (2024)
- [Partial Decoder Attention Network with Contour-weighted Loss Function for Data-Imbalance Medical Image Segmentation](https://mlsanity.com/paper/2601.14338/) (2026)

## Abstract (arXiv)

This paper presents an effective and general data augmentation framework for medical image segmentation. We adopt a computationally efficient and data-efficient gradient-based meta-learning scheme to explicitly align the distribution of training and validation data which is used as a proxy for unseen test data. We improve the current data augmentation strategies with two core designs. First, we learn class-specific training-time data augmentation (TRA) effectively increasing the heterogeneity within the training subsets and tackling the class imbalance common in segmentation. Second, we jointly optimize TRA and test-time data augmentation (TEA), which are closely connected as both aim to align the training and test data distribution but were so far considered separately in previous works. We demonstrate the effectiveness of our method on four medical image segmentation tasks across different scenarios with two state-of-the-art segmentation models, DeepMedic and nnU-Net. Extensive experimentation shows that the proposed data augmentation framework can significantly and consistently improve the segmentation performance when compared to existing solutions. Code is publicly available.