Source-linked AI summary
Population Based Augmentation: Efficient Learning of Augmentation Policy Schedules
Daniel Ho, Eric Liang, Ion Stoica, Pieter Abbeel, Xi Chen
TL;DR
Effective augmentation policy search is difficult because the candidate space is large, while AutoAugment’s computation is impractical for ordinary users. PBA uses Population Based Training to learn nonstationary augmentation schedules instead of fixed policies, matching AutoAugment across CIFAR-10, CIFAR-100, and SVHN with far less compute. The approach is open source, though its search hyperparameters and schedule duration leave scope for further improvement.
Problem
Choosing effective augmentation policies from a large operation space is challenging, and AutoAugment is computationally infeasible for ordinary users.
Method
PBA uses Population Based Training to jointly optimize child models and epoch-dependent augmentation parameters, producing a schedule rather than a fixed policy.
Results
PBA reaches similar final performance to AutoAugment across network models and leads slightly on PyramidNet and Wide-ResNet-28-10 while performing comparably on Shake-Shake models.
Takeaways & Limitations
PBA provides a fast, efficient augmentation-search formulation that can produce competitive schedules and is released as open-source code.
Takeaways & Limitations
The search may improve with more careful tuning of the distribution for the maximum number of augmentation functions applied per batch.
Abstract
from arXiv · showhide
A key challenge in leveraging data augmentation for neural network training is choosing an effective augmentation policy from a large search space of candidate operations. Properly chosen augmentation policies can lead to significant generalization improvements; however, state-of-the-art approaches such as AutoAugment are computationally infeasible to run for the ordinary user. In this paper, we introduce a new data augmentation algorithm, Population Based Augmentation (PBA), which generates nonstationary augmentation policy schedules instead of a fixed augmentation policy. We show that PBA can match the performance of AutoAugment on CIFAR-10, CIFAR-100, and SVHN, with three orders of magnitude less overall compute. On CIFAR-10 we achieve a mean test error of 1.46%, which is a slight improvement upon the current state-of-the-art. The code for PBA is open source and is available at https://github.com/arcelien/pba.
1. Introduction
Data augmentation can improve generalization, but selecting effective policies is difficult and AutoAugment requires substantial computation. PBA addresses this by learning augmentation schedules efficiently while maintaining competitive performance.
- Data augmentation increases image diversity and regularizes neural networks against overfitting.
- PBA matches AutoAugment’s CIFAR-10 accuracy across network models while using 1,000x fewer GPU hours.At $1.5 per GPU hour, the caption estimates about $7.5 for PBA versus $7,500 for AutoAugment; the same scaling holds on SVHN.
- Selecting augmentation functions is important because redundant or overly aggressive transformations can slow training and introduce dataset bias.
- PBA learns an epoch-dependent augmentation schedule with orders of magnitude less compute than fixed-policy search.The schedule defines the augmentation policy for each training epoch rather than applying identical transformations throughout training.
- PBA learns a robust CIFAR-10 augmentation policy in five hours on one NVIDIA Titan XP GPU.The authors describe this pre-computation cost as marginal relative to training large CIFAR-10 networks for several days.
- The authors release open-source code intended to let ordinary workstation users experiment with augmentation search and operations.
2. Background
Prior work established augmentation as a useful regularizer and automated policy search as a route to stronger image models. AutoAugment searches a large fixed-policy space using reinforcement learning, whereas PBA learns a time-varying schedule with PBT.
- Data augmentation: Common transformations such as cropping, flipping, rotation, scaling, and translation generate additional training samples and improve generalization with limited data.
- Automated search: Automated augmentation research has used manifold learning, Bayesian optimization, generative adversarial networks, and image-combination methods.
- PBA: PBA generates one augmentation function f(x, t), where t is the current epoch, instead of AutoAugment’s ensemble of fixed-policy functions.
- AutoAugment: AutoAugment trains a controller RNN with reinforcement learning to output augmentation policies that maximize validation accuracy.
- AutoAugment: AutoAugment policies contain five sub-policies, each with two operations and associated parameters, with one sub-policy randomly selected per batch.The final policy concatenates five best-performing policies for 25 sub-policies in total.
- AutoAugment: AutoAugment evaluates 15,000 sampled policies on a reduced-data Wide-ResNet-40-2 child model trained for 120 epochs.CIFAR-100 uses an augmentation policy transferred from CIFAR-10.
3. Population Based Augmentation
PBA formulates augmentation search as hyperparameter schedule learning and uses Population Based Training to optimize augmentation parameters jointly with child models. Its smooth, time-varying representation supports efficient exploration despite a larger schedule space.
- Why Augmentation Schedules?: PBA searches for an augmentation schedule over training epochs rather than a fixed policy.The paper argues that schedules can be more efficient to search because useful augmentations may differ between early and late training.
- Learning a Schedule: Population Based Training trains a parallel population, then clones strong models and perturbs hyperparameters of weaker models through exploit-and-explore steps.Because weights are cloned rather than reinitialized, total computation scales with training one model times the population size.
- Learning a Schedule: PBA optimizes probability and magnitude schedules for two instances of each augmentation operation while keeping the operation identities fixed.
- Policy template: PBA first shuffles operations and applies them sequentially until a sampled limit of zero, one, or two operations is reached.The limit is sampled with probabilities 0.2, 0.3, and 0.5, respectively.
- Policy Search Space: The policy template uses 15 AutoAugment operations, excluding SamplePairing, with two magnitude-probability tuples per operation and 60 hyperparameters.
- Policy Search Space: PBA’s scheduled search space contains approximately 1.75 × 10^61 possibilities for a single operation, compared with 2.8 × 10^32 for AutoAugment.
- PBT procedure: PBT evaluates trials after each epoch, makes them ready after three epochs, and applies truncation selection to the bottom 25% of the population.
4. Experiments and Analysis
The experiments compare PBA with state-of-the-art and random-search baselines, examine whether nonstationary schedules matter, and test computational efficiency and hyperparameter sensitivity. PBA is competitive across datasets and models while using substantially less computation, and schedule structure improves accuracy over fixed or shuffled alternatives.
- Accuracy: PBA leads AutoAugment slightly on PyramidNet and Wide-ResNet-28-10 and performs comparably on Shake-Shake models.
- Learned schedules: PBA schedules emphasize Cutout, Posterize, Invert, Equalize, and AutoContrast throughout CIFAR-10 training, while all augmentations appear at least moderately at some point.
- Accuracy: CIFAR-100 results remain competitive with AutoAugment and significantly outperform the baseline and Cutout-only training.
- Computational cost: PBA uses 3,200 training epochs versus AutoAugment’s approximately 1.8 million, reducing compute by over 500x.
- Computational cost: More than 250 random trials are needed for expected child accuracy to approach a single 16-trial PBA run, giving PBA over an order of magnitude speedup.
- Does having a schedule matter?: Fixed-policy training degrades accuracy by approximately 10% on average, while shuffled schedules also perform significantly worse than the full schedule.
5. Conclusion
PBA is introduced as a fast, efficient method for learning state-of-the-art augmentation policy schedules, with implementation available as open source.
- PBA learns state-of-the-art augmentation policy schedules quickly and efficiently.
- PBA is designed for straightforward implementation within any Population Based Training framework.
- The authors release PBA code as open source.
A. PBA Scalability with Compute
PBA scalability depends on population size and model choice: 16 WRN-40-2 models perform best, while ResNet-20 uses about half the compute with somewhat lower accuracy.
- A population size of 16 WRN-40-2 models performs best for learning an effective schedule.More than 16 trials did not help, while fewer than 16 reduced performance.
- Population size controls both search-space exploration and PBA’s computational overhead.
- ResNet-20 requires about half the compute of WRN-40-2 but does not achieve as high test accuracy.The results were relatively close, and the smaller model’s training accuracy plateaus faster.
B. Model Hyperparameters
The paper specifies WideResNet-40-2 training hyperparameters for schedule discovery and final evaluation, while Figure 7 illustrates operation settings on an example image.
- Table 5 lists hyperparameters for WideResNet-40-2 schedule discovery and final-model training.
- Figure 7 displays each augmentation operation with its name, probability, and magnitude value.
C. SVHN discovered schedule
The SVHN schedule is visualized through example-image augmentations and parameter evolution, with several operations remaining frequent across epochs and all operations appearing at some stage.
- SVHN discovered schedule: Cutout, Translate Y, Shear X, and Invert appear with high probability across all epochs.
- SVHN discovered schedule: All augmentations appear with reasonable probability at some point in the schedule.
- SVHN discovered schedule: Table 4 reports test error during PBT search and subsequent schedule evaluation across population sizes and models.
- SVHN discovered schedule: Table 5 provides evaluation hyperparameters for CIFAR-10, CIFAR-100, and Reduced-CIFAR-10, with limited tuning beyond specified models.
- SVHN discovered schedule: Figure 8 visualizes the evolution of PBA operation parameters, averaging the two parameter entries for each operation.