ScaLA: Speeding-Up Fine-tuning of Pre-trained Transformer Networks via Efficient and Scalable Adversarial Perturbation

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 address this challenge , we propose ScaLA , a scalable and robust method for largebatch 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 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 × 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× faster than the baselines . 1 INTRODUCTION . We have seen exponential growth in DL model size since the debut of the transformer network ( Vaswani et al. , 2017 ) . For example , while BERTbase has around 100M parameters , the model size has increased to multi-billion parameters such as Megatron-LM ( 8B ) ( Shoeybi et al. , 2019 ) , T5 ( 11B ) ( Raffel et al. , 2019 ) , Turing-NLG ( 17B ) ( tur ) , and with GPT-3 hitting a new staggering 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 a 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 is to increase the batch size in the optimization algorithm in order to leverage multi-GPU training ( Li et al. , 2020 ; Liu et al. , 2019 ; Huang et al. , 2019 ; Shazeer et al. , 2018 ; Shoeybi et al. , 2019 ; Rajbhandari et al. , 2019 ) . By increasing the batch size , a mini-batch of sizeB can be divided across more workers ( GPUs ) , where the gradients are computed locally on each worker using back-propagation 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 et al. , 2020 ) . 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 new ideas more rapidly . 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-batch optimization for transformer networks . LAMB is a variant of Adam ( Kingma & Ba , 2015 ) that applies layer-wise 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 et al. , 2019a ) 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 to 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 as 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 generalizability 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 a 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 generalizability in large batch task-specific fine-tuning . ( 4 ) We evaluate our approach against BERT ( Devlin et al. , 2019 ) and RoBERTa ( Liu et al. , 2019 ) and show that ScaLA obtains significant improvements over the state-of-the-art 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 18× speedups on 2 NVIDIA DGX-2 nodes over the baseline fine-tuning time and is up to 9.8× 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 parallelizable 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 difficult ( Keskar et al. , 2017 ; Hoffer et al. , 2017 ; Nado et al. , 2021 ) . Furthermore , prior works mostly focus on reducing the pre-training time ( You et al. , 2019a ; Zhang & He , 2020 ; Gong et al. , 2019 ; Clark et al. , 2020 ) instead of the adaptation time at the fine-tuning stage . While many researchers and practitioners focus on how to reduce the pre-training time , few attention has been paid to accelerate the fine-tuning stage , which gradually becomes a 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 is to speed-up task-specific fine-tuning without hurting model accuracy . Different from pre-training , the fine-tuning stage often employs a much smaller batch size ( e.g. , log2B = { 4 , 5 } ) than pre-training ( e.g. , log2B ≥ 10 ) ( Devlin et al. , 2019 ; Liu et al. , 2019 ) . The small batch size results in inefficient data parallelism ( i.e. , sub-optimal computation-communication ratio ) , making it difficult for fine-tuning to benefit from multi-GPU training . Moreover , a common understanding is that small-batch sizes provide implicit regularization effects ( e.g. , from gradient noise ) that help improve generalization of downstream tasks . In contrast , our goal is to speed-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 a 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 et al. , 2020 ; Liu et al. , 2020 ; Yao et al. , 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 Zhu 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 GAME-THEORETIC METHOD VIA ADVERSARIAL PERTURBATION ORACLE . Formulation : Let X denote the parameter space and Y denote the data ( mini-batch/sample ) 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 : min x∈X Eξ∼Q [ g ( x , ξ ) ] = min x∈X Eξ∼Q [ f ( x , ξ ) + λr ( x ) ] = min x∈X max y∈Y Eξ∼Q [ f ( x , ξ ) + λr ( x , y ) ] : = min x∈X max y∈Y Eξ∼Q [ f ( x , y , ξ ) ] ( 1 ) where g : X× Y→ R denotes the robust training objective , f : X× Y→ R denotes the standard training objective , f : X × Y × Y → R denotes the augmented objective , r : X → R denotes a deterministic regularization term on the parameters controlled by a strength factor λ ∈ ( 0 , ∞ ) , r : X→ R denotes the augmented regularization and ξ denotes samples drawn fromQ ( for simplicity , we slightly abuse the notation in using ξ to denote the random variable , e.g . Eξ [ g ( x , ξ ) ] , or its empirical realizations , e.g . 1K ∑K k=1 g ( x , ξk ) 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 a 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 a 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 in a large data batch size to undermine the performance of the transformer parameter configuration chosen by the first player . 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 r captures the prediction deviation from the perturbation . In a given round of the game , with respect to the first player ’ s proposal , let Φ denote the transformer network under consideration ( specifically , Φ is BERT in this paper ) and ξ be a large batch of data sampled from Q . We construct a label for the second player as γ : = Φ ( x , ξ ) . Next , for classification tasks , we choose r to be the symmetric KL divergence ( Jiang et al. , 2020 ) , i.e. , r ( x , y ) : = KLsym ( γ , Φ ( x , y ) ) . We use symmetric KL divergence to measure the distributional divergence to generate adversarial perturbation . For regression tasks , we choose r to be the squared loss , i.e. , r ( x , y ) : = ( γ − Φ ( x , y ) ) 2 . In practice , we add an ` ∞ constraint on y , which is achieved by simple clipping with a radius of ω ( projection ) . Intuitively , a large r corresponds to a situation wherein the transformer is highly sensitive to a given perturbation in the input , suggesting that the model parameters are close to a 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 . Lexicon encoder + sentence inputs Input embeddings Adv . noise Transformer encoder Context embeddings Output layer Minimize loss Maximize prediction deviation ( Delayed PGA-1 ) Noise generator Figure 1 : The architecture of the proposed method . Figure 2 : Time breakdown with and with PGA-1 . Figure 3 : Impact of perturbation steps . Inner Maximization : For any given outer step t , let xt denote the parameter proposed by the first player . Since the exact inner maximization in Equation equation 1 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. , yτ+1 = Πω ( yτ + ρτ∇yr ( xt , y ) ) where ρτ for τ ∈ [ T ] is the step size sequence and Π projects the result of the gradient ascent update into an ` ∞ ball of diameter 2ω around the original input embeddings , ξ , considered by the first player . 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 xit+1 = x i t−ηtν ( ‖xit‖ ) ∇̂ixg ( x ) /‖∇̂ixg ( x ) ‖ , ∀i ∈ [ h ] where i denotes the ith-layer of the transformer . The normalized gradient descent mitigates issues due to exploding gradients . The learning rate sequence ηt , ∀t ∈ [ T ] is scaled by a clipping function ν ( c ) : = max ( L , min ( c , U ) ) 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 a batch size B distributed across P workers in order to obtain a noisy gradient estimate ∇̂ixg ( x ) at epoch t. Computational Cost : Given that the primary interest of using a large batch size is to improve hardware efficiency , we are motivated to look carefully into the computational cost of adversarial perturbation . Adversarial perturbation requires an extra PGA inner loop that standard training does not have . Figure 2 provides the time breakdown of adversarial training using PGA with T = 1 ( denoted as PGA-1 ) . 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-propagation 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 2 times , even with T =1 . This motivates us to look at 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 et 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 NLP 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 fine-tuning 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 , the numerical convergence in the model ’ s parameter space is of primary importance rather than the numerical convergence in the data space , i.e. , the maximization is an auxiliary procedure that augments the training phase to 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 to get away with inexact deterministic maximization . Therefore , we apply PGA-1 in our large-batch optimization scheme , given that it produces sufficiently good solutions while being much more computationally efficient . Delayed Perturbation Injection : Given that even PGA-1 still adds an overhead factor of 2 , we are motivated to investigate how useful adversarial perturbations are in the initial phase of largebatch optimization . We conduct additional experiments to measure the final accuracy corresponding to starting from a regular fine-tuning and then enabling PGA-1 for t ≥ ts where ts ∈ [ 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 , at initialization , the model ’ s parameters are relatively far from their final values and are less likely to get stuck at local minima . Therefore the adversarial perturbations generated in the initial training iterations are quite different from the perturbations towards the end of training because they would not maximize the adversarial loss in Equation 1 . We remark that a similar phenomenon has been observed in computer vision tasks ( Cai et al. , 2018 ; Ye et al. , 2020 ; Gupta et al. , 2020 ) . 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 1 ) , whose convergence rate is characterized in Theorem 1 . Algorithm 1 , ScaLA , Scalable Large-batch Adversarial Perturbation 1 : Input : Epochs T , delay ts , perturbation ( inner ) step size ρ , clipping radius ω , regularization strength λ , ( outer ) learning rate η 2 : Output : h-layer transformer model Φ with converged robust parameters x : = xT 3 : for t ∈ [ T ] do . Loop through epochs 4 : for worker p ∈ [ P ] do . In parallel across homogeneous workers 5 : for mini-batch ξp ∼ Q do . Subsample BP data instances on each worker 6 : r ( xt ) ← 0 , γ ← Φ ( x , ξp ) , select y0 . Initialize regularization and label 7 : if t ≥ ts then . Check delay condition 8 : y1 ← Πω ( y0 + ρ∇yr ( xt , y ) ) . Perform adversarial perturbation with PGA-1 9 : r ( xt ) ← KLsym ( γ , Φ ( xt−1 , y1 ) ) . Calculate the adversarial regularization 10 : g ( xt , ξp ) ← f ( xt−1 , ξp ) + λr ( xt ) . Calculate the augmented loss 11 : ∇xg ( xt , ξp ) ← Backward pass on Φ . Compute local gradients using accumulation 12 : ∇̂xg ( xt ) ← 1B ∑P p=1∇xg ( xt , ξp ) . Gradient averaging using all-reduce 13 : xit ← xit−1 − ηtν ( ‖xit‖ ) ∇̂ixg ( xt ) ‖∇̂ixg ( xt ) ‖ , ∀i ∈ [ h ] . Update model parameters Theorem 1 ( Complexity of Algorithm 1 ; Informal – Details in Appendix D ) . Consider the problem in Equation 1 . Let ts = 0 . Setting the outer learning rate as η = O ( 1/ √ T ) and scaling batch size as b = O ( T ) , for Algorithm 1 , we have E [ ‖∇g1/2α ( x ) ‖2 ] ≤ O ( + κα/ √ T ) where x is the estimator obtained from running T steps of Algorithm 1 and picking xt uniformly at random for t ∈ [ T ] . Here , is the error due to the approximate inner maximization oracle , α characterizes the smoothness of f ( x , . ) , g1/2α is the Moreau-envelope of g and κα = maxi αi/mini αi .

Summary: 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.