Source-linked AI summary
ActiveAugment: Online Active Learning for Augmentation Selection in Deep Learning
Noah Videcrantz, Mostafa Mehdipour Ghazi
TL;DR
Deep learning often needs more labelled data, while conventional augmentation is static, expensive to search, or model-agnostic. ActiveAugment selects model-fragile views online using uncertainty and feature discrepancy, then trains with classification and supervised contrastive objectives. Across eight datasets and varied architectures and training regimes, it outperforms standard augmentation baselines, especially at low labelling budgets and on medical imaging, while assuming candidate transformations are label-preserving.
Problem
Deep learning requires large labelled datasets, but augmentation is commonly static or model-agnostic despite models’ changing vulnerabilities and limited labelling budgets.
Method
ActiveAugment scores candidate augmented views per sample using predictive uncertainty and feature discrepancy, then jointly trains classification and supervised contrastive objectives.
Results
ActiveAugment consistently outperforms AutoAugment, RandAugment, and TrivialAugment across datasets, architectures, and training regimes, with the strongest gains at low labelling budgets and in medical imaging.
Takeaways & Limitations
Online augmentation selection adapts to the evolving model state and shows strong cross-domain adaptability, while selection strategy directly shapes generalisation.
Takeaways & Limitations
The candidate space assumes, rather than automatically verifies, that predefined transformations preserve labels, and uses three discrete intensity levels during selection.
Abstract
from arXiv · showhide
Data augmentation is a cornerstone of deep learning pipelines, yet existing strategies treat it as a static, model-agnostic preprocessing step, either relying on expensive dataset-specific policy search or applying transformations uniformly at random, regardless of what the model has already learned. We introduce ActiveAugment, a unified framework that treats augmentation selection as an online active learning problem. For each training minibatch, ActiveAugment generates a pool of candidate augmented views and scores each candidate using a combination of the model's predictive uncertainty and the feature discrepancy induced by the augmentation. The augmentation under which the current model is most fragile is selected per sample, and the model is then trained with a joint supervised classification and supervised contrastive objective that enforces intra-class invariance to the selected augmentations while maintaining inter-class separation. We evaluate ActiveAugment on eight benchmark datasets spanning natural and medical imaging, using CNN and transformer architectures across three training regimes (training from scratch, full fine-tuning, and linear probing), and comparing eight active selection strategies for augmentation scoring. ActiveAugment outperforms AutoAugment, RandAugment, and TrivialAugment under controlled augmentation shifts across all domains and budgets, with the most pronounced gains at low labelling budgets. On medical imaging datasets, where data is scarce and domain shift relative to natural-image pretrained models is large, ActiveAugment achieves higher test F1 than all baselines, demonstrating strong cross-domain adaptability. Our analysis reveals that the augmentation selection policy evolves meaningfully during training and that strategy choice has a direct impact on generalisation. Code is available at: https://github.com/noahvide/ActiveAugment.
1 Introduction
ActiveAugment addresses the mismatch between static augmentation pipelines and models whose vulnerabilities change during training by selecting augmentations online. It combines active scoring with contrastive learning and shows consistent gains across broad evaluation settings, especially under low labelling budgets and in medical imaging.
- Motivation: Limited labelled data and static augmentation policies motivate a data-efficient approach that adapts augmentation to the model’s current state.Existing pipelines may be computationally expensive, model-agnostic, or increasingly misaligned as model invariances evolve.
- ActiveAugment: ActiveAugment selects candidate augmented views online using predictive uncertainty and augmentation-induced feature discrepancy, prioritising transformations under which the model is most fragile.The framework separates selection from learning and scores candidates for each training minibatch.
- ActiveAugment: The joint objective pairs active augmentation selection with supervised contrastive learning to enforce intra-class invariance while maintaining inter-class discriminability.This connects the selected, challenging views to the representation-learning objective rather than treating augmentation as preprocessing alone.
- Evaluation design: Eight active-learning strategies are translated into online augmentation-scoring strategies, enabling analysis of their selection behaviour and performance consequences.The strategies span uncertainty-based, diversity-based, and hybrid approaches.
- Evaluation design: Evaluation spans eight datasets, two backbone architectures, and three training regimes, with consistent gains over AutoAugment, RandAugment, and TrivialAugment.The strongest gains occur at low labelling budgets and in medical imaging, where domain-adaptive augmentation selection is most beneficial.
2 Related Work
Prior augmentation methods expand training data through heuristics, search, or randomised policies, while active learning selects informative labelled instances under budget constraints. ActiveAugment connects these ideas by selecting informative augmented views and pairing them with supervised contrastive learning.
- Data augmentation: Data augmentation increases effective dataset size and can encourage reliance on robust features rather than surface-level memorisation.The approach has developed from geometric and photometric transformations to methods such as Cutout and Mixup.
- Automated augmentation: AutoAugment searches dataset-specific policies expensively, whereas RandAugment reduces search to two hyperparameters but remains dataset-specific and model-agnostic.These methods differ in search cost while retaining static policy selection.
- Hardness and active selection: MaxUp and AugMax optimise augmentation hardness, while ActiveAugment scores candidate views using predictive informativeness and relative representation shift.ActiveAugment therefore differs in both its online-selection framing and its use of active-learning acquisition strategies.
- Active learning: Active learning selects informative instances under labelling-budget constraints, using uncertainty and diversity as complementary considerations.Uncertainty-based methods target low-confidence samples, while diversity-based methods avoid redundant selections.
- Connection to ActiveAugment: Active-learning informativeness measures can be applied to augmented views as well as novel images, motivating their direct translation into augmentation selection.ActiveAugment combines this selection idea with supervised contrastive learning, which treats same-class augmented views as positives and different-class views as negatives.
3 Method
ActiveAugment selects, for each sample, the augmentation that most challenges the current model, then trains on it with supervised classification and contrastive objectives. Selection combines predictive uncertainty with feature discrepancy, while the contrastive term promotes within-class invariance and between-class separation.
- Augmentation Space: ActiveAugment defines low, medium, and high intensity levels for parameterised augmentations, with raw strengths sampled within fixed ranges.Candidate transformations and intensity ranges are set before training; the resulting discrete space contains approximately |A|×3 candidates, excluding intensity-free augmentations.
- Active Invariance Learning: The method separates online selection from learning: it identifies each sample’s most non-invariant domain-valid augmentation and enforces invariance specifically to those selections.This avoids applying invariance uniformly across augmentations that may expose model fragilities to different degrees.
- Augmentation Scoring: Each candidate is scored using predictive uncertainty and feature discrepancy, measuring instability in prediction and representation displacement from the clean sample.The discrepancy is normalised by the mean discrepancy for the same augmentation to reduce domination by globally large or small embedding norms.
- Theoretical Motivation: The selection criterion is motivated by augmentation-shift generalisation bounds: discrepancy estimates feature-distribution shift, while predictive entropy targets updates with larger expected loss reduction.The greedy rule a_i* = argmax_k s_ik is presented as an approximation to jointly maximising bound tightening.
- Online Selection Algorithm: For every minibatch, the algorithm samples an intensity, generates one candidate view per augmentation, selects the highest-scoring candidate per sample, and trains on original-plus-selected views.The expanded minibatch updates both encoder and classifier parameters.
- Training Objective: Training combines cross-entropy on original and selected augmented samples with SupCon over the expanded minibatch, using cosine similarity and a balance parameter λ > 0.The objective is described as preventing feature collapse, encouraging intra-class invariance, and maintaining inter-class separation; eight active-learning strategies are implemented for augmentation scoring.
4 Experiments and Results
Across controlled low-label experiments, ActiveAugment generally matches or outperforms established augmentation baselines, with especially clear advantages on overall robustness and medical imaging. Its selection behavior changes with domain and training, favoring perturbations that expose model fragilities.
- Natural-image results: ActiveAugment matches or outperforms all baselines on overall natural-image test sets, with the clearest and most consistent gains at low labelling budgets.Its advantage shrinks at larger budgets, while CIFAR-100 shows a visible margin across the full budget range.
- Natural-image results: On clean natural-image tests, augmentation improves CIFAR-10 and CIFAR-100, while ActiveAugment remains competitive with state-of-the-art methods across budgets.For STL-10 and MNIST, augmentation provides no advantage over no augmentation under the employed transformations.
- Robustness analysis: ActiveAugment’s overall-test gains primarily reflect improved robustness to the evaluated augmentations rather than solely improved clean-image accuracy.When transformations are held out during training, it remains competitive or improves over baselines, but gains are smaller than for observed transformations.
- Selection-strategy analysis: BADGE, Margin, Entropy, and Coreset outperform BAL, Least Confident, and TypiClust on MNIST, with the largest gap under Flip.The strongest strategies approach near-perfect F1 at moderate budgets, while weaker strategies lag throughout training.
- Selection-strategy analysis: Feature discrepancy drives non-uniform random selection, with Flip selected at 22.5% versus an expected 16.7% and Gaussian Blur at 13.4%.Strategies favoring Flip, Gaussian Noise, and Rotation generally outperform those preferring Gaussian Blur, which comparatively weakly challenges learned representations.
- Policy evolution: Selection confidence rises monotonically during training, while Flip increases from below 20% to 30–35% selection frequency by the end.Early preferences for low-intensity Gaussian Noise and Rotation give way to Brightness and Contrast, whereas several augmentation settings are virtually never selected.
- Medical-image results: ActiveAugment is consistently among the top-performing methods on clean medical-image tests across four datasets and all budget levels.The authors attribute the stronger clean-set advantage to adaptation beyond natural-image policies; on BRISC it reaches approximately 0.95 macro-F1 at the highest budget and remains competitive on low-budget ISIC-2019.
- Medical-image results: Gaussian Noise is selected as frequently as Flip by several medical-image strategies, and its performance gap is most pronounced on FETAL-PLANES.This contrasts with natural images and indicates that noise-based perturbations are particularly discriminative in the medical domain.
5 Conclusion
ActiveAugment reformulates augmentation selection as online active learning, scoring model-fragile views and training with classification plus supervised contrastive objectives. It consistently outperforms established augmentation baselines, especially at low labelling budgets and under medical-image domain shift.
- ActiveAugment scores augmented views per sample using predictive uncertainty and normalized feature discrepancy, then enforces invariance to selected model-fragile augmentations.The training objective jointly combines classification and supervised contrastive learning.
- ActiveAugment consistently outperforms AutoAugment, RandAugment, and TrivialAugment across datasets, backbones, and training regimes, with the largest gains at low labelling budgets.
- On medical imaging datasets, ActiveAugment adapts to larger domain shifts by learning augmentation policies from the evolving model state.The paper identifies this cross-domain adaptability as a practical advantage.
- Selection policies evolve during training, and strategies favoring Flip and Gaussian Noise generalize better than strategies favoring Gaussian Blur.The analysis links strategy choice directly to generalization and provides guidance for selecting active scoring strategies.