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: The paper argues that redundant or "useless" parameters are not an axiom we should take for granted, but rather a symptom of current optimization settings. The authors propose an adaptive learning-rate schedule that specifically aims to eliminate redundant parameters. Through extensive experiments on fine-tuning transformers, they show that indeed this method (SAGE) does reduce redundancy and also slightly improves results.