\section{Introduction}

Data augmentation is an effective way to improve the performance of CNN models for image recognition tasks, particularly when its policy is optimized for the target dataset. Conventional data augmentation for images consists of image transformation operations, such as random cropping and flipping, and color enhancing including modification of color intensities \cite{Krizhevsky2012,He2016b}. However, designing good data augmentation strategies requires profound understanding of the target data and operations.
For example, CutOut \cite{DeVries2017} randomly erases a patch region of each image and is known to improve performance on the CIFAR-10 dataset, but is also reported to degrade the performance on other datasets, e.g., ImageNet \cite{Lopes2019}.

Therefore, automatically designing effective augmentation strategies according to target data and tasks is desirable to improve the performance of image recognition models. One approach to augment existing data is to generate new data samples by interpolating several training images \cite{Hauberg2015,Devries2017,Zhang2018} or by using conditional generative models \cite{Antoniou2018,Baluja2018}. However, this approach requires a large amount of labeled data \cite{Rather2017} and sometimes fails to improve the performance \cite{Ravuri2019}, even if powerful conditional generative models are used.
On the other hand, some methods improve the performance by efficiently selecting effective combinations of image transformation operations from exponentially large candidate pools \cite{Rather2017,Cubuk2018a,Zhang2020}. In particular, AutoAugment\xspace \cite{Cubuk2018a} and its family \cite{Ho2019,Lim2019,Hataya2019a,Lin2020,Cubuk2019,Berthelot2019a,Lin2020} optimized combinations of operations to improve validation performance and achieved state-of-the-art results. 

These methods involve a bi-level optimization: the inner process optimizes parameters of a CNN on training data using a given combination of operations, and the outer process optimizes the combination of operations to maximize the validation performance. Particularly, the inner loop, i.e., training of a CNN, is usually expensive. Therefore, prior works use proxy tasks that adopt small subsets of training datasets or small models \cite{Cubuk2018a,Ho2019,Lim2019,Hataya2019a} or reduce the search space \cite{Cubuk2019,Berthelot2019a,Lin2020} to keep the entire training feasible. 
Such approaches may result in sub-optimal solutions. We will further review and formalize this problem in \Cref{sec:generalizing_dao}.

In this paper, we tackle the original bi-level optimization problem directly without using proxy tasks or reducing the search space. We propose Meta Approach to Data Augmentation Optimization\xspace (MADAO\xspace), which optimizes CNNs and augmentation policies simultaneously by using gradient-based optimization. Here, policies are updated so that they directly increase CNNs' performance. 
Na\"ively applying gradient-based optimization to this bi-level optimization requires differentiation through the inner optimization process \cite{Finn2017b} or computation of the inverse Hessian matrix \cite{Bengio2000}, both of which suffer from large space complexity. These problems are fatal because data augmentation optimization needs to handle large networks, e.g., CNNs for ImageNet. We bypass these issues by using the implicit gradient method with Neumann series approximation. 
Thanks to these approximations, MADAO\xspace is simple with little overhead as shown in \Cref{ls:bilevel_optimization}. Notably, this simplicity allows MADAO\xspace to scale to problems of ImageNet size, which has been nearly impossible for existing bi-level optimization methods \cite{hutter19}.

We empirically demonstrate that MADAO\xspace learns effective data augmentation policies and achieves performance comparable or even superior to existing methods on benchmark datasets for image classification: CIFAR-10, CIFAR-100, SVHN and ImageNet, as well as four fine-grained datasets. In addition, we show that MADAO\xspace improves the performance of self-supervised learning on ImageNet. All of the reported results have been achieved \textit{without using dataset-specific configurations}.


\newcommand{\highlig}[1]{\colorbox[rgb]{1.0, 0.8, 0.8}{#1}}
\begin{listing}[t]
\centering

\begin{minted}[%
linenos,
xleftmargin=20pt,
breaklines,
frame=lines,
mathescape,
escapeinside=||,
]{python}
cnn.initialize() # parameterized by ${\bm \theta}$
|\highlig{policy.initialize()}| # parameterized by ${\bm \phi}$
for epoch in range(num_epochs):
    for train_data, val_data in data_loader:
        for i in range(num_inner_iters):
            input = |\highlig{policy(}|train_data[i]|\highlig{)}| # augment data by policy
            criterion = cnn.train(input) # referred to $f$ in the text
            cnn.update(criterion)
        vcriterion = cnn.val(val_data) # referred to $g$ in the text
        |\highlig{policy.update(vcriterion)}|
\end{minted}
\vspace{-0.8\baselineskip}
\caption{Our proposed method MADAO\xspace in a Python-like pseudo code. MADAO\xspace optimizes \mintinline[]{python}{policy} as well as \mintinline[]{python}{cnn} by gradient descent with \textit{slight modifications} (\highlig{\texttt{highlighted}}) and \textit{little overhead} to the standard optimization protocol.}
\label{ls:bilevel_optimization}
\end{listing}

\section{Generalizing Data Augmentation Optimization} \label{sec:generalizing_dao}

In \Cref{sub:optimization_da}, we describe the preliminaries of data augmentation optimization of AutoAugment\xspace family, and then review prior methods by generalizing the problem in \Cref{sub:generalization}.

\subsection{Designing Data Augmentation Space}\label{sub:optimization_da}

Let us define a set of input images $\mathcal{X}$ and a set of operations $\mathcal{S}$ consisting of data augmentation operations such as rotation and color inversion. In the AutoAugment\xspace family, each image ${\bm{x}}\in\mathcal{X}\subset[0, 1]^D$ is augmented by an operation $O: \mathcal{X}\to\mathcal{X}$ with a probability of $p_O\in[0, 1]$ and a magnitude of $\mu_O\in[0, 1]$ as illustrated in \Cref{fig:subpolicies}. The magnitude parameter can correspond for example to the degree of rotation, while some operations, such as inversion, have no magnitude parameter. By applying $K$ consecutive augmentation operations, each image results in $2^K$ possible images, that is, the number of images virtually increases. This formulation makes the size of the search space $(\lvert\mathcal{S}\rvert\times[0, 1]\times[0, 1])^{K}$, where $\lvert\mathcal{S}\rvert$ is the size of the operation set.

Operations and accompanied parameters need to be selected so that they minimize the validation criterion, such as the error rate. Usually, this selection is performed heuristically \cite{Krizhevsky2012,He2016b}. However, \citeauthor{Cubuk2018a} showed that data-driven optimization surpasses handcrafted selection \cite{Cubuk2018a}.

\begin{figure}
    \centering
    \includegraphics[width=0.65\linewidth]{subpolicy2}
    \vspace{-0.7\baselineskip}
    \caption{In AutoAugment\xspace family, a \mintinline[]{python}{policy} consists of operations. Each operation, e.g., $\texttt{Rotate}$ or $\texttt{Invert}$, augments an image with a probability of $p$ and a magnitude of $\mu$. As can been seen, performing multiple operations virtually increases the diversity of images. The operations are selected from given operation sets. Our proposed method MADAO\xspace learns to select suitable operations and their parameters, and simultaneously optimize a CNN by using gradient descent. This figure shows a 2-stage augmentation.}
    \label{fig:subpolicies}
\end{figure}

\subsection{Generalizing  AutoAugment Family}\label{sub:generalization}


Let ${\bm \theta}\in\mathbb{R}^M$ denote parameters of a CNN model and ${\bm \phi}\in\mathbb{R}^N$ denote parameters of a policy for augmentation, i.e., selection of operations from the operation set and their accompanied parameters $\{(\mu_O, p_O); O\in\mathcal{S}\}$. Let the empirical risk be $f({\bm \theta}, {\bm \phi}; \mathcal{D}_T)$ and the validation criterion be $g({\bm \theta}; \mathcal{D}_V)$. $\mathcal{D}_T$ and $\mathcal{D}_V$ are training and validation datasets. 
In the case of classification task, $\mathcal{D}\coloneqq\{({\bm{x}}_i, y_i)\}\in\mathcal{X}\times \{1, 2, \dots, C\}$, where $C$ is the number of classes.

Optimization of data augmentation policy in AutoAugment\xspace family methods can be generalized as

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:bi_level_optimization}
    \argmin_{{\bm \phi}}g(\argmin_{{\bm \theta}}f({\bm \theta}, {\bm \phi}; \mathcal{D}_T); \mathcal{D}_V),
\end{equation}

\noindent that is, optimizing CNNs on training data with policies that minimize validation criteria on validation data.


Na\"ively solving this bi-level minimization problem takes a long time because CNN training $\min_{{\bm \theta}}f({\bm \theta}, {\bm \phi})$ is costly and the number of possible combinations of augmentation operations and their parameters is infeasibly large. Therefore, prior works tried to alleviate this problem in several ways. One direction is to reduce the search space over augmentation policies, i.e., dimension and range of the parameter ${\bm \phi}$ \cite{Cubuk2019,Berthelot2019a,Lin2020}. For example, RandAugment\xspace \cite{Cubuk2019} randomly samples operations from the operation set $\mathcal{S}$ and shares $\mu_{O}$ among all operations. This reduction changes the outer problem in \Cref{eq:bi_level_optimization} from $\argmin_{\bm \phi} g$ to $\argmin_\mu g$, which makes it possible to use a simple searching process, such as grid search. OHL-AutoAug \cite{Lin2020} enables online searching using policy gradient \cite{Williams1992} by restricting the search space only to a limited range.

On the other hand, some methods use proxy tasks that approximate \Cref{eq:bi_level_optimization} to obtain (sub-) optimal policy ${\bm \phi}'$ to reduce the searching time of the inner optimization. The obtained policy ${\bm \phi}'$ is then used to train a CNN as $\min_{{\bm \theta}}f({\bm \theta}, {\bm \phi}')$ \cite{Cubuk2018a,Ho2019,Lim2019,Hataya2019a} in an ``offline'' manner. 
For instance, AutoAugment\xspace employs a proxy task $f'$ that approximates the original inner problem $f$ as

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:proxy_task}
     {\bm \phi}'=\argmin_{{\bm \phi}}g(\argmin_{{\bm \theta}'}f'({\bm \theta}', {\bm \phi}'; \mathcal{D}'); \mathcal{D}_V),
\end{equation}

\noindent with a smaller dataset $\lvert\mathcal{D}'\rvert\ll\lvert\mathcal{D}_T\rvert$ and a smaller network $\dim{\bm \theta}'\ll\dim{\bm \theta}$ for efficiency. The outer problem $g$ is optimized by black-box optimization techniques, such as reinforcement learning \cite{Cubuk2018a}. Fast AutoAugment\xspace \cite{Lim2019} and Faster AutoAugment\xspace \cite{Hataya2019a} approximate \Cref{eq:bi_level_optimization} as minimizing distance of distributions between augmented images and original images without directly minimizing $f$, which allows faster searching.

To summarize, prior methods indirectly solve the bi-level optimization problem in \Cref{eq:bi_level_optimization}, as displayed in \Cref{tab:ourmethod}. We instead propose to tackle this problem directly.


\begin{table}[t]
    \centering
    \caption{Comparison of prior methods and our proposal, MADAO\xspace. \textbf{MADAO\xspace can efficiently search \textit{full space} for policies on \textit{full datasets} with \textit{full CNN models}}, e.g., on ImageNet with ResNet-50.}
    \label{tab:ourmethod}
    \vspace{-0.5\baselineskip}
    \begin{tabular}{lcc}
       \toprule
           Method                            & Direct Inner Problem             & Direct Outer Problem \\  \midrule
           AutoAugment\xspace \cite{Cubuk2018a}          &                                  & \checkmark           \\
           Fast AutoAugment\xspace \cite{Lim2019}        &                                  & \checkmark           \\
           \midrule
           OHL-AutoAug \cite{Lin2020}        & \checkmark           & \\
           RandAugment\xspace \cite{Cubuk2019}           & \checkmark           & \\ 
           \midrule
           MADAO\xspace (ours)            & \checkmark           & \checkmark\\
       \bottomrule
    \end{tabular}
\end{table}

\section{Method} \label{sec:method}

In this paper, we propose to directly optimize the bi-level optimization problem in \Cref{eq:bi_level_optimization}. In other words, we optimize CNNs and augmentation policies \textit{simultaneously}, i.e., in an \textit{online} manner, without reducing the search space or using proxy tasks. With this simultaneous optimization, policies are expected to augment images according to the learning state of CNN models. \Cref{ls:bilevel_optimization} is a simple depiction of our approach, which we call Meta Approach to Data Augmentation Optimization\xspace, or in short, MADAO\xspace.

\subsection{Optimizing Policies by Gradient Descent}


MADAO\xspace directly optimizes the bi-level problem \Cref{eq:bi_level_optimization} using gradient descent. To this end, we assume that  $f$ and $g$ are differentiable w.r.t.\xspace ${\bm \theta}$. Taking $C$-category classification as an example, cross entropy $-\mathbb{E}_{{\bm{x}}_i, y_i}\log [{\bm{h}}_{{\bm \theta}}({\bm{x}}_i)]_{y_i}$ can be used as $f$ and $g$, but error rate $\mathbb{E}_{{\bm{x}}_i, y_i}\mathbbm{1}(\argmax {\bm{h}}_{{\bm \theta}}({\bm{x}}_i) \neq y_i)$ cannot. Here, $\mathbbm{1}$ is the indicator function, and ${\bm{h}}_{{\bm \theta}}: \mathbb{R}^D\to\mathbb{R}^C$ is a CNN with a softmax output layer.

Gradient-based optimization of \Cref{eq:bi_level_optimization} requires $\nabla_{{\bm \phi}}g$ for iterative updating. Since the data augmentation implicitly affects the validation criterion, in other words, data augmentation is not used for validation, we obtain

\vspace{-0.6\baselineskip}
\begin{equation}\label{eq:grad_phi_g}
    \frac{\partial g}{\partial {\bm \phi}}=\frac{\partial g}{\partial {\bm \theta}}\frac{\partial {\bm \theta}}{\partial {\bm \phi}}.
\end{equation}

Because of the requirement of $g$, $\nabla_{{\bm \theta}}g$ can be obtained. On the other hand, the exact computation of $\nabla_{{\bm \phi}}{\bm \theta}$ has a large space complexity, as we will describe in \Cref{eq:unroll_steps}. Yet, if this gradient was available, the policies could be optimized by gradient descent.

\subsection{Approximating Gradients of Policy and Inverse Hessian}

To obtain $\nabla_{{\bm \phi}}{\bm \theta}$ without suffering from a large space complexity, we can use the Implicit Function Theorem. If there exists a fixed point $({\bm \theta}^{\star}, {\bm \phi}^{\star})$ that satisfies $\nabla_{{\bm \theta}}f({\bm \theta}^{\star}, {\bm \phi}^{\star})=0$, then there exists a function $\hat{\theta}$ around ${\bm \phi}^{\star}$ such that $\hat{{\bm \theta}}({\bm \phi}^\star)={\bm \theta}^{\star}$. If this condition holds, we also obtain

\vspace{-0.6\baselineskip}
\begin{equation}\label{eq:ift}
    \frac{\partial \hat{{\bm \theta}}}{\partial {\bm \phi}} = -\left( \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \theta}^{\top}} \right)^{-1} \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \phi}^{\top}}\Big|_{{\bm \theta}^\star, {\bm \phi}^\star}.
\end{equation}

We can obtain an approximated gradient using this property. Unfortunately, $M=\dim{\bm \theta}$ is usually large; therefore, computing the inverse Hessian matrix $(\nabla^2_{{\bm \theta}}f)^{-1}$ is prohibitively expensive as it usually requires $\mathcal{O}(M^3)$ computations. To avoid computing the inverse Hessian matrix, we use iterative methods based on the Neumann series, which boast better scalability than conjugate gradients in various problems \cite{Koh2017,Liao2018,Lorraine2019}.

The Neumann series ${\bm{I}}+{\bm{A}}+{\bm{A}}^2+\dots=\sum_i^J {\bm{A}}^j\to({\bm{I}}-{\bm{A}})^{-1} ~(J\to\infty)$ holds with a given squared matrix ${\bm{A}}$ if $\|{\bm{A}}\|<1$. Using this property, \Cref{eq:grad_phi_g} can be approximated with a positive integer $J$ as

\vspace{-0.75\baselineskip}
\begin{equation}
    \frac{\partial g}{\partial {\bm \phi}} = 
   -\frac{\partial g}{\partial {\bm \theta}} 
    \left( \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \theta}^{\top}}\right)^{-1} 
    \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \phi}^{\top}}
    \approx
   -\frac{\partial g}{\partial {\bm \theta}}
    \sum_{j=0}^{J} \left({\bm{I}}-\frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \theta}^{\top}} \right)^j
    \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \phi}^{\top}}.
\end{equation}

We regularize the norm by simply introducing a scalar $\alpha\in\mathbb{R}^+$ as ${\bm{I}}-\alpha\nabla_{{\bm \theta}}^2 f$ as \cite{Lorraine2019}. We use $\alpha=10^{-3}$ with $J=5$ in the experiments. This Neumann series approximation can also help us avoid explicitly storing the Hessian matrix, whose space complexity is $\mathcal{O}(M^2)$, by using Hessian-vector products derived from the following identity:

\vspace{-0.75\baselineskip}
\begin{equation}
    {\bm{v}}^{\top}\left(\frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \theta}^{\top}}\right)
   =\frac{\partial}{\partial{\bm \theta}}\left({\bm{v}}^{\top}\frac{\partial f}{\partial{\bm \theta}}\right)\hphantom{xxxx}(\forall{\bm{v}}\in\mathbb{R}^M).
\end{equation}

This right-hand side can be used instead of storing the Hessian matrix and only has the space complexity of $\mathcal{O}(2M)$. 

\subsection{Differentiable Data Augmentation}\label{sub:dda}

To differentiate through ${\bm \phi}$, MADAO\xspace adopts the differentiable data augmentation pipeline following \cite{Hataya2019a}. As described in \Cref{sub:optimization_da}, each image ${\bm{x}}$ is augmented with an operation $O$ with a magnitude of $\mu_O$ and a probability of $p_O$, which can be written as

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:operation_process}
    {\bm{x}} \mapsto 
    \begin{cases}
        O({\bm{x}}; \mu_O)  &\text{ with probability of } p_O \\
        {\bm{x}}            &\text{ with probability of } 1-p_O.
    \end{cases}
\end{equation}

This right-hand side can be written as $bO({\bm{x}}; \mu_O)+(1-b){\bm{x}}$ with a binary stochastic variable $b\sim\text{Bern}(b; p_O)$. Although the original Bernoulli distribution $\text{Bern}(b; p_O)$ is not differentiable w.r.t.\xspace $p_O$, Gumbel trick \cite{Jang2016} relaxes this restriction to enable backpropagation to update $p_O$. Similarly, some color-enhancing operations are non-differentiable w.r.t.\xspace the magnitude $\mu_O$, so the straight-through estimator \cite{Bengio2013} is used for such operations. We clamp $\mu_O$ and $p_O$ by a sigmoid function to limit their range to [0, 1].

MADAO\xspace uses a different method to select operations compared to Faster AutoAugment\xspace in order to accelerate training.
MADAO\xspace selects operations using categorical distribution parameterized by a weight parameter ${\bm \pi}\in[0, 1]^{\lvert\mathcal{S}\rvert}$, where $\sum_i\pi_i=1$. Since the original categorical distribution is non-differentiable as Bernoulli distribution, we use Gumbel-softmax with a temperature of $\tau\in\mathbb{R}^{+}$. This distribution, referred to as $\text{RelaxCat}({\bm \pi}; \tau)$, samples one-hot-like vectors as $\tau\longrightarrow 0$. Using this distribution, an operation is selected and applied as

\vspace{-1.2\baselineskip}
\begin{align}
    {\bm{x}} &\mapsto \frac{u_i}{\text{SG}(u_i)}O_i({\bm{x}}; \mu_{O_i}, p_{O_i}), \label{eq:transformation} \\
    i &= \argmax{\bm{u}}, \\
    {\bm{u}} &\sim \text{RelaxCat}({\bm \pi}; \tau).
\end{align}

Here, $\text{SG}$ is the stop gradient operation, and thus, $\frac{u_i}{\mathrm{SG}(u_i)}=1$ so that the transformation \Cref{eq:transformation} keeps the range in $(0, 1)$.
Different from this approach, Faster AutoAugment\xspace applies all operations and takes the weighted sum of the outputs to approximate this selection.

\subsection{Connection to Gradient-based Hyperparameter Optimization}

As can been seen, \Cref{eq:bi_level_optimization} is a hyperparameter optimization (HO) problem to neural networks. Traditional HO methods, such as grid search, random search \cite{Bergstra2012} and Bayesian optimization \cite{Snoek2012}, have poor scalability to increasing dimensionality of hyperparameters \cite{hutter19}. For that reason, gradient-based HO attracts attention. From HO viewpoint, policy parameters ${\bm \phi}$ are hyperparameters.

As shown in \Cref{eq:grad_phi_g}, we need to obtain $\nabla_{{\bm \phi}}{\bm \theta}$ to optimize the outer problem $g$.  The inner optimization process (\Cref{ls:bilevel_optimization} Line 5-8) can be rewritten as 

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:unrolling}
{\bm \theta}_T
={\bm \theta}_{T-1}-\eta\nabla_{{\bm \theta}} f({\bm \theta}_{T-1}, {\bm \phi})
=\cdots=
{\bm \theta}_0-\eta\nabla_{{\bm \theta}}\sum_{t=0}^{T-1} f({\bm \theta}_{t}, {\bm \phi})
\eqqcolon {\bm \theta}_T({\bm \theta}_0, {\bm \phi})
\end{equation}

\noindent after $T$ SGD steps with learning rate of $\eta$. One approach to obtain $\nabla_{{\bm \phi}}{\bm \theta}$ is to unroll the steps in \Cref{eq:unrolling} as  \cite{Maclaurin2015,Franceschi2018,Shaban2018}:

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:unroll_steps}
    \nabla_{{\bm \phi}}{\bm \theta}\approx\nabla_{{\bm \phi}}{\bm \theta}_T({\bm \theta}_0, {\bm \phi}).
\end{equation}

This unrolling requires to cache ${\bm \theta}_0, \dots, {\bm \theta}_{T-1}$, and thus, the space complexity becomes $\mathcal{O}(TM)$, which might be prohibitive for large neural networks, while MADAO\xspace requires $\mathcal{O}(2M)$.

Alternatively, implicit gradient yields $\nabla_{{\bm \phi}}{\bm \theta}$ as explained in \Cref{eq:ift}. This approach needs the inverse Hessian matrix $(\nabla^2_{{\bm \theta}}f)^{-1}$ \cite{Bengio2000}, but this computation is infeasible for modern neural networks with millions of parameters. Iterative methods, such as conjugate gradient \cite{Do2009,Domke2012,Pedregosa2016a,Rajeswaran} or Neumann series approximation \cite{Lorraine2019}, effectively compensate for this issue by approximating this inverse matrix in gradient hyperparameter optimization. Such iterative approximation methods using the Neumann series are also used in approximating influence function \cite{Koh2017} and enabling RNNs to handle long sequences \cite{Liao2018}. We exploit the knowledge from these prior works and adopt the implicit gradient method with Neumann series approximation to efficiently handle large-scale datasets and CNNs.

\section{Experiments and Results} \label{sec:experiments}

This section describes the empirical results of our proposed method in supervised learning for image classification and self-supervised learning tasks. For classification tasks, we use CIFAR-10, CIFAR-100 \cite{Krizhevsky2009}, SVHN \cite{Netzer2011} and ImageNet (ILSVRC-2012) \cite{Russakovsky2015}. In addition, we also used fine-grained classification datasets: Oxford 102 Flowers \cite{Nilsback08}, Oxford-IITT Pets \cite{parkhi12a}, FGVC Aircraft \cite{maji13} and Stanford Cars \cite{Krause2013}. For the self-supervised learning task, we used ImageNet. In all experiments except those on ImageNet, we set 10 \% of the original training data aside as validation data $\mathcal{D}_V$ and report error rates on test data. For ImageNet, we used 1 \% of the training data as validation data. Note that this data split means that we use less training data than prior works and that changes the performance of baseline models. More details about experiment configurations are in \Cref{ap:exp_details}.

We used 14 operations for augmentation: \texttt{ShearX}, \texttt{ShearY}, \texttt{TranslateX}, \texttt{TranslateY}, \texttt{Rotate}, \texttt{Invert}, \texttt{AutoContrast}, \texttt{Equalize}, \texttt{Solarize}, \texttt{Color}, \texttt{Posterize}, \texttt{Contrast}, \texttt{Brightness} and \texttt{Sharpness} (see \Cref{ap:operations} for more details.) To make these operations differentiable, we implemented them using PyTorch \cite{NIPS2019_9015} and kornia \cite{eriba2019kornia}. We scale the magnitude of operations from $0$ to $1$ so that the magnitude of $0$ means no change and the magnitude of $1$ corresponds to the strongest level of the particular augmentation operation. We initialized the parameters of magnitude and probability with $\text{sigmoid}(0.5)\approx0.62$ and the parameters for operation selection to be equal. That means that MADAO in its initial state is nearly equivalent to RandAugment\xspace with the magnitude of $0.62$. We set the number of augmentation stages $K=2$ in all experiments below (see \Cref{fig:subpolicies}).

For training, we used SGD with a momentum for CNN optimization and RMSprop optimizer with learning rate of $10^{-2}$ for policy optimization, following \cite{Lorraine2019}. \Cref{eq:ift} requires the existence of fixed points that satisfies $\nabla_{{\bm \theta}}f({\bm \theta}, {\bm \phi})=0$. Following \cite{Lorraine2019}, we assume that $s$ iterations, which corresponds to \mintinline[]{python}{num_inner_iters} in \Cref{ls:bilevel_optimization}, makes the parameters hold the condition. To further encourage parameters to satisfy the condition, we also used warm up, i.e., training CNNs without augmentation for the first $w$ epochs. We set $s=30$ and $w=20$ by default. Validation criterion $g$ is the same loss function that is used as the training loss function $f$, e.g., cross-entropy, for simplicity.


Training of WideResNet 28-2 on CIFAR-10 for 200 epochs with MADAO\xspace takes 1.7 hours, while training without MADAO\xspace takes 1.0 hours in our environment. As regards memory consumption, training with MADAO\xspace requires 3.38 GB of GPU memory, while training without MADAO\xspace requires 1.74 GB in this setting. We tune the magnitude parameter of RandAugment\xspace using three runs of random search and report the test error rate for the run with the lowest validation error rate. 

\subsection*{CIFAR-10, CIFAR-100 and SVHN}

\Cref{tab:results_small} presents test error rates on CIFAR-10, CIFAR-100 and SVHN with various CNN models: WideResNet-28-2, WideResNet-40-2, WideResNet-28-10 \cite{Zagoruyko} and ResNet-18 \cite{He2016b}. We show the average scores of three runs. For comparison, we present the results of RandAugment\xspace and the default augmentation: Cutout \cite{DeVries2017}, random cropping, and random horizontal flipping (except SVHN), following \cite{Cubuk2019}, which we refer to as Baseline. For CIFAR-10 and CIFAR-100, we also present the results of AutoAugment\xspace, which are trained on the split training set using the policy provided by the authors \footnote{\url{https://github.com/tensorflow/models/tree/master/research/autoaugment}}. RandAugment\xspace and AutoAugment\xspace are selected here as representative methods that limit the search space and use proxy tasks, as discussed in \cref{sub:generalization}.

As can been seen, MADAO\xspace achieves performance superior to baseline methods, which demonstrates the effectiveness of directly solving the bi-level problem rather than using proxy tasks as AutoAugment\xspace or a reduced search space as RandAugment\xspace.

\subsection*{ImageNet}

\Cref{tab:results_imagenet} shows the results on ImageNet with ResNet-50 \cite{He2016b} trained in supervised and self-supervised learning manners. In the supervised learning setting, we trained models for 180 epochs. In addition, to showcase the effectiveness of MADAO\xspace in other tasks than supervised learning, we apply MADAO\xspace to contrastive self-supervised learning \cite{he2019moco} for 100 epochs of training. We report the results of the linear classification protocol.

Most importantly, MADAO\xspace can scale to an ImageNet-size problem without using proxy tasks or reducing search space. We hypothesize that the slightly inferior performance of MADAO\xspace compared to RandAugment\xspace could be attributed to the fact that ImageNet has 1,000 categories and a single policy might be insufficient for such diverse data. We also believe that class- or instance-conditional data augmentation might be required to further improve performance, which we leave as an open problem.

\subsection*{Fine-grained Classification}

To showcase the ability of MADAO\xspace to augment images according to target datasets, we conducted experiments on four fine-grained datasets with ResNet-18 \cite{He2016b}. \Cref{tab:results_finedgrained} shows the average test error rates over three runs. These datasets are nearly ten to twenty times smaller than CIFARs and SVHN, yet MADAO\xspace drastically improves the performance without using dataset-specific hyperparameters. 
This improvement emphasizes the importance of searching for good policies according to full target datasets from the full search space. MADAO\xspace can capture the characteristics of each fine-grained dataset and generate tailored policies for each dataset from a large search space.

\begin{table}[tb]
    \centering
    \caption{Test error rates on CIFAR-10, CIFAR-100 and SVHN. For AutoAugment${}^{\ast}$, we train CNN models with the official AutoAugment\xspace policy for CIFAR-10 provided by the authors. \textbf{MADAO\xspace uses the same hyperparameters for all configurations and achieves superior performance to baselines, which demonstrates the effectiveness of our method.}}
    \vspace{-0.5\baselineskip}
    \label{tab:results_small}
    \begin{tabular}{llcccc}
       \toprule
       Dataset   &  Model            & Baseline  & AutoAugment${}^{\ast}$   & RandAugment\xspace      & MADAO\xspace (ours) \\ 
       \cmidrule(r){1-2} \cmidrule(rl){3-6}
       CIFAR-10  &  WideResNet 28-2  & 4.8       &    4.7              & 4.7          &  4.6          \\
                 &  WideResNet 40-2  & 4.5       &    4.3              & 4.3          &  4.2          \\
                 &  WideResNet 28-10 & 3.4       &    3.1              & 3.1          &  3.1          \\
                 &  ResNet-18        & 5.2       &    4.5              & 4.5          &  4.4          \\
       \cmidrule(r){1-2} \cmidrule(rl){3-6}
       CIFAR-100 &  WideResNet 28-2  & 24.8      &   24.0              & 24.0         &  23.9           \\
       \cmidrule(r){1-2} \cmidrule(rl){3-6}               
       SVHN      &  WideResNet 28-2  & 2.9       &    N/A              & 2.6           & 2.5             \\
       \bottomrule                                        

    \end{tabular}
\end{table}


\begin{table}[tb]
    \centering
    \caption{ImageNet test error rates (Top-1/Top-5) on supervised and self-supervised learning.}
    \vspace{-0.5\baselineskip}
    \label{tab:results_imagenet}
    \begin{tabular}{llccc}
       \toprule
       Task                 &  Model            & Standard & RandAugment\xspace       & MADAO\xspace (ours) \\ 
       \cmidrule(r){1-2} \cmidrule(rl){3-5}
       Supervised Learning       &  ResNet-50        & 23.7/ 6.9  &  22.5/ 6.4    & 23.0/ 6.6             \\
       \midrule
       Self-Supervised Learning  &  ResNet-50        & 42.1/19.4 &  39.6/17.6   & 39.6/17.7            \\
       \bottomrule
    \end{tabular}
\end{table}

\begin{table}[tb]
    \centering
    \caption{Test error rates on fine-grained datasets. MADAO\xspace outperforms RandAugment\xspace with a large margin. \textbf{Note that we use the same hyperparemeters as in the experiments in \Cref{tab:results_small}.}}
    \vspace{-0.5\baselineskip}
    \label{tab:results_finedgrained}
    \begin{tabular}{ccccccc}
       \toprule
       Dataset   &  Model            & \# Classes  & \# Training Set & Standard & RandAugment\xspace       & MADAO\xspace (ours) \\ 
       \cmidrule(r){1-4} \cmidrule(l){5-7}
       Flower    &  ResNet-18        & 102         &  2,040          & 10.7      &  \zero8.5   &  \zero7.5             \\
       Pet       &  ResNet-18        & \zero37     &  3,680          & 15.0      &  13.4       &  12.5                 \\
       Aircraft  &  ResNet-18        & \zero70     &  6,667          & 12.8      &  \zero9.7   &  \zero8.7            \\
       Car       &  ResNet-18        & 196         &  8,144          & 14.2      &  11.3       &  10.8                 \\
       \bottomrule
    \end{tabular}

\end{table}

\section{Analysis}

\subsection{How Policies Develop during Training}


In \Cref{fig:policy_development}, we present the development of policies during training on fine-grained datasets. As can be seen, each dataset has its specific operations that are selected, which could be thought as reflecting the characteristics of each dataset. Besides, the way of selection changes according to the learning phase. Note that the first 20 epochs are set to warm-up and the augmentation policy parameters are not updated. We show further observations in \Cref{ap:additional_experiments}.

\begin{figure}[tb]
    \centering
    \includegraphics[width=0.9\linewidth]{policy_development.pdf}
    \vspace{-1\baselineskip}
    \caption{The development of operation selection probabilities at the first augmentation stage during training on fine-grained datasets for 200 epochs. Note that we set the first 20 epochs as a warm-up period during which parameters are not updated.}
    \label{fig:policy_development}
\end{figure}

\begin{wrapfigure}[11]{r}{0.4\linewidth}
    \centering
    \vspace{-3.5\baselineskip}
    \includegraphics[width=\linewidth]{ablations.pdf}
    \vspace{-2\baselineskip}
    \caption{The relationship between the number of inner steps and the final test error rates on CIFAR-10 with WideResNet-28-2. We present the mean and standard deviation over three runs.}
    \label{fig:ablations}
\end{wrapfigure}

\subsection{How Inner Steps Affect the Performance}\label{sub:num_inner_optim}

\Cref{fig:ablations} presents test error rates with various numbers of inner update steps $s$ using WideResNet-28-2 on CIFAR-10. CNNs yield the best performance when $s=30$ and the results indicate that there exists a trade-off between ``exploration and exploitation'' of obtained policies: a small number of inner steps might not correctly evaluate the current policies, while running a large number of inner steps might fail to explore better strategies. Importantly, unrolled-based implementations would require to store $s$ model caches, which is infeasible for $s=30$ with modern CNNs. On the other hand, MADAO\xspace can efficiently handle a large $s$.

\section{Conclusion}

In this paper, we have proposed MADAO\xspace, a novel approach to optimize an image recognition model and its data augmentation policy simultaneously. To efficiently achieve this goal, we use the implicit gradient method with Neumann series approximation. As a result, the overhead of MADAO\xspace to the standard CNN training, w.r.t., time and memory, is marginal, which enables ImageNet-scale training. Empirically, we demonstrate on various tasks that MADAO\xspace achieves superior performance to prior works without restricting search space or using sub-optimal proxy tasks.

Data augmentation is known to boost the performance in various visual representation learning settings, such as semi-supervised learning \cite{Berthelot2019a,Suzuki2020}, domain generalization \cite{Volpi2018}, and self-supervised learning \cite{Chen2020}. We believe that our method can be introduced into these representation learning methods and efficiently enhance their performance, as we showcased with self-supervised learning in this paper.

\clearpage
\makeatletter
\if@preprint
   
\else
\section*{Potential Impacts of Our Work}

MADAO\xspace boosts performance of image recognition models with little overhead and minimal hyperparameter tuning. These advantages help to reduce energy consumption, which is a known issue of large-scale hyperparameter optimization, such as neural architecture search \cite{strubell2019}. MADAO\xspace showcases superior performance, especially on fine-grained datasets. This property enables further application of image recognition models to various domains, but it may also help in misuse of image recognition.

\fi
\makeatother


\printbibliography



