Concurrent Adversarial Learning for Large-Batch Training
Authors: Yong Liu
Xiangning Chen
Minhao Cheng
Cho-Jui Hsieh
Yang You
Venue: ICLR 2022 Poster (2022)
Keywords: Distributed Machine Learnig
Large-Batch Training
Adversarial Learning

Abstract: Large-batch training has become a commonly used technique when training neural networks with a large number of GPU/TPU processors. As batch size increases, stochastic optimizers tend to converge to sharp local minima, leading to degraded test performance. Current methods usually use extensive data augmentation to increase the batch size, but we found the performance gain with data augmentation decreases as batch size increases, and data augmentation will become insufficient after certain point. In this paper, we propose to use adversarial learning to increase the batch size in large-batch training. Despite being a natural choice for smoothing the decision surface and biasing towards a flat region, adversarial learning has not been successfully applied in large-batch training since it requires at least two sequential gradient computations at each step, which will at least double the running time compared with vanilla training even with a large number of processors. To overcome this issue, we propose a novel Concurrent Adversarial Learning (ConAdv) method that decouple the sequential gradient computations in adversarial learning by utilizing staled parameters. Experimental results demonstrate that ConAdv can successfully  increase the batch size on both ResNet-50 and EfficientNet training on ImageNet while maintaining high accuracy. In particular, we show ConAdv along can achieve 75.3\% top-1 accuracy on ImageNet ResNet-50 training with 96K batch size, and the accuracy can be further improved to 76.2\% when combining ConAdv with data augmentation. This is the first work successfully scales ResNet-50 training batch size to 96K. 

Published as a conference paper at ICLR 2022
CONCURRENT ADVERSARIAL LEARNING FOR LARGE- BATCH TRAINING
Yong Liu', Xiangning Chen', Minhao Cheng', Cho-Jui Hsieh?, Yang You' "Department of Computer Science, National University of Singapore
DDepartment of Computer Science, University of California, Los Angeles DDepartment of Computer Science and Engineering, Hong Kong University of Science and Technology (1 (liuyong, youy J coop. nus. edu. sg, ( (xiangning, chohsieh)@cs www..llaaaa. mi ninhaocheng@ustt hk
ABSTRACT
Large-batch training has become a widely used technique when training neural networks with a large number of GPU/TPU processors. As batch size increases, stochastic optimizers tend converge to sharp local minima, leading to degraded test performance. Current methods usually use extensive data augmentation to in- crease the batch size as a remedy, but we found the performance brought by data augmentation decreases with the increase of batch size. In this paper, we propose to leverage adversarial learning to increase the batch size in large-batch training. Despite being natural choice for smoothing the decision surface and biasing towards a flat region, adversarial learning has not been successfully applied large-batch training since it requires at least two sequential gradient computations at each step. To overcome this issue, we propose a novel Concurrent Adversarial Learning (ConAdv) method that decouples the sequential gradient computations in adversarial learning by utilizing stale parameters. Experimental results demon- strate that ConAdy can successfully increase the batch size of ResNet-50 training on ImageNet while maintaining high accuracy. This is the first work that success- fully scales the ResNet-50 training batch size to 96K.
INTRODUCTION
As larger datasets and bigger models being proposed, training deep neural networks has become quite time-consuming. 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 TPUs. With the developments of high performance computing clusters, (e.g.. Google and NVIDIA build high performance clusters with thousands of TPU or GPU chips), how to fully utilize those computing resources accelerate the training process has become important research topic.
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 parameter update. Training with hundreds or thousands of processors with data parallelism is thus equivalent running stochastic gradient optimizer (e.g.. SGD Adam) with very large batch size, also known as large batch training. For example, Google and NVIDIA show that by increasing the batch size 64k on ImageNet, they can finish 90-epochs 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 avail- able? Large batch training often faces two challenges. First, given fixed number of training epochs, increasing the batch size implies reducing the number training iterations. Even worse, has been observed that large-batch training often converges to solutions with bad generalization performance (also known as sharp local minima) (Keskar et al., 2017), possibly due to the lack of inherent noise in each stochastic gradient update. Although this problem can be partially mitigated by using differ- ent 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 distributed batch normaliza-

---

Published as a conference paper at ICLR 2022
tion and maixed-precisie training, to further scale the training of ResNet-50 4096 v3 TPU chips. However, can just expand the batch size to 64k (Kumar et al., 2021; Ying et al., 2018).
To mitigate the generalization gap brought by large-batch training. data augmentation has become an indispensable component. For instance, researchers at Facebook use augmentation to scale the training of ResNet-50 to 256 NVIDIA P100 GPUs with batch size of 8k on ImageNet (Goyal et al., 2017). You et al. (2018) also use data augmentation to expand the batch size to 32k on 2048 KNL. nodes. However, in this paper we observe that when batch size is large enough (i.e., larger than 32k), data augmentation 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, which finds a perturbation within a bounded set around each sample to train the model. Previous works find that adversarial training can lead to significant decrease in the curvature of the loss surface and make the network behave more "linearly" in small batch size cases, which could be used as way to improve generalization (Xie et al., 2020; Moosavi-Dezfooli et al., 2019). However, adversarial training has not been used large-batch training since requires a series of sequential gradient computations within each update to find an adversarial example. Even when conducting only one gradient ascent to find the adversarial example, adversarial training requires two sequential gradient computations (one for adversarial example and one for weight update) that cannot be parallelized. Therefore, even with infinite computing resources, adversarial training is at least two times slower than standard training and increasing the batch size cannot compensate for that.
To resolve this issue and make adversarial training applicable for large-batch training, we propose novel Concurrent Adversarial Learning (ConAdv) algorithm. We show that by allowing the com- putation of adversarial examples using stale weights, the two sequential gradient computations in adversarial training can be decoupled, leading to fully parallelized computations at each step. As 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 ConAdy is better choice than existing augmentations.
Our main contributions are listed below:
This is the first work showing that 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 com- putations 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 ConAdy achieves almost identical performance the original adversarial training. We also provide theoretical analysis for ConAdy.
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, ConAdy alone achieves 75.3% accuracy when using 96K batch size. Further, the accuracy will rise to 76.2% when combined with data augmentation. This the first method scaling ResNet-50 batch size to 96K with accuracy matching the MI.Prf standard (75.9%), while previous methods fail to scale beyond 64K batch size.
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 an extremely large batch is difficult to converge and has a general- ization gap (Keskar et al., 2017; Hoffer et al., 2017). Therefore, related work starts to carefully fine-tune the hyper-parameter. to bridge the gap, such as learning rate, momentum (You et al., 2018; Goyal et al., 2017; Li, 2017; Shallue et al., 2018; Xue et al., 2021; Lou et al., 2021). Goyal et al. (2017) try to narrow the generalization gap with the heuristics of learning rate scaling. However, there is still big room to increase the batch size. Several recent works try to use adaptive learning rate to reduce the tinettunng of
and further scaling the batch size to larger value (You et al., 2018: 1: landola et al., 2016; Codreanu et al., 2017; Akiba et al., 2017; 7; Jia al., 2018; Smith
2

---

Published as a conference paper at ICLR 2022
et al., 2017; Martens & Grosse, 2015; Devarakonda et al., 2017: Osawa et al., 2018: You et al., 2019; Yamazaki et al., 2019; Liu et al., 2022). You et al. (2017) propose Layer-wise Adaptive Rate Scaling (LARS) for better optimization and scaling to the batch size of 32k without performance penalty ImageNet. In addition, related work also tries to bridge the gap from the aspect of augmentation. Goyal al. (2017) use data augmentation Yao et al. (2)18a) propose an adaptive batch size method based on Hessian information to gradually increase batch size during training and use vanilla adversarial training regularize against the sharp minima. However, the process of adversarial training is tinecconuming and they just use the batch size of 16k in the second half of the training process (the initial batch size is 256). How further accelerate the training process based on adversarial training and reduce its computational burden is still an open problem.
further scale the training of ResNet-50 ImageNet.
2.2 ADVERSARIAL LEARNING
Adversarial training has shown great success improving the model robustness through collecting adversarial examples and injecting them into training data (Goodfellow et al., 2015; Papernot al., 2016; Wang et al., 2019). Madry al. (2017) formulates it into min-max optimization framework as follows:
min E (e,, ()) max L(0,, x + 5,)1.
(1)
[0
where ID = ((r;;;:)) denotes training samples and x; E Rdd 3/: E {l, Zj. 5 is the adversarial perturbation, ||- llp denotes some Ly-norm distance metric, 02 the parameters of time and Z is the number of classes. Goodfellow et al. (2015) proposes FGSM to collect adversarial data, 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 generate adversarial examples. Recently, several papers (Shafahi et al., 2019; Wong al., 2020; Andriushchenko & Flammarion, 2020) aim to improve the computation overhead brought by adversarial training. Specifically, FreeAdv (Shafahi et al., 2019) tries update both weight parameter & and adversarial example at the same time by exploiting the correlation between the gradient to the input and the model weights. Similar Free-adv, Zhang al. (2019) further restrict most of the forward and backpropagation within the first layer to speedup computation. Wong et al. (2020) finds the overhead could be further reduced by using single-step FGSM with random initialization. While these works 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 computations to further boost efficiency, which is more suitable for large-batch training. Recently, several works (Xie et al., 2020; Cheng et al., 2021; Chen serve
et al., 2021; Mei et al., 2022) show that the adversarial example benefit the clean accuracy in the small batch size setting. However, whether adversarial training can improve the performance of large-batch training is still open problem.
augmentation to
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, 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. 5.9% accuracy on ImageNet.
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 batch training. Then we discuss the bottleneck of adversarial training in distributed systems and propose a novel Concurrent Adversarial Learning (ConAdv) method for large-batch training.
3

---

Published as a conference paper at ICLR 2022
8--
The updsle process stop
Step
The updari process of slepp
Worker 2
Werkor
Worker
Loca Dato D
Looal Oata 13
Local Dasa 1D
Lock Dats D
Trainin Dusa TD 69
Traning Data B
(t)
Figure 1: (a) Distributed Adversarial Learning (DisAdv), (b) Concurrent Adversarial Learning (ConAdv). To ease the understanding, we just show the system including two workers.
3.1 DOES DATA AUGMENTATION IMPROVE THE PERFORMANCE OF LARGE-BATCH TRAINING?
Data augmentation can usually improve the generalization of models and is a commonly used tech- nique to improve the batch size limit in large-batch training. To formally study the effect of data augmentation in large-batch training, we train ResNet-50 using ImageNet (Deng et al., 2009) by AutoAug (AA) (Cubuk et al., 2019). The results shown in Figure reveal that although AA helps improve generalization under batch size S 64K, the performance gain decreases batch size in- creases. Further, it could lead a negative effect when the batch size is large enough (c.g.. 128K or 256K). For instance, the top-1 accuracy is increased from 76.9% to 77.5% when using AA Ik batch size. However, it decreases from 73.2% to 72.9% under data augmentation when the hatch size is 128k and drops from 64.7% to 62.5% when the batch size is 256k. The main reason is that the augmented data increases the diversity of training data, which leads to slower convergence when using fewer training iterations. Recent work tries to concat the original data and augmented data to jointly train the model and improve their accuracy (Berman et al., 2019). However, we find that con- cating them will hurt the accuracy when batch size is large. Therefore, we just use the augmented data to train the model. The above experimental results motivate us to explore a new method for large batch training.
3.2 ADVERSARIAL LEARNING IN THE DISTRIBUTED SETTING
Adversarial learning can be viewed as a way automatically conduct data augmentation. Instead of defining fixed rules to augment data, adversarial learning conducts gradient- based adversarial attacks to find adversarial examples. As re- sult, adversarial learning leads to smoother decision bound- ary (Karimi et al., 2019; Madry et al., 2017). which often comes with flatter local minima ( (Yao al., 2018b). Instead of solving the original empirical risk minimization problem, adversarial learning aims to solve a min-max objective that minimizes the loss under the worst case perturbation of sam- plcs within a small radius. In this paper, since our main goal is to improve clean accuracy instead of robustness, con- sider the following training objective that includes loss on both natural samples and adversarial samples:
80.0
RooNar-t0
S0-AAA
Top-1 Accrery
75.0
72.5
70.0
67.5
62.5
1k
4k
8k
16k
30k
64k
96k
128k
256x
Batch Size
Figure Augmentation Analysis
min E I~D [C00 ( M1) + max C(0 + 8, 8,91),, ||
where C is the loss function and e represents the value of perturbation. Although many previous work in adversarial training focus on improving the trade-off between accuracy and robustness (Shafahi et al., 2019; Wong al., 2020), recently Xie et al. (2020) show that using split BatchNorm for adversarial and clean data can improve the test performance on clean data. Therefore, we also adopt this split Norch approach.
4

---

Published as a conference paper at ICLR 2022
Sequence
Concurrent
Betth
Algorithm 1 ConAdv
for 1 = 1, T do for v; B do
First Gradient
gr. (8t)
9a (b))
Compute Loss: L(0 V1) using main BN,
# (0.)
(0,)
(h,; (0,-), using adv BN,
(0, BBg [(0 Vh)+ Egg (2:(6t-)):))
Concurrent Training (Independet)
gc (0,)
Second Gradient
2: (8(rr)
Minimize the L (0,) and obtain of (0c) end for
for 2; E B 12,++ do Calculate adv gradient (2,)))n B Obtain adv examples (2)(0)), ys) end for end for
Update
å (0))
Sequence 2
E Bat
Update
Batee
Concurrent 2
Aggregate: 9e (o.) = Σ (0.)
Update weight 8,+1 parameter sever
Figure 3: Vanilla Training and Concurrent Train- ing
For Distributed Adversarial Learning (DisAdv), training data ID is partitioned into N local dataset D& and DD k=1 D For worker k, we firstly sample mini-batch data (clean data) B from the local dataset DD at each step t. After that, each worker downloads the weights 6g from parameter sever and then uses O, to obtain the adversarial gradients 인º (0.) V C(0e; Ti, y:) on input example ar E B Noted that we just use the local loss E B(ss.an)= ((0;; Ni; ys) to calculate the adversar- ial gradient of (0,) rather than the global loss E( x;, Vi), since we aim to reduce the communication cost between workers. In addition,
use 1-step Project Gradient Descent (PGD) to calculate 2 (0,) = + a V, Vi) to approximate the optimal adversarial example x;- Therefore, we can collect the adversarial mini-batch B% = { (0)),::) and use both the clean example (æi, y:) B and adversarial example (#7 (0:), 31) B% to update the weights 0,, More specially, we use main BatchNorm to calculate the statics clean data and auxiliary BatchNorm to obtain the statics of adversarial data.
We show the workflow of adversarial learning on distributed systems (DisAdv) as Figure 1, and more importantly, we notice that it requires two sequential gradient computations at each step which is cone-consuming and, thus, not suitable for large-batch training. Specifically, we firstly need to compute the gradient 9² (6.) to collect adversarial example .² After that, we these examples to update the weights 0r, which computes the second gradient. In addition, the process of collecting adversarial example द and use 2 update the model are tightly coupled, which means that each worker cannot calculate local loss -( L(e,; 24, Vi) and B(=..a))>> (ff;; 2 yt) to update the weights 8r, until the total adversarial examples & are obtained.
3.3 CONCURRENT ADVERSARIAL LEARNING FOR LARGE-BATCE TRAINING
As mentioned in the previous section, the vanilla DisAdv requires two sequential gradient computa- tions at each step, where the first gradient computation to obtain 4 based on C(O,, 2i, y:) and then compute the gradient of L(O,, 2f, y4) update 0,. Due to the sequential update nature, this overhead cannot be reduced even when increasing the number processors - even with an infinite number of processors, the speed of two sequential computations will be twice of one parallel update. This makes adversarial learning unsuitable for large-batch training. In the following, we propose simple but novel method to resolve this issue, and provide theoretical analysis the proposed method.
Concurrent Adversarial Learning (ConAdv) As shown in Figure 3, our main finding is that if we use stale weights (0,-r) for generating adversarial examples, then two sequential computations can
5

---

Published as a conference paper at ICLR 2022
be de-coupled and the parameter update step run concurrently with the future adversarial example generation step.
Now formally define the ConAdv procedure. Assume is sampled at iteration t, instead of the current weights Oz. we use stale weights 0,-, (where T is the delay) to calculate the gradient and further obtain approximate adversarial example (0(-r))
go (0:-r) V. Te, Vi),
ft (0t-r) = I; + a 9o (0
In this way, we can obtain the adversarial sample 2, ( through stale weights before updating the model each step Therefore, the training efficiency can be improved. The structure of ConAdv is shown in Figure 1: At each step I, each worker l can directly concatenate the clean mini-hatch data and adversarial mini-batch data to calculate the gradient 9. (0+) and update the model. That is because the system has obtained the approximate adversarial example 2; based the stale weights 0 before iteration t.
In practice, we т 1 the adversarial examples è, is computed at iteration - 1. Therefore, each iteration will compute the current weight update and the adversarial examples for the next batch:
0c+1 0, + Vo BB,
[(), (4, + EELYN 20 1s)),
(4)
(0,) - - aa V,L(O,; Mi), where (II, E Bo
(5)
where Ba, = - B denotes clean mini-batch of all workers and B. = U a.t represents adversarial mini-batch of all workers. These two computations can be parallelized so there is no longer two sequential computations at each step. In large-batch setting when the number of workers reaches the limit that each batch size can use, ConAdv is similarly fast as standard optimizers such as SGD or Adam. The pseudo code of proposed ConAdv is shown in Algorithm 1.
3.4 CONVERGENCE ANALYSIS
In this section, we will show that despite using the stale gradients, ConAdv still enjoys nice con- vergence properties. For simplicity, will use L(0, 22) shorthand for C(0; M4) and || indicates the le norm. We define the optimal adversarial example as x; = arg max L(0,, x;)) In order to present our main theorem, will need the following assumptions.
Assumption 1. The function [(0, .r) satisfies the Lipschitzian conditions:
[v,((0;;)) V ((0(();))) S Lue ||0 0211. [V ve[(O2; - Loo 01 021l. [Vo((0;)) - (0)))))))) < La - -22l. [V)L(0;)) - (.,(:::))V(((((;0.) _ Luel || - 22 (6)
Assumption 2. -(0,x) is locally p-strongly concave in X; fo- ||05 xelloo S for all i [n]. for for any 21, x2 E X.
11 ||]
C( mn) S C(,, x22 + (V,((O, x2), x22)
x2ll
2
Assumption 2 can be verified based on the relationship between robust optimization and distribu- tional robust optimization in (Sinha et al., 2017: Lee & Raginsky, 2017).
Assumption 3. The concurrent stochastic gradient ğ(O.) is bounded by the constant M:
E (VoC(O;;
(0)().)1 s M.
(8)
Assumption 4. Suppose (0,)
I E (Z(), + L(B,, g(0,)
28 Ei (Vol(i)) + (oL(0r,x))) and E[g(O.)) VLD(O.), where [B] represents batch
1 Σ
size. . The variance of g(O)) is bounded by o2:
Elllg(0.) V (0.))
Based on the above assumptions, we can obtain the upper bound between original adversarial exam- ple (0.) and concurrent adversarial example x (0.r)), where is the delay time.

---

Published as a conference paper at ICLR 2022
Lemma 1. Under Assumptions 1 and 2, we have
Lag |1x;; x (0,-))l 1 - -10, 0
nTT M. 14
(10)
Lemma illustrates the relation between (0.) and x (0,-)) which bounded by the delay т. When the delay is small enough, x, (0.->) can be regarded as an approximator of x; (0z). We now establish the convergence rate as the following theorem.
Theorem 1. Suppose Assumptions 1, 2, 3 and 4 hold. Let loss function (D().) [C(0,; aq , Ue) + [(0;; Di, y:)) and ã; (0:--) be the X-solution of x (0t-r): : {xī (0--))-
4a)) S X. Under Assumptions and for the concurrent stochastic gradient @(0). the step size of outer minimization is set to np =1 min(1/L, / / A/LoTT) Then the oulput of Algorithm satisfies:
- 2
T-1 I
LA
L 2
LM
A3 $V l
/
T
20
(1I)
where L Las +
Lie Low 2μ
Our result provides formal convergence rate of ConAdv and it can converge to a first-order sta- tionary point at sublinear rate up to precision of 2
222 which is related to т. In practice we use the smallest delay 1 as discussed in the previous subsection.
1 Lis
(TML (
4 EXPERIMENTAL RESULTS
4.1 EXPERIMENTAL. SETUP
Architectures and Datasets. We select ResNet default architectures. More specially, the mid-weight version (ResNet-50) to evaluate the performance of our proposed algorithm. The dataset we used in this paper is ImageNet-1k, which consists of 1. 1.28 million images for training and 50k images for testing. The convergence baseline of ResNet-50 in MLPerf is 75.9% top-1 accuracy in 90 epochs (i.e. ResNet-50 version 1.5 (Goyal et al., 2017)).
Implementation Details. We use TPU-v3 for all our experiments and the same setting as the base- line. We consider 90-epoch training for ResNet-50. For data augmentation, mainly consider AutoAug (AA). In addition, we use LARS (You et al., 2017) to train all the models. Finally, for adversarial training, we always use 1-step PGD attack with random initialization.
Throughuu (imagessms)
200
CanAdv
CoeAdiy Bosciie
60
40 Broughmann
100
30
10
16k.
Betthh
ResNet-50
ResNet-50 Limit
Figure 4: (a): throughput scaling up batch size for ResNet-50, (b): throughtput when the number of processors reach the limit that each batch size can use for ResNet-50.
4.2 IMAGENET TRAINING WITH RESNET
We train Re with ConAdy and compare it with vanilla training and DisAdy. The exper- imental results of scaling up batch size in Table 1 illustrates that ConAdy obtain the similar

---

Published as a conference paper at ICLR 2022
accuracy compared with DisAdv and meanwhile speed up the training process. More specially, can find that the top-1 accuracy of all methods are stable when the batch size is increased from 4k to 32k. After that, the performance starts to drop, which illustrates the bottleneck of large-batch training. However, ConAdy can improve the top-1 accuracy and the improved performance is stable DisAdv does when the batch size reaches the bottleneck (such as 32k, 64k, 96k), but AutoAug gradually reaches its limitations. For instance, the top-1 accuracy increases from 74.3 to 75.3 when using ConAdv with a batch size of 96k and improved accuracy is 0.7%, 0.9% and 1.0% for 32k, 64k and 96k. However, AutoAAg cannot further improve the top-1 accuracy when the batch size 96k. The above results illustrate that adversarial learning can successfully maintain a good test performance in the large-batch training setting and can outperform data augmentation.
Table I: Top-1 accuracy for ResNet-50 on ImageNet
Method
1k
4k
8k
16k
32k
64k
96k
ResNet-50
76.9
76.9
76.6
76.6
76.6
75,3
74.3
ResNet-50+AA
77.5
77.5
77.4
77.1
76.9
75.6
74.3
ResNet-50+DisAdv
77.4
77.4
77.4
77.4
77.3
76.2
75.3
ResNet-50+ConAdy
77.4
77.4
77.4
77.4
77.3
76.2
75.3
In addition, Figure 4(a) presents the throughput (images/ms) on scaling up batch size. We can observe that ConAdy can further increase throughput and accelerate the training process. To obtain accurate statistics of BatchNorm, we need to make sure each worker has at least 64 examples to calculate them (Normal Setting). Thus, the number of cores is [Batch Size 64]. For example, we use TPU v3-256 to train DisAdv when batch size is 32k, which has 512 cores (32k/64-512) As shown in Figure 4(a), the throughput of DisAdv increases from 10.3 on 4k to 81.7 on 32k and CondAdv achieve about 1.5x speedup compared with DisAdv, which verifies our proposed ConAdv can maintain the accuracy of large-batch training and meanwhile accelerate the training process.
To simulate the speedup when the number of workers reach the limit that each Batch Size can use, we use a large enough distributed system train the model with the batch size of 512, 1 and 2k on TPU v3-128. TPU v3-256 and TPU v3-512, respectively. The result is shown in Figure 4(b), can obtain that ConAdv can achieve about 2x speedup compared with DisAdv. Furthermore, this scenario we observe ConAdy can achieve the similar throughput Baseline (vanilla ResNet-50 training). For example, compared with DisAdv, the throughput increases from 36.1 to 71.3 when using ConAdy with a batch size of 2k. In addition, the throughput is 75.7, which illustrates that ConAdy can achieve a similar speed as baseline. However, ConAdy can expand to larger batch size than baseline. Therefore, ConAdv can further accelerate the training of deep neural network.
4.3 IMAGENET TRAINING WITH DATA AUGMENTATION
To explore the limit of our method and evaluate whether adversarial learning can be combined with data augmentation for large-batch training, we further apply data augmentation into the proposed adversarial learning algorithm and the results are shown in Table 2. We can find that ConAdy can further improve the performance of large-batch training on ImageNet when combined with Autoaug (AA). Under this setting, we can expand the batch size to more than 96k, which can improve the algorithm efficiency and meanwhile benefit the machine utilization. For instance, for ResNet, the top-1 accuracy increases from 74.3 to 76.2 under 96k when using ConAdy and AutoAug.
Table Top-1 accuracy with AutoAug on ImageNet
Method
1k
4k
8k
16k
32k
64k
96k
ResNet-50
76.9
76.9
76.6
76.6
76.6
75.3
74.3
ResNet-50+AA
77.5
77.5
77.4
77.1
76.9
75.6
74.3
ResNat-004CoAAdv4AA
78.5
78.5
78.5
78.5
78.3
77.3
76.2
4.4 TRAINING TIME
The wall clock training times for ResNet-50 are shown Table 3. We can find that the training time gradually decreases with batch size increasing. For example, the training time of ConAdy decreases
8

---

Published as a conference paper at ICLR 2022
from 1277s 227s when scale the batch size from 16k to 64k. addition, we can find that DisAdv need about 1.5x training time compared with vanilla ResNer-50 but ConAdv can effi- ciently reduce the training of DisAdv to level similar to vanilla ResNet. For instance, the train- ing time DisAdv reduced from 1191s 677s when using ConAdv. Noted that we don't report the clock time for vanilla ResNet-50 at 64k since the top-1 accuracy is below the MLPerf standard 75.9%. The number of machines required to measure the maximum speed at 96K exceeds our current resources. The comparison on 32k and 64k also can evaluate runtime improvement.
Table Training Time Analysis
Method
16k
32k
64k
ResNet-50
1194s
622s
DisAdv
1657s
1191s
592s
ConAdv
1277s
677s
227s
4.5 GENERALIZATION GAP
To the best of our knowledge, theoretical analysis of generalization errors in large- batch setting is still an open problem. However, we empirically found that our method can successfully reduce the generalization gap in large-batch training. The experimental results in Table 4 indicate that ConAdv can narrow the generalization gap. For example, the generalization gap is 4.6 for vanilla ResNet-50 at 96k and ConAdv narrows the gap to 2.7. In addition, combining ConAdy with AutoAug. the training accuracy and test accuracy can further increase and meanwhile maintain the similar generalization gap.
Table 4: Generalization Gap of Large-Batch Training on ImageNet-Ik
Vanilla ResNet-50
ConAdv
ConAdy AA
16k
32k
64k
96k
16k
32k
64k
96k
16k
32k
64k
96k
Training Accuracy
Test Accuracy
77.3
76.2
Generalization Gap
4.8
2.9
3.5
2.0
2.3
2.2
4.6 ANALYSIS OF ADVERSARIAL PERTURBATION
Adversarial learning calculates adversarial perturbation input data to smooth the decision boundary and help the model converge to the flat minima. In this section, we analyze the effects of different perturbation values for the performance of large-batch training ImageNet. The analysis results are illustrated Table It presents that we should increase the attack intensity as the batch size increasing. For example, the best attack perturbation value increases from (32k) (96k) for ResNet-50 and from 8 (16k) to 12 (64k). In addition, we should increase the perturbation value when using data augmentation. For example, the perturbation value should be for the original ResNet-50 but be when data augmentation is applied.
Table 5: Experiment Results (Top-1 Accuracy) when useing Different Adversarial Perturbation.
Method
Batch Size
p-0
p=4
p=5
p-10
p=12
ResNet-50 ConAdv
32x
76.8
77.2
77.3
77.4
77.3
77.3
77.3
77.3
77.3
77.3
77.2
77.2
ResNet-50 ConAdv AA
32x
77.8
78.0
78.1
78.1
78.0
78.3
78.2
78.2
78.2
78.2
78.2
78.1
ResNet-50 ConAdv
64K
75.7
76.2
76.3
76.3
76.4
76.7
76.4
76.4
76.4
76.4
76.4
76.3
ResNet-50 ConAdy AA
64x
76.8
77.0
76.8
77.0
77.1
77.1
77.2
77.4
77.2
77.1
77.1
77.1
ResNet-50- ConAdy
96x
74.6
75.1
75.1
75.1
75.3
75.1
75.1
75.1
75.3
75.2
75.1
75.1
ResNet-50- ConAdv AA
96x
73.8
75.9
75.8
76.0
76.0
76.0
76.0
76.1
76.2
76.2
76.0
76.0
5 CONCLUSIONS
We firstly analyze the effect of data augmentation for large-batch training and propose a novel dis- tributed adversarial learning algorithm to scale to larger batch size. To reduce the overhead of adversarial learning, we further propose novel concurrent adversarial learning to decouple the two sequential gradient computations in adversarial learning. We evaluate proposed method on ResNet. The experimental results show that our proposed method beneficial for large- batch training.

---

Published as a conference paper at ICLR 2022
6 ACKNOWLEDGEMENTS
We thank Google TFRC for supporting us get access to the Cloud TPUs. We thank CSCS (Swiss National Supercomputing Centre) for supporting us to get access the Piz Daint supercomputer. We thank TACC (Texas Advanced Computing Center) for supporting to get access to the Longhorn supercomputer and the Frontera supercomputer. We thank Lux Provide (Luxembourg national super- computer HPC organization) for supporting to get access to the Melu supercomputer. CJH and XC are partially supported by NSF under IIS- ASS-2008133 173, IIS- 2048280 and by Army Research Laboratory under agreement number W91 10111-20022-11
ETHICS STATEMENT
We do not have any potential ethics issues in this paper. We hope to propose a novel distributed adversarial learning algorithm to accelerate large-batch batch training.
REPRODUCIBILITY STATEMENT
we list our main hyperparameters for large-batch training in Appendix A.4 (Table 6). For experimen- tal details, we introduce our experiment settings in section 4, such the dataset, model architecture, data augmentation, optimizer and so on.
REFERENCES
Takuya Akiba, Shuji Suzuki, and Keisuke Fukuda. Extremely large minibatch sgd: Training resnet- 50 on imagenet 15 minutes. arXiv preprint arXiv:1711.04325, 2017.
Maksym Andriushchenko and Nicolas Flammarion. Understanding and improving fast adversarial training. In Advances in Neural Information Processing Systems, 2020.
Maxim Berman, Hervé Jégou, Andrea Vedaldi, Iasonas Kokkinos, and Matthijs Douze. Multigrain: unified image embedding for classes and instances. arXiv preprint arXiv: 19Xiv:1902.05509 2019.
Xiangning Chen, Cihang Xie, Mingxing Tan, Zhang, Cho-Jui Hsieh, and Boqing Gong. Robust and accurate object detection via adversarial learning. In Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition (CVPR), pp. 16622-16631, June 2021.
Minhao Cheng, Zhe Gan, Yu Cheng, Shuohang Wang, Cho-Jui Hsieh, and Jingjing Liu. Adversarial masking: Towards understanding robustness trade-off for generalization, 2021. URL https: / //openzeview..o..mmo watt/foruai um? TXJ
Valeriu Codreanu, Damian Podareanu, and Vikram Saletore. Scale out for large minibatch sgd: Residual network training on imagenet- lk with improved accuracy and reduced time to train. arXiv preprint arXiv: 1tX.:x:711.04400 2017.
Ekin D. Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoe V. Le. Autoaugment: Learning augmentation strategies from data. In IEEE Conference on Computer Vision and Pattern Recognition, 2019.
Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-FFii ImageNet: A Large-Scale Hierarchical Image Database. In IEEE Conference on Computer Vision and Pattern Recognition, 2009.
Aditya Devarakonda, Maxim Naumov, and Michael Garland. Adabatch: Adaptive batch sizes for training deep neural networks. arXiv preprint 171vv2171..02029, 2017.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of
deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Conference of the North American Chapter of the Association for Compu- fational Linguistics: Human Language Technologies, NAACL-HILT, 2019.
10

---

Published as a conference paper at ICLR 2022
Ian Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. In International Conference on Learning Representations, 2015.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv:1706.002: meee:/1066,,7, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In IEEE Conference Computer Vision and Pattern Recognition, 2016.
Elad Hoffer, Iay Hubara, and Daniel Soudry. Train longer, generalize better: closing the general- ization gap in large batch training of neural networks. arXiv preprint 0rXiv:/705.08744 2017.
Forrest N landola, Matthew W Moskewicz, Khalid Ashraf, and Kurt Keutzer. Firecaffe: near- linear acceleration of deep neural network training on compute clusters. IEEE Conference on Computer Vision and Pattern Recognition, 2016.
Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang. Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang. Liwei Yu, et al. Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes. arXiv preprint arXiv: 1205, 2018.
Hamid Karimi, Tyler Derr, and Jiliang Tang. Characterizing the decision boundary of deep neural networks. arXiv preprint arXiv: 19/2. / 1460, 2019.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference Learning Representations, 2017.
Sameer Kumar, Yu Wang, Cliff Young, James Bradbury, Naveen Kumar, Dehao Chen, and Andy Swing. Exploring the limits of concurrency in ml training on google tpus. Machine Learning and Systems, 2021.
Jaeho Lee and Maxim Raginsky. Minimax statistical learning with wasserstein distances. arXiv 2017.
preprint
Mu Li. Scaling distributed machine learning with system and algorithm co-design. PhD thesis, PhD thesis, Intel, 2017
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-aware minimization. arXiv preprint arXiv:2203.02711 2022.
Yuxuan Lou, Fuzhao Xue, Zangwei Zheng, and Yang You. Sparse-mip: A fully-mip architecture with conditional computation. arXiv preprint arXiv:2109.02008 2021.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. arXiv preprint arXiv:1706.06083, 2017.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factoree approximate curvature. In International conference on machine learning. PMLR, 2015.
Peter Mattson, Christine Cheng, Cody Coleman, Greg Diamos, Paulius Micikevicius, David Patter- son, Hanlin Tang, Gu-Yeon Wei, Peter Bailis, Victor Bittorf, et al. MIperf training benchmark. arXiv preprint arXiv: 1910.01500, 2019.
Jieru Mei, Yucheng Han, Yutong Bai, Yixiao Zhang, Yingwei Li, Xianhang Li, Alan Yuille, and Cihang Xie. Fast advprop. In International Conference Learning Representations, 2022. URL https: http///pppeeeeeeeeeeeeeeee net/ forum? hidhhhhsssssaaaae..
Seyed-Mohsen Moosavi-Dezfoole Alhussein Fawzi, Jonathan Uesato, and Pascal Frossard. Robust- ness via curvature regularization, and vice versa. IEEE Conference on Computer Vision and Patlern Recognition, 2019.
11

---

Published as a conference paper at ICLR 2022
Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Second-order optimization method for large mini-batch: Training resnet-50 on imagenet in 35 epochs. arXiv preprint arXiv:/811.1201 2018.
Nicolas Papernot, Patrick McDaniel, and Ian Goodfellow. Transferability in machine learning: from phenomena to black-box attacks using adversarial samples. arXiv preprint arXiv: 1rXiv:6605.0727,, 2016.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenA/ blog, 2019.
Ali Shafahi, Mahyar Najibi, Amin Ghiasi, Zheng Xu, John Dickerson, Christoph Studer, Larry Davis, Gavin Taylor, and Tom Goldstein. Adversarial training for free! arXiv preprint arXiv:1904. /2843, 2019.
Christopher Shallue. Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism neural network training. arXiv preprint 181 1.03600, 2018.
Aman Sinha, Hongseok Namkoong, and John Duchi. Certifiable distributional robustness with principled adversarial training. arXiv preprini wwXiv:/710.105777 2017.
Samuel L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe V Le. Don't decay the learning rate, increase the batch size. arXiv preprint arXiv: 1,00489. 2017.
Yisen Wang, Xingjun Ma, James Bailey, Jinfeng Yi, Bowen Zhou, and Quanquan Gu. On the convergence and robustness of adversarial training. In ICMI,, volume 1, pp. 2, 2019.
Eric Wong, Leslie Rice, and J Zico Kolter. Fast is better than free: Revisiting adversarial training.
arXiv preprint wrXiv:2001.0399. 2020.
Cihang Xie, Mingxing Tan, Boqing Gong, Jiang Wang, Alan L Yuille, and Quoc Le. Adversar- ial examples improve image recognition. In IEEE Conference Computer Vision and Pattern Recognition, 2020.
Fuzhao Xue, Ziji Shi, Futao Wei, Yuxuan Lou, Yong Liu, and Yang You. Go wider instead of deeper. arXiv preprint arXiv: 10..2 2021.
Masafumi Yamazaki, Akihiko Kasagi, Akihiro Tabuchi, Takumi Honda, Masahiro Miwa, Naoto Fukumoto, Tsuguchika Tabaru, Atsushi Ike, and Kohta Nakashima. Yet another accelerated sgd: Resnet-50 training on imagenet in 74.7 seconds. arXiv preprint arXiv: /2650, 2019.
Zhewei Yao, Amir Gholami, Daiyaan Arfeen, Richard Liaw, Joseph Gonzalez, Kurt Keutzer, and Michael Mahoney. Large batch size training of neural networks with adversarial training and second-order information. arXiv preprint arXiv: (at::01100.02 2018a.
Zhewei Yao, Amir Gholami, Qi Lei, Kurt Keutzer, and Michael W Mahoney. Hessian-based analysis of large batch training and robustness to adversaries. arXiv preprint 1802.0880000244 2018b.
Chris Ying, Sameer Kumar, Dehao Chen, Tao Wang, and Youlong Cheng. Image classification at supercomputer scale. arXiv preprint / 1.06992, 2018.
Yang You, Igor Gitman, and Boris Ginsburg. Scaling sgd batch size to 32k for imagenet training. arXiv preprint arXiv:
2017.
Yang You, Zhao Zhang, Cho-Jui Hsich, James Demmel, and Kurt Keutzer. Imagenet training in minutes. In International Conference Parallel Processing, 2018.
Yang You, Jonathan Hseu, Chris Ying, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large- batch training for Istm and beyond. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 2019.
Dinghuai Zhang, Tianyuan Zhang, Yiping Lu, Zhanxing Zhu, and Bin Dong. You only propagate once: Accelerating adversarial training via maximal principle. arXiv preprint arXiv: 1905. 005700877, 2019.
12

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1 THE PROOF OF LEMMA 1:
This completes the proof. The proof is inspired by Sinha et al. (2017); Wang al. (2019). Under Assumptions and 2, we have
Tup0 [0)((()) (0,-t)l1
Tzg8 01--11
110:
TTM
(12)
/
where x; (0a) and x; (0t-r) denote the adversarial example of x; calculated by 0, and stale weight 0,-,, respectively.
Proof:
According to Assumption 2, we have
S [(1,, (0))) + (V,,(0,, x1 -0, (0.))
"
2 | (0t-r) n (8:))1
(13)
$ 2 (0:))
H
2 | (0(-)) - =0;
In addition, we have
2(0,,2;(0.) & x5 (V==)(0) (0(-))
M || -2, -1)
(14)
2
Combining (13) and (14), we can obtain:
< (V L,,,,, (0))), x2 (0,-))
I (V L,,,,,, (0.-r)) V.L(0, x, (0)=))))) (0,)) (0r-r)
< Vn (0t-r)) V,0(O, -rI (0+=+))111= 2/ (0.
$ Lug 10: - O 01-111122 - x =0(-,))
(15)
where the second inequality due to (V.[(e,-,, 27 (0.-,)) (0.) (01-r)) < O. the third inequality holds because CauchySchwarz inequality and the last inequality follows from Assumption Therefore,
Las 110 01-rll
[x: - (0,-
S
11
Lug
Σ
(0--jt1-00--)) - 0z-,))l
jE(1.T]
(16)
Lug Σ
gg:=;(
μ
jE(1.7]
Lz0
M
where the second inequality follows the calculation of delayed weight, the third inequality holds
because the difference of weights is calculated with gradient 9 [1, 7]) and the last inequality holds follows Assumption 3.
13

---

Published as a conference paper at ICLR 2022
Thus,
Lux [0, x/ (0.-r))l
17M M
(17)
16
This completes the proof.
A.2
Lemma 2. Under Assumptions and 2, we have (D is L-smooth where L Les + Lee Low i.e., for any 0, and O2, we can say
2p
(VoLD(O.) VoLD(0 & L1[0: - f6--|
(18)
L
(p(+)) = Lo(0rrr) (CD(8)-))) - -b-r) +
t ーー 110
0e-r1l
(19)
Proof:
Based on Lemma I, we can obtain:
(0,) - x (0t-r))l s Lzd
10, 0--
(20)
We can obtain for i E [n]]
$ VoL(O,, (0:)) VoL x (((--))
||VoL x (0--)) - VoL 1-r (00000r,)0 J.j (02-+))l
LLoo - 0.- + Loz |, (0t) - x (0t-))l
(21)
Luo 110 0,-r|
<Loo|10: 01-7 Lgz
Law )110, )100
=(Loo + Los
0---11
11
where the second inequality holds because Assumption 1, the third inequality holds follows Lemma 1.
Therefore,
1 "
[V,( - VLD(0r-))l22 S.
Σ
v +
2n
VoL(e,, - VoC(0r-,
2n
(22)
+ 2n
VeL(O,, VoL x (01-)))
1=1
1 LLo + Loz Lee Lx0) ||B: 0
< Loell0e - 0 |l
2
(Lee + L
Lew)ll8, - -0,-71
24
This completes the proof.
14

---

Published as a conference paper at ICLR 2022
A.3
Lemma 3. Let (.( be the X-solution of x (0z): (x) (0.) - (0.), &; S X. Under
Assumptions and 2, for the concurrent stochastic gradient g(0), have
Les
X
11g(0))-g((4))) -((00-)) -
(23)
2
Proof:
llo(0.) - -0(0==>)] =
VoL(O,,, + (2)
BEIB
(V ( 1.) +
1
(Vo[(0,, (0:)) VoL(0,, (0c-,)))
B
B
VoL(O,, 2 (0.)) a0
2
iE|8I
<
Σ
Loz (1)) - (0,--))
2|
sE|8|
(24)
Løx [0,) x1 (0c_r) x (0t-r) (0c=)ll
2|3
iE|B]
<
(Loz || 1 (0t) x (0. )l + Lozz || (0:-r) -; (0,-0=))ll)
2|
kE|8I
Loz (0,) (0.--)) Loz (0,-7) (02-))))
2|
iE|B)
(L M Tuad
Σ
+ -
2|
kE|B
Let (0(--)) be the X-approximate of x;7 (0:-r), we can obtain:
- (0
(25)
In addition, we can obtain:
- - 5o
(26)
Combining 25 and 26, we have:
- (,(0))))2(((-)))) ||
(27)
Based on Assumption have
(V, 1,000000100000000000000000000010 (28)
Combining 28 with 27, we obtain:
- <
(29)
15

---

Published as a conference paper at ICLR 2022
Therefore, have
| (0(-)) x(0--)) S 1
(30)
Thus, we can obtain
Loe
(17 M Lua
ll((,)) - 0(0e-r)l1
+
(31)
2
1
This completes the proof.
A.3 THE PROOF OF THEOREM 1:
Suppose Assumptions 2, 3 and 4 hold. Let A = CD(0)) - ming (D)),, (C) =
-i=ı (VL(0,, 24, vi) + VL 25 9:)). If the step size of outer minimization is set to 1)
7 = min A/ where L Leg + Lue Loz. Then the output of Algorithm satisfies: 2μ
LA
L
TM
A
ElJ B
200
-
(32)
T 1==
T
Lμp.
V 11
where = Lo baa eep + Loo)-
Proof:
LD(0e+1) < (o,)))) + (V.D(0))) (11) - 0²) + 10++1 0.11
CD(()) // (1))) Ln²
110(0.)11 + (VVLD(()) VCD(O)) g(0.)
2
CD(()) 11 Ln
- L() vLD(0) ((0.)
2
Li7 (1)(0.)) - (o.(()))
- 7/11 L1
[V) (n))))=++((1 10 -9(0.)
2
+ +n - LN) (VLD (0.), (L) g(,)) + Lnp 2 l0(0)) - ((0)) + +(0)) (LD(1)))
$ (D) " /VL (0.) + "
//2 Ln)lg((e)) ((0))
+1(11 - (LD 01)) (LD(O)) g(0()) + (I1(()) o(0) 2 (lg(0.) VLD
=D)(0) (VD (1+ Lm)ll((0)) o(
+ + (L) (L)( ,())) Lo7 VLD
(33)
Taking expectation both sides of the above inequality conditioned 0,, we can obtain:
16

---

Published as a conference paper at ICLR 2022
E[Cn(0++1) - Ln(0.))0,) - - ( + (1 + Lnji
Lox
(nTM La
Luo
)2 + Lnfo
2
+ (
(1 Lay (MTT
Lux
1) + LnI²0²
2 ( + nL (1 (1+ (nT
Luo
)2 + Lif
8
M
(34)
where we used the fact that Elg(6:)] VCp(0,) Assumption 2, Lemma and Lemma 3. Taking the sum of (34) over t = O,... T 1, we obtain that:
T
Lue n (I+ M 1 +
A
Σ
1
Σ o²0²
1=0
1=0
(35)
Choose y = minm1/L, / riox) where A (D((o) (D(0r) and Loo +
Lee
Lozs we can 2p
show that:
LA
L
TMLa
5 12
( s 201
V
(36)
T
2
Lu
A.4 HYPERPARAMETERS
HYPERPARAMETERS:
More specially, our main hyperparameters are shown in Table 6.
Table 6: Hyperparameters of ResNet-50 on ImageNet
32k
64k
96k
Peak LR
35.0
41.0
43.0
Epoch
90
90
90
Weight Decay
5B-4
5B-4
5B-4
Warmup
40
41
41
LR decay
POLY
POLY
POLY
Optimizer
LARS
LARS
LARS
Momentum
0.9
0.9
0.9
Label Smoothing
0.1
0.1
0.1
17


Reviews:
[{"review_id": "ytYmG-4q5uq", "reviewer": "Reviewer_nTet", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "My main concerns are the scalability of the proposed augmentation strategy and the gap between theorems and experiments.\n\nDetails:\n1. The experiments mainly focus on ResNet-50/ImageNet training. How to set a large batch size for a new training problem? Whether 96k, 64k or 32k batch size is suitable for other problems?\n2. The authors actually prove a non-asymptotic convergence rate and I also notice that LARS is used for the proposed methods in experiments. Such a gap may lead to many unknown problems. Furthermore, according to my experience, LARS is sensitive to the hyper-parameters and NaN often occurs for large batch training.\n", "summary_of_the_paper": "This paper proposes a novel strategy which adding adversarial data for training to improve the performance of large batch training. Experiments also verify its efficiency.", "summary_of_the_review": "Above all, the authors should improve the theory for a better understand of the adversarial data. It would be better if the authors propose the strategy on how to choose a suitable large batch size for a given problems.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "XiT8Ds2q-X9", "reviewer": "Reviewer_Zv8z", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strength\n\n- Well-written. The motivation is clearly depicted and so is the thought process to the proposed solution.\n\n- The proposed approach is simple yet novel and useful. As demonstrated in the experiments, using stale gradients can achieve a great speed up over the DisAdv baseline with no noticeable accuracy differences.\n\n## Weaknesses\n\n- It seems that the local machine has to support two parallel models. From what I understand, Fig 1b demonstrates the case where the local machine supports a stale model and a current model so both can be executed by different workers in parallel. If a model is so large that a single machine can't fit two models, then it creates extra communication overhead with other workers for communicating stale weights and adversarial examples. Long story short, using stale weights as a solution might lead to additional communication overhead that might be an issue depending on the setting of the training system.", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, I think this paper is well-written and it provides a novel, simple, yet useful algorithm for further increasing the batch size for large-batch training.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "O_WF9i-BzeT", "reviewer": "Reviewer_M9ZZ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is clearly motivated and well written. To the best of my knowledge, the findings regarding the adversarial training for improving generalization performance in **large-batch training** is new. Although the proposed method is simple and straightforward, and the theoretical analysis is easy to conduct, the empirical results are nice. \n\nThe authors conducted various experiments, in which many empirical findings are quite interesting. However, the results lack of reasonable interpretations. For example, \n1.\tAbout the perturbation magnitude. From the experiment result in Table 5, it shows we should increase the attack intensity as the batch size increasing and augmentation used. Could you elaborate a little bit why should be this? It is a well known conclusion that there exists accuracy-robustness trade-off when using adversarial training, which means that larger perturbation might ruin the clean accuracy. This appears contrary to your finding. It is a really interesting point. \n2.\tAnother important side regarding the training details is that whether the pure adversarial training helps, i.e. not mixing the clean data and adversarial data for training the networks. I think the authors should provide this result to deepen the understanding the role of adversarial examples for generalization. \n", "summary_of_the_paper": "This manuscript empirically show that adversarial training in large-batch training scenario has better performance than traditional data augmentation. And furthermore, the authors proposed a simple method to conduct adversarial example generation and gradient computation w.r.t. to weights concurrently to accelerate the adversarial training in distributed setting. The key strategy is to use staled weights to generate adversarial examples, and then decoupled the bi-level optimization. ", "summary_of_the_review": "New empirical findings regarding the advesarial training helps generalization in large-batch training scenarios, but some results lack of deeper understanding. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Thank you for your submission to ICLR.  This paper presents a straightforward but reasonable approach to (slightly) improving the performance of large-batch training via adversarial training.  The basic approach is to apply (small epsilon) adversarial training, shown to help performance in small-batch settings, but accelerate the method using stale parameters to allow for parallel computation of the perturbations.  This speeds up adversarial training while improving performance, enough to enable it to be more effective than existing techniques for this large-batch setting.

The reviewers are not entirely in agreement about this paper, but I personally found the objections of the reviewer to be fairly generic, and not really addressing the core contributions of the paper.  However, I also felt that the overall contribution of this work seems somewhat incremental, using a not-particularly-unexpected result (that we can use stale gradients for this form of adversarial training) to achieve moderate speedup in what ultimately seems like one point in hyperparameter space.

That all being said, though, clearly the authors are working within standard benchmark frameworks, and "simple" algorithms here are indeed a positive rather than a negative.  So I am inclined to slightly recommend the paper for acceptance.
Decision: Accept (Poster)