# Adaptively Weighted Data Augmentation Consistency Regularization for Robust Optimization under Concept Shift

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 57 references that are in the map.
- First posted on arXiv 4 Oct 2022.

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

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [FixMatch: Simplifying Semi-Supervised Learning with Consistency and Confidence](https://mlsanity.com/paper/2001.07685/) (2020)
- [Unsupervised Feature Learning via Non-Parametric Instance-level Discrimination](https://mlsanity.com/paper/1805.01978/) (2018)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [MixMatch: A Holistic Approach to Semi-Supervised Learning](https://mlsanity.com/paper/1905.02249/) (2019)
- [Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2105.05537/) (2021)
- [Temporal Ensembling for Semi-Supervised Learning](https://mlsanity.com/paper/1610.02242/) (2016)
- [Regularization With Stochastic Transformations and Perturbations for Deep Semi-Supervised Learning](https://mlsanity.com/paper/1606.04586/) (2016)
- [Learning Models with Uniform Performance via Distributionally Robust Optimization](https://mlsanity.com/paper/1810.08750/) (2018)
- [Not All Samples Are Created Equal: Deep Learning with Importance Sampling](https://mlsanity.com/paper/1803.00942/) (2018)
- [On The Power of Curriculum Learning in Training Deep Networks](https://mlsanity.com/paper/1904.03626/) (2019)
- [Learning with Pseudo-Ensembles](https://mlsanity.com/paper/1412.4864/) (2014)
- [Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach](https://mlsanity.com/paper/1610.03425/) (2016)
- [Deep learning with noisy labels: exploring techniques and remedies in medical image analysis](https://mlsanity.com/paper/1912.02911/) (2019)
- [Transformation Consistent Self-ensembling Model for Semi-supervised Medical Image Segmentation](https://mlsanity.com/paper/1903.00348/) (2019)
- [Data augmentation using learned transformations for one-shot medical image segmentation](https://mlsanity.com/paper/1902.09383/) (2019)
- [Stochastic Gradient Descent, Weighted Sampling, and the Randomized Kaczmarz algorithm](https://mlsanity.com/paper/1310.5715/) (2013)
- [Deep Semantic Segmentation of Natural and Medical Images: A Review](https://mlsanity.com/paper/1910.07655/) (2019)
- [Online Batch Selection for Faster Training of Neural Networks](https://mlsanity.com/paper/1511.06343/) (2015)
- [Iterative Learning with Open-set Noisy Labels](https://mlsanity.com/paper/1804.00092/) (2018)
- [Stochastic Optimization with Importance Sampling](https://mlsanity.com/paper/1401.2753/) (2014)
- [Unified Focal loss: Generalising Dice and cross entropy-based losses to handle class imbalanced medical image segmentation](https://mlsanity.com/paper/2102.04525/) (2021)
- [Variance Reduction in SGD by Distributed Importance Sampling](https://mlsanity.com/paper/1511.06481/) (2015)
- [Combo Loss: Handling Input and Output Imbalance in Multi-Organ Segmentation](https://mlsanity.com/paper/1805.02798/) (2018)
- [Semi-Supervised Medical Image Segmentation via Learning Consistency under Transformations](https://mlsanity.com/paper/1911.01218/) (2019)

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

- [Discrepancies are Virtue: Weak-to-Strong Generalization through Lens of Intrinsic Dimension](https://mlsanity.com/paper/2502.05075/) (2025)
- [Data Augmentation: A Fourier Analysis Perspective](https://mlsanity.com/paper/2606.24418/) (2026)
- [Randomized Dimension Reduction with Statistical Guarantees](https://mlsanity.com/paper/2310.01739/) (2023)

## Abstract (arXiv)

Concept shift is a prevailing problem in natural tasks like medical image segmentation where samples usually come from different subpopulations with variant correlations between features and labels. One common type of concept shift in medical image segmentation is the "information imbalance" between label-sparse samples with few (if any) segmentation labels and label-dense samples with plentiful labeled pixels. Existing distributionally robust algorithms have focused on adaptively truncating/down-weighting the "less informative" (i.e., label-sparse in our context) samples. To exploit data features of label-sparse samples more efficiently, we propose an adaptively weighted online optimization algorithm -- AdaWAC -- to incorporate data augmentation consistency regularization in sample reweighting. Our method introduces a set of trainable weights to balance the supervised loss and unsupervised consistency regularization of each sample separately. At the saddle point of the underlying objective, the weights assign label-dense samples to the supervised loss and label-sparse samples to the unsupervised consistency regularization. We provide a convergence guarantee by recasting the optimization as online mirror descent on a saddle point problem. Our empirical results demonstrate that AdaWAC not only enhances the segmentation performance and sample efficiency but also improves the robustness to concept shift on various medical image segmentation tasks with different UNet-style backbones.