Concurrent Adversarial Learning for Large-Batch Training

1 INTRODUCTION . With larger datasets and bigger models proposed , training neural networks has become quite timeconsuming . For instance , training BERT ( Devlin et al. , 2019 ) takes 3 days on 16 v3 TPUs . GPT-2 ( Radford et al. , 2019 ) contains 1,542M parameters and requires 168 hours of training on 16 v3 TPU chips . This also leads to the developments of high performance computing clusters . For example , Google and NVIDIA build high performance clusters with thousands of TPU or GPU chips . How to fully utilize those computing resources for machine learning training thus becomes an important problem . Data parallelism is a commonly used technique for distributed neural network training , where each processor computes the gradient of a local batch and the gradients across processors are aggregated at each iteration for a parameter update . Training with hundreds or thousands of processors with data parallelism is thus equivalent to running a stochastic gradient optimizer ( e.g. , SGD or Adam ) with a very large batch size , also known as large batch training . For example , Google and NVIDIA showed that by increasing the batch size to 64k on ImageNet , they can finish 90-epoch ResNet training within one minute ( Kumar et al. , 2021 ; Mattson et al. , 2019 ) . But why can ’ t we infinitely increase the batch size as long as more computing resources are available ? Large batch training often faces two challenges . First , under a fixed number of training epochs , increasing the batch size implies reducing number of training iterations . Even worse , it has been observed that large-batch training often converges to solutions with bad generalization performance ( also known as sharp local minima ) , possibly due to the lack of inherent noise in each stochastic gradient update . Although this problem can be partially mitigated by using different optimizers such as LARS ( You et al. , 2017 ) and LAMB ( You et al. , 2019 ) , the limit of batch size still exists . For instance , Google utilizes several techniques , such as distributed batch normalization and Mixedprecision training , to further scale the traing of ResNet-50 on 4096 v3 TPU chips . However , it can just expand the batch size to 64k ( Kumar et al. , 2021 ; Ying et al. , 2018 ) . Data augmentation has become an indispensable component of large-batch training pipeline . For instance , researchers at Facebook use augmentation to scale the training of ResNet-50 to 256 NVIDIA P100 GPUs with a batch size of 8k on ImageNet ( Goyal et al. , 2017 ) . You et al . also use data augmentation to expand the batch size to 32k on 2048 KNL nodes ( You et al. , 2018 ) . However , in this paper we find that when batch size is large enough ( i.e. , larger than 32k ) , the increased diversity in augmented data will also increase the difficulty of training and even have a negative impact on test accuracy . This motivates us to study the application of adversarial training in large-batch training . Adversarial training methods find a perturbation within a bounded set around each sample to train the model . Previous work finds the adversarial training would lead to a significant decrease in the curvature of the loss surface and make the network more ” linear ” in the small batch size case , which could be used as a way to improve generalization ( Xie et al. , 2020 ; Moosavi-Dezfooli et al. , 2019 ) . However , adversarial training has not been used in large-batch training since it requires a series of sequential gradient computations within each update to find an adversarial example . Even when conducting only 1 gradient ascent for finding adversarial examples , adversarial training requires 2 sequential gradient computations ( one for adversarial example and one for weight update ) that can not be parallelized . Therefore , even with infinite computing resources , adversarial training is at least 2 times slower than standard training and increasing the batch size can not compensate for that . To resolve this issue and make adversarial training applicable for large-batch training , we propose a novel Concurrent Adversarial Learning ( ConAdv ) algorithm for large-batch training . We show that by allowing the computation of adversarial examples using staled weights , the two sequential gradient computations in adversarial training can be decoupled , leading to fully parallelized computations at each step . As a result , extra processors can be fully utilized to achieve the same iteration throughput as original SGD or Adam optimizers . Comprehensive experimental results on large-batch training demonstrate that ConAdv is a better choice than existing augmentations . Our main contributions are listed below : • This is the first work showing adversarial learning can significantly increase the batch size limit of large-batch training without using data augmentation . • The proposed algorithm , ConvAdv , can successfully decouple the two sequential gradient computations in adversarial training and make them parallelizable . This makes adversarial training achieve similar efficiency with standard stochastic optimizers when using sufficient computing resources . Furthermore , we empirically show that ConAdv achieves almost identical performance as the original adversarial training . We also provide a theoretical analysis on ConAdv . • Comprehensive experimental studies demonstrate that the proposed method can push the limit of large batch training on various tasks . For ResNet-50 training on ImageNet , ConAdv alone achieves 75.3 % accuracy when using 96K batch size . Further , the accuracy will rise to 76.2 % when combined with data augmentation . This is the first method scaling ResNet-50 batch size to 96K with accuracy matching the MLPerf standard ( 75.9 % ) , while previous methods fail to scale beyond 64K batch size . 2 BACKGROUND . 2.1 LARGE-BATCH TRAINING . Using data parallelism with SGD naturally leads to large-batch training on distributed systems . However , it was shown that extremely large batch is difficult to converge and has a generalization gap ( Keskar et al. , 2017 ; Hoffer et al. , 2017 ) . Therefore , related work start to carefully fine-tune the hyper-parameters to bridge the gap , such as learning rate , momentum ( You et al. , 2018 ; Goyal et al. , 2017 ; Li , 2017 ; Shallue et al. , 2018 ) . Goyal et al . try to narrow the generalization gap with the heuristics of learning rate scaling ( Goyal et al. , 2017 ) . However , there is still big room to increase the batch size . Several recent works try to use adaptive learning rate to reduce the fine-tuning of hyper-parameters and further scaling the batch size to larger value ( You et al. , 2018 ; Iandola et al. , 2016 ; Codreanu et al. , 2017 ; Akiba et al. , 2017 ; Jia et al. , 2018 ; Smith et al. , 2017 ; Martens & Grosse , 2015 ; Devarakonda et al. , 2017 ; Osawa et al. , 2018 ; You et al. , 2019 ; Yamazaki et al. , 2019 ) . You et al . propose Layer-wise Adaptive Rate Scaling ( LARS ) ( You et al. , 2017 ) for better optimization and scaling to the batch size of 32k without performance penalty on ImageNet . Ying et al . use LARS optimizer to train ResNet-50 on TPU Pods in 2.2 minutes . In addition , related work also try to bridge the gap from aspect of augmentation . Goyal et al . use data augmentation to further scale the training of ResNet-50 on ImageNet ( Goyal et al. , 2017 ) . Yao et al . propose an adaptive batch size method based on Hessian information to gradually increase batch size during training and use vanilla adversarial training to regularize against the sharp minima Yao et al . ( 2018a ) . However , the process of adversarial training is time consuming and they just use the batch size of 16k in the second half of training process ( the initial batch size is 256 ) . How to further accelerate the training process based on adversarial training and reduce its computational burden is still an open problem . 2.2 ADVERSARIAL LEARNING . Adversarial training has shown great success on improving the model robustness through collecting adversarial examples and injecting them into training data Goodfellow et al . ( 2015 ) ; Papernot et al . ( 2016 ) ; Wang et al . ( 2019 ) . Madry et al . ( 2017 ) formulates it into a min-max optimization framework as follows : min θ E ( xi , yi ) ∼D [ max||δ||p∈ L ( θt , x+ δ , y ) ] , ( 1 ) where D = { ( xi , yi ) } ni=1 denotes training samples and xi ∈ Rd , yi ∈ { 1 , ... , Z } , δ is the adversarial perturbation , || · ||p denotes some Lp-norm distance metric , θt is the parameters of time t and Z is the number of classes . Goodfellow et al . proposes FGSM to collect adversarial data Goodfellow et al . ( 2015 ) , which performs a one-step update along the gradient direction ( the sign ) of the loss function . Project Gradient Descent ( PGD ) algorithm Madry et al . ( 2017 ) firstly carries out random initial search in the allowable range ( spherical noise region ) near the original input , and then iterates FGSM several times to generate adversarial examples . Recently , several papers Shafahi et al . ( 2019 ) ; Wong et al . ( 2020 ) ; Andriushchenko & Flammarion ( 2020 ) aim to improve the computation overhead brought by adversarial training . Specifically , FreeAdv Shafahi et al . ( 2019 ) tries to update both weight parameter θ and adversarial example x at the same time by exploiting the correlation between the gradient to the input and to the model weights . Similar to Free-adv , Zhang et al . Zhang et al . ( 2019 ) further restrict most of the forward and back propagation within the first layer to speedup computation . FastAdv Wong et al . ( 2020 ) finds the overhead could be further reduced by using single-step FGSM with random initialization . While these work aim to improve the efficiency of adversarial training , they still require at least two sequential gradient computations for every step . Our concurrent framework could decouple the two sequential gradient computation to further boost the efficiently , which is more suitable for large-batch training . Recently , several works Xie et al . ( 2020 ) ; Cheng et al . ( 2021 ) ; Chen et al . ( 2021 ) show that the adversarial example can serve as an augmentation to benefit the clean accuracy in the small batch size setting . However , whether adversarial training can improve the performance of large-batch training is still an open problem . 2.3 MLPERF . MLPerf Mattson et al . ( 2019 ) is an industry-standard performance benchmark for machine learning , which aims to fairly evaluate system performance . Currently , it includes several representative tasks from major ML areas , such as vision , language , recommendation . In this paper , we use ResNet-50 He et al . ( 2016 ) as our baseline model and the convergence baseline is 75.9 % accuracy on ImageNet . 3 PROPOSED ALGORITHM . In this section , we introduce our enlightening findings and the proposed algorithm . We first study the limitation of data augmentation in large-batch training . Then we discuss the bottleneck of adversarial training in distributed system and propose a novel Concurrent Adversarial Learning ( ConAdv ) method for large-batch training .

Summary: This paper presents a simple algorithm named ConAdv to incorporate adversarial training into the large-batch training setting such that one can further increase the batch size without harming too much accuracy while maintaining the high utilization of the hardware. The core idea is to use adversarial training to improve the accuracy for large batch training and at the same time use stale weights to allow parallel computation of the adversarial example and the normal gradient updates. With his simple yet novel approach, the paper has demonstrated good performance on ImageNet with batch size as large as 96k while maintaining the accuracy above MLPerf's 75.9.