NO PARAMETERS LEFT BEHIND: SENSITIVITY GUIDED ADAPTIVE LEARNING RATE FOR TRAINING LARGE TRANSFORMER MODELS

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

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

heading: 1 INTRODUCTION
text: 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).

heading: 2 PRELIMINARY
text: We briefly review the sensitivity of the parameters and adaptive learning rate methods.

heading: 2.1 SENSITIVITY OF THE PARAMETERS
text: 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.

heading: 2.2 ADAPTIVE LEARNING RATE METHODS
text: 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.

heading: 3 METHOD
text: 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.

heading: 4 EXPERIMENTS
text: We evaluate SAGE on widely used benchmarks for natural language understanding (NLU), neural machine translation (NMT), and image classification.

heading: 4.1 NATURAL LANGUAGE UNDERSTANDING
text: Model and data. We evaluate the fine-tuning performance of the pre-trained language models, BERT-base (Devlin et al., 2018) and RoBERTa-large (Liu et al., 2019b), on the General Language
Understanding Evaluation (GLUE, Wang et al. (2018)) benchmark. GLUE contains nine NLU tasks, including textual entailment, question answering, sentiment analysis, and text similarity. Details about the benchmark are deferred to Appendix A.1.1.
Implementation Details. We implement our method using the MT-DNN code-base1. We follow the suggested training and hyper-parameters settings from Liu et al. (2020). Specifically, we adopt Adam and Adamax (Kingma & Ba, 2014) with corrected weight decay (Loshchilov & Hutter, 2018) as the baseline optimizer and we set β = (0.9, 0.999). We use a linear-decay learning rate schedule, and we apply SAGE to both Adam and Adamax.
We select learning rates in range of {1, 2, 3, 5, 8}× {10−5, 10−4}. We select β0 in range of [0.6, 0.9] with an increment of 0.05. Other training details are reported in Appendix A.1.2.
Main results. Table 2 and Table 3 show the evaluation results on the GLUE benchmark. The dev results are averaged over 5 different random seeds, and all gains are statistically significant2. We select the best single task model for test evaluation.
Our method gains 1.4 on dev and 1.1 on test of the task-average score on BERT-base. In large datasets, i.e., MNLI (392K) and QNLI (108K), SAGE improves around 0.5 points. In small datasets, i.e., RTE (2.5K) and CoLA (8.5K), we obtain more than 2 points of improvements. Such observations indicate that SAGE is very effective on the small datasets. Furthermore, SAGE improves upon RoBERTa-large by 0.6 average scores, suggesting SAGE can still achieve significant improvements for larger and more adequately pre-trained models than BERT-base.

heading: 4.2 NEURAL MACHINE TRANSLATION
text: Model and Data. We evaluate SAGE on the Transformer-base NMT models (Vaswani et al., 2017) using two widely used NMT datasets, IWSLT’14 De-En (Cettolo et al., 2015)3 and WMT’16 En-De
1https://github.com/namisan/mt-dnn 2The dev results on RoBERTa-large are averaged over 3 different random seeds. All results have passed a
paired student t-test with p-values less than 0.05. The detailed statistics are summarized in Appendix A.1.3. 3https://wit3.fbk.eu/
(Bojar et al., 2016)4. IWSLT’14 De-En is a low-resource dataset, which contains 160K sentence pairs. WMT’16 En-De is a rich-resource dataset, which contains 4.5M sentence pairs. Dataset and pre-processing details are deferred to Appendix A.2.1.
Implementation Details. We implement the algorithms using the fairseq code-base and follow the training and hyper-parameters settings from Ott et al. (2018; 2019). Specifically, we adopt the inverse square root learning rate schedule and we employ Adam (Kingma & Ba, 2014) as the optimizer with β = (0.9, 0.98). We apply SAGE to the same setting.
We select learning rates in range of {5, 7} × 10−5 ∪ {1, 2} × 10−4 and select β0 in range of {0.5, 0.6, 0.7, 0.8, 0.9}. Comprehensive training details are reported in Appendix A.2.2. Main results. Table 4 shows the BLEU scores on the IWSLT’14 De-En and the WMT’16 En-De test set, where SAGE improves around 0.6 and 0.4 points, respectively. This suggests that other than fine-tuning, SAGE can also improve the generalization of trained-from-scratch models in both low-resource and rich-resource settings.

heading: 4.3 IMAGE CLASSIFICATION
text: Model and data. We evaluate SAGE using Vision Transformer models (ViT) on the CIFAR100 (Krizhevsky et al., 2009) and ILSVRC-2012 ImageNet dataset (Deng et al., 2009). Specifically, we evaluate the fine-tuning performance of the ViT-base and ViT-large pre-trained using ImageNet-21k, a superset of ImageNet dataset with 21k classes and 14M images. Data and pre-processing details are deferred to Appendix A.3.1.
Implementation details. All experiments follow the suggested training configuration of Dosovitskiy et al. (2020) and a jax-implemented code base 5. We adopt SGD as the baseline optimizer with a momentum factor 0.9. We fine-tune the models for 100K steps for CIFAR100, and 200K steps for ImageNet. We select learning rates in range of {0.02, 0.05, 0.08, 0.1} and select β0 in range of {0.85, 0.90, 0.95}. Comprehensive training details are reported in Appendix A.3.2. Main results. Table 5 shows the evaluation results on CIFAR100 and ImageNet. SAGE outperforms baselines by a significant margin. This demonstrates that SAGE is quite general, and can be applied to various tasks (e.g., NLP and computer vision) and optimizers (e.g., Adam, Adamax and SGD).

heading: 5 ANALYSIS
text: We verify that SAGE leads to more sufficient training (Section 5.1), better generalization performance (Section 5.2), and is complementary to existing state-of-the-art regularization methods (Section 5.3). We also provide ablation studies in Appendix A.4.4.

heading: 5.1 SAGE LEADS TO MORE SUFFICIENT TRAINING
text: Recall that SAGE adjusts the learning rate for each parameter according to two factors: the sensitivity of parameters and the local temporal variation of sensitivity. By inspecting these factors, we verify that SAGE leads to more sufficient training.
The sensitivity distribution is more concentrated. Figure 2 shows the sensitivity distribution of parameters in the SAGE optimized models and the baseline models. We select the hyper-parameters that yield the best generalization performance on the BERT-base model, and we evaluate the sensitivity of each parameter using the entire training set. See Appendix A.4.2 for implementation details.
4http://data.statmt.org/wmt16/translation-task/ 5https://github.com/google-research/vision_transformer
We observe that the sensitivity distribution exhibits a lower variance in the SAGE optimized models than the baseline models. This suggests that the sensitivity of parameters becomes more concentrated. In other words, the amount of each parameter’s contribution is more balanced, and the model is more sufficiently trained.
Even the most redundant parameters contribute to the model performance. Recall that sensitivity is a type of importance score in pruning, which is a straightforward approach to measure each parameter’s contribution. Therefore, we conduct an unstructured, one-shot pruning experiment on the fine-tuned BERT-base models. Specifically, we remove up to 40% parameters6 with the lowest sensitivity scores and evaluate the pruned models’ performance. We average the results over 5 models trained with different random seeds. Figure 3 Upper shows the generalization performance of the pruned models. To ease the comparison, Figure 3 Lower shows the change in generalization performance with respect to the un-pruned models.
We have the following observations:
• The pruning performance of the SAGE optimized models remains higher than that of the baseline models (Figure 3 Upper).
• Even the most redundant parameters in the SAGE optimized models makes contributions (Figure 3 Lower). When there are over 80% of weights remaining, the pruning performance of the baseline models is comparable or even superior than their un-pruned alternatives. In contrast, the performance of the SAGE optimized models consistently deteriorates. This suggests that the most redundant parameters in the baseline models fail to contribute, while those in the SAGE optimized models are trained more sufficiently and are able to make contributions.
6Embedding weights are excluded.
Sensitivity is a reliable indicator of redundancy. We visualize the local temporal variation (Figure 4) to verify that sensitivity indeed becomes a more reliable indicator of redundancy in SAGE than in the baselines. We track the variation for all parameters in the BERT-base model at each iteration, and we evaluate the variation based on the current mini-batch of training data. See Appendix A.4.2 for implementation details.
We observe that the local temporal variation in SAGE remains lower or decreases faster than in the baselines for all tasks. For example, the variation in the baseline approach remains large in QNLI. In contrast, the variation in SAGE decreases, suggesting the sensitivity indeed stabilizes and becomes a reliable indicator of redundancy.

heading: 5.2 SAGE LEADS TO BETTER GENERALIZATION PERFORMANCE
text: We verify that SAGE leads to better generalization performance through inspecting the learning curves, decision boundary and hyper-parameter search space.
Learning Curves. Figure 6 shows the training loss, validation loss, learning rate, and sensitivity score obtained by fine-tuning BERT-base on SST-2. All experiment details are deferred to Appendix A.4.3. We have two major observations: 1) SAGE’s validation loss descends faster and SAGE is less prone to overfitting. This observation suggests that SAGE has a regularization effect and reduces the model variance. 2) SAGE’s variance of the sensitivity score becomes lower through training, aligning with our observation in Figure 2. This suggests that SAGE gives rise to a more balanced and sufficient training. Both observations agree with our initial motivation (Figure 1) that
redundancy elimination can lead to better generalization.
Hyper-parameter Study. Figure 7 shows the validation accuracy heatmap obtained by fine-tuning BERT-base on the RTE dataset. We plot the accuracy obtained by training with different learning rates,
Adam’s βs and SAGE’s β0s. We can observe that SAGE consistently achieves a better generalization performance within a larger region of hyper-parameter search space under different β0s. We also provide a hyper-parameter study for more datasets in Appendix A.4.5.
Decision Boundary. Figure 5 shows the decision boundary predicted with Adam and SAGE on the Spiral dataset. Specifically, we train a multi-layer perceptron with 3 hidden layers, each with a hidden dimension of 100. The decision boundary predicted with SAGE is smoother and has a larger margin than with Adam, suggesting SAGE produces a better generalized model.

heading: 5.3 COMBINE WITH STATE-OF-THE-ART METHODS
text: We further show that SAGE is complementary to existing state-of-the-art regularization methods. Specifically, we apply SAGE to SMART (Jiang et al., 2019), a state-of-the-art smoothness-inducing adversarial regularization method. As shown in Table 6, SAGE can further improve upon SMART, suggesting the two techniques are complementary.

heading: 6 DISCUSSION
text: SAGE is complementary to Adaptive Gradient Methods. Our proposed method and the mainstream adaptive gradient methods (e.g., Adam and AdaGrad) are for fundamentally different purposes. The mainstream adaptive gradient methods aim to improve optimization by adapting to the optimization landscape, while SAGE aims to improve generalization by eliminating the weight redundancy. The quantities of our interest (i.e., Eq. (2) and Eq. (3)) are related to the weight redundancy. They are not directly related to the moduli of the objective function, e.g., smoothness, curvature (which are of the interests for optimization). As shown in our experiments (See Section 4), we do not observe any conflicts between the two methods, as SAGE improves the model generalization performance when being combined with several adaptive gradient methods (e.g., Adam).
Redundant Weights vs. Insufficiently Trained Weights. Lottery Ticket Hypothesis (Frankle & Carbin, 2018) suggests that, in a randomly initialized network, there exists a well-initialized subnetwork, which outperforms any other subnetworks and matches the full model’s performance. This suggests the rest parameters contribute marginally to the model performance. Although the initialization of these parameters may not be satisfactory, SAGE provides them sufficient training so that they can learn to contribute.

heading: 7 CONCLUSION
text: We begin with a hypothesis that the redundant parameters can become useful if they are sufficiently trained by desirable optimization strategies. We verify this hypothesis by proposing an adaptive learning schedule – SAGE, which excites the under-fitted parameters to reduce redundancy and regularize the well-fitted parameters to prevent overfitting. We demonstrate that SAGE can benefit model generalization in a wide range of tasks and strengthen various types of optimizers.

Review summaries:
This paper focuses on the parameter redundancy issue in large transformer architectures. Instead of pruning redundant parameters, it strengthens training them to make them contribute better performance. To this end, it proposes an adaptive learning rate algorithm SAGE, which automatically scales the learning rate for each parameter based on its sensitivity. The sensitivity is approximated by the dot product between parameters and their gradients. An exponential moving average is used to track the sensitivity scores to reduce uncertainty in mini-batch training. The algorithm is applied to fine-tuning pre-trained transformer models in benchmarks for natural language understanding (NLU), neural machine translation (NMT), and image classification.

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.

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. 

The paper proposes the SAGE method. The idea is that neural networks have redundant parameters. Some approaches will prune these parameters which has the effect of not decreasing the performance but to decrease the number of parameters. The paper studies if it is possible to learn better these parameters in order to make them more useful for the network. The paper proposes a method that will allow to train differently the parameters of a network in order to have a better use of the weights. The paper is evaluated in NLP and image classification with transformers.