FreeAugment: Data Augmentation Search Across All Degrees of Freedom

TL;DR

FreeAugment addresses the data augmentation search bottleneck by jointly optimizing all four degrees of freedom—policy depth, transformation types, order, and magnitudes—through a fully differentiable, bilevel framework. It introduces learnable depth via $Gumbel$-$Softmax$, permutation learning via $Gumbel$-$Sinkhorn$, and differentiable magnitude sampling with reparameterization, enabling end-to-end gradient-based search without heuristics. Empirical results across CIFAR-10/100, ImageNet-100, and DomainNet show state-of-the-art or competitive performance, validating the method's generality and domain transferability. The approach reduces manual tuning, mitigates transformation repetition, and scales to deeper policies, offering a practical path to automatic, task-tailored augmentation policies.

Abstract

Data augmentation has become an integral part of deep learning, as it is known to improve the generalization capabilities of neural networks. Since the most effective set of image transformations differs between tasks and domains, automatic data augmentation search aims to alleviate the extreme burden of manually finding the optimal image transformations. However, current methods are not able to jointly optimize all degrees of freedom: (1) the number of transformations to be applied, their (2) types, (3) order, and (4) magnitudes. Many existing methods risk picking the same transformation more than once, limit the search to two transformations only, or search for the number of transformations exhaustively or iteratively in a myopic manner. Our approach, FreeAugment, is the first to achieve global optimization of all four degrees of freedom simultaneously, using a fully differentiable method. It efficiently learns the number of transformations and a probability distribution over their permutations, inherently refraining from redundant repetition while sampling. Our experiments demonstrate that this joint learning of all degrees of freedom significantly improves performance, achieving state-of-the-art results on various natural image benchmarks and beyond across other domains. Project page at https://tombekor.github.io/FreeAugment-web

Figures (10)

• Figure 1: FreeAugment differentiable search space over four degrees of freedom (DoF). The search space over policy depth, transformations type, order, and magnitudes is constructed in a fully differentiable manner and these are learned end-to-end.
• Figure 2: Illustration of FreeAugment's policy depth search mechanism (Left) and learned depth distribution along the search (Right), where the last epoch represents the resulting probability distribution over policy depth.
• Figure 3: Learning the transformation types and order as the distribution over permutations of transformations (Top Middle) and sampling from it (Bottom). The learning dynamics of a column (Top Right) and a row (Top Left) of the matrix along the search are shown, with the final values embedded in the matrix.
• Figure 4: The dynamics of the learnable magnitude ranges of transformations at the first augmentation layer during the search over the CIFAR10 datasets with WRN-40-2. The lower and upper bounds of the uniform distribution are presented with solid lines and with their mean value with a dashed line.
• Figure 5: The benefits of a variable policy depth. Top-1 test accuracy on CIFAR100 with WRN-40-2 obtained by FreeAugment compared to fixed-depth policies. Each point quantifies the average accuracy over 5 random seeds, with the shaded areas indicating the 95% confidence intervals. The horizontal line marks the performance of FreeAugment of learned variable policy depth.