AdaAugment: Adaptive Data Augmentation

• AdaAugment is a tuning-free, reinforcement learning based framework that dynamically selects augmentation strength per sample based on the model's training state.
• It employs a dual-model architecture combining a CNN target network and an actor–critic policy network to adapt augmentation magnitude and balance underfitting and overfitting risks.
• Experimental results demonstrate consistent gains on benchmarks like CIFAR-10, CIFAR-100, and ImageNet, with faster convergence and improved feature clustering.

AdaAugment is a tuning-free, reinforcement learning–based adaptive augmentation framework developed to address the longstanding challenge of unregulated augmentation variability in deep learning pipelines. Traditional data augmentation (DA) methods, while increasing data diversity, often employ random or fixed-magnitude transformations throughout the entirety of training, disregarding the evolving capacity of the target network. This indiscriminate variability can lead to either underfitting—when large augmentations early in training misalign with immature features—or overfitting—when small or ineffective augmentations late in training fail to diversify robust representations. AdaAugment eliminates the need for heuristic tuning of augmentation strengths by dynamically selecting, for each sample and at every training iteration, the augmentation magnitude that is most appropriate for the current training status of the model (&&&0&&&).

1. Motivation and Problem Definition

Most DA approaches, including Cutout, Mixup, and AutoAugment, apply either uniformly random or pre-specified magnitudes of augmentation. Although this fosters sample diversity, empirical (Yang et al., 2024) and theoretical studies (Bashir et al., 2020) demonstrate that such augmentation practices create a persistent mismatch between data variability and model learning state.

During initial training epochs, networks possess limited feature extraction capabilities, and high-magnitude augmentations can corrupt sample integrity, impeding learning and causing underfitting. Conversely, in later epochs, models become sensitive to minor distributional peculiarities unless augmentation intensity increases accordingly, which if mishandled, precipitates overfitting. The fundamental need is an augmentation strength selection process adaptable in real time, tightly coupled with model status, and individualized per sample—a requirement AdaAugment directly fulfills by formulating augmentation selection as a sequential decision problem amenable to reinforcement learning.

2. Dual-Model Architecture

AdaAugment implements a dual-model architecture consisting of a target network $f_\theta$ and a policy network (parameterized by actor $\pi_\phi$ and critic $V_\psi$), co-optimized in training:

• Target Network $f_\theta$: Standard CNN backbone (e.g., ResNet, WRN), which ingests adaptively augmented samples $\tilde x = e(m, x)$, where $e$ denotes a chosen augmentation operation and $m \in [0,1]$ controls strength. The network is trained by minimizing cross-entropy $\mathcal{L}(f_\theta(\tilde x), y)$.
• Policy Network ($\pi_\phi, V_\psi$): Consists of a lightweight MLP or convolutional encoder. It receives the state vector composed of feature maps from: (i) the non-augmented input $S_{\rm none}$, and (ii) the adaptively augmented input $S_{\rm ada}$, both extracted by a frozen portion of $f_\theta$. The actor outputs a distribution over possible augmentation magnitudes for each sample, and the critic estimates the expected cumulative reward.

Per mini-batch, the actor selects magnitudes, constructs augmented samples, computes three distinct losses (none, full, adaptive), and derives an immediate reward, used to jointly update both networks via actor–critic optimization.

3. Reinforcement Learning Formulation and Reward Design

AdaAugment formulates the per-sample augmentation selection task as a Markov Decision Process (MDP):

• State ($s_t$): Tuple $(S_{\rm none}, S_{\rm ada})$, encoding the current sample's non-augmented and adaptively augmented feature maps.
• Action ($a_t = m_t$): Continuous magnitude $m \in [0,1]$.
• Reward ($r_t$): A curriculum-inspired function balancing underfitting and overfitting risk, based on three losses:

$$L_{\rm none} = \mathcal{L}(f_\theta(x), y), \quad L_{\rm full} = \mathcal{L}(f_\theta(e(1, x)), y), \quad L_{\rm ada} = \mathcal{L}(f_\theta(e(m, x)), y)$$

The reward combines two terms:

$$r =\lambda (L_{\rm full} - L_{\rm ada}) + (1-\lambda) (L_{\rm ada} - L_{\rm none}), \quad \lambda \in [0,1]$$

$\lambda$ is linearly annealed from 1 to 0 throughout training, thus adapting the policy’s focus from avoiding underfitting to addressing overfitting.

The objective is to maximize the expected discounted return: $$J(\phi) = \mathbb{E}_{\tau \sim \pi_\phi}\left[\sum_{t=0}^{T} \gamma^t r_t\right], \quad \gamma=0.99$$ Actor–critic (A2C) algorithms are employed, using the advantage $A_t = r_t + \gamma V_\psi(s_{t+1}) - V_\psi(s_t)$ for gradient updates to both policy and value networks.

4. Adaptive Magnitude Selection Mechanism

The action space for each sample comprises a continuous magnitude $m$ in $[0, 1]$, representing the interpolation between no augmentation ($m=0$) and maximal preset augmentation ($m=1$) for a chosen transformation. At each batch:

1. The actor network processes the sample states and outputs $m_i$ for each.
2. Augmentation $e(m_i, x_i)$ is applied, creating adaptively perturbed inputs.
3. The resulting losses and reward inform the next policy update.
4. This closed-loop process enables AdaAugment to regulate the augmentation intensity per sample, matching the model’s real-time training status.

5. Training Algorithm and Operational Protocol

AdaAugment is tuning-free, with all hyperparameters fixed:

• Discount factor: $\gamma=0.99$
• Curriculum weight: $\lambda$, linearly annealed from 1 to 0
• Joint optimization of target and actor–critic networks using SGD (with prescribed learning rates and batch sizes)

The procedure for each epoch involves:

1. Mini-batch sampling.
2. Extracting states for each sample.
3. Sampling adaptive magnitudes via the actor.
4. Performing augmentations and computing relevant losses and rewards.
5. Backpropagating actor, critic, and target network losses, updating parameters accordingly.

The policy network adds minimal computational and parameter overhead: +0.15M parameters (+1.3%) for ResNet-18, +0.19M (+0.5%) for WRN-28-10, and <0.5 GPU hours additional cost per run.

6. Experimental Results on Benchmark Datasets

Extensive evaluation demonstrates that AdaAugment delivers consistent generalization improvements over both fixed and dynamic DA baselines:

Dataset/Model	AdaAugment	Best Prior DA	Baseline
CIFAR-10/ResNet-18	96.75%	96.51% (AutoAugment)	95.28%
CIFAR-100/WRN-28-10	83.23%	82.90% (Fast-AutoAug)	78.96%
Tiny-ImageNet/ResNet-18	71.25%	70.03% (DADA)	61.38%
ImageNet/ResNet-50	78.2%	77.9% (CutMix)	77.1%

Further empirical findings:

• Convergence: AdaAugment achieves faster test error reduction, particularly evident on CIFAR-10, and retains performance advantages after each scheduled learning-rate drop.
• Transferability: Pre-training on CIFAR-100 and fine-tuning on CIFAR-10 yields a transferred accuracy of 92.82% (compared to 92.55% for the baseline).
• Clustering Effects: On MNIST (demonstrated with t-SNE embeddings), the Dunn Index improves by 108% relative to baseline, indicating increased intra-class compactness and inter-class separation.

7. Theoretical Insights and Empirical Analysis

Under the regime $L_{\rm none} < L_{\rm ada} < L_{\rm full}$, the losses incurred by non-augmented and maximally augmented samples proxy overfitting and underfitting risks, respectively. The curriculum schedule ($\lambda$), by continuously interpolating between these extremes, operationalizes a form of curriculum learning: initially, it enforces small augmentations to avoid underfitting immature models (high similarity, low diversity), and gradually transitions to larger perturbations as learning progresses, counteracting overfitting tendencies in mature networks.

Observed training dynamics confirm that AdaAugment automatically increases augmentation diversity with epoch progression, as measured by declining sample similarity and rising diversity indices (see Fig. 3 in the source). These effects are accompanied by significant improvements in cluster structure of learned representations.

In summary, AdaAugment's principal contribution resides in its reinforcement learning–driven, per-sample adaptive control of augmentation magnitude, delivered in a tuning-free fashion. This strategy tightly synchronizes augmentation variability with model development stage, yielding systematic reductions in both underfitting and overfitting risks, and producing consistent, architecture-agnostic improvements in generalization accuracy (Yang et al., 2024).