SAFLEX: Self-Adaptive Augmentation via Feature Label Extrapolation

TL;DR

SAFLEX is an efficient, self-adaptive data augmentation method that learns to optimize augmented sample weights and labels, improving model robustness across diverse datasets and tasks with minimal additional computational cost.

Contribution

We introduce SAFLEX, a novel bilevel optimization-based approach that enhances existing augmentation pipelines by learning sample weights and labels, reducing noise and errors effectively.

Findings

Effective across natural, medical, and tabular data

Improves few-shot learning and out-of-distribution generalization

Seamlessly integrates with popular augmentation methods

Abstract

Data augmentation, a cornerstone technique in deep learning, is crucial in enhancing model performance, especially with scarce labeled data. While traditional techniques are effective, their reliance on hand-crafted methods limits their applicability across diverse data types and tasks. Although modern learnable augmentation methods offer increased adaptability, they are computationally expensive and challenging to incorporate within prevalent augmentation workflows. In this work, we present a novel, efficient method for data augmentation, effectively bridging the gap between existing augmentation strategies and emerging datasets and learning tasks. We introduce SAFLEX (Self-Adaptive Augmentation via Feature Label EXtrapolation), which learns the sample weights and soft labels of augmented samples provided by any given upstream augmentation pipeline, using a specifically designed efficient bilevel optimization algorithm. Remarkably, SAFLEX effectively reduces the noise and label errors of the upstream augmentation pipeline with a marginal computational cost. As a versatile module, SAFLEX excels across diverse datasets, including natural and medical images and tabular data, showcasing its prowess in few-shot learning and out-of-distribution generalization. SAFLEX seamlessly integrates with common augmentation strategies like RandAug, CutMix, and those from large pre-trained generative models like stable diffusion and is also compatible with frameworks such as CLIP's fine-tuning. Our findings highlight the potential to adapt existing augmentation pipelines for new data types and tasks, signaling a move towards more adaptable and resilient training frameworks.