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

TL;DR

This paper introduces AdaWAC, an adaptive data augmentation regularization method that dynamically balances supervised and unsupervised learning to improve robustness and performance in medical image segmentation under concept shift.

Contribution

The paper proposes a novel adaptive reweighting algorithm, AdaWAC, that effectively exploits label-sparse samples by balancing regularization and supervised loss, with theoretical convergence guarantees.

Findings

AdaWAC improves segmentation accuracy across various tasks.

The method enhances robustness to concept shift.

Empirical results show increased sample efficiency.

Abstract

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.