Local Patch AutoAugment with Multi-Agent Collaboration

Despite the remarkable performance of deep learning models in computer vision tasks, they often suffer from overfitting. Data augmentation (DA) as an effective technique has been proved to improve the generalization ability of deep learning models (Shorten & Khoshgoftaar, 2019). Previous works (Cubuk et al., 2020; Gontijo-Lopes et al., 2020) indicate that the main benefit of DA arises from increasing the diversity of images. Popular techniques, such as rotation, flipping, color transformation, have been performed as commonly used augmentation methods. Recently, thanks to the skillful design of human experts, many DA methods (e.g., Cutout (DeVries & Taylor, 2017), Mixup (Zhang et al., 2017) and CutMix (Yun et al., 2019)) have been proposed and show significant performance. However, these manually designed methods require additional human prior knowledge on the dataset and sometimes they are limited to certain datasets and target tasks. Naturally, automatically finding DA methods from data have emerged to overcome the limitations of dependence on manual cumbersome exploration. Some works use generative adversarial networks (GANs) to directly generate training data (Shrivastava et al., 2017; Tran et al., 2017).

Furthermore, recent studies aim to automate the search for augmentation policies that choose the optimal transformations for training images. AutoAugment (AA) (Cubuk et al., 2018) adopts a controller to generate an augmentation policy which is used to train a proxy network, then gets the validation accuracy as the reward signal to update the controller using reinforcement learning. Unfortunately, the evaluation of thousands of policies makes AA computationally expensive. Therefore, multiple automated DA approaches focus on reducing the huge complexity and have achieved great progress. For example, PBA (Ho et al., 2019) employs hyperparameter optimization, Fast AA (Lim et al., 2019) uses a density matching algorithm and Adversarial AA (Zhang et al., 2019b) proposes an adversarial framework to jointly optimize the target network and the augmentation network. However, these automated augmentation methods perform the search at the image level, i.e., they use the same policy on the whole image. It inevitably ignores the diversity of different regions in an image, which limits the diversity of data increased by data augmentation, and sometimes causes the damage of critical semantic information. In contrast, our method takes diversity in local regions and contextual relationship into account, aiming to search for augmentation policies for multiple regions and achieve the joint optimal DA effect across the entire image through MARL.

Our proposed method is conceptually orthogonal to most region-based DA methods where DA transformations are performed in a non-automated way. For example, RandomErasing (Zhong et al., 2020), CutMix (Yun et al., 2019) perform cropping or replacement operation on a randomly selected rectangle region. Some works further exploit class activation map (Zhou et al., 2016) or saliency map (Zhou et al., 2015) to select representative regions which are augmented by a randomly selected DA operation (e.g., SaliencyMix (Uddin et al., 2020), SnapMix (Huang et al., 2020)) and KeepAugment (Gong et al., 2020)). Yet our proposed method automatically searches for the augmentation transformations based on the given datasets and tasks.

The most significant characteristic of MARL is the cooperation between agents (Tampuu et al., 2017; Ma & Cameron, 2008; Foerster et al., 2017) which is distinct from directly applying reinforcement learning (RL) algorithm to multi-agent systems. Due to the limited observation and action of a single agent, cooperation is necessary in the reinforced multi-agent system to achieve the common goal. Compared with independent agents, cooperative agents can improve the efficiency and robustness of the model (Neto, 2005; Zhang et al., 2019a; Busoniu et al., 2008). Many vision tasks use MARL to interact with the public environment to make decisions, with the goal of maximizing the expected total return of all agents, such as image segmentation (Liao et al., 2020; Han et al., 2019; Ma et al., 2020), image processing (Furuta et al., 2019).

We first introduce the preliminaries of reinforcement learning (RL). RL models the decision-making problem as a Markov decision process (MDP) which is presented with a tuple $({\mathcal{S}},{\mathcal{A}},{P},R,\gamma,T)$. In RL framework, given the state $s\in{\mathcal{S}}$, the agent takes an $a\in{\mathcal{A}}$ according to its policy ${\pi}(a\left|{s)}\right.:{\mathcal{S}}\times{\mathcal{A}}\to[0,1]$ and then receives a reward $r:{\mathcal{S}}\times{\mathcal{A}}\to\mathbb{R}$. The environment moves to the next state with a transition function denoted as ${P}:{\mathcal{S}}\times{\mathcal{A}}\times{\mathcal{S}}\to[0,1]$. $\gamma\in(0,1]$ is a discount factor and $T$ is a time horizon. The agent aims to maximize the long-term reward $R$ over $T$ steps to learn the optimal policy ${{\pi^{*}}}$.

Furthermore, multi-agent reinforcement learning (MARL) considers a group of $N$ agents, denoted as $\mathcal{N}$, operating cooperatively in a shared environment towards a common goal. It can be formulated as a multi-agent MDP (MAMDP) (Boutilier, 1996) represented with a tuple $({\mathcal{S}},\{{\mathcal{O}_{i}\}{{{}^{{N}}_{i=1}}}},\{{\mathcal{A}_{i}\}{{{}^{{N}}_{i=1}}}},{P},R,\gamma,T)$. Here, $\mathcal{S}$ describes the shared state space and $\{\mathcal{O}_{1},\cdot\cdot\cdot,\mathcal{O}_{{N}}\}$ is a set of observations for agents. In MARL configuration, each agent receives a private observation correlated with part of state, i.e., $o_{i}:\mathcal{S}\to\mathcal{O}_{i}$. According to the global state $s$ and its observation $o_{i}$, the agent takes its action ${a_{i}}\in{\mathcal{A}_{i}}$ based on its policy ${\pi_{i}}{{(a_{i}}{|o_{i},s)}}$, $\mathcal{A}_{i}$ is the action space for the $i$-th agent and $\bm{\mathcal{A}}=\{\mathcal{A}_{1},\cdot\cdot\cdot,\mathcal{A}_{N}\}$ denotes the joint action space. Then, the joint action $\bm{a}={a}_{1}{\times}\cdot\cdot\cdot{\times}{a}_{{N}}$ produces next state according to transition function ${P}:{\mathcal{S}}\times\bm{\mathcal{A}}\times{\mathcal{S}}\to[0,1]$ and the environment gives the team reward $r:{\mathcal{S}}\times\bm{\mathcal{A}}\to\mathbb{R}$ to agents $\mathcal{N}$. The objective of each agent is to maximize the total accumulative reward $R$ to cooperatively learn the globally optimal policy $\bm{\pi^{*}}=\{{\pi{{}_{1}^{*}}},\cdot\cdot\cdot,{\pi{{}^{*}_{{N}}}}\}$ that consists of the optimal policy of each agent. In this paper, as mentioned before, we employ MARL to search optimal policy for each patch, and achieve the optimal effectiveness of DA across the entire image to further improve the performance of the target network as possible.

In our proposed Patch AutoAugment (PAA), we formulate the task of policy search for the patches as a cooperative multi-agent decision-making problem and adopt multi-agent reinforcement learning (MARL) to solve it. In the following, We clarify the detailed formulation $i$) the state, observation and action modeling for the policy, $ii$) an effective team reward function design and $iii$) the detailed MARL algorithm for policy learning) of Patch AutoAugment.

Policy Modeling. As illustrated in Figure 2, given the original input batch $\bm{x}$ and the corresponding label $\bm{y}$ (i.e., $\bm{\{x,y\}}=\{x_{j},y_{j}\}{{}_{j=1}^{b}}$ and $b$ is the batch size), we divide an image $x_{j}$ into ${N}$ equal-sized and non-overlapping patches, denoted as $x_{j}=\{{\mathcal{P}{{}_{j}^{i}}\}{{}_{i=1}^{{N}}}}$ where ${\mathcal{P}{{}_{j}^{j}}}$ is the $i$-th patch of the image $x_{j}$. In our proposed method, we aim to search the augmentation policy for each patch. Therefore, in MARL formulation, the augmentation policy of the patch ${\mathcal{P}{{}^{i}}}$ is controlled by an agent $i\in\mathcal{N}$ and we detail the state, observation and action for the augmentation policy as below.

State. As aforementioned, the selection of augmentation operation for a patch is closely bound up with the contextual relationship between regions. Therefore, the augmentation policy needs to perceive the image semantics and we take the deep features of the whole image extracted by a backbone (e.g., ImageNet pre-trained ResNet-18 (He et al., 2016)) as the global state $s$ which is visible to all agents.

Observation. Apart from capturing state (i.e., the global information), the agent only use their own observation (i.e., the local information) which is invisible to other agents. The augmentation policy seeks to choose augmentation operation based on the content of the patch and the observation is generally part of the state. Considering all these factors, we utilize the deep features of the $i$-th patch as observation $o_{i}$, which are extracted by the same backbone as the state extractor.

Action. The augmentation policy is responsible for choosing which transformations to apply from pre-defined operations. Following the previous automated DA methods (Cubuk et al., 2018; Lim et al., 2019), we define the fifteen operation functions (i.e., ShearX/Y, TranslateX/Y, Rotate, Invert, Equalize, Solarize, Posterize, Contrast, Color, Brightness, Sharpness, RandomErasing, Cutout, Mixup, Cutmix) to construct the action space $\mathcal{A}$. Given the state $s$ and the observation $o_{i}$, according to policy ${\pi_{i}}{{(a_{i}}{|o_{i},s)}}$, each agent $i$ determines an action $a_{i}(\cdot)\in\mathcal{A}$, which is the operation performed on the patch. Each operation is associated with two hyperparameters: probability $p$ and magnitude $m$. In order to dramatically and effectively reduce the action space for augmentation policy, similar with (Cubuk et al., 2020; Ho et al., 2019; Lim et al., 2019), we take a fixed probability and magnitude schedule. Among them, the probability of applying the operation is sampled from the uniform distribution (i.e., $p_{i}\sim U(0,1)$. Following (Ho et al., 2019), we employ the same linear scale to be the magnitude schedule. In summary, the processed $i$-th patch in image $x_{j}$ is denoted as $\widetilde{P}{{}_{j}^{i}}=a_{i}({P}{{}_{j}^{i}})$ with the probability $p_{i}$ and the magnitude of $a_{i}(\cdot)$ is $m_{i}$, otherwise $\widetilde{P}{{}_{j}^{i}}={P}{{}_{j}^{i}}$.

Note that most operations are label-invariant, except Mixup (Zhang et al., 2017) and CutMix (Yun et al., 2019) are label-disturbing operations that combine different patches as well as their labels. We take Mixup as an example, then $\mathcal{\widetilde{P}}{{}_{j}^{i}}=\lambda\mathcal{{P}}{{}_{j}^{i}}+(1-\lambda)\mathcal{{P}}{{}_{t}^{i}}$ with probability $p_{i}$, where $\mathcal{{P}}{{}_{t}^{i}}$ is a patch from another image $x_{t},t\neq{j}$, $\lambda\sim Beta(\alpha,\alpha)$, for $\alpha\in(0,\infty)$, and the one-hot label is modified as $\widetilde{y}_{j}=\frac{\lambda}{N}y_{j}+\frac{(1-\lambda)}{N}y_{t}$. Until all patches are processed by the corresponding operations chosen by the augmentation policies, we obtain the image augmented by our PAA, denoted as $\widetilde{x}_{j}=\{\mathcal{\widetilde{P}}{{}_{j}^{1}},\cdot\cdot\cdot,\mathcal{{\widetilde{P}}}{{}_{j}^{{N}}}\}$, and the final label $\widetilde{y}_{j}$. More operation details are shown in Appendix A.

Reward Function. The reward function is of importance to guide the agents to learn so that they follow desired behaviors. The previous work, Adversarial AA (AdvAA) (Zhang et al., 2019b), attempts to increase the training loss of the target network to generate harder augmentation policies and explore the weakness of the target network. Inspired by AdvAA, we reformulate the reward design appropriately under our configuration. In our proposed PAA, the common objective of all agents is to improve the performance of the mainstream target task through enhancing the benefits of DA. Therefore, we compare the feedback of target network $\phi(\cdot)$ on the augmented data processed by our proposed PAA $\widetilde{\bm{x}}$ with the original data $\bm{x}$ and take their difference on the training losses as the reward for the policy in MARL, as in Eq. (1):

$${r}=l(\phi(\bm{x}),\bm{y})-l(\phi(\bm{\widetilde{x}}),\bm{\widetilde{y}}),$$

(1)

where $\bm{x}$ and $\bm{y}$ denote raw inputs and labels in supervision tasks. In our PAA model, all agents are encouraged to cooperate to achieve the common goal, thus we adopt the team reward function design (i.e., the shared reward mechanism in MARL) as Eq. (1) for all agents to make the joint augmentation policy achieve the optimal effectiveness.

Policy Learning. Here, we introduce the training for the augmentation policies mentioned above. Considering that the action space is discrete, we adopt multi-agent Advantage Actor-Critic algorithm (Lowe et al., 2017) to learn the augmentation policies and encourage the coordination behaviors. In MARL, the framework of centralized training with decentralized execution (Foerster et al., 2018; Rashid et al., 2018) is generally adopted. More concretely, each agent $i$ has an actor which is to learn discrete policy $\pi_{i}(a_{i}|o_{i},s)$ and agents share a common critic which aims to estimate the value of global state $V^{\bm{\pi}}(s)$. And we use the centralized critic to train decentralised actors. Here, we reformulate it appropriately for our task. We model the centralized action-value Q function that takes the actions of all agents in addition to state information $s$ and outputs the Q-value for the team, formulated as:

$$Q^{\bm{\pi}}(s,\bm{a})={\mathbb{E}_{\pi}}[{R_{t}}|s,a_{1},\cdot\cdot\cdot,a_{N}],$$

(2)

where $\bm{a}$ is the joint action of all agents $\bm{a}=\{a_{i},\cdot\cdot\cdot,a_{N}\}$ and ${R_{t}}=\sum\nolimits_{l=0}^{T}{{\gamma^{l}}{r_{t+l}}}$ is the long-term discounted reward. Then, the advantage function on the augmentation policy is given as follows:

$$A^{\bm{\pi}}(s,\bm{a})=Q^{\bm{\pi}}(s,\bm{a})-V^{\bm{\pi}}(s).$$

(3)

And we use $\varphi$ to denote the critic network parameters. We take the square value of the advantage function $A^{\pi}$ as the loss function to update $\varphi$:

$$L(\varphi)={\left(A^{\bm{\pi}}(s,\bm{a})\right)^{2}}.$$

(4)

Besides, to further achieve the ability to cooperate, similar to (Foerster et al., 2018; Rashid et al., 2018), the parameters of the actor networks of all agents are shared, denoted as $\theta$. In addition, the loss function for updating $\theta$ is defined as:

$${L}(\theta)=-\log{\pi_{{\theta}}}(\bm{a}|s)A^{\bm{\pi}}(s,\bm{a}).$$

(5)

Framework Summary. In this part, we summarize the overall framework of our proposed PAA. As shown in Figure 2, PAA first divides an image into a grid of patches. Then, we use a feature extractor to obtain the deep features of the whole image as the state. Each agent draws its individual observation, i.e., the deep features of the patch. According to the global state (i.e., the semantics of the entire image) and the local observation (i.e., the content of the patch), the actor networks output the augmentation operations of patches to further construct the joint operation map performed on the whole image. The augmented mini-batch is processed by our proposed PAA and then we input it to the target task network for parameters updating. Moreover, the feedback of the target network is used as the team reward signal to update the policy network.

In this section, to study the effectiveness of Patch AutoAugment (PAA), we experiment with core image classification and fine-grained image recognition tasks. Exactly, we focus on CIFAR-10, CIFAR-100 (Krizhevsky et al., 2009) and ImageNet (Deng et al., 2009) datasets as well as three fine-grained object recognition datasets, i.e., CUB-200-2011 (Wah et al., 2011), Stanford Cars (Krause et al., 2013) and FGVC-Aircraft (Maji et al., 2013). We describe the datasets in detail in Appendix B. We compare PAA with baseline pre-processing, Cutout (DeVries & Taylor, 2017), Mixup (Zhang et al., 2017), Cutmix (Yun et al., 2019), AutoAugment (AA) (Cubuk et al., 2018) and Fast AutoAugment (FastAA) (Lim et al., 2019). The baseline follows (Zoph et al., 2018; Yamada et al., 2018; Gastaldi, 2017): standardizing the data, horizontally flipping with 0.5 probability, zero-padding and random cropping. More details about our policy network architectures and model hyperparameters are supplied in Appendix C and D, respectively. Moreover, we set the number of patches $N$ to $4$ for CIFAR tasks, and $N=16$ for ImageNet together with fine-grained recognition datasets. To ensure the reliability of our experiments, we run each experiment four times using different random seeds.

Classification Results on CIFAR-10 and CIFAR-100. For CIFAR-10 and CIFAR-100, we examine on Wide-ResNet-28-10 (WRN-28-10) (Zagoruyko & Komodakis, 2016), Shake-Shake (SS) (Gastaldi, 2017) and Pyramid-Net+ShakeDrop (Pyramid+SD) (Han et al., 2017; Yamada et al., 2018) models. The results are reported in Table 1, which shows the proposed approach consistently outperforms several state-of-the-art DA methods. We observe that the improvement of performance is relatively slight, due to the small image size of CIFAR which is $32\times{32}$. In the following, we further apply our proposed PAA on datasets with larger image sizes and other networks.

Classification Results on ImageNet. As shown in Table 2, we evaluate our method on ResNet-50 and ResNet-200 (He et al., 2016) backbone on ImageNet, and our PAA significantly improves the performance of the target networks. The results further demonstrate that our proposed method is an effective DA technique for consistent and expressive benefits for datasets with larger image sizes.

Effectiveness of Fine-grained Classification. Furthermore, we evaluate our proposed method on fine-grained image recognition tasks. According to previous work (Du et al., 2020; Chen et al., 2019), we take ResNet-50 and ResNet-101 as the backbones. The results are shown in Table 4, which illustrates that the performance of PAA is consistently better than other methods and PAA has achieved remarkable performance on these challenging fine-grained tasks.

In this section, in order to further demonstrate the performance of PAA in terms of complexity, we compare the policy search time and training time of PAA with AA (Cubuk et al., 2018), FastAA (Lim et al., 2019) and AdvAA (Zhang et al., 2019b), as illustrated in Table 3. As shown in Table 3, compared to the previous works, PAA requires the fewest total computational resources, and the search time is almost negligible. As for the parameters, the total number of PAA model parameters (about 0.23M) is less than 1$\%$ of the target network (e.g., ResNet50: about 25.5M).

In summary, the main reasons to reduce the computational cost lie in three aspects: 1) The augmentation policy network is jointly optimized with the target network, similar to (Zhang et al., 2019b; Lin et al., 2019). Namely, our proposed method searches for policies in an online manner, obviating thousands of policies validation on a small proxy network and the requirement for retraining the target network. Besides, we use fixed schedules for the two corresponding hyperparameters (i.e., probability and magnitude) of each transformation to effectively reduce the search space, which makes it easier to search for effective policies. By these means, PAA compresses most of the policy search time 2) As mentioned before, a MARL algorithm is adopted, in which all agents parallelly learn the augmentation policies, to reduce the training time. 3) Compared with the previous DA methods using image-by-image sequential transformations, we perform parallel transformations on tensor. Specifically, we pick the patches in a batch performing the same operation to reconstruct a new tensor. And we use Kornia¹¹1Kornia (Riba et al., 2020) is a differentiable computer vision library for PyTorch. We use it to accelerate augmentation operation on tensors. to realize tensor transformations on GPU to further reduce the processing time which is included in the training time.

In this section, we study the effectiveness of MARL and also discuss the design of patch numbers, through several ablation experiments.

Performance of random policies. We randomize the augmentation policy for each patch, dubbed Patch RandAugment and compare it with our proposed PAA. To be specific, the augmentation policy is randomly selected from predefined transformations with uniform probability. As shown in Table 5, Patch RandAugment (PRA) leads to considerable appreciable improvements. However, our Patch AutoAugment with meaningful guidance significantly surpasses the random patch policies.

Performance of the policies searched using single-agent RL. Furthermore, we directly use the same single-agent RL to search for each patch’s policy, termed as Patch SingleAutoAugment (PSAA) and compare it to our PAA, where the search task is a cooperative multi-agent problem. In short, PSAA ignores the contextual relationship between patches, where patches are non-cooperative. The results (see Table 5) indicate that taking the contextual information into account with cooperative RL-agents has improved the joint effectiveness of DA.

Impact of the number of patches. We explore the effect of the number of patches on the target model performance. As shown in Table 6, we respectively set the number of patches $N=\{4,16,64,256,1024\}$. The results show that on ImageNet / CUB-200-2011, when patch numbers $N$ increase, the performance accuracy first increases and then decreases. In addition, when $N=16$, PAA achieves the best performance. We analyze that the small number of patches may cause PAA to be unable to effectively explore local diversity, and its advantages would be limited. In particular, when $N=1$, PAA degenerates into image-wise automated DA. In contrast, under too larger values (e.g., the extreme case is to search for DA policy for each pixel), the local semantic consistency is broken and the benefit brought by the consideration of contextual relationship between patches is gradually overtaken.

Policy Visualization. Grad-CAM (Selvaraju et al., 2017) is used to localize the important regions in the image. Therefore we can calculate the importance score of each patch, and we divide the importance scores into four bins, i.e., very important, important, normal and not important. Then, we categorize patches into four groups and draw four stacked area charts showing the percentages of operations selected by PAA augmentation policies over time. We take ResNet-50 backbone trained on CUB-200-2011 as an example.

Since our proposed PAA searches for patch policies in an online manner, the strategies change dynamically over time as shown in Figure 4. At the beginning of the training process, the selected actions are messy since the MARL network is in the exploratory stage. At the tail end of the training, the target network has converged, causing the percentage of all operations to be almost the same and the percentage to flatten out.

In addition, we have some interesting findings which may provide some insights to the DA community. $i$) The optimal augmentation strategies of patches vary by their important levels. Therefore, it is necessary to take the image content into account when performing data augmentation. $ii$) In the middle of the training process, different types of patches prefer to select different augmentation operations. Concretely, as illustrated in Figure 4, for the important patches, color transformation is mostly picked. For the unimportant patches, RandomErasing and Cutout are usually chosen by PAA. Important patches commonly take along semantic information that is related to mainstream tasks. It is better to choose the mild transformations (e.g., color) for them, which can effectively protect the semantic information from being damaged. In contrast, unimportant patches typically carry unexpected features (Singla et al., 2021) which are causally unrelated to the desired class. Severe transformations (e.g., Cutout, RandomErasing) could be chosen for them, which introduce noise and disturbance to reduce the impact of unexpected features on the target network learning.

Grad-CAM Visualization. Here, we adopt Grad-CAM (Selvaraju et al., 2017) to visualize the learned features to intuitively show the impact of PAA, as shown in Figure 4. We take the ResNet-50 as the backbone which is trained with dataset processed by 1) Baseline 2) AutoAugment and 3) Patch AutoAugment, respectively. We observe that the model trained with PAA focus on more task-related areas rather than spurious correlations (e.g., the branch where the bird stands) or overemphasized features (e.g., birds’ claws). More visualization results are provided in the Appendix E.