Local Patch AutoAugment with Multi-Agent Collaboration

1 INTRODUCTION Data Augmentation ( DA ) has been widely used to alleviate the overfitting risk in training deep neural networks by appropriately enriching the diversity of training data ( Shorten & Khoshgoftaar , 2019 ; Gontijo-Lopes et al. , 2020 ) . Notable DA methods improve the performance and robustness of the neural networks , such as rotation , Mixup ( Zhang et al. , 2017 ) and Cutmix ( Yun et al. , 2019 ) . But these approaches are typically handcraft and require human prior knowledge , which causes weak transferability of DA across different datasets . To relieve the dependence on manual design and further explore more adaptive augmentation , AutoAugment ( AA ) ( Cubuk et al. , 2018 ) , as a new DA paradigm , is proposed to automate the search of the optimal DA policies ( i.e. , DA operation , probability and magnitude ) from the training dataset . To be specific , AA trains a proxy model with the augmentation policy generated by a controller , which is updated through reinforcement learning using validation accuracy as the reward signal . In spite of the superior performance of AA , its optimization procedure is computationally intensive due to the need to evaluate thousands of policies . Therefore , to reduce the search costs , multiple automated DA approaches ( Lim et al. , 2019 ; Ho et al. , 2019 ; Hataya et al. , 2019 ; Zhang et al. , 2019b ; Lin et al. , 2019 ) have been proposed . For example , ( Lim et al. , 2019 ) employs density matching as a search method to accelerate the policy search , and ( Zhang et al. , 2019b ) introduces adversarial learning to organize the target network training and augmentation policies search in an online manner . Yet the aforementioned automated DA methods all search for policies at the image level . They ignore the exploration of diversity in local regions , which may result in insufficient diversity of dataset and limit the benefits of DA ( Gontijo-Lopes et al. , 2020 ) . In addition , due to this coarse-grained augmentation , they may damage key semantic information and introduce ambiguity into the training process ( see a simple example in Figure 1 row 2 , column 2 ) . With those in mind , it is necessary to automatically search for the optimal augmentation policies for different regions by taking regional diversity into account . One straightforward idea is to directly apply image-wise automated DA methods in different regions . But such an intuitive solution ignores the contextual relationship between regions , which may lead to the non-globally optimal effectiveness of DA policies across the entire image . Besides , it may encounter an extremely high computational cost due to the need of optimizing multiple policies for regions respectively . In this paper , we propose a new approach , named Patch AutoAugment ( PAA ) ( see the last row of Figure 1 ) , to address the above-mentioned problems . We first divide an image into a grid of patches to increase the flexibility of representations for different regions and take “ patch ” as the basic control unit . Then , we model the search for the augmentation policies of patches as a fully cooperative multiagent task , and we leverage a multi-agent reinforcement learning ( MARL ) algorithm where agents cooperatively learn the policies . Specifically , based on the content of each patch and the semantics of the entire image , the agent searches for a policy in terms of choosing which transformation to apply out of the pre-defined DA operations . To encourage our policy networks to adaptively learn the beneficial augmentation policies for the target network , we use the feedback of the target network as the team reward signal to guide the policy networks to learn on the fly . All agents wind up benefiting from two mechanisms ( i.e. , parameter sharing and centralized training with decentralized execution ) in MARL . In this way , all agents collaboratively and parallelly learn policies to further achieve the joint optimal DA policy across the whole image and alleviate the computational cost . In summary , our contributions can be summarized as follows : 1 ) We pinpoint that exploring diversity in local regions is important for automated learned DA approaches . To our best knowledge , we are the first to propose a more fine-grained automated DA approach , that searches for the optimal policies for the patches according to the content of the patch and the semantics of the entire image . 2 ) To further achieve the joint optimal policy across the image , we model the DA policy search of patches as a fully cooperative multi-agent task , and adopt a multi-agent reinforcement learning algorithm for Patch AutoAugment by considering the contextual relationship between the patches . 3 ) Our visualization results provide some insights to the DA community on which augmentation operation should be chosen for patches with different content during the whole training process . 2 RELATED WORK . 2.1 DATA AUGMENTATION . 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 im- prove 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 ) , 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 , Adv AA ( Zhang et al. , 2019b ) proposes an adversarial framework to jointly optimize the target network and the augmentation network , and DADA ( Li et al. , 2020 ) and DDAS ( Liu et al. , 2021a ) target relaxing the discrete DA policy selection to a differentiable problem . Besides , RandAugment ( Cubuk et al. , 2020 ) removes the separate search on a proxy task in the offline automated DA methods ( e.g. , AA ) while also outperforming them . However , these automated DA 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 relationships into account to search for augmentation policies for multiple regions . 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 , Hard chips ( Hong et al. , 2019 ) extracts the object patches from an object pool and pastes them to an image to train the detector , and SECOND ( Yan et al. , 2018 ) performs DA by pasting the objects into the 3D point cloud . 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 ) ) , AugMix ( Hendrycks et al. , 2019 ) , KeepAugment ( Gong et al. , 2020 ) and TrivialAugment ( Müller & Hutter , 2021 ) ) . Yet our proposed method automatically searches for the augmentation transformations based on the given datasets and tasks . 2.2 MULTI-AGENT REINFORCEMENT LEARNING . 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 ) algorithms 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 ) . 3 PROPOSED METHOD . As above mentioned , we aim to search augmentation policies for the patches to explore more augmentation diversity . Learning the optimal DA policies often needs to interact with the target network , which aligns with the nature of RL ( i.e. , RL agents learn to make decisions through interacting with the environment ( Sutton & Barto , 2018 ; Arulkumaran et al. , 2017 ) ) . In addition , it is of great necessity to take the patch content and the contextual relationship between patches into account . Therefore , we use multi-agent reinforcement learning ( MARL ) algorithm . Specifically , the search for patch policies is based on the content of the patch together with the semantics of the image , and policies are encouraged to coordinate to achieve the joint optimal DA policy across the whole image . In this section , we first describe the preliminaries of MARL , then elaborate on our augmentation policy formulation and modelling . Furthermore , we summarize the framework of Patch AutoAugment . 3.1 PRELIMINARIES OF MULTI-AGENT REINFORCEMENT LEARNING . 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 ( S , A , P , R , γ , T ) . In RL framework , given the state s ∈ S , the agent takes an a ∈ A according to its policy π ( a |s ) : S × A → [ 0 , 1 ] and then receives a reward r : S × A → R. The environment moves to the next state with a transition function denoted as P : S × A × S → [ 0 , 1 ] . γ ∈ ( 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 π∗ . Furthermore , multi-agent reinforcement learning ( MARL ) considers a group of N agents , denoted as 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 ( S , { Oi } Ni=1 , { Ai } Ni=1 , P , R , γ , T ) . Here , S describes the shared state space and { O1 , · · · , ON } is a set of observations for agents . In the MARL configuration , each agent receives a private observation correlated with part of state , i.e. , oi : S → Oi . According to the global state s and its observation oi , the agent takes its action ai ∈ Ai based on its policy πi ( ai|oi , s ) , Ai is the action space for the i-th agent and A = { A1 , · · · , AN } denotes the joint action space . Then , the joint action a = a1× · · · ×aN produces next state according to transition function P : S ×A×S → [ 0 , 1 ] and the environment gives the team reward r : S ×A → R to agents N . The objective of each agent is to maximize the total accumulative reward R to cooperatively learn the globally optimal policy π∗ = { π∗1 , · · · , π∗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 .

Summary:   The paper proposes an evolution of the traditional   pipeline of image data augmentation used to reduce ML   model overfitting.    Instead of applying transformations such as shear,   rotate, CutOut, etc. at the image level, the proposed   technique divides the images into a fixed grid and   applies a potentially different transformation to each   cell. The problem of selecting a transform for each cell   is cast as a multi-agent RL (MARL) task, and the agents   learn as the main network trains within a (multi-agent)   Advantage Actor Critic framework.    The agents use a shared reward mechanism, with the reward   defined as the difference between the loss on the   augmented sample and the original loss.    The regular grid is fixed for a dataset and   hyperparameters like the magnitude of the augmentations   follow a fixed schedule; the agents only pick _which_   augmentation to apply on a patch.    Experiments performed on CIFAR-{10,100}, ImageNet,   CUB-200-2011, Stanford Cars, and FGVC-Aircraft show   relatively small but very consistent improvements in   terms of image classification accuracy. Different design   choices (MARL vs. single-agent vs. random, grid size,   etc.) are ablated and discussed in detail.