Beyond Cropping and Rotation: Automated Evolution of Powerful Task-Specific Augmentations with Generative Models

TL;DR

EvoAug tackles the problem of learning task-specific data augmentations in the presence of rich generative models by automatically discovering augmentation policies via an evolutionary search. It leverages a diverse operator set that includes classical transforms, diffusion-based conditioning via ControlNet, and NeRF-based 3D rotations, organized as augmentation trees to encode hierarchical transformations. The main contributions are the augmentation-tree framework, novel fitness strategies for low-data and one-shot settings, and empirical evidence showing improvements on fine-grained, few-shot and one-shot classification across multiple datasets. The approach reduces the syn-to-real gap in synthetic augmentations and demonstrates that learned generative augmentations can outperform traditional baselines in scarce-data regimes, with implications for robust model training in data-scarce domains.

Abstract

Data augmentation has long been a cornerstone for reducing overfitting in vision models, with methods like AutoAugment automating the design of task-specific augmentations. Recent advances in generative models, such as conditional diffusion and few-shot NeRFs, offer a new paradigm for data augmentation by synthesizing data with significantly greater diversity and realism. However, unlike traditional augmentations like cropping or rotation, these methods introduce substantial changes that enhance robustness but also risk degrading performance if the augmentations are poorly matched to the task. In this work, we present EvoAug, an automated augmentation learning pipeline, which leverages these generative models alongside an efficient evolutionary algorithm to learn optimal task-specific augmentations. Our pipeline introduces a novel approach to image augmentation that learns stochastic augmentation trees that hierarchically compose augmentations, enabling more structured and adaptive transformations. We demonstrate strong performance across fine-grained classification and few-shot learning tasks. Notably, our pipeline discovers augmentations that align with domain knowledge, even in low-data settings. These results highlight the potential of learned generative augmentations, unlocking new possibilities for robust model training.

Figures (7)

• Figure 1: Example image augmentations using our pipeline. Classical augmentations include color jitter, rotation, and random cropping. Canny, color, depth, and segment use existing image information to steer a ControlNet diffusion model. NeRF uses a zero-shot NeRF to perform a 3D rotation.
• Figure 2: Mutation and Crossover for Augmentation Trees
• Figure 3: Tree Learning Pipelines. (a) K-Fold applies when there is more than one example per class. (b) We can measure cluster quality for the 1-shot case. (c) We can duplicate the image and assume the problem to be 2-shot instead of 1-shot (d) We can simply use training loss, though it is risky to assume that lower train loss equates to better performance.
• Figure 4: Success Case: ViT-B/16-Flowers
• Figure 5: Failure Case: ResNet50-Flowers