ScaLA: Speeding-Up Fine-tuning of Pre-trained Transformer Networks via Efficient and Scalable Adversarial Perturbation
Authors: Minjia Zhang
Niranjan Uma Naresh
Yuxiong He
Venue: ICLR 2022 Submitted (2022)
Keywords: Efficient Training Methods
Large Batch Optimization
Transformer Networks
BERT

Abstract: The size of transformer networks is growing at an unprecedented rate and has increased by three orders of magnitude in recent years, approaching trillion-level parameters. To train models of increasing sizes, researchers and practitioners have employed large-batch optimization to leverage massive distributed deep learning systems and resources. However, increasing the batch size changes the training dynamics, often leading to generalization gap and training instability issues that require extensive hyperparameter turning to maintain the same level of accuracy. In this paper, we explore the steepness of the loss landscape of large-batch optimization and find that it tends to be highly complex and irregular, posing challenges to generalization. To address this challenge, we propose ScaLA, a scalable and robust method for large-batch optimization of transformer networks via adversarial perturbation. In particular, we take a sequential game-theoretic approach to make large-batch optimization robust to adversarial perturbation, which helps smooth the loss landscape and improve generalization. Moreover, we perform several optimizations to reduce the computational cost from adversarial perturbation, improving its performance and scalability in the distributed training environment.
We provide a theoretical convergence rate analysis for ScaLA using techniques for analyzing non-convex saddle-point problems. Finally, we perform an extensive evaluation of our method using BERT and RoBERTa on GLUE datasets. Our results show that our method attains up to 18 $\times$ fine-tuning speedups on 2 DGX-2 nodes, while achieving comparable and sometimes higher accuracy than the state-of-the-art large-batch optimization methods. When using the same number of hardware resources, ScaLA is 2.7--9.8$\times$ faster than the baselines.

Under review as a conference paper at ICLR 2022
SCALA: SPEEDING-UP FINE-TUNING OF PRE- TRAINED TRANSFORMER NETWORKS VIA EFFICIENT AND SCALABLE ADVERSARIAL PERTURBATION
Anonymous authors Paper under double-bling blind review
ABSTRACT
The size of transformer networks is growing at an unprecedented rate and has increased by three orders of magnitude in recent years, approaching the trillions. To train models of increasing sizes, researchers and practitioners have employed large-batch optimization to leverage massive distributed deep learning systems and resources. However, increasing the batch size changes the training dynamics, often leading to generalization gap and training instability issues that require extensive hyperparameter turning to maintain the same level of accuracy. In this paper, we explore the steepness of the loss landscape of large-batch optimization and find that it tends to be highly complex and irregular, posing challenges to generalization, To propose ScaLA, scalable and robust method for large-
address this challenge, batch optimization of transformer networks via adversarial perturbation. Moreover, we perform several optimizations to reduce the computational cost of performing the adversarial perturbation, thereby improving its performance and scalability in the distributed training environment. We provide theoretical convergence rate analysis for ScaLA using techniques for analyzing non-conver saddle-poing problems. Finally, we perform an extensive evaluation of our method using BERT and RoBERTa GLUE datasets. Our results show that our method attains up to 18 lne-tuning speedups on 2 DGX-2 nodes, while achieving comparable and sometimes higher accuracy than the state-of-the-an large-batch optimization methods. When using the same number of hardware resources, ScaLA is 2.7-9.8x faster than the baselines.
INTRODUCTION
We have seen exponential growth in DL model size since the debut of the transformer net- work VVaswani et al. 2017). For example, while BERT base has around 100M parameters, the model size has increased to multi-billion parameters such as Megatron-LM (8B) (Shoeybi et al.) stagger-
2019). T5 (1 (1 (Raffel et al. 2211)) Turing-NI-G (17B) (tur), and with GPTT- hitting ing record of 175B parameters. With the three orders of magnitude growth, these large models also have powered accuracy breakthroughs in many challenging Natural Language Processing (NLP) tasks such as the General Language Understanding Evaluation (GLUE) benchmark (Wang et al. 2019a). Recent studies show that the performance of these models continues to scale with their sizes Kaplan et al. 2020 As result, we expect that the model size would continue to grow in the future.
To accelerate the training speed of large models, the most common way increase the batch size in the optimization algorithm in order to leverage multi-GPU training (Li et al. 2020 et al.) 2019 Huang et al. 2019, Shazeer et al. 2018 (hoeybi al., 2019 Rajbhandari al. 2019). By increasing the batch size, mini-batch of size B can be divided across more workers (GPUs), where the gradients are computed locally on each worker using pook-propggttion and then aggregated. Furthermore, most of the operations used in transformer networks are highly optimized in modern linear algebra frameworks on GPUs and can scale to larger batch sizes without significantly increasing the time per step ( (Wang et al. 2019c Kaplan al.) (020). If researchers can train each neural network with more GPUs and increased throughput, then it makes it possible for them to achieve better results by training even larger models, using larger datasets and exploring ideas rapidly.

---

Under review as a conference paper at ICLR 2022
However, changing the batch size is not always straightforward, as it often impacts the training dynamics. You et al. propose LAMB (You et al. 2019a) to exploit large- optimization for transformer networks. LAMB is variant of Adam (Kingma & Ba 2015) that applies layer-wist normalization before applying each gradient update, which has been used to successfully train BERT on 1024 TPU chips in 76 minutes. Despite showing promising results, prior work (You al.] 2011a) primarily focuses on pre-training. On the other hand, the fine-tuning stage starts to become a bottleneck (e.g., it takes tens of hours fine-tune RoBERTa-large on MNLI (Wang et al. 2019a)) and becomes more expensive as model size increases. If we can speed-up pre-training by increasing batch sizes, why do we not also increase the batch size during fine-tuning in the interest of making fine-tuning more efficient well? Contemporary experience is that fine-tuning with large batch sizes is harder to train, often reaching lower accuracy than the baseline accuracy using small batch sizes.
To address these challenges, we develop new approaches to improving the scalability and generaliz- ability in fine-tuning pre-trained transformer networks by making the following contributions: (1) We present an adversarial perturbation based large batch optimization algorithm ScaLA (Scalable Large-batch Adversarial Perturbation) for training transformer networks, in the distributed training setting. We show how adversarial perturbation helps improve the generalization and more importantly how we reduce the cost of injecting adversarial perturbations to improve computational efficiency. (2)
We also present theoretical convergence rate analysis using techniques for analyzing non-convex saddle-point problems. (3) We conduct evaluation on a wide range of natural language understanding (NLU) tasks and assess the impact of adversarial perturbations on both the scalability and the general- izability in large batch task-specific fine-tuning, (4) We evaluate our approach against BERT (Devlin) al. 2019) and RoBERTa (Liu et al.) 2019) and show that ScaLA obtains significant improvements over the state Mfate-of-the-arr algorithms,, such as LAMB, for the large batch optimization of fine-tuning tasks. Concretely, while LAMB leads to 1 point accuracy drop on average (e-g.. GLUE) as we increase the batch size, our approach achieves the same and sometimes higher accuracy (up to 0.9 points) after drastically increasing the batch size. Furthermore, with our cost-efficient optimizations, ScaLA achieves up to 18x speedups on 2 NVIDIA DGX-2 nodes over the baseline fine-tuning time and is up to 9.8x faster than the baseline when using the same number of GPUs.
2 BACKGROUND AND RELATED WORK
Despite the great success of pre-trained transformer networks such as BERT (Devlin et al., 2019). a big challenge. in general, comes from the training efficiency even with self-attention and par- allelizable recurrence (Vaswani et al. 2017), and high-performance hardware (Jouppi et al. 2017). training transformer networks can still take a significant amount of time. One effective approach to reducing training time is through data parallelism (Devlin et al.) 2019 Liu et al. 2019: Shoeybi et al., 2019) which motivates studies on large-batch stochastic non-convex optimizations for transformer networks (You et al. 2019a). These studies have raised concerns with respect to its convergence, generalizability. and training stability by observing that training with a large batch could be diffi- cult (Keskar et al., 2017: Hoffer al., [017]] Nado et al., (021). Furthermore, prior works mostly focus on reducing the pre-training time (You et al. 2019a] [Zhang & He. 2020;/Gong et al. 2019
Clark al. 2020 instead of the adaptation time the fine-tuning stage.
While many researchers and practitioners focus on how to reduce the pre-training time, few attention has been paid accelerate the fine-tuning stage, which gradually becomes bottleneck as model sizes increase (e.g., it takes tens of hours to fine-tune MNLI on RoBERTa-large). Therefore, in contrast to previous works, our goal in this paper to speed-up task-specific fine-tuning without hurting model accuracy. Different from pre-training, the fine-tuning stage often employs much smaller batch size (e.g., log, B {4, 5)) than pre-training (e.g.. log, B > 10) (Devlin et al. 2019 Liu et al. 2019). The small batch size results in inefficient data parallelism (i.e., sub-optimal compuration-commniicaa
ratio), making it difficult for fine-tuning to benefit from multi-GPU training. Moreover, common understanding is that small-batch sizes provide implicit regularization effects (c.g., from gradient noise) that help improve generalization of downstream tasks. In contrast, our goal is to -peed-up the fine-tuning process with large batch sizes while preserving model accuracy.
On a separate line of research, adversarial training was first proposed in the computer vision literature to improve model's robustness against adversarial attacks (Goodfellow et al. 2015) Madry et al. 2018). Recently, there has been some work that shows that adversarial training helps improve model generalizability (Cheng et al., 2019 Wang et al. 2019b: Jiang al. 2020
Liu al. 2020 Yao et al.
2

---

Under review as a conference paper at ICLR 2022
2018a Zhu et al., 2020). However, very few works examine large-batch optimization of transformer networks and NLP tasks with adversarial perturbations from a computational efficiency and scalability perspective. The work most similar to ours is Zhu et al. (2020). who study adversarial training for NLP tasks by accumulating the gradient of the parameters from each of the ascent steps and updates the parameters only once after K inner ascent steps with the accumulated gradients. Unlike ZZuu et al., (2020), we consider accelerating the fine-tuning speed by parallel adversarial training and by adjusting the number of inner maximization steps, which offers much higher speedups.
3 THE PROPOSED METHOD
In this section, we present a principled method for large batch optimization that is highly scalable while maintaining the quality of the solutions as measured by task-appropriate accuracy metrics.
3.1 A SEQUENTIAL GAMORTHEOR METHOD VIA ADVERSARIAL PERTURBATION ORACLE
Let X denote the parameter space and denote the data saini-batchssample) space and Q denote a distribution supported on Y. To improve the generalizability of transformer fine-tuning while retaining the scalability, we augment the usual stochastic optimization objective by constructing an adversarial (Keskar et al.) 2017 Madry et al.] 2018) regularization. In particular, we solve the following robust optimization problem, which is a stochastic minimax (Lin et al.) 2020) optimization problem augmented with a regularization term involving a deterministic adversarial perturbation, instead of vanilla risk minimization:
Formulation:
min Eg [)] = min Eg + Ar(x)]
#EX
#EX
min max R Re- + Ar(x, = min max B Re zeX yEY
(1)
=EX vEY
+ R denotes the robust training objective, X R denotes the standard training objective, / X x V x V + R denotes the augmented objective, r X - R denotes a deterninistic regularization term on the parameters controlled by a strength factor À E (0, oo), r: X IR denotes the augmented regularization and & denotes samples drawn from Q (for simplicity, we slightly abuse the notation in using & to denote the random variable, e.g. [x or its E empirical realizations, e.g. K Σi= g(xr Ex) for any K; the meaning is clear from the context). The overall (outer) training objective involves a minimization problem in the parameter space while being stochastic with respect to the data space. The adversarial regularization (inner) term is deterministic maximization problem operating in the data space conditioned on a fixed parameter configuration. We wish to emphasize that this formulation is a two-player sequential (Jin et al.. 2020), not simultaneous, game wherein the goal is to optimize a transformer network that is robust to adversarial perturbation. In a given round, the first player (associated with the outer minimization) proposes parameter configuration, and the second player (associated with the inner maximization) responds with a penalty to capture the effect of label errors due to perturbations a large data batch size to undermine the performance of the transformer parameter configuration chosen by the first player.
where g
Practical Considerations: Language expressions are quite sensitive to individual words or clauses, where perturbations against those would likely generate incorrect or biased training data with wrong labels (Zhang & Yang 2018). Following prior success in applying adversarial training to NLP models (Miyato et al.] 2017: Zhu et al. 2020), we apply perturbations to the continuous word embeddings instead of directly to discrete words or tokens. The term I captures the prediction deviation from the perturbation. In given round of the game, with respect the first player's proposal, let & denote the transformer network under consideration (specifically, $ is BERT in this paper) and E be a large batch of data sampled from Q. We construct a label for the second player as y : 4(z, xx E). Next, for classification tasks, we choose to be the symmetric KL divergence (Jiang al. 2020), i.e., r(x, y) == KI.ym(( (7, (xx, v)). We use symmetric KL. divergence to measure the distributional divergence to generate adversarial perturbation. For regression tasks, we choose to be the squared loss, i.c., r(x, y) := (7 - (z )) In practice, we add an loo constraint on y, which is achieved by simple clipping with radius of (projection). Intuitively, large corresponds to situation wherein the transformer is highly sensitive to a given perturbation in the input, suggesting that the model parameters close to sharp minimum. Augmenting the original training objective with r makes the first player incur an additional penalty if the outer minimization solution veers closer to sharp minima, thereby encouraging flatter solutions and better generalizability.
3

---

Under review as a conference paper at ICLR 2022
Minimize loss
Output layer
Maximize
Exocuitin time (mz))
-00
120
202
Reseiie
Niversarial
33000
30000
Training Accuracy
V1.0
915
prediction
25000
Context embeddings
deviation
Training time
90.5
Transformer encoder
(Delayed
30.0
Accil
PGA-1]
Input embeddings
Adv. nolse
10000
80.5
Lexicon encoder
Noise generator
30-
5000
senteence inpuls
25.5
Figure I: The architecture of the pro- Figure 2: Time breakdown Figure 3: Impact of perturba- with and with PGA-1.
posed method.
tion steps.
Inner Maximization: For any given outer step f, let Ie denote the parameter proposed by the first player. Since the exact inner maximization in Equation equation is intractable for non-convex models such as transformers, we adopt truncated methods as in prior works. Specifically, we use Projected Gradient Ascent (PGA) (Madry et al. 2018 Jiang et al. 2020) to solve this problem, i.e. Ur+1 II. (y+ + P-V, v,V (I2,)) where Pr for the result of the gradient ascent update into an loo ball of diameter 2w around the original input embeddings, ç. considered by the first player.
[T] is the step size sequence and II projects
Outer Minimization via Groupwise Adaptive Learning Rates: Inspired by prior works that stabilize large-batch training and improve its convergence quality, we employ algorithms with groupwise adaptive learning rates (You et al., 2019a) to solve the outer minimization. Specifically, we solve the minimization problem in Equation equation 1 by ri+1 = .r} --i-ne
(x))), Vi E [h] where i denotes the lay- of the transformer. The normalized gradient descent mitigates issues due to exploding gradients. The learning rate sequence Th. Vt E [T] is scaled by a clipping function v(c) == max C, mi)(c, where L U (e.g., L = 0 and U 10), which ensures the norm of the update is of the same order as that of the weights. Note that we use gradient averaging on &, i.e., gradient accumulation and all-reduce, over batch size B distributed across P workers in order to obtain a noisy gradient estimate (zg(2)) at epoch
Computational Cost: Given that the primary interest of using large batch size is to improve hardware efficiency, we are motivated look carefully into the computational cost of adversarial perturbation. Adversarial perturbation requires an extra PGA inner loop that standard training does not have. Figure provides the time breakdown of adversarial training using PGA with T (denoted as PGA-T), PGA-1 performs the perturbation and takes approximately the same time as making three forward passes (Fwd) through the network. This is because one step of PGA requires to make one forward and backward pass (Bwd) over the entire network. The backward pass of adversarial training takes roughly twice the amount of time as the standard backward step because the back- socck-ropaggtion is triggered twice to calculate the perturbation noise and the gradients. The time spent on the optimizer step function (step) remains the same. In total, adversarial training slows down training by at least times, even with T=1. This motivates us to look the effectiveness of different perturbation steps as well as the usefulness of perturbation from the initial epochs in practice.
Impact of Perturbation Steps, T: Prior works often do multiple gradient computation steps (T > 1) and take several times longer training time to produce adversaries (Madry et al.) 2018 Zhu al. 2020), because their focus is not on computational efficiency. Subsequently, researchers presented Curriculum Adversarial Training (CAT) (Cai et al., 2018) and Annealing based Adversarial Training (Amata) (Ye et al.) 2020). which progressively increase the perturbation with various strengths, cutting the adversarial training cost while maintaining good accuracy. To investigate how CAT and similar methods affect large-scale NLPP problems involving transformers, we evaluate the final fine-tuning accuracy and training cost of QNLI, varying the number of perturbation steps T and report the results in Figure 3 Interestingly, although using a large T helps to produce stronger adversaries, we find that this does not lead to improved mun-ttnnng accuracy, despite the fact that the training overhead still increases almost linearly. In fact, the best accuracy is achieved with T = 1.
Our hypothesis to this phenomenon is below. The model has two components, namely, the parameter space and data space. First, unlike the minimization in the parameter space, which is stochastic, the maximization in the data space is deterministic. Second, with respect to the testing phase,
4

---

Under review as a conference paper at ICLR 2022
the numerical convergence in the model's parameter space is of primary importance rather than the numerical convergence in the data space, i.c., the maximization is an auxiliary procedure that augments the training phase make the parameter space "aware" of effects of the batch size across epochs. Due to these two points, at a certain epoch, for a given batch, the marginal utility of an additional PGA step is low, and we are able get away with inexact deterministic maximization. Therefore, we apply PGA-l in our large-batch optimization scheme, given that it produces sufficiently good solutions while being much more computationally efficient.
Given that even PGA-I still adds an overhead factor of 2, we are motivated investigate how useful adversarial perturbations are in the initial phase of large- batch optimization. We conduct additional experiments to measure the final accuracy corresponding to starting from a regular fine-tuning and then enabling PGA-1 for 2 l, where is [T]. Our observation is that enabling PGA-1 from the beginning does not offer much improvement in accuracy, whereas adversarial perturbation becomes more potent as the model begins to stabilize towards the end of training. Intuitively, this makes sense because generally, initialization, the model's parameters are relatively far from their final values and are less likely to get stuck at local minima. quite different
Delayed Perturbation Injection:
Therefore the adversarial perturbations generated in the initial training iterations from the perturbations towards the end of training because they would not maximize the adversarial loss in Equation [ We remark that similar phenomenon has been observed in computer vision tasks (Cai et al. 2018 Ye et al.) 2020) Gupta et al. (020). We show that it is possible to delay the injection of adversarial perturbations for large-batch optimization of transformers for NLP tasks.
The Algorithm, ScaLA: Combining the formulation with the above investigations, we construct our distributed large-batch transformer fine-tuning algorithm, named ScaLA (Algorithm) I whose convergence rate is characterized in Theorem |
Algorithm 1, ScaLA, Scalable Large- batch Adversarial Perturbation
Input: Epochs T. delay t. perturbation (inner) step size p. clipping radius w, regularization strength A, (outer) learning rate 1)
2: Output: h-layer transformer model & with converged robust parameters F == XT
3: for E [T] do
> Loop through epochs
4: for worker p E [P] do
> In parallel across homogencous workers Subsample E data instances on each worker >> Initialize regularization and label > Check delay condition Perform adversarial perturbation with PGA-1 > Calculate the adversarial regularization
5: for mini-batch Sp ~ Q do
6:
r(x;)) + 0, <- (xx, Ep). select yo
7:
it 2 1, then
8:
y1 (yo pV rzr y))
9:
r(x)) t KLaym (1,,(,11111
10:
g( f 1 5x) + Ar(xz)
Calculate the augmented loss Backward pass & Compute local gradients using accumulation
11:
V,g(z,1
(z) {p)
12:
V (.g(zz) + B Ep= Vag(ze, Ep)
b Gradient averaging using all-reduce
13:
# - n ((==(==))))) [[io(z)]]"
(ig(s)) Vi [A]
Update model parameters
Theorem (Complexity Algorithm D Informal Details Appendix D
Consider the problem
in Equation |
Let t. = O. Setting the outer learning rate as n O (11/v1 I/VTT and scaling batch
size as b (T)) for Algorithm we have E (1V So (E( Ka/VT) where is the estimator obtained from running T steps of Algorithm / and picking uniformly at random for t E (). Here, e the error due the approximate inner maximization oracle, a characterizes the smoothness of (x, -), 91/2a is the Moreau-envelus of g and Ko. max; ai/ min; os.
4 EVALUATION
We evaluate the effectiveness of ScaLA in training transformer networks set of NLP tasks.
Hardware: We study the efficiency of computation using 2 NVIDIA DGX-2 nodes. Each node consists of 16 NVIDIA V100 GPUs. The nodes connected with InfiniBand using 648-port port Mellanox MLNX- OS CS7500 switch. We vary the number of workers (i.e., from GPU to 32 GPUs) in the experiments evaluate the scalability. Software: We Py Distributed Data Parallel to
5

---

Under review as a conference paper at ICLR 2022
scale the fine-tuning from single GPU to multiple GPUs, which uses all-reduce (Proficz 2018 to compute the average of the gradients across all workers for a parameter update. We use NCCL V2.4 as the underlying all-reduce implementation. Model/Dataset: We use pre-trained BERT hase model released by HuggingFace (Wolf et al.) 2020). We use the GLUE benchmark (Wang et al. 2019a), which is collection of sentence sentence-pair natural language understanding tasks inc including question answering, sentiment analysis, and textual entailment. Given that large models normally require days or even weeks (e.g., RoBERTa) trained from scratch, it is not cost-efficient to evaluate from scratch using large-scale pretraining. Therefore, experiments presented in this section focus on fine-tuning pre-trained models, e.g., BERT (Devlin et al., 2019) and RoBERTa (Liu et al. 2019). and exclude fine-tuning tasks that have very small datasets (e.g.,CoLA, RTE). We report the details about the hyperparameters in Appendix
4.1 MAIN RESULTS- TRAINING TIME ACCELERATION AND ACCURACY IMPROVEMENT
We first compare the following schemes: (1) Single GPU + SB: This is the existing PyTorch implementation of Transformer fine-tuning from HuggingFace (HF), using small batch (SB) sizes (e.g.. 32). Our fine-tune tuning results on single GPU achieves higher accuracy than what was reported by (Devlin et al. 2019), (2) Multi-GPU + SB: This is multi-node multi-GPU PyTorch fine-tuning implementation using DistribuutedDattaaaallee (Li et al. 2020), and (3) Multi-GPU + L.B + ScaL.:: This is approach described in Algorithm| [ using large minibatches (LB), e.g., IK. Table 2 shows results on MNLI, QNLI, QQP, and SST2, which are larger datasets and less sensitive to random seeds. X 9 refers to P., nodes each with P, GPUs for a total of P P. homogencous workers (e.g.. 32 GPUs on 2 NVIDIA DGX-2 nodes). For a fair comparison, we reproduce BERT and RoBERTa baseline. Our reproduced baseline achieves the same slightly higher accuracy than the originally reported results in (Devlin et al.) 2019) and (Liu et al., 2019). We now discuss our results and observations.
SSTZ.
Table 1:
The training time and accuracy results on GLUE benchmark. Results show that ScaLA is able to achieve the same average accuracy as the baseline while providing up to 18x speedups than single GPU, and up to 9.8x speedups when using the same amount of hardware resources.
MNLI-m
QNLI
QQP
SST-2
BERT base
nxg Steps
Time
Aec. Steps
Time
Steps
Time
Acc/Fl
Steps Time Ace.
Avg.
Devlin et al. 2019
84.4
88.4
92.7
Baseline (B=32)
1x1
32 73632
19635
84.8
19644
5535
90.6
68226
16494
91/88.0
12630 2736
93.1
89.4
Bascline (B=32)
2x16
32
73632
8848
84.8
19644
2408
90.6
68226
11311
91/88.0
12630
1494
93.1
89.4
ScaL. (B=IK) 2x16
IK 2301
1323
85.1
615
432
90.0
2133
4229
90.9/87.7
396
151
93.5
89.4
MNLI-m
QNLI
QQP
SST-2
RoBERTa
nxg bsz
Avg.
Steps
Time Acc.
Steps
Time
AAce.
Steps
Time
Acc/F1
Steps Time Acc.
Liu ct al. 2020
90.2
94.7
92.27-
96.4
Baseline (B=32)
1x1 32 776333
1964414188
94.7
40926
12630)4440
96.4
Baseline (B=32)2x16
32
73632
18114
19644
4842
94.7
68226
1661492.0/89
12630,
96.4 92.5
Scal.A (B=1K) 2x TK 2301
3363
90.9
615
1168
95.1
2133
2404 92.3/89.8
396 401
96.7 92.9
Improving scalability with ScaLA: Compared with single-GPU training, baseline fine-tuning multi-GUU leads to only modest training speedup improvements, e.g., with 1.5 2.4x faster training speed for both BERT and RoBERTa, even with 32x more compute resources. The speedup is limited because of the small mini- (e.g.. 32) used for fine-tuning, which do not provide sufficient workload fully utilize the underlying hardware (e.g., many cores stay idle). Thus, communication overhead becomes the dominant part, and fine-tuning tasks struggle to obtain speedups with more than 4 GPUs (see detailed analysis results later). In contrast, ScaLA achieves up to 18x speedups with 32 GPUs. When using the same number of GPUs (c.g., 32), ScaLA is 2.7-9.8x 8x faster than Baseline (B- = 32). The speedup is significant because (1) large batches prevent hardware from under-utilizatie and enables processing more samples per second; (2) it takes fewer iterations to process an epoch, hence the reduced all-reduce operations to exchange gradients and the overall increased
ratio; (3) PGA-1 significantly reduces the cost to generate adversarial perturbations; and (4) With delayed perturbation injection, PGA-1 1 is only added to later iterations instead of throughout the entire training process, further reducing the training cost. Finally, ScaLA obtain the speedups while achieving the same accuracy (88.4 vs. 88.4) 4) average accuracy for BERT and higher accuracy (92.9 vs. 92.5) for RoBERTa as the baselines.

---

Under review as a conference paper at ICLR 2022
Closing generalization gap with ScaLA; We also compare alternative methods that perform large- batch optimizations: (1) Multi-GPUL + LB: This configuration uses large mini-batches (c.g., IK), and applies Souristic-based scheduling rule (e.g., square root) and grid search for learning rates; (2) Multi-GPU + L.B + LAMB: Applies the large-batch optimizer LAMB (You et al., 2019a) to fine-tuning tasks. First, compared with the baseline, the accuracy of Multi-GPU LB drops by close to point (88,4 4 vs. 89,4, and 92.1 vs. 92.9) in average and close to 2 points for some tasks (c.g., QQP on BERT), indicating that is challenging to obtain on-par accuracy with large-batch optimizations for fine-tuning tasks. Second, LAMB leads to only marginal improvements (88.6 vs. 88.4, and 92.1 vs. 92.1) than the baseline and is 0.8 points lower than the small-batch baseline. Since LAMB performs no explicit regularization to model complexity, it may still lead to overfitting on downstream tasks. With ScaLA, we are able to close the generalization gap from large-batch optimization (89.4 vs. 89.4, and 92.5 vs. 92.9) and achieve 0.8 points higher accuracy (89.4 vs. 88.6, 92.9 vs. 92.1) than LAMB on both BERT and RoBERTa. ScaLA improves generalizability because it introduces adversarial perturbations, which serves regularizer. By training the network to be robust to such perturbations, the model loss landscape is smoothed out, leading to improved generalization.
Table 2: The comparison results of the GLUE benchmark. ScaLA outperforms aarge-baccl optimiza- tion baselines by achieving higher accuracy after training the same number of samples.
MNLI-m
QNLI
QQP
SST-2
BERT
nxg
Batch size
Steps
Time
Acc.
Steps
Time
Acc. Steps
Time
Acc/F1
Steps
Time
Acc.
Avg.
Bascline (B=IK) K
2x16
IK
2301
1148
84.3
615
349
89.3 2133
2892 89.6/86.1
396
134
93
88,4
LAMB (B=1K)
2x16
IK
2301
1180
84.1
615
359
89.6
2133
2978
90.5/87.0
396
139
92.4
88.6
ScaLA (B=IK)
2x16
IK
2331
1323
85.1
613
432
90.0 2133
422942222990 90.9787.7
396
151
93.5
89.4
MNLI-m
QNLI
QQP
SST-2
RoBERTaagge
nXg
Batch
S12G
Steps
Time
Acc.
Steps
Time Acc. Steps Time Acc/Fl
Steps
Time Acc.
Avg.
Baseline (B=IKg K)
2x16 IK
2301
2514
90.1
615
936 94.3 2133 1874 91.7/89.1
396
317
95.9
92.1
LAMB (B= K)
2x
IK
2.301
2646
90.5
615
973 94.5 2133 1998 97.3788.5
396
324
96.2
92.1
ScaLA (B=IK)
2x16
TK
2301
3363
90.9
615
1168
95.1
2133 2404 92.3/89.8
396
401
96.7
92.9
4.2 EXPERIMENT - ANALYSIS RESULTS
Ablation analysis: In this part, we study the importance of components in ScaLA. We set t, to which denotes as w/o Delaying PGA-I. We replace the outer minimization to use ADAM (Kingma & Ba 2015), which is noted as w/o Groupwise LRR. We set to O which denotes as w/o PGA-1. The
results are reported in Table 3
Table 3: Ablation study of ScaLA using BERT base on GLUE tasks.
MNLI-m
QNLI
QQP
SST-2
Avg.
Time
Acc.
Time
Acc.
Time
Aco/F1
Time
Acc.
BERT
19635
84.8
5535
90.6
16494
91/88.0
2736
93.1
89.4
Scal.A AA
1323
85.1
432
90
4229
90.9/87.7
151
93.5
89.4
w/o Delaying PGA-1
2503
85.2
726
90.2
6407
91.3/88.3
272
93.1
89.5
w/o Groupwise LR
1290
85.0
422
89.9
4212
90.7/87.6
146
93.0
89.2
w/o PGA-1
1180
84.1
359
89.6
2978
90.5/87.0
139
92.4
88.6
The results show that the removal of either design element would result in a performance drop. For example, removing groupwise learning rates leads to 0.2 points accuracy drop (89.2 vs. 89.4), while completely removing PGA-1 leads to 0.8 points accuracy drop (88.6 vs. 89.4). This result demonstrates that these two components are complementary to each other. If we perform PGA-1 without delayed injection, the average accuracy increases by 0.1 point (89.5 vs. 89.4), but the execution time is increased significantly. This difference will be particularly felt during pre-training or training larger models. With our delayed PGA-1 approach, we save the training time by 1.5-1.9x while retaining high accuracy.
Scalability analysis varying GPUs. We carry out scalability test by varying the number of GPUs from to 32, with and without communication. We choose a batch size 32, and divide the samples among P GPUs. If the Per-GPU batch size (e.g.. 16) is larger than the maximum admissible per-GPU batch size (e.g.. 8), we use local gradient accumulation (Goyal et al. 2017) to avoid running out of memory. Figure shows the scalability results. For batch size 32, the training time decreases when the number of workers increases. However, it quickly plateaus and increases slightly after 4 GPUs. This is because small-batch leads more frequent communication among workers. As result, the
7

---

Under review as a conference paper at ICLR 2022
86.0
le5
20000
83.6
81.4
81.5
8/18
81.5
82.9
80.0
32.7
32.7
32.7 35.1
84.5
Treining time (s]
22000
B-1K
B-32 ino commi
1021
8
83
83
83
88.5
84.0
Senvall
6
IK (no co mml
2048
78.5
81.0
02.4
83.0
83.5
63.7 83.7
84.1 04.1 83.6 81.8
68.0
5.0
32830
62.5
5900
73.0
79.8
81.2 82.3
83.0
83.7
62.6
82.0
2.5
8192
666 0
/6.8 /9 800
#1
HCI
833. ls
833 5. 83 b.
833 4
#2
3
AGPU:
16
16-05
Jes
05
5a-08
To-05
10:04
20-04
94-0A
10-05
Se 2e-05
81.0
Be
Bwik (SCALA)
(a) Scalability
(b) Generalizability
(c) Sharpness
Figure 4: Scalability and generalizability results by fine-tuning BERT base using ADAM with different batch sizes on the MNLI task from the GLUE benchmark.
communication overhead dominates the total execution time and hinders the scalability of multi-GPU training. In contrast, by increasing the batch size, the training time keeps decreasing the number of GPUs increases because large batch not only reduces the number of expensive all-reduce operations for exchanging gradients (i.e., increased cumputation-v--oommnn training throughput per GPU due to increased computation granularity. Figure 3 shows the scalability comparison on SST-2 after optimizations. While the speedup still plateaus at 4 GPUs with a small batch size (c.g., B 32), the four large-batch configurations are able to scale well up to 32 GPUs and take a similar amount of time with 32 GPUs. ScaLA scales better than ScaLA without delaying PGA-1, and achieves a much faster training speed, especially in the 1-16 GPU range.
ratio) but also increases the
Batch sizes vs. learning rates for fine-tuning tasks. Figure 4b reports the learning rate scaling effects fine-tuning Transformer networks. We observe that the learning rate scales roughly with the the square root of the increase of the mini-batch size (That said, the Appendix (rroviidde analysis to show that the best learning rates do not always follow sqrt rule). This is consistent with prior findings in pre-training y You et al. 2019a) and simplifies the hyperparameter tuning effort for Scal.A.
Curvature analysis. Prior work (Keskar al. 2017) correlates the low generalization with sharp minima (which are characterized by a positive curvature of large magnitude in the parameter space). To verify this hypothesis, we quantitatively measure the steepness of loss landscape by loading the checkpoint of a converged model and computing the curvature, i.c., properties of the second derivative of the model, with respect to its parameters, for fixed batch of samples. In particular. following (aao et al. 2018b), for model (cr), we compute the largest eigenvalue of the model's Hessian, Xmax V [x)]. using the Hessian- vector product primitive and the power method. We use the largest eigenvalue as a measure of sharpness since the corresponding (top) eigenvector characterizes the direction of the largest change in gradient at a given point in the parameter space. From Figure 4c the largest eigenvalue of the model trained with large batch (e.g.. IK) is much larger (e.g.. 2.6x) with higher deviations (e.g.. 3.9x) compared to the small batch baseline. Thus, we empirically find that large-batch optimization makes the loss landscape of the model more prone to ill-conditioning and less robust to perturbation, which may explain the loss in generalization. We then measure the steepness of the loss landscape again after applying ScaLA. As shown in Fig. 4c the largest eigenvalue of the model becomes much smaller (6.9x) with lower deviations with ScaLA and is slightly better than the small batch baseline, which a strong indication that our approach enforces the smoothness of the model that leads to the accuracy improvement.
Comparison with random noise. We have performed additional experiments by creating pertur- bations through adding Gaussian noise to the embeddings. Table 5 that random noise indeed can improve the accuracy for MNLI-m (84.3 vs. 84.5), QNLI (89.3 vs. 89.4), and QQP (90.3/87.0 vs. 99.6/86.1 1) over the baseline, but it also leads to worse results on SST-2 (93. vs. 92.6). Compared with Scal.A, .A, random noise consistently falls behind S in its ability to reduce the generalization error on all tested tasks and is on average 0.7 points lower than ScaLA (88.7 vs. 89.4).
Perturbations via ground-truth vs. label probability. We also create one-hot labels and use those to generate perturbations instead of using label probability generated by the network. Table shows that using label probability (LP) consistently leads to higher accuracy than using the ground-trun -truh (GT), e.g., 89.4 vs. 89.0 on average. Label probability leads to better generalization, probably because it provides a better measurement of the adversarial direction, which is the direction in the input space in which the label probability of the model is most sensitive small perturbations.
8

---

Under review as a conference paper at ICLR 2022
Table Evaluation results on alternative methods generate per- Table Comparison results turbations using random noise, ground-truth, and label probability. with FreeLb.
Model
MNLI-m
QNLI
QQP
SST-2
Avg
MNLI-m
SST-2
Baseline
84.3
89.3
89.6/86.1
88.4
Acc.
Time
Acc.
Time
Gaussian noise
84.5
89.4
90.3/87.0
92.6
88.7
Baseline
84.8
8848
93.1
2736
ScaLA (GT)
84.1
89.6
90.7/87.6
93.2
89.0
Freel.b
85.1
3773
93.3
389
ScaLA (LP)
85.1
90
90.9/87.7
93.5
89.4
ScaLA
85.1
1323
93.5
151
Comparison with FreeLb, We also compare ScaLA with FreeL.. The original FreeLb does not support multi-node training. We extend it with PyTorch DDP to train in multi-node distributed training environment. We observe that although both FreeLb and ScaLA achieve similar accuracy, ScaLA is much faster than FreeLb. ScaLA is faster because FreeLb still performs multiple ascent steps to calculate adversaries cross the entire training. In contrast, ScaLA takes several optimizations to reduce the adversarial perturbation cost well leveraging group-wise adaptive learning rates to enable training with larger batch sizes, which improves the computational efficiency.
4000
3500
80.0
80.5
Training time is)
3000
-Q-
LAMI (0=1.)
84
84.8
85.2
98.1
63.0
68.4
2900
ScaLA we calay (1u020)
SealA (5-23))
Test accuraic
04-
1K
04.5
1K
02.3
02
98.5
80
1500
64.7
90.4
Baseline (B=2K)
K
84.0
84.1
84.5
1K
91.8
91.9
98.2
500
ScaLA (8-2K)
86.3
80.2
#GPUs
tepochs
8
12
Baselina
LAMB
ScaL.
80.0
Baselino
LAMB
ScaLA
01.0
Figure 5: Comparison of Figure 6: Comparison of scalability using different test accuracy by training the large-batch optimization bascline longer. methods on SST-2.
(a) MNLI-m
(b) SST-2
Figure 7: Comparison of accuracy under even larger batch sizes.
Train longer, generalize better? Despite improved training speed and accuracy, one may still wonder whether the generalization gap can be mitigated by training the downstream tasks longer. (Figure&sshow the comparison results between Scal.A and the baseline on batch size of 2K. Scal obtains an accuracy of 85.2 after 6 epochs of training, whereas the baseline has difficulty to reach 84 after training twice longer (c.g., 12 epochs). ScaLA achieves better accuracy because it explicitly penalizes model weights from getting stuck sharp minima, leading to better generalizability.
Generalizability under different batch sizes. In this part, we evaluate how different batch sizes affect the generalizability of fine-tum fine-tuning transformer tasks. Figure the results on MNLI-m and SST-2. We make two major observations: (1) The accuracy tends to drop as the batch size increases.
(2) While both the bascline and LAMB suffer from significant accuracy drop by drastically increasing the batch size (e.g., from 32 to 8K), Scal.A is able to mitigate the generalization gap and consistently achieves higher accuracy than the baseline (e.g., 84.4 vs. 83.5 for MNLI, and 92.6 vs. 91.3 for SST-2 at batch size 8K) and LAMB (e.g., 84.4 vs. 83.9 for MNLI, and 92.6 vs. 91.7 for SST-2 at batch size 8K). These results indicate the benefit of ScaLA is maintained by further increasing the batch size, which could bring even greater speedups when increasing the data parallelism degree.
CONCLUSIONS AND FUTURE DIRECTIONS
In this paper, we study how to add adversarial perturbations to improve the scalability and gener- alizability of training large transformer networks in a principled manner. We introduce ScaLA, scalable and generalizable large-batch optimization method using adversarial perturbations. The experiment results show that ScaL.A obtains up to 18x speedups on fine-tuning transformer networks and outperforms the state-of-the-an large-batch optimization methods in accuracy. Despite offering great speedups without losing accuracy, ScaLA limited that it currently has only been evaluated against fine-tuning tasks of pre-trained transformer networks due to the expensive pre-training cost. Furthermore, it has not been tested on emerging pre-trained transformer networks for computer vision tasks, such as ViT, so it is unclear about the effectiveness of ScaLA on those tasks. We plan to explore how to generalize our proposed method to improve the generalization of the pre- pre-training of large-scale language models as well as tasks in other domains in future work.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Turing-NLG:
A
17-billion 17-billion-parameerr
language
model by
Microsoft.
https
www mi microsoft . com en research blog
htt.::pp-p..llooopppaaaaeeeeaaaa Accessed: 19-May-2020.
Qi-Zhi Cai, Chang Liu, and Dawn Song. Curriculum adversarial training. In Jérôme Lang (ed.), Proceedings of the Twenty-Seventl International Joint Conference on Artificial Intelligence, IJCAI 2018, July 13-19, 2018, Stockholm, Sweden, pp. 3740-3747. ijcai.org, 2018.
Yong Cheng, Lu Jiang, and Wolfgang Macherey. Robust neural machine translation with doubly adversarial inputs. In Anna Korhonen, David R. Traum, and Lluís Marquez (eds.), Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2019, Volume Linguistics, 2019.
Long Papers, pp. 4324 4333. Association for Computational
Kevin Clark, Minh-Thang Luong, Quoe V. Le, and Christopher D. Manning. ELECTRA: pre-training text encoders as discriminators rather than generators. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Damek Davis and Dmitriy Drusvyatskiy. Stochastic subgradient method converges at the rate o(k{-1/4) weakly convex functions. arXiv preprint 10XX:08802.0288, 2018.
Damek Davis and Dmitriy Drusvyatskiy, Stochastic model-based minimization of weakly convex functions. SIAM Journal on Optimization, 29(1):207-239 2019.
Jacob Devlin, Ming-Wast Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. NAACL-HLT 2019). pp. 4171-41888 2019.
Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal Optimization, 23(4):2341-2366 2013.
Linyua Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tie-Yan Liu. Efficient training of BERT by progressively stacking. In Proceedings of the 36th International Conference on Machine Learning, JCML 2019, 9-15 June 2019, Long Beach, California, USA, pp. 2337-2346, 2019, URL mlr. wwwps/vv7..... https///9///ggg9../ooooo
http::/procceeding
Ian J. Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Priya Goyal, Piotr Dollár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch. Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: training imagenet in 1 hour. CoRR, abs/1 706.02677, 2017.
Sidharth Gupta, Parijat Dube, and Ashish Verma. Improving the affordability of robustness training for dnns. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR Workshops 2020, Seattle, WA, USA, June 14-19, 2020, pp.
IEEE, 2020.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the general- ization gap in large batch training of neural networks. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 1731-1741, 2017.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Advances Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, pp. 103-112, 2019.
10

---

Under review as a conference paper at ICLR 2022
Prateck Jain and Purushottam Kar. Non-convex optimization for machine learning. arXiv preprint arXiv:1712.07897. 2017.
Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao. SMART: robust and efficient tin-ttning for pre-trained natural language models through principled regu- larized optimization. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel R. Tetreault (eds.), Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020. pp. 2177-2190. Association for Computational Linguistics, 2020.
Chi Jin, Praneeth Netrapalli, and Michael Jordan. What is local optimality in noncomvex-nonconeee minimax optimization? International Conference on Machine Learning, pp. 4880 4889. PMLR, 2020.
Norman P. Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, AI Borchers, Rick Boyle, Pierre-luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, C. Richard Ho. Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Daniel Killebrew, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le, Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross, Amir Salck, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang. Eric Wilcox, and Doe Hyun Yoon. In- in-Datacenter Performance Analysis of Tensor Processing Unit. In Proceedings of the 44th Annual International Symposium on Computer Architecture, ISCA '17, pp. 121 2017. ISBN doi: 10. N0.1145/3079856.3000 1145/30 3080246
URL http: doi acm. /00
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, abs/2001.08361 2020. UtLpttps::/ /ttpais////aavv.e/eeeeeeee org/ /atg/bb//// 2001 08361
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings, 2017. URL https:h http:::s///pperevieeee net forum?idi H1 HloyR1Ygg
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In 3rd Interna- tional Conference on Learning Representations, ICLR 2015, 2015.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, and Soumith Chintala. Pytorch distributed: Experiences on accelerating data parallel training. Proc. VLDB Endow., 13(12):3005-300 3018, 2020.
Tianyi Lin, Chi Jin, and Michael Jordan. On gradient descent ascent for nonconvex-conee minimax problems. In International Conference on Machine Learning, pp. 6083-6093 PMLR, 2020.
Xiaodong Liu, Hao Cheng. Pengcheng He, Weizhu Chen, Yu Wang, Hoifung Poon, and Jianfeng Gao. Adversarial training for large neural language models. CoRR, abs/2004.08994, 2020.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized BERT pretraining approach. CoRR, abs/1907.11699 1692, 2019.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver. BC, Canada, April 30- May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018.
11

---

Under review as a conference paper at ICLR 2022
Takeru Miyato, Andrew M. Dai, and lan J. Goodfellow. Adversarial training methods for semi- supervised text classification. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
Zachary Nado, Justin Gilmer, Christopher Shallue, Rohan Anil, and George E. Dahl. A large batch optimizer reality check: Traditional, generic optimizers suffice across batch sizes. CoRR, hass/102200336,, 2021.
Jerzy Proficz.
Improving all-reduce collective operaionss for imbalanced process arrival patterns. The Journal of Supercomputing, 74(7):3071-3092, 2018.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. CoRR, 0bs/1910.10688 2019.
10683
Hassan Rafique, Mingrui Liu, Qihang Lin, and Tianbao Yang. Weakly-convex-come min-mat max optimization: provable algorithms and applications in machine learning. Optimization Methods and Software, pp. 1-35, 2021.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimization towards training A trillion parameter models. CoRR, abs/1910.02054, 2019.
Ralph Tyrell Rockafellar. Convex analysis. Princeton university press, 2015.
Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young. Ryan Sepassi, and Blake A. Hechtman. Mesh-tensorflow: Deep learning for supercomputers. In Advances in Neural Informa- tion Processing Systems 31: Annual Conference Neural Information Processing Systems 2018, pp. 10435-10444, 2018
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro. Megatron-Im: Training multi-biliion parameter language models using model parallelism. CoRR, abs//900900053, 2019.
Leslie N. Smith. A disciplined approach to neural network hyper-parameters: Part learning rate, batch size, momentum, and weight decay. CoRR,
2018. URL http
7 Tarxii. mrxxi..org/abs/100333
Samuel L. Smith, Pieter-Jan Kindermans, Chris Ying, and Quoc Le. Don't decay the learning rate, increase the batch size. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30- May 3, 2018, Conference Track Proceedings. OpenReview.net. 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, pp. 6998-6008, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In 7th International Conference on Learning Representations, 2019a.
Dilin Wang, Cheng Gong, and Qiang Liu. Improving neural language modeling via adversarial training. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th Inter- national Conference Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Pruceedings of Machine Learning Research, pp. 6555-6565. PMLR, 2019b.
Yu Wang, Gu- Yeon Wei, and David Brooks. Benchmarking tpu, gpu, and CPU platforms for deep learning. CoRR, adts/1907.10ll 10701, 2019c.
12

---

Under review as a conference paper at ICLR 2022
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma. Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin I Lhoest, and Alexander M. Rush. Transformers: State-of-the-art natural language processing. In Qun Liu and David Schlangen (eds.), Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, EMNLP 2020 Demos, Online, November 16-20, 2020, pp. 38-45, Association for Computational Linguistics, 2020.
Zhewei Yao, Amir Gholami, Qi Lei, Kurt Keutzer, and Michael W. Mahoney. Hessian-based analysis of large batch training and robustness to adversaries. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 4954 4964, 2018a.
Zhewei Yao, Amir Gholami, Qi Lei, Kurt Keutzer, and Michael W. Mahoney. Hessian-based analysis of large batch training and robustness to adversaries, In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 4954-4964, 2018b.
Nanyang Ye, Qianxiao Li, Xiao Yun Zhou, and Zhanxing Zhu, Amata: An annealing mechanism for adversarial training acceleration. CoRR, abs/2012.08 12, 2020.
Yang You, Jing Li, Jonathan Hseu, Xiaodan Song, James Demmel, and Cho-Jui Hsieh. Reducing BERT pre-training time from 3 days to 76 minutes. CoRR, abs/1904.00062. 2019a.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 1904.00602 2019b.
Dongxu Zhang and Zhichao Yang. Word embedding perturbation for sentence classification. CoRR, abs/1804.08 166, 2018
Minjia Zhang and Yuxiong He. Accelerating training of transformer- language models with progressive layer dropping- In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
Chen Zhu, Yu Cheng, Zhe Gan, Siqi Sun, Tom Goldstein, and Jingjing Liu, Freelb: Enhanced adversarial training for natural language understanding. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, OpenReview.net, 2020.
13

---

Under review as a conference paper at ICLR 2022
A
ADDITIONAL RESULTS
In the part, we present results that are not included in the main text due to the space limit.
A.1
THE USEFULNESS OF PERTURBATION THE INITIAL EPOCHS
85.50
93.4
85.25
93.2
Test Test accuracy
85.00
Test accuracy
84.75
92.8
84.50
84.25
84.00
92.4
83.75
92.2
83.50
0
3
5
0
3
5
#Epochs
#Epochs
(a) MNLI-m
(b) SST-2
Figure 8: Accuracy results from delaying the injection of PGA-1 different epochs.
In Section 3.T. we mention that no adversaries are needed at the initial epochs of fine-tuning. To verify,, conduct experiments measure the final accuracy corresponding to starting from regular training and switching to PGA-1 after t, epochs, where t, E [T'. Figure S shows that enabling PGA-1 from the very beginning does not offer much improvement accuracy. However, we delay the adversarial perturbation, the model accuracy starts to increase. This is because, at initialization, the model's parameters are relatively far from their final values and less likely to get stuck local minima. By delaying the injection of adversarial perturbations, we observe improved test accuracy fine-tuning tasks. However, it also seems that such perturbation should be injected too late, which may inadvertently affect the accuracy. It is possible that a more advanced method to adaptively choose the value of f, desired. However, given that (1) the primary focus of this work is to demonstrate that it is possible and effective to leverage adversarial training for ange-batch optimization of transformer networks and (2) the search space of quite small for fine-tuning tasks, we leave this as an interesting research question for future exploration.
B HYPERPARAMETERS
For all configurations, we fine-tune against the GLUE datasets and the maximum number of epochs to 6. We use a linear learning rate decay schedule with a warm-up ratio of 0.1. For ScaLA, we set À = 1, perturbation clipping radius w 10-5, step size p = 10-4 and (,=(335). These values worked well enough that we did not feel the need to explore more. For fairness, we perform grid search of learning rates in the range of {1c-5, 3e-5, 5e-5, 7 9e-5, le-4, 3e-4] for small batch sizes and (5.6e-5, 8e-5, le-4, 1.7e-4, 2.4c-4, 2.8e-4, 4e-4, 5.6e-4, le-3} } for large batch sizes. We keep the remaining hyperparameteres unchanged.
C LEARNING RATE SCALING FOR FINE-TUNING TASKS
In this part, we investigate how large-batch optimization affects the generalizability in fine-tuning transformer networks. As there are various heuristics for setting the learning rates (Smith et al. 2018) Goyal et al.) 2017 Smith 2018 You et al., 2019a), and because few work studies the learning rate scaling effects on fine-tuning Transformer networks, we perform a grid search on learning rates (le-4, 3,-4,5e-4, 7e-4, 9e-4, le-3, 3--3) and batch sizes (1K, 2K, 4K, 8K} while keeping the other hyperparameters the same to investigate how ScaLA affects the hyperparameter tuning effort for large batch optimizations.
Table below shows the results of using the square root scaling rule to decide the learning rates for large batch sizes accuracy results with tuned learning rate results, without and with ScaL.A A.. The
14

---

Under review as a conference paper at ICLR 2022
first row represents the best accuracy found through fine-tuning with small batch size 32. The next two rows correspond to fine-tuning with batch size 1024 using tuned learning rates vs. using the scaling rule. The last two rows represent tine-tuning using ScaLA with batch size 1024, also using tuned learning rates vs. the scaling rule. Even with square-root scaling, the large-batch baseline still cannot reach the small-batch accuracy (88.7 vs. 89.4). Moreover, although tuning the learning rates lead to better results on some datasets such as MNLI-m (84.9 vs, 85.1) and SST-2 (92.9 vs. 93.5), the square-root scaling rule leads better results on other tasks such as QNLI (90.8 vs. 90) and QQP (91 1//888.4 vs. 90.9/87.7). So the best learning rates on fine-tuning tasks are not exactly sqrt. However, given that ScaLA with square-root learning rate scaling achieves on average better results than the grid search of learning rates (89.4 vs. 89.7), we suggest to use sqrt scaling for learning rates to simplify the hyperparameter tuning effort for ScaL.A
Table Evaluation results on hyperparameter tuning vs. using square-root learning rate scaling.
MNLI-m
QNLI
QQP
SST-2
Avg
Bsz=32 (tuned, baseline)
84.8
90.6
91/88
93.1
89.4
Bsz=1022 (tuned, baseline)
84.3
89.3
89.6/86.1
93
88.5
Bsz= 1024 (scaling rule, baseline)
83.9
89.2
90.6/87.4
92.5
88.7
Bsz= 1024 (tuned, Scal.A)
85.1
90
90.9/87.7
93.5
89.4
Bsz= (scaling rule, ScaLA)
84.9
90.8
91,4/88.4
92.9
89.7
D CONVERGENCE ANALYSIS
In this section, we provide the formal statements and detailed proofs for the convergence rate. The convergence analysis builds on techniques and results in (Davis & Drusvyatskiy 2018 You et al. 2019b). We consider the general problem of a two-player sequential game represented as nonconvex- nonconcave minimax optimization that is stochastic with respect to the outer (first) player playing r E X while sampling & from Q and deterministic with respect to the inner (second) player playing y E Y, i.e.
min max maxfff=lff y. 6)] == min Eg~ol9(x,
(2)
Since finding the Stackelberg equilibrium, i.e. the global solution to the saddle point problem. is NP-hard, we consider the optimality notion of a local minimax point (Jin et al. 2020). Since maximizing over y may result non-smooth function even when f smooth, the norm of the gradient is not particularly suitable metric track the convergence progress of an iterative minimax optimization procedure. Hence, we use the gradient of the Moreau envelope (Davis & Drusvyatskiy 2019) the appropriate potential function. Let E 184, The μ-Moreau envelope for function g: x + R is defined as (x) == min; g(z) + Em= 1 Z 11' 2²1 Another reason for the choice of this potential function is due to the special property (Rockafellar 2015) of the Moreau envelope that if its gradient [x)) almost vanishes at x, such x close to stationary point of the original function g.
We assume that X x= - X is partitioned into h disjoint groups i.e., in terms of training a neural network, we can think of the network having the parameters partitioned into (hidden) layers. The measure Q characterizes the training data. Let fx, y) denote the noisy estimate of the true gradient Vzf(z, v). We assume that the noisy gradients are unbiased, i... E V.f e)] f.. y). For each group E [h]. we make the standard (groupwise) boundedness assumption (Ghadimi & Lan 2013) on the variance of the stochastic gradients, i.c.,
Assumptions:
E y) Vz. e)(l < of. Vi E [h]. We assume that f(x, y) has Lipschitz continuous gradi- ents. Specifically, let f(r y) be a-smooth in x where a := (a1..... a1) denotes the h-dimensional vector of (groupwise) Lipschitz parameters, i.e., [ (af vill I as|l= aflll max, di.
Vi E [h] and &a, vy E X, E W. Let Kn min,
Super-scripte are used to index into a vector (i denotes the group index and j denotes an element in group i). For any E R, the function v R + [C, clips its values, i.e., v(c) == max(C. min (c)) where < U. Let II-II, Il-IIa and I| denote the fg, l1, and fo norms. We assume that the true gradients are bounded, i..., [V, 1l 1 G.
First, we begin with relevant supporting lemmas. The following lemma characterizes the convexity of an additive modification of
15

---

Under review as a conference paper at ICLR 2022
Lemma 1 ((Lin et al.
2020 Jin et al., 2020 Rafique et al. 2021 Let g(z) maxy f(x, v)
with f being a-smooth in where E R. is the vector of groupwise Lipschitz parameters. Then, g(x) + Ei= 2 ||='||
convex in x.
The following property of the Moreau envelope relates it to the original function.
Lemma 2 ((Rockaelllar 2015)). Let be defined as in Lemma Let 3 = arg ming g(ā)
Ei=1 2p - - x²1 Then, (x))())] & implies | S || and mins ||A]| S with h E 0g where Ög denotes the subdifferential of
We now present the formal version of Theorem Theorem Note that Lemma facilitates giving the convergence guarantees in terms of the gradient of the Moreau envelope. Recall that & E [T] denotes the epochs corresponding to the outer maximization. Without loss of generality, we set the delay parameter for injection of the adversarial perturbation in Algorithm as t, = O. Here, we assume that the PGA provides an e-approximate maximizer.
Theorem 2 (Groupwise outer minimization with e-approximate inner maximization oracle). Let us define relevant constants as := (91/2a (Za) Emmin g(x))) being the optimality gap due
to initializution, Rig
max, og being the condition number, V (x,y)) S G being gradient min mi being the variance term, C,U being clipping constants such that _ and scaling batch size as
bound, Z := max
U. For the outer optimization, setting the learning rate as n uuT
b 16T 2² z²
we have
4
2ka DG
V 21/2a a
[1]
S 4e|la +
(3)
VT
where E is the estimator obtained from running steps of Algorithm 7 and picking 2t uniformly at random for E [T].
Proof. In this proof, for brevity, we define the vector V. := [I, v). i.e., the gradient of the objective with respect to v; evaluated at the outer step l. Since evaluating gradients using mini-hatches produces noisy gradients, we use V to denote the noisy version of a true gradient V. i.e. V V + A for noise vector A. For any outer step f, have f(z,,)) g(x,)) where ÿ is an e-approximate maximizer. For any I E X, using the smoothness property (Lipschitz gradient) of f. we have
g(x) ( 2 (II
h
2 f(r vr) +
V
x1)
0
2
1
2
1=1
1=1
2 2
Σ
Vi
21)
04 |
(4)
2
Let (x, == g(=) + Σか 2p
xz
2 1 Recall that the μ-Morreu envelope for g is defined as
gu (I) = min, (x 2) and its gradient the groupwise proximal operator given by V (=)]
(xt arg ming1 (2r, 2))
Mx arg ming* (x,
Now, let 3 arg min, 2 (2e, x) arg min, (9(x) Σhe asllzt - 2² 2T(12). Then, plugging in the update rule for at step in terms of quantities step t, using the shorthand of =(1[=E|l)
16

---

Under review as a conference paper at ICLR 2022
and conditioning the filtration up to time t, we have
(xc+1) S g(te) +
Σ a 0/ll//111 - 211?
5g(0,) +
0/
==
V
5g(ā,) + Σ
|||/
Σ
20
1/
+
am (r/)
h
S (11/2((z)) +
Σ
20
1
x
+
Σ
ao (v/)²
|1
1=1
1=1
s (1/2 + 21)
(3) x2,"
am (v:)
|
i=1
h
S (1z) + 217/
a.V
Σ
x 1
di
|
+ 27)
x
(
2و. + Σ
a.
-=1
j=1
1=1
a.
011/22 (1)) + 211
Σ
-
|
d:
V
Ai
+ 27p
Σ
| +
|
x
(
+
am (v:)²
|=1
17

---

Under review as a conference paper at ICLR 2022
S
(z1) + 2n1
a
< ( - xi)
d;
14 + A x
-
+ 21)t
[ + IIII
(V.
i=1
h
+
a:T (11)²
Bg
(11) + 2mpel max
((Gr) (ay)6) +
Σ
h
|
=fll?
[
2
(3) -dd-xt" 2.:2)
-
(V
|
+ 21x Σ
u
max
+
anh (vi)²
(Vit)
V + الl
i=1
i=1
S 91/22((z) + 21d max
Oi |ill
((कt)) g2( +
0z | - ==il?²
2
(3 - x
[ +
V 2
V v )
2nh
Σ
a max
(Vi')
V + A:l
i=1
+
an (vt)²
(5)
S $1122a (2x) + 21jil max
a;
(((t)) - gg(0)) + +
ai |1 - =x/1?) 2
|i:|
1=1
max (3/ - xد² (Vi2)
K44 V + اللهه )
- 21h
a
| 11 +
+ Σ
aani (1i)
| + Ailه
1=1
h
Eg
S (x+) + 2m max a;
((G)) (a0)) +
04 |8 =1'1?
1 |
2
h
(3 - =
21
Σ
ane max
(!A!) - [V: + +لل1 - [a:il + Σ (vi)
(Vi)))
1=1
1=1
a; 2 (x)) + 2ndll max |iill
(((e) (a0g) + c
Σ
04
|z x'd1?
2
1=1
(3)
4nj
a,vi max
|لنم| + Σ (1²)
(Vi2)
T-1 1
91/2a (IT)
S
91 (5o) + 22
it max
g(5x) - g(xz)) +
04 |1 - -x;11
|
2
1=1
T-1
تي - x
max
|A:||
+
on (
(v-²)
1=0
t=0 1=1
where we have used Hölder's inequality along with bound equation 4 in E,. Cauchy-Schward inequality in Eg, triangle inequality in Eg, telescoping sum in E4. Rearranging and using 7)t = 1j in
18

---

Under review as a conference paper at ICLR 2022
Eg along with Hölder's inequality,
h
20j2 (91/2a 91/2a (zo))
_
max
0g
g(fc)) - g((z)) +
Σ
04
13 =11??
Vill
2
1=0
(ot (V-))
)
T-1 h
a max
+
1
Σ Σ
xx(2))²
u 0=
211
1=0 1=1
2nH (01/2a(UT) 91/2a (20))
R6
max 1,t
g(0e) - ((x))) + Σ
2
|
-
-
xil??
1=
1=1
2
aivi max j
(o) 2.
|بإا|
+
T-1
Σ
as(vi)?
u
(Vi))
211
1=
t=0
Dividing by T and rearranging,
h
g((r)) g
Σ
0i
|
xrll?) )
Se- 2n (91/2a(1T) - (20))
T
2
1=1
2
T
(3 - It
2.5
aivi max
المه
UCT
(V:2)
(=0
7
(v))2
21ÇT
i=1
1=
where we define Ç == max,
|!V!|| Defining D : (91/2(
E(min (())) and taking expec-
tation with respect to & on both sides, we have
T-1
T Σ E g(xz) g(5e)
Σ 22110 ||
-
=i(l?)
D set 2nRCC
2
2L
T-1 h
(5) x1 25.)
nMaalll
E E| +
a max
uCT
(V!')
2ç
1=0 i=1
Eg D
se
2TJUCT
2C T-1
h
(5) xx 0, + nllallll
a max
UCT 1=0 i=1
(V/') vb
2
F7
D
S
c
2TUUÇT
2C||al11 max
(7) - 11.)) -2)))
0;
ucVb
(V:')
2
Eg e
D
2CIla|113 2
(alla
(6)
2TUCT
usVb
2ç
where we have used the assumption on the variance of stochastic gradients in Eg, Hölder's inequality in Ey and we define Z maxi.tt
(of a og in Eg: b denotes batch size. Now, we lower bound
(VP
19

---

Under review as a conference paper at ICLR 2022
the left hand side using the convexity of the additive modification of g-
g(xx) g(()) Σ [ 13'-x'11
h
h
2 g((z) +0 g( - Σ
a
||
i=1 -2/12-
0g
x;
2
iml
h
2tll? ))
g(xx))
||
x//1?
min
g
+
-
+ 2
Σ im
1
i=1
a
1|00
3-
4a
1
2
4cts
1 maxx Og
/
(zxe)l
where we have used the expression for the gradient of the Moreau envelope in Eg. Combining the inequalities from Equation equation 7 and Equation equation we have
1
T-1
TM
2CZ ucvb
(21/2a(z=))1 S E
+
||a||a
T
4 max;
24UÇT
26
t=0
and batch size as b 16T
Setting the learning rate as
uVT
U12
T-1
2D max; as
Σ T 1=0
E (1V9222 < 4e maxa; +
CVT
Now, to simplify S. using the inequality that max:(a;: by) 2 mink, Ok,, mink, bR, for two finite sequences {a, b) with positive values, along with the bounded gradients assumption, we have
T
2DG max; a; 4c|lalloc + 28a DG
Σ 1=0
E [1V91/2 (Z)11) < max as +
T
VT min, as
VT
where Ka.
max
min;
In analyzing inexact version, as in Theorem 2 we assumed the availability of an adversarial oracle. Next, we open up the adversarial oracle characterize the oracle-free complexity. In order do this, we will assume additional properties about the function f as well as our deterministic perturbation space, Y, Y, Vt E [T'. Note that, for any given Vr E Y,, VT E T. We recall the following guarantee for generalized non-convex projected gradient ascent.
Lemma 3 (JJiin & Kar (2017)). For every t, Let f(r,, satisfy restricted strong convexity with parameter C and restricted strong smoothness with parameter S over non-convex constraint set with < S/C
C11== 1 y < (xx,y) f(I,, 2) (V= 2), y 2) s $11== - yl2 for y, We. For any given let the PGA-T algorithm Yr +- 11- + pV (x v)] be executed with step size
ie, C
1/S. Then after at most T = 20-5 log !) steps, f(z,, ((1),V)) 2 max, f(11,
Using Theorem 2 and Lemma (together with the additional restricted strong convexity/smoottness assumptions), we have the following theorem on the full oracle-free rates for Algorithm |
Theorem (Groupwise outer minimization with inner maximization using projected gradient ascent).
Setting the inner iteration count as T O x log ||aاا and the outer iteration count as
T 16k. D'g² for combined total of O log )) adaptive adversarial iterations, Algorithm ! achieves E (T) S.
20


Reviews:
[{"review_id": "r5e3T13kz8", "reviewer": "Reviewer_jy9c", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "Pros:\n1. Firstly, the paper is generally well written and easy to follow.\n2. The proposed method contains several techniques to reduce the computational overhead from adversarial training, leading to significant training time savings.\n2. The experimental results seem satisfactory. The convergence is faster compared to the baseline. \n\nCons:\n1. Firstly, the core contribution, using adversarial training to improve generalization and convergence under large-batch training, has been adopted in NLP (FreeLb), which limits the novelty of the proposed paper. (A minor point, the comparison in Table 5 does not seem fair; ScaLA uses LARS optimizer for large-batch training, while FreeLb does not. This optimizer is not part of the technical contribution, and should be kept the same). \n2. FreeLb can also be used for training from scratch, which is another important setting. It is unclear whether ScaLA can also be applied to this setting.\n3. The authors argued that ScaLA can help to escape sharp minimum. Other techniques like the SAM optimizer [a] also have the same effect and proved to be working for ViTs [b]. It would be great to compare against these methods.\n4. It is a bit unclear to me what are the training settings for each method. It seems that due to the different convergence speeds, different methods use different numbers of training epochs/iterations. How do you choose when to halt training to report the results (e.g., in Table 1)?\n\n\n[a] Foret et al., Sharpness-aware minimization for efficiently improving generalization\n[b] Chen et al., When Vision Transformers Outperform ResNets without Pre-training or Strong Data Augmentations\n", "summary_of_the_paper": "In this paper, the authors propose ScaLA to speed up the fine-tuning of large pre-trained transformer language models. Specifically, ScaLA employs adversarial training to solve the worse convergence in large-batch training. Several techniques are proposed to reduce the computational overhead. Experiments on the GLUE benchmark show faster convergence compared to existing methods.", "summary_of_the_review": "Generally, I feel this paper presents an interesting solution to solve the convergence and generalization problem of large-batch fine-tuning. However, the proposed method is somewhat similar to existing work, which limits the novelty. I would like to hear back from the authors for the final opinion.", "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": "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": {}}, {"review_id": "L7iRd1XBvlH", "reviewer": "Reviewer_io2b", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "The main idea of this paper is to introduce adversarial perturbations to improve the convergence rate of large batch fine-tuning of sequence models. The proposed idea is very simple and effective. In this paper, the authors demonstrate a nearly three times speedup in BERT and RoBERTa fine-tuning on the GLUE dataset. Although the observations made in the paper such as only one PGD step is sufficient for generating perturbation, the novelty of the idea is limited. As mentioned in the paper, both aspects of the idea have been proposed and validated for training convolutional neural networks. Therefore, a throughout evaluation of sequence models at different scale on several NLP datasets is critical. My major concerns are listed blow.\n\nmajor:\n1) Empirical results on other datasets: Since the goal of the paper is to speed up finetuning, it is important to validate whether the proposed method can be generalized to different datasets such as SQuAD v2.\n\n2) Empirical results on large models: The paper is motivated by the increasing cost of finetuning large sequence models. However, the paper only includes results on BERTBase model. It is important to include more results on large sequence models such as BERTLarge, Google T5 Large [1], and GPT-2/3 models. \n\n3) Concurrent adversarial learning: A recent work in [2] on improving the convergence rate of very large small batch training also suggests adding adversarial perturbations. More interestingly, they retain a delayed model to generate adversarial perturbations, which makes the process of training and generating adversarial perturbations completely parallel.\n\nminor:\n1) caption of Figure 2: with and with\n\n[1] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer\n[2] Concurrent Adversarial Learning for Large-Batch Training", "summary_of_the_paper": "The authors propose introducing adversarial perturbation in the late phase of finetuning of sequence models to enable large mini-batch. The proposed method can speed up the finetuning by almost three times.", "summary_of_the_review": "The authors proposed to add adversarial perturbation to enable large mini-batch finetuning of  sequence models. However, the empirical study is limited to BERTBase model on GLUE datasets.", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "G4AYkDCRFkM", "reviewer": "Reviewer_FdJg", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "**Strong Points:**\n- This paper is well organized and it's good to see that we can also use large batch training for the fine-tuning process.\n- The authors did plenty of ablation studies to break down the impact of each component in ScaLA, and it's interesting that using delayed perturbation just has minor side effects.\n\n**Weak Points and Questions:**\n- From Table 1, ScaLA even matches or surpass the performance of small batch training, I'm curious about the performance of similar adversarial training algorithm when applied on small batch fine-tuning.\n- In Algorithm 1, should line 9 also be different for classification and regression task?\n- It would be better if the author can show the results of baselines when using the same fine-tuning time. Table 1 shows that ScaLA performs similar to small batch baselines but is faster, Table 2 shows that ScaLA achieves higher accuracy but is slower than large batch baselines. I would like to know the result of baseline under the same training cost (e.g., using a batch size of 512).", "summary_of_the_paper": "This paper proposes an adversarial training algorithm to speed up the fine-tuning of large Transformer models while retaining strong performance. The authors also incorporate several tricks such as delayed perturbation to lower the computational overhead.", "summary_of_the_review": "In summary, this paper is well-written and i can see the efforts of the authors to make it rigorous and well-organized, the ablation studies and analysis are also very comprehensive. But I'm not sure about the novelty of this paper, ScaLA to me is more like a combination of several methods, incorporating adversarial perturbation and layer-wise optimizer. Moreover, the benefits of adversarial training on normal batch size in both NLP and CV domains are well-studied, so I think the novelty is restricted. I vote for a weak accept for the written quality and the rigorous analysis of this paper.", "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": {}}]

Author rebuttal:

Meta-review: The submission considers a method involving adversarial training to speed up the fine-tuning of large pre-trained transformer language models. Reviewers consider it to be a borderline paper.  Many suggestions are made by the reviewers which will help improve the presentation and substance and make it more useful for the community.
Decision: Reject