DADA: Differentiable Automatic Data Augmentation

Let $\mathcal{O}$ be the set of all image pre-processing operations in our search space. Then the candidate $N$ sub-policies $\mathcal{S}$ are all the combinations of $k$ operations in $\mathcal{O}$. To determine which sub-policies should be selected, we sample the sub-policy from a Categorical distribution $p(c|\pi)$ (where we can sample a one-hot random variable) with probabilities $\pi$. Probabilities are computed as a softmax over parameters $\alpha=\alpha_{1:N}$ defining preference for sub-policies:

After $\alpha$ optimized, we can select the sub-policies with the top probabilities as the final DA policy $P$. Therefore, to optimize the DA policy, our task becomes optimizing the parameters $\alpha$ of the sub-policy sampling distribution.

Given a particular sub-policy choice, the constituent operations are executed or not according to a Bernoulli distribution sample. We start from the simple example: a sub-policy $s$ contains only one operation $O_{1}^{s}$. Then, the image operation $O$ (we omit the superscript and subscript of $O_{1}^{s}$ for simplicity) with application probability $\beta$ and magnitude $m$ can be represented:

To generalize Eq. (4) to a sub-policy with multiple operations, we formulate the composition of $O_{i}$ and $O_{i-1}$ in Eq. (2) as below:

where $b_{i-1}\sim\text{Bernoulli}(\beta_{i-1})$, $b_{i}\sim\text{Bernoulli}(\beta_{i})$.

With the above formulations, the DA policy can be sampled from a joint distribution $p(c,b|\alpha,\beta)=p(b|\beta,c)p(c|\alpha)$ (we use $p(c|\alpha)$ to describe Eq. (3)). Therefore, our objective can be represented as:

where $\mathcal{L}_{w}(c,b)$ is the validation-set loss achieved by the network (with weights $w$) which is trained using the DA policy {$c,b$} on training set. Unlike AutoAugment which uses validation accuracy as reward, we use the validation loss to facilitate gradient computation.

The Categorical distribution is not differentiable w.r.t. the parameter $\alpha$, therefore we use the Gumbel-Softmax reparameterization trick to achieve the differentiable relaxation. This reparameterization is also used in network architecture search, e.g SNAS [26] and GDAS [5]. With the Gumbel-Softmax reparameterization, Eq. (3) becomes:

where $g=-\log(-\log(u))$ with $u\sim\text{Uniform}(0,1)$, and $\tau$ is the temperature of softmax function. In our implementation, the straight-through gradient estimator is applied: the backward pass uses differentiable variables as Eq. (7), the forward pass uses discrete variables as shown:

Similar to the Categorical distribution, the Bernoulli distribution is not differentiable w.r.t. $\beta$. We apply the same reparameterization trick to the Bernoulli distribution:

Similar to Eq. (8), we only execute the operation if $b>0.5$ but backpropagate using the gradient estimated by Eq. (9).

Different from optimizing the discrete distribution parameters, we optimize augmentation magnitudes by approximating their gradient. Since some operations (e.g. flipping and rotation) in our search space are not differentiable w.r.t. the magnitude, we apply the straight-through gradient estimator [1] to optimize the magnitude. For an image $\hat{x}=s(x)$ augmented by sub-policy $s$, we approximate the gradient of magnitude (of the sampled operations) w.r.t. each pixel of the augmented image:

Then the gradient of the magnitude w.r.t. our objective $L$ can be calculated as:

Here, we consider how to estimate the gradient of parameters of distribution $p(q|\theta)$, which can represent either $p(c|\alpha)$ or $p(b|\beta)$ in our algorithm. Unlike the Gumbel-Softmax estimator, the gradient of RELAX cannot simply be achieved by the backward of loss function. Furthermore, the RELAX estimator even does not require the loss function $\mathcal{L}$ to be differentiated w.r.t. the distribution parameters. It means we do not need to apply continuous relaxation the forward stage like Eq. (3) or Eq. (4). Instead, the RELAX estimator requires a differentiable neural network $c_{\phi}$ which is a surrogate of loss function, where $\phi$ are the parameters of neural network $c$. To make $c_{\phi}$ be differentiable w.r.t. the distribution parameters, we need the same Gumbel-Softmax reparameterized distribution $p(z|\theta)$ for $p(q|\theta)$, where $\theta$ is the distribution parameter and $z$ is a sampled continuous variable. $p(z|\theta)$ can be either the Gumbel-Softmax reparameterized distributions of $p(c|\alpha)$ or $p(b|\beta)$. Then, in the forward stage, the DA policy can also be sampled from $z\sim p(z|\theta)$ but only forward the policy using deterministic mapping $q=H(z)=b\sim p(b|\theta)$ to guarantee the probability distribution curve is the same. Furthermore, since RELAX applies the control variate at the relaxed input $z$, we also need a relaxed input conditioned on the discrete variable $q$, denoted as $\tilde{z}\sim p(z|q,\theta)$. Then the gradient estimated by RELAX can be expressed as:

To achieve a small variance estimator, the gradient of parameters $\phi$ can be computed as Eq. (13), which can be jointly optimized with the parameters $\theta$.

As shown in Eq. (3), the sub-policies can be sampled from the Categorical distribution. To apply the RELAX estimator, we first sample $z$ from the relaxed Categorical distribution $\text{RelaxCategorical}(\alpha,\tau)$ from Eq. (7), but we only forward the sub-policy with $c=\text{one\_hot}(\text{argmax}(z))$ in Eq. (8). We further sample $\bar{z}$ conditioned on variable $c$ to control the variance:

As shown in Eq. (4), the parameters $b$ can be sampled from Bernoulli distribution. To adapt to the RELAX gradient estimator, we also utilize the Gumbel-Softmax reparameterization trick for $p(z|\beta)$ like Eq. (9), but only forward the operation with $b=\mathbb{I}(z>0.5)$, where $\mathbb{I}$ is the indicator function. To control the variance, we sample $\tilde{z}$ conditioned as the sampled parameters $b$:

As for the gradient of magnitudes, we approximate them following Eq. (10). Since we need to estimate the gradient of the parameters of joint distribution $p(c,b|\alpha,\beta)$, we feed $c_{\phi}$ with the relaxed variables $c$ and $b$ in the same time for the surrogate of loss function.

For the search space, we follow AA for a fair comparison. Specifically, our search space contains the same $15$ data augmentation operations as PBA [10], which is also the same as AA and Fast AA except that the SamplePairing [11] is removed. Following AA, our sub-policy consists of two DA operations ($k=2$), and the policy consists of $25$ sub-polices.

Following [3, 10, 15], we search the DA policies on the reduced datasets and evaluate on the full datasets. Furthermore, we split half of the reduced datasets as training set, and the remaining half as validation set for the data augmentation search. In the search stage, we search the policy parameters for $20$ epochs on the reduced datasets. We use the Adam optimizer for the policy parameters $d=\{\alpha,\beta,m,\phi\}$ optimization with learning rate $\eta_{d}=5\times 10^{-3}$, momentum $\beta=(0.5,0.999)$ and weight decay $0$. For the optimization of neural network parameters, we use momentum SGD as optimizer, with the same hyper-parameters as evaluation stage except the batch size and the initial learning rate. In the search stage, we only apply the data augmentation policy to training examples, and set the batch size to $128$ for CIFAR-10 and $32$ for other datasets. The initial learning rate is set according to the batch size by the linear rule. We set $\tau$ to $0.5$ and use a two layer fully connected neural network with $100$ hidden units for $c_{\phi}$. Parameters $\alpha$ are initialized to $10^{-3}$. Parameters $\beta$ and magnitudes $m$ are all initialized to $0.5$.

We use the official publicly available code of Fast AA to evaluate the searched DA policies for a fair comparison. Following Fast AA [15] and AA [3], we use the same hyper-parameters for policy evaluation: weight decay, learning rate, batch size, training epoch.

Both CIFAR-10 and CIFAR-100 have $50,000$ training examples. Following [3, 10, 15], we conduct DA optimization on the reduced CIFAR-10 dataset ($4,000$ randomly selected examples) and evaluate the trained policies on the full CIFAR-10 test set. [3, 10, 15] use the discovered policies from CIFAR-10 and evaluate these policies on CIFAR-100. Since our DADA is much more efficient than other methods, thus, we conduct both the search (using a reduced dataset of $4,000$ randomly selected examples) and evaluation on CIFAR-100. Following [3, 10, 15], we search the DA policy using a Wide-ResNet-40-2 network  [29] and evaluate the searched policy using Wide-ResNet-40-2 [29], Wide-ResNet-28-10 [29], Shake-Shake [6] and PyramidNet+ShakeDrop [27].

From the results in Table 1(b) and 2, we can see that DADA requires significantly less computation than the competitors, while providing comparable error rate. For example, we require only $0.1$ GPU hours for policy search on reduced CIFAR-10, which is at least one order of magnitude faster than AA ($50,000\times$) and Fast AA ($35\times$). Similar to CIFAR-10, we achieve very competitive performance on CIFAR-100 yet with much less searching cost. Despite the lower error rates of OHL AA, OHL AA is not directly comparable to other methods since it uses a larger and dynamic search space.

To verify the generalization ability of our search algorithm for different datasets, we further conduct experiments with a larger dataset: SVHN. SVHN dataset has $73,257$ training examples (‘core training set’), $531,131$ additional training examples, and $26,032$ testing examples. Following AA and Fast AA, we also search the DA policy with the reduced SVHN dataset, which has $1,000$ randomly selected training samples from the core training set. Following AA, PBA and Fast AA, we evaluate the learned DA policy performance with the full SVHN training data. Unlike AA, we use the Wide-ResNet-28-10 [29] architecture in the search stage and evaluate the policy on the Wide-ResNet-28-10 [29] and Shake-Shake (26 2x96d) [6]. Our results are shown in Table 3. As shown in Table 3, our DADA achieves similar error rate to PBA, slightly worse than AA and Fast AA. However, we only use $0.1$ GPU hours in the search stage.

Finally, we evaluate our algorithm on the ImageNet dataset. We train the policy on the same ImageNet subset as Fast AA, which consists of $6,000$ examples from a randomly selected $120$ classes. We use ResNet-50 [9] for policy optimization and report the performance trained with the full ImageNet dataset. As shown in Table 1(b) and Table 4, we achieve very competitive error rate against AA and Fast AA while requiring only $1.3$ GPU hours in the search stage – compared to $15000$ and $450$ hours respectively for these alternatives. Again, OHL AA  [16] uses a larger and dynamics search space, thus, it is not comparable to other methods.

One technical contribution of this paper is the derivation of a RELAX estimator for DA policy search, which removes the bias of the conventional Gumbel-Softmax estimator. To evaluate the significance of this contribution, we conduct the search experiments for both estimators with the same hyper-parameters on the CIFAR-10 dataset. As we can see from Fig. 3(a), the policy found using our RELAX estimator achieves better performance on CIFAR-10 compared with Gumbel-Softmax estimator.

We further evaluate the performance when we train the DA policy on the full dataset rather than on the reduced one, noting that this is feasible for the first time with DADA, due to its dramatically increased efficiency compared to alternatives. We conduct DA policy search on both the reduced and full CIFAR-10 dataset with Wide-ResNet-40-2. As we can see from Table 6, the policy searched on the full dataset works better than that on the reduced one evaluated on CIFAR-10.

We finally bring together some of these results and compare the speed accuracy trade-off provided by DADA, Fast AA and AA in Fig. 3(b) for CIFAR-10 on WRN-40-2 architecture. Fast AA does not benefit from DA policy on the full CIFAR-10 dataset. However, DADA provides an excellent tradeoff, especially at low resource operating points.