No Parameters Left Behind: Sensitivity Guided Adaptive Learning Rate for Training Large Transformer Models

Recent research has shown the existence of significant redundancy in large Transformer models . One can prune the redundant parameters without significantly sacrificing the generalization performance . However , we question whether the redundant parameters could have contributed more if they were properly trained . To answer this question , we propose a novel training strategy that encourages all parameters to be trained sufficiently . Specifically , we adaptively adjust the learning rate for each parameter according to its sensitivity , a robust gradient-based measure reflecting this parameter ’ s contribution to the model performance . A parameter with low sensitivity is redundant , and we improve its fitting by increasing its learning rate . In contrast , a parameter with high sensitivity is well-trained , and we regularize it by decreasing its learning rate to prevent further overfitting . We conduct extensive experiments on natural language understanding , neural machine translation , and image classification to demonstrate the effectiveness of the proposed schedule . Analysis shows that the proposed schedule indeed reduces the redundancy and improves generalization performance . 1 INTRODUCTION . Large-scale Transformer models have achieved remarkable success in various fields . Performance of these models scales with their number of parameters , which can be up to hundreds of millions , e.g. , BERT ( Devlin et al. , 2018 ) , DeBERTa ( He et al. , 2020 ) , GPT-3 ( Brown et al. , 2020 ) . Recent research , however , has shown the existence of significant redundancy in the Transformer models ( Michel et al. , 2019 ; Fan et al. , 2019 ; Wang et al. , 2019 ; Chen et al. , 2020 ; Sanh et al. , 2020 ) . For example , Sanh et al . ( 2020 ) removes around 90 % of the parameters , and the models exhibit only a marginal performance drop . The existence of redundancy can hurt the model performance . Recent works have demonstrated that the removal of the redundant parameters can lead to better generalization performance , a phenomenon observed in both small-scale models ( Mozer & Smolensky , 1989 ; Rasmussen & Ghahramani , 2001 ; Grünwald & Grunwald , 2007 ) and large-scale Transformer models ( Bartoldson et al. , 2019 ; Voita et al. , 2019 ; Hou et al. , 2020 ; Liang et al. , 2021 ) . As illustrated in Figure 1 , with up to 20 % of the parameters pruned , the generalization performance boosts up to 1 % . As a result , we aim to improve model generalization through redundancy elimination . However , the existence of redundancy has long been regarded as inevitable . The common belief is that , in each network , there always exists a set of parameters “ born ” to be useless ( Frankle & Carbin , 2018 ; Liu et al. , 2018 ) . Following this belief , pruning , where redundant parameters are directly zeroed out , becomes one of the most widely adopted solutions to redundancy elimination . However , we ask a critical question here : Are these parameters really redundant , or just insufficiently trained by commonly used training strategies ? Our question is motivated by empirical observations , which show that training strategies indeed play a role in causing redundancy . For example , different learning rates ( Table 1 ) , random seeds and optimizers ( Morcos et al. , 2019 ) can produce models with similar performance but different sets of redundant parameters . This suggests that the redundancy of parameters depends on the training strategy : A training strategy often prefers specific parameters and provides them with sufficient training . In contrast , the other parameters receive insufficient training and become under-fitted . As a result , these parameters become redundant , such that they fail to contribute to the generalization and prevent the model from achieving its ideal performance . Therefore , we hypothesize that with a desirable training strategy , these redundant parameters can receive more sufficient training and become useful ultimately . We verify the hypothesis by proposing a novel training strategy , which encourages all parameters to be trained sufficiently . Throughout the training process , we simultaneously excite the under-fitted parameters to reduce redundancy and regularize the well-fitted parameters to prevent overfitting . More specifically , we propose an adaptive learning rate schedule – SAGE ( Sensitivity-guided Adaptive learninG ratE ) , where each parameter learns at its own pace guided by its sensitivity . Sensitivity originated in model pruning , where it is used to measure the redundancy of the parameters ( Molchanov et al. , 2016 ; 2019 ; Theis et al. , 2018 ; Lee et al. , 2018 ; Ding et al. , 2019 ) . In pruning literature , parameters with low sensitivity are considered redundant . Since a redundant parameter could be insufficiently trained and under-fitted , we promote its training by increasing its learning rate . In contrast , for a parameter with high sensitivity , i.e. , it is considered sufficiently trained and well-fitted , we slow down its training by decreasing its learning rate to prevent overfitting . Moreover , we introduce a local temporal variation of the sensitivity as a second factor to further guide the learning rate . The local temporal variation essentially measures the uncertainty of sensitivity , which mainly comes from two sources : ( 1 ) The sensitivity can have large variance due to data sampling . This is because during training , the sensitivity is evaluated using a randomly sampled mini-batch instead of all the training data . ( 2 ) The sensitivity of a parameter may not be stable and can vary drastically among iterations , which introduces extra uncertainty . We define the local temporal variation of a parameter as the absolute difference between its sensitivity and an exponential moving average of its sensitivity from all previous iterations . A large local temporal variation implies high uncertainty in the sensitivity at the current iteration , and therefore it is not yet a reliable indicator of redundancy . Accordingly , we should avoid significantly decreasing its learning rate even though its sensitivity at the current iteration might be large . Therefore , we eventually require the overall learning rate schedule for each parameter to be proportional to the ratio between the local temporal variation and the sensitivity . This can effectively account for the uncertainty issue in sensitivity . We conduct experiments on a wide range of tasks and models to demonstrate the effectiveness of SAGE . In natural language understanding , the fine-tuning performance of BERT-base ( Devlin et al. , 2018 ) and RoBERTa-large ( Liu et al. , 2019b ) improves 1.4 and 0.6 task-average score on the dev set of the GLUE benchmark ( Wang et al. , 2018 ) , respectively . Furthermore , SAGE improves neural machine translation performance using Transformer-base ( Vaswani et al. , 2017 ) on two datasets , suggesting it also benefits training-from-scratch . SAGE also boost the image classification accuracy on ImageNet dataset ( Deng et al. , 2009 ) with Vision Transformer models ( Dosovitskiy et al. , 2020 ) . Furthermore , our experiments demonstrate SAGE is complementary to various types of optimizers , e.g. , SGD ( Robbins & Monro , 1951 ) , Adam , and Adamax ( Kingma & Ba , 2014 ) . Moreover , we observe several favorable proprieties of SAGE . First , it leads to balanced and sufficient training on all parameters and produces a better-generalized model . Second , SAGE is complementary to state-of-the-art training methods . Specifically , we show that SAGE achieves better performance on GLUE when combined with adversarial regularization ( Jiang et al. , 2019 ) . 2 PRELIMINARY . We briefly review the sensitivity of the parameters and adaptive learning rate methods . 2.1 SENSITIVITY OF THE PARAMETERS . The sensitivity of a parameter essentially approximates the change in the loss magnitude when this parameter is completely zeroed-out ( LeCun et al. , 1990 ; Mozer & Smolensky , 1989 ) . If the removal of a parameter causes a large influence on the loss , then the model is sensitive to it . More specifically , we define a deep neural network with parameters Θ = [ θ1 , ... , θJ ] ∈ RJ , where for j = 1 , ... , J , θj ∈ R denotes each parameter . We further define Θj , −j = [ 0 , ... , 0 , θj , 0 , ... , 0 ] ∈ RJ . We denote the loss of the model as L ( Θ ) , and the gradients of the loss with respect to Θ as ∇ΘL ( Θ ) . The sensitivity of the j-th parameter is defined as the magnitude of the gradient-weight product : Ij = |Θ > j , −j∇ΘL ( Θ ) | . ( 1 ) This definition is derived from the first-order Taylor expansion of L ( · ) with respect to θj at Θ . Specifically , Ij approximates the absolute change of the loss given the removal of θj : Θ > j , −j∇ΘL ( Θ ) ≈ L ( Θ ) − L ( Θ−Θj , −j ) . The sensitivity was originally introduced for model pruning ( Molchanov et al. , 2016 ; 2019 ; Theis et al. , 2018 ; Lee et al. , 2018 ; Ding et al. , 2019 ; Xiao et al. , 2019 ) , and it was commonly used as an “ importance score ” for model weights . The parameters with high sensitivity are of high importance and should be kept ( Lubana & Dick , 2020 ) . Parameters with low sensitivity are considered redundant , and they can be safely pruned with only marginal influence on the model loss . 2.2 ADAPTIVE LEARNING RATE METHODS . Adaptive learning rate methods adjust the learning rate of each individual parameter based on the training progress . Most of these methods focus on adapting the training to the optimization landscape , e.g. , AdaGrad ( Duchi et al. , 2011 ) , AdaDelta ( Zeiler , 2012 ) , RMSProp ( Hinton et al. , 2012 ) , Adam ( Kingma & Ba , 2014 ) and RAdam ( Liu et al. , 2019a ) . Their purpose is to make the model converge faster to the first-order stationary solutions . Specifically , these methods prefer updating the weights with smaller second-order moments , as the loss function is generally flat along directions corresponding to such weights . There are also some adaptive learning rate methods focusing on the perspective of improving model generalization ( Loshchilov & Hutter , 2018 ; Foret et al. , 2020 ) . For example , AdamW ( Loshchilov & Hutter , 2018 ) propose to decouple the weight decay and gradient update to avoid regularizing weights that have larger gradient magnitudes with a weaker strength . 3 METHOD . We introduce our proposed adaptive learning rate schedule , SAGE . Our method customizes a specific learning rate for each parameter at each iteration . A parameter ’ s learning rate at a certain iteration is determined by two factors : sensitivity and its local temporal variation . Sensitivity of the parameters . At the t-th iteration , following Eq . ( 1 ) , we define the sensitivity of θ ( t ) j as I ( t ) j = |Θ ( t ) > j , −j∇Θ ( t ) L ( Θ ( t ) ) | , ( 2 ) which reflects the influence of removing θ ( t ) j in the model loss . In previous literature , θ ( t ) j is considered redundant when I ( t ) j is small . In contrast , we hypothesize that θ ( t ) j is just insufficiently trained and under-fitted , and can become less redundant when receiving further training . Local temporal variation . Recall that the sensitivity measure involves excessive uncertainty , which comes from : ( 1 ) Sensitivity is measured based on a randomly sampled mini-batch of the training data at each iteration , which leads to a large variance ; ( 2 ) Sensitivity can be unstable and vary drastically , as changes of the model introduce extra uncertainty to the measure . One way to measure the uncertainty of sensitivity of θj is the absolute change of sensitivity , i.e. , |I ( t ) j − I ( t−1 ) j | . Such a quantity often has a large variance in practice . Therefore , we propose to keep track of an exponential moving average of I ( t ) j as Î ( t ) j = β0Î ( t−1 ) j + ( 1− β0 ) I ( t ) j , where Î ( 0 ) j = 0 and β0 ∈ ( 0 , 1 ) is a hyper-parameter . Based on Î ( t ) j , we measure the uncertainty of the j-th parameter ’ s sensitivity using the local temporal variation defined as : U ( t ) j = |I ( t ) j − Î ( t ) j | . ( 3 ) We remark that a large U ( t ) j implies that there exists high uncertainty in I ( t ) j , and therefore it is not yet a reliable indicator of the redundancy of θ ( t ) j . Algorithm . We denote the learning rate at the t-th iteration as η ( t ) under the original schedule . Then the sensitivity-guided learning rate for the j-th parameter at the t-th iteration can be computed as η ( t ) j = η ( t ) · U ( t ) j + Î ( t ) j + = η ( t ) · |I ( t ) j − Î ( t ) j |+ Î ( t ) j + , ( 4 ) where 0 < 1 prevents zero learning rate and zero denominator . Algorithm 2 shows the SAGE algorithm for SGD , and extensions to other algorithms , such as Adam ( Kingma & Ba , 2014 ) , are straightforward ( Appendix A.4.1 ) . In Eq . ( 4 ) , we place Î ( t ) j in the denominator , as one of our goals is to encourage all parameters to be sufficiently trained . If Î ( t ) j is small , we promote its training by increasing its learning rate . If Î ( t ) j is large , we slow down its training to prevent overfitting by decreasing its learning rate . We place U ( t ) j in the numerator to measure the uncertainty in the sensitivity . A large U ( t ) j implies I ( t ) j is not yet a reliable indicator of the redundancy in θ ( t ) j . We thus avoid significantly decreasing its learning rate . Algorithm 1 SGD-SAGE ( denotes Hadamard product and denotes Hadamard division ) Input : Model parameters Θ ∈ RJ ; Data D ; Learning rate schedule η ( · ) ; Total training iteration T ; Moving average coefficient β0 . 1 : Initialize Î ( 0 ) = 0 ∈ RJ . 2 : for t = 1 , ... , T do 3 : Sample a minibatch b ( t ) from D. 4 : Compute gradient ∇Θ ( t ) L ( b ( t ) , Θ ( t ) ) . 5 : I ( t ) = |Θ ( t ) ∇Θ ( t ) L ( b ( t ) , Θ ( t ) ) | . 6 : Î ( t ) = β0Î ( t−1 ) + ( 1− β0 ) I ( t ) . 7 : U ( t ) = |I ( t ) − Î ( t ) | . 8 : Θ ( t+1 ) = Θ ( t ) − η ( t ) ( U ( t ) + ) ( Î ( t ) + ) ∇Θ ( t ) L ( b ( t ) , Θ ( t ) ) . 9 : end for Computation and memory usage . SAGE adds a marginal cost to computation and memory usage . At each iteration , we only perform an extra element-wise multiplication between the weight matrix and the corresponding gradient matrix obtained through back-propagation . The only memory cost is to store the exponential moving average of sensitivity .

Summary: This paper proposes a method for scaling the learning rate during training in order to encourage all parameters in a neural net to be fully used. Specifically, the updates for parameters are scaled in inverse proportion to how much they affect the loss (their sensitivity); the scaling factor also depends on how stable the estimate of sensitivity is, so that high-sensitivity parameters whose role changes rapidly aren't down-weighted as much. This technique is shown to improve the performance of Transformer models on several different problems, even when used in combination with other adaptive learning rate methods (Adam, Adamax). Analysis experiments verify that the method has the intended effect: compared to the baseline, more parameters have higher sensitivity, and pruning is less effective.