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: This paper proposed a fine-grained automated data augmentation approach, Patch AutoAugment (PAA), which tries to increase diversity in local regions by divide an image into a grid of patches and search for the joint optimal augmentation policies for the patches. The proposed PAA considers the task as a multi-agent reinforcement learning problem,  and adopt a multi-agent reinforcement learning algorithm to automatically search for the optimal augmentation policies by considering the contextual relationship between the patches. They verify the proposed method on many classification and fine-grained recognition dataset(CIFAR10, CIFAR-100, ImageNet, CUB-200-2011, Stanford Cars and FGVC-Aircraft). The experiments show a good result and visualization results provide some insights that the PAA  help the target network to localize more class-related cues.