No Parameters Left Behind: Sensitivity Guided Adaptive Learning Rate for Training Large Transformer Models
Authors: Chen Liang
Haoming Jiang
Simiao Zuo
Pengcheng He
Xiaodong Liu
Jianfeng Gao
Weizhu Chen
Tuo Zhao
Venue: ICLR 2022 Poster (2022)
Keywords: Training Large Transformer Models
Reducing Model Redundancy
Parameter Sensitivity
Adaptive Learning Rate Method
Model Generalization
Model Pruning

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.

Published as a conference paper at ICLR 2022
No PARAMETERS LEFT BEHIND: SENSITIVITY GUIDED ADAPTIVE LEARNING RATE FOR TRAINING LARGE TRANSFORMER MODELS
Chen Liang*, Haoming Jiang* Simiao Zuo Pengeheng He*, Xiaodong Liu",
Jianfeng Gaod, Weizhu Chen' & Tuo Zhao* Georgia Institute of Technology, f Amazon, Miirosooft Azure AI, Microsoft Research [cliang73, jianghm, s simiaozus tourzhao) @gatech. @gatechh edu (penhe, xiaodl, jEgao, wzchen @microsoft.com com
ABSTRACT
Recent research has shown the existence of significant redundancy in large Trans- former models. One can prune the redundant parameters without significantly sacrificing the generalization performance. However, question whether the redundant parameters could have contributed more if they were properly trained. To answer this question, propose novel training strategy that encourages all parameters to trained sufficiently. Specifically, we adaptively adjust the learning for each parameter according its sensitivity, a robust gradient-based measure reflecting this parameter's contribution to the model performance. A parameter with low sensitivity redundant, and we improve its fitting by increasing its learning rate. In contrast, parameter with high sensitivity is well-trained, and we regularize by decreasing its learning rate prevent further overfitting. We conduct extensive experiments natural language understanding. neural machine translation, and image classification demonstrate the effectiveness of the pro- posed schedule. Analysis shows that the proposed schedule indeed reduces the redundancy and improves generalization performance
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 hundreds of millions. e.g., BERT (Devlin et al., 2018). DeBERTa (He et al.] (022), GPT-3 (Brown et al.) (020)). Recent research, however, has shown the existence of significant redundancy
the Transformer models (Michel et al. 2019; Fan et al. 2019 Wang al. 2019:) Chen et al. [020 |aanh al. (022)). For example, Sanh et al. (2020) removes around 90% of the parameters, and the models exhibit only marginal performance drop.
The existence of redundancy can hurt the model performance. Re- cent works have demonstrated that the removal of the redundant parameters can lead to better generalization performance, a phe-
70.0
Accurcyy
nomenon observed in both small-scale models (Mozer & Smolensky
67.5
1989
Rasmussen & Ghahramani 2001 Grünwald & Grunwald
2007 and large-scale Transformer models (Bartoldson et al.) 2019
62.5
Voita et al. 2019 Hou et al.) 2020 Liang et al.) (2021). As illus- trated in Figure with up 20% of the parameters pruned, the
60.0
100
$0
generalization performance boosts up to 1%.
Percentage of Weight Remaining (%)
As a result, we aim to improve model generalization through re- dundancy elimination. However, the existence of redundancy has long been regarded 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., (008).. Following this belief, pruning, where redundant parameters becomes one of the most widely adopted solutions to redundancy elimination. However, ask critical question here:
Figure 1: Validation results
of fine-tuning BERT-base at different sparsity levels the RTE dataset (Wang et al. 2018) in Liang al. (2021). Solid black curve represents the full model performance.
directly zeroed out,
Our code has been released httpt:tp//// httpp:://gttttt.... com/cliacccccccccao iang htt.eeinnnn1113/3531 453/SAGH

---

Published as a conference paper at ICLR 2022
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 D. random seeds and optimizers (Morcos 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 result, these parameters become redundant, such that they fail contribute to the generalization and prevent the model from achieving its ideal performance. Therefore, we hypothesize that with desirable training strategy, these redundant parameters can receive more sufficient training and become useful ultimately.
We verify the hypothesis by proposing novel training strat- egy, which encourages all parameters to be trained sufficiently.
OVLP Among
Avg OVLP
Throughout the training process, we simultaneously excite the under-fitted parameters to reduce redundancy and regularize the well-fitted parameters to prevent overfitting.
2 Models
59.8%
3 Models
46.5%
5 Models
35.7%
More specifically, propose adaptive learning rate sched- ule SAGE (Sensitivity-guidee Adaptive learninG ratE), where each parameter learns at its own pace guided by its sensitivity. Sensitivity originated in model pruning, where is used to mea- sure the redundancy of the parameters (Molchanov et al. 2016)]
Table 1: Percentage of overlapping between the 30% most redundant parameters
BERT---aus models fine-tuned using (1, 5, 8, 10, 20) x 10-5 as learning rates on SST-2.
2019: Theis et al. 2018; Lee et al.] 2018) Ding et al. 2019).
In pruning literature, parameters with low sensitivity are considered redundant. Since redundant parameter could be insufficiently trained and under-fitted learning rate. In contrast, for a parameter with high sensitivity, i.c., it considered sufficiently trained and well-fitted we slow down its training by decreasing its learning rate to prevent overfitting.
we promote its training by increasing its
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 data sampling. This is because during training. the sensitivity is evaluated using randomly sampled mini-batch instead of all the training data. (2) The sensitivity of parameter may not stable and can vary drastically among iterations, which introduces extra uncertainty. We define the local temporal variation of 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 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 pro- portional 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 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-fromm from-scratchh. SAGE also boost the image classification accuracy on ImageNet dataset (Deng et al. (2099 with Vision Transformer models (Dosovitskiy et al. (020). Furthermore, our experiments demonstrate SAGE is complementary to various types of optimizers, 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 all parameters and produces better-generalie model. Second, SAGE is complementary to state-of-the-ar training methods. Specifically, we show that SAGE achieves better performance on GLUE when combined with adversarial regularization (Jiang et al. 2019).
2

---

Published as a conference paper at ICLR 2022
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). the removal of 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 θ = [01,..., 0y] R. where for = 1,....J, 0 E R denotes each parameter. We further define + j.-; = (o...., 0, 0, 0, .... Of E R We denote the loss of the model as L()), and the gradients of the loss with respect to as VeL(O). The sensitivity of the j-th parameter is defined as the magnitude of the gradient-weight product:
I, = -=
(1)
This definition is derived from the first-order Taylor expansion of IL..) with respect to 0g at O. Specifically, 1 approximates the absolute change of the loss given the removal of 0y:
(o = L( L( z.-1)-
The sensitivity was originally introduced for model pruning (Molchanov et al. 2016 2019; Theis et al. 2018 Lee 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, c-g.. AdaGrad (Duchi et al. 2011). AdaDelta (Zeiler 2012). RMSProp (Hinton al.. 2012), Adam(Kingma & Ba| 2014) and RAdam (Liu et al. 2019a). Their purpose is to make the model converge faster 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] 22018: Foret et al. 2220).. 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.
METHOD
We introduce our proposed adaptive learning rate schedule, SAGE. Our method customizes a specific learning rate for each parameter at each iteration. parameter's learning rate at certain iteration is determined by two factors: sensitivity and its local temporal variation.
Sensitivity of the parameters. At the t-th iteration, following Eq. 8/0) as
C.. we define the sensitivity of
(c) 1 - 1.= e (1)))
(2)
which reflects the influence removing g the model loss. previous literature, 8() is considered redundant when r(()) is small. In contrast, we hypothesize that o(0) 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 randomly sampled mini-batch of the training data at each iteration, which leads to 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 0, is the absolute change of sensitivity, i.c., 1750) 7/ gt--)). Such quantity often has large variance in practice. Therefore, we propose keep
3

---

Published as a conference paper at ICLR 2022
track of exponential moving average of //0
7 Bo7, + (1-B3)1
where / and 8 E (0, 1) a the j-th parameter's sensitivity using the local temporal variation defined as:
Based on 7 we measure the uncertainty of
Ug 17, 1.
(3)
We remark that a large U/ implics that there exists high uncertainty in 1700, and therefore is not yct a reliable indicator the redundancy of 0/0))
Algorithm. We denote the learning rate at the 1-th iteration as r) 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) +
14, (0)
1
77(0
17
(4)
7 70t+c +
7 +
where 0 < < prevents zero learning rate and zero denominator. Algorithm shows the SAGE algorithm for SGD, and extensions to other algorithms, such Adam (Kingma & Ba 2014). are straightforward (appediixAA.4.11.
In Eq. (4, we place 7) in the denominator, as one of our goals is to encourage all parameters to be sufficiently trained. If 75 is small, we promote its training by increasing its learning rate. If 7(0) is large, we slow down its training to prevent overfitting by decreasing its learning rate.
We place U/ in the numerator to measure the uncertainty in the sensitivity. A large U/!) implies y(() is not yet a reliable indicator of the redundancy in ( We thus avoid significantly decreasing its leaming rate.
Algorithm 1 SGD-SAGE ( denotes Hadamard product and @ denotes Hadamard division)
Input:
Model parameters E RR;; Data D; Learning rate schedule 17):; Total training iteration T: Moving average coefficient S6. R.V.
1: Initialize 7
2: for port=1
do 3: Sample minibatch b(() from D.
4: Compute gradient Vow L(b(), ((1)..
((t) ()(+) @ Lb(it, (())).
5:
700 ft + (())
6:
(1(0) (1(0) |.
7:
(++1)) = (t) (0 + 2 (10) +c) Low)
8:
9: end for
Computation and memory usage. SAGE adds a marginal cost to computation and memory usage. At each iteration, only perform extra element-wise multiplication between the weight matrix and the corresponding gradient matrix obtained through back pock-propagttion The only memory cost is store the exponential moving average of sensitivity.
4 EXPERIMENTS
We evaluate SAGE on widely used benchmarks for natural language understanding (NLU), neural machine translation (NMT), and image classification.
4.1 NATURAL LANGUAGE UNDERSTANDING
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. GLÜE contains nine NLU tasks, including textual entailment, question answering, sentiment analysis, and text similarity. Details about the benchmark are deferred to Appendix A.I.I
4

---

Published as a conference paper at ICLR 2022
Implementation Details. We implement method using the MT-DNN code-base We follow the suggested training and hyper-parameters parameters settings from Liu et al. (2020). Specifically, we adopt Adam and Adamax (Kingma & Ba 2014) with corrected weight decay (Loshchilov & Hutter 2018)) the baseline optimizer and we set B (0.9, 0.999). We use a linear-decay learning rate schedule, and apply SAGE to both Adam and Adamax.
(roumLiu
We select learning rates in range of 2,,2,3, 8) x (10-5, 10-4). We select in in range of [0.6, 0.9) with increment of 0.05. Other training details reported in Appendix A..12
Main results. Table 2 and Table show the evaluation results on the GLUE benchmark. The dev
results averaged over 5 different random seeds, and all gains are statistically significant We select the best single task model for test evaluation.
Model
Optimizer
CoL.A
SST-2
QNLI
QQP
MNI.Im/mmm
Average
Acc
Acc/Fi
Mec
Acc
PS Curr
Are
Acc/F1
Ace
Score
Devlin al
2018
-886.7
92.7
88.4
+/-
84.4/-
BERT BASE
Adam Adam-SAGE
8/.1/89.0
54.7
92.9
89.2/88.8
90.9/88.1
84.5/84.4
81.5
73.3
87.0/90.9
60.3
93.5
90.3/89.9
91.7
91.2/88.1
84.7/84.8
84,0
Adamax
69.2
86.2/90.4
57.8
92.9
89.7/89.2
91.2
90.9/88.0
84,5/84.4
82.8
Adamax SAGE
74.0
87.3/91.0
59.7
93.8
90.3/89.8
91.8
91.2/88.2
85.0/85.2
84.2
Liu et al.
2019b
86.6
-/90,9
68.0
96.4
92.4/-
94.7
92.2/-
90.2/902
RoBERTa
Adamax
86.6
90.4/93.1
67.5
96.4
92.4/92.2
94.7
92.1/89.5
90,4/90.3
88.7
Adamaa-AAGE SAGE
87.8
91.5/93.9
68.7
96.7
92.7/92.4
94.9
92.2/89.4
90.8/90.4
89.3
Table 2: Single task fine-tuning dev results on GLUE. All results are from our implementations. denotes missing results.
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 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.
RTE
MRPC
CoLA
SST-2 Acc
STS-B Corr
QNLI Acc
QQP
MNLI-m/um Average Score
Acc
F1
Mec
Ave
BERTRASE (Delin al. 2018
4
88.9
52.1
93.5
85.8
90.5
71.2
84.6/83.4
79.6
BERTBASE: Adamax
88.6
93.4
71.1
84.7/83.6
79.9
BERT Ademax-SAGE
69.8
89.7
54.5
94.1
87.1
90.8
71.3
84.9/83.8
80.7
Table 3: Single task fine-tuning test results from the GLUE evaluation server.
Model
Optimizer
IWSL1'14 De-En
WMT'16 En-De
Adam
34.5
TransformeruASE
Adam-SAGE
35.1
27.7
Table 4: Neural machine translation BLEU scores on test set. All results are from our implementation.
Model
Optimizer
CIFAR100
ImageNet
SGD*
91.97
81.28
VIT-B/32
SGD-SAGE
92.68
81.72
SGD"
93.04
80.99
VIT-L/32
SGD-SAGE
93.74
81.90
Table 5: Image classification test accuracy. Results with are from Dosovitskiy et al. (2020). ViT- B/32 and ViT-L/32 each denotes ViT-base and ViT-large model with 32 x 32 input patch size.
hhnnt://ghhhhh.m///mmmammt com/namisan/mttdmm 'The dev results on RoBERTa-large are averaged over different random seeds. All results have passed paired student t-test with p-values less than 0.05. The detailed statistics summarized in Appendix A.1.3
5

---

Published as a conference paper at ICLR 2022
4.2 NEURAL MACHINE TRANSLATION
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 al. 2015 and WMT 16 En-De (Bojar et al. 2016 IWSLT" 14 De-En is low-resource dataset, which contains 160K sentence pairs. WMT 16 En-De is rich-resource dataset, which contains 4.5M sentence pairs. Dataset and pre-processing details are deferred to Appendix [A.2.]
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 employ Adam (Kingma & Ba 2014) as the optimizer with 3 = (0.9, 0.98). We apply SAGE to the same setting.
We select learning rates in range of {5,7) 10 5 1 (1.2) x 10-4 and select 3n 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 ceined-frrm-scceee models both low-resource and rich-resource settings.
4.3 IMAGE CLASSIFICATION
Model and data. We evaluate SAGE using Vision Transformer models (ViT) the CIFARI00 (Krizhevsky et al. 2009) and ILSVRC-2012 ImageNet dataset (Deng et al. (009). Specifically, we evaluate the -tnningning performance of the ViT-base and ViT-large pre-trained using ImageNet-2IK. a superset of ImageNet dataset with 21 classes and 14M images. Data and pre-processing details deferred to Appendix A.3.
Implementation details. All experiments follow the suggested training configuration of Dosovitskiy al. (2020) and jax-implemented code base [6] We adopt SGD the baseline optimizer with a momentum factor 0.9. We fine-tune the models for 100K steps for CIFAR 100, and 200K steps for ImageNet. We select learning rates in range of {0.02, 0.05, 0.08, 0.1} and select (9 in range of {0.85, 0.90, 0.95J. Comprehensive training details are reported in Appendix A.3.2
Main results. Table 5 shows the evaluation results on CIFAR 100 and ImageNet. SAGE outperforms baselines by 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).
ANALYSIS
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
5.1 SAGE LEADS TO MORE SUFFICIENT TRAINING
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 shows the sensitivity distribution of parameters in the SAGE optimized models and the baseline models. We select the 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.
We observe that the sensitivity distribution exhibits 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.
"htts::/witt.fkkkcc "htps://gituub..o//oogggeee httptr///thia..ooooooooorererres
16/ranslation--tskk

---

Published as a conference paper at ICLR 2022
Distribution of Sensitivity
Schedule Adam
1.5
1.5
Adan SAGE
1.0-
0.5
0.0
MNLI
QNLI
SST-2
MRPC
Figure 2: The sensitivity distribution of the BERT-base models fine-tuned on GLUE tasks. Note that we drop some outliers to ease visualization.
Even the most redundant parameters contribute to the model performance. Recall that sensitivity is type of importance score in pruning, which is 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 40% parameters with withh 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 Lower shows the change in generalization
of the pruned models. To ease the comparison, Figure 3
performance with respect to the un-pruned models.
SST-2
MRPC
RTE
Schedule
96-
80
Adam
90.0
Adam-SAGi
75
87.5
92-
$5.0
90-
65-
85-
100
90
70
100
90
60
100
90
70
Pereentage of Weight Remaining (%)
Percentage of W Wiight Remaining (%6
Porcentage Weight Remaining (%)
24
5
2.5
Change in Accumiyy
0-
0.0
0-
2.5
5-
-4
-10
-6-
-10.0
-15
100
100
90
80
70
100
70
Percentage Weight Remmining (%)
Percentage of Weight Remaining (%)
Percentage Weight Remaining (%6
Figure 3: Upper: Model generalization performance at different pruning ratios; Lower: Change in generalization performance with respect to the full model. Pruning is conducted on the fine-tuned BERT-base models.
We have the following observations:
The pruning performance of the SAGE optimized models remains higher than that of the baseline models (Figure B Upper).
Even the most redundant parameters in the SAGE optimized models makes contributions (Figure 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.
Sensitivity is 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 the baselines. We track the variation for all parameters in the BERT-base model at each iteration, and
"Embedding weights are excluded.
7

---

Published as a conference paper at ICLR 2022
evaluate the variation based the current mini-bater of training data. See Appendix A.4.2 for implementation details.
Local Tempora) Variation
Schedale
Adbm-AAGE
30
10.
MNLI
QNLI
Local Temporal Vaciation
0.0
0.5
1G
0.5
2.0
2.5
SST-2 Number of Training Updates (K)
MRPC
Number of Training Updates (K)
Figure The local temporal variation of sensitivity (with 30 0.7) during training.
We observe that the local temporal variation SAGE remains lower or decreases faster than in the baselines for all tasks. For example, the variation in the baseline approach remains large in QNLL In contrast, the variation in SAGE decreases, suggesting the sensitivity indeed stabilizes and becomes a reliable indicator of redundancy.
5.2 SAGE LEADS TO BETTER GENERALIZATION PERFORMANCE
We verify that SAGE leads to better generalization performance through inspecting the learning curves, decision boundary and hyper-parameter search space.
Learning Curves. Figure 0 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 s valida- tion loss descends faster and SAGE is less prone overfitting. This observation suggests that SAGE has a regularization effect and reduces the model vari- ance. 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 D that redundancy elimination can lead to better generalization.
Adam
Adam-SAGE
Figure 5: Decision boundary predicted on the Spiral dataset. The white curve on Adam- SAGE corresponds the decision boundary of Adam.
0.5
Schadole
0.5
0.08
Adm SAGE
Adam
Validation Loss
Looming Rate
0.07-
Trining
0.3
6.4
Sauiiviity
0.06
0.2
0.0001
5000
5000
Training Step
Trining Step
Trinnnn Step
Training Stop
Figure Learning curves obtained by fine-tuning BERT-base on SST-2 dataset.
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 Es and SAGE's Sys. We can observe that SAGE consistently achieves a better generalization performance within a larger region of seper-parameter search space under different Bns. We also provide param-er study for more datasets in Appendix A.4.5
8

---

Published as a conference paper at ICLR 2022
091
3.0 0..0900
0.8
6.75
63
60
50
60
40
044
15
125
1o-5
Ie55
1o-5
fe-S
Learning Kate
Learning Adan-SAGE,
Learning Adam-SAGE,
Adum
Adam-SAGE; 8%
Figure 7: Validation accuracy obtained by fine-tuning BERT-base on RTE dataset with wide range of tlyper-praamters
Decision Boundary. 3 the decision boundary predicted with Adam and SAGE on the Spiral dataset. Specifically, train multi-layer perceptron with 3 hidden layers, each with hidden dimension of 100. The decision boundary predicted with SAGE is smoother and has larger margin than with Adam, suggesting SAGE produces better generalized model.
5.3 COMBINE WITH STATE-OF-THIE-AS METHODS
We further show that SAGE complementary to existing state-of-the-art regularization methods. Specifically, we apply SAGE to SMART (Jiang et al. 2019). state-of-the-ar smoothness-inde -indusns---ninnnggg adversarial regularization method. As shown in Table 6 SAGE can further improve upon SMART, suggesting the two techniques are complementary.
Model
Optimizer
RTE
MRPC
CoLA
SST-2
STS-B
QNLI
QQP
MNLI-m/mm
Average Score
BERTBAS DASE
Adamax
69.2
86.2/90.4
57.8
92.9
89.7/89.2
91.2
90.9/88.0
84.5/84.4
82.8
SMARTBAS:
Adamax
72.5
87.7/91.4
59.5
93.5
90.0/99.6
91.9
91.7/88.9
85.2/85.7
84.1
SAGE
Table 6: Single task fine-tuning dev results on GLUE
DISCUSSION
SAGE complementary 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 improve optimization by adapting to the optimization landscape, while SAGE aims to improve generalization by eliminating the weight redundancy. The quantities our interest (i.e., Eq. (2) and Eq. (3) are related to the weight redundancy. They are not directly related 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 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 be satisfactory, SAGE provides them sufficient training so that they can learn to contribute.
7 CONCLUSION
We begin with hypothesis that the redundant parameters can become useful if they 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 wide range of tasks and strengthen various types of optimizers.

---

Published as a conference paper at ICLR 2022
REFERENCES
Roy Bar-Haim, Ido Dagan, Bill Dolan, Lisa Ferro, and Danilo Giampiccolo. The second PASCAL recognising textual entailment challenge. In Proceedings of the Second PASCAL Challenges Workshop on Recognising Textual Entailment, 01 2006.
Brian R Bartoldson, Ari S Morcos, Adrian Barbu, and Gordon Erlebacher. The generalization-stablliy tradeoff in neural network pruning. arXiv preprint arXiv: 0et.:/996.0378 2019.
Luisa Bentivogli, Ido Dagan, Hoa Trang Dang, Danilo Giampiccolo, and Bernardo Magnini. The fifth pascal recognizing textual entailment challenge. In In Proc Text Analysis Conference (TAC 2009.
Ondřej Bojar, Rajen Chatterjee, Christian Federmann, Yvette Graham, Barry Haddow, Matthias Huck, Antonio Jimeno Yepes, Philipp Koehn, Varvara Logacheva, Christof Monz, et al. Findings of the 2016 conference machine translation. In Proceedings of the First Conference on Machine Translation: Volume 2. Shared Task Papers, pp. 131-198, 2016.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14 141,5, 2020.
Daniel Cer, Mona Diab, Eneko Agirre, lñigo Lopez-Gazpio, and Lucia Specia. Semeval-2017 task 1: Semantic textual similarity multilingual and crosslingual focused evaluation. In Proceedings of the 11th International Workshop on Semantic Evaluation (SemEval-2017), pp. 1-14, 2017.
Mauro Cettolo, Jan Nichues, Sebastian Stüker, Luisa Bentivogli, Roldano Cattoni, and Marcello Federico. The iwslt 2015 evaluation campaign. In /WSLT 2015, International Workshop on Spoken Language Translation, 2015.
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. The lottery ticket hypothesis for pre-trained bert networks. arXiv preprint arXiv:2007.1222. 2020.
Ido Dagan, Oren Glickman, and Bernardo Magnini. The pascal recognising textual entailment challenge. In Proceedings of the First International Conference on Machine Learning Chal- lenges: Evaluating Predictive Uncertainty Visual Object Classification, and Recognizing Textual Entailment, MLCW'OS 05, pp. 177-190, Berlin, Heidelberg, 2006. Springer- Verlag. ISBN 3-540- 33427-0, 978-3-540-33427-99
doi: 10.1797/11336799 9. URL http: www.xxxdoo....
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810. 04805, 2018.
Xiaohan Ding, Guiguang Ding, Xiangxin Zhou, Yuchen Guo, Jungong Han, and Ji Liu. Global sparse momentum sgd for pruning very deep neural networks. arXiv preprint arXiv: 12778, 2019.
William B Dolan and Chris Brockett. Automatically constructing a corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (/WP2005), 2005.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint 1rXiv:2010.11929, 2020.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: 1909. 1556, 2019.
10

---

Published as a conference paper at ICLR 2022
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010001412, 2020.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural
networks. arXiv preprint arXiv: a 000118033..33 2018.
Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan. The third PASCAL recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL Workshop Textual Entailment and Paraphrasing, pp. 1-9, Prague, June 2007. Association for Computational Linguistics. URL http://wwwwc.. www www.wwwwwwbb................wwww http//aahoooooooooooaaaa
https:
Peter D Grünwald and Abhijit Grunwald. The minimum description length principle. 2007.
Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding-enhanee bert with disentangled attention. arXiv preprint N.XYv:2006.03654 2020.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on, 14(8), 2012.
Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. Dynabert: Dynamic bert with adaptive width and depth. arXiv preprini arXiv:2004.04000 04037, 2020.
Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao. Smart: Robust and efficient fine-tuning for pre-trained natural language models through principled regu- larized optimization. arXiv preprint
2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun, John S Denker, and Sara A Solla. Optimal brain damage. In Advances in neural information processing systems, pp. 598-6098 605, 1990.
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr. Snip: Single-shot network pruning based connection sensitivity. arXiv preprint arXiv: arXiv:/810.00 1810.02340, 2018
Chen Liang, Simiao Zuo, Minshuo Chen, Haoming Jiang, Xiaodong Liu, Pengcheng He, Tuo Zhao, and Weizhu Chen. Super tickets in pre-trained language models: From model compression to improving generalization. arXiv preprint arXiv:2 12002, 2021.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning and beyond. arXiv preprini arXiv: 1e0e:990800.000 03265, 2019a.
Xiaodong Liu, Yu Wang, Jianshu Ji, Hao Cheng, Xueyun Zhu, Emmanuel Awa, Pengcheng He, Weizhu Chen, Hoifung Poon, Guihong Cao, et al. The microsoft toolkit of multi-task deep neural networks for natural language understanding. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics: System Demonstrations, pp. 118-122 126, 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 1692, 2019b.
approach. arXiv preprint
Zhuang Liu, Mingjie Sun, Tinghui Zhou, Gao Huang, and Trevor Darrell. Rethinking the value of network pruning, arXiv preprint 1rXiv:1810.05270 2018
Ilya Loshchilov and Frank Hutter. Fixing weight decay regularization in adam. 2018.
Ekdeep Singh Lubana and Robert P Dick. A gradient flow framework for analyzing network pruning. arXiv preprint arXiv:2009. 1839, 2020.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? arXiv preprint arXiv: 1905. 10650, 2019.
11

---

Published as a conference paper at ICLR 2022
Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convolutional neural networks for resource efficient inference. arXiv preprint arXiv:1611.06440 2016.
Pavlo Molchanov, Arun Mallya, Stephen Tyree, luri Frosio, and Jan Kautz. Importance estimation for neural network pruning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 1224-112722 2019.
Ari S Morcos, Haonan Yu, Michela Paganini, and Yuandong Tian. One ticket to win them all: general- izing lottery ticket initializations across datasets and optimizers. arXiv preprint arXiv: 1906.0277.. 2019.
Michael C Mozer and Paul Smolensky. Skelctonization: A technique for trimming the fat from network via relevance assessment. In Advances in neural information processing systems, pp. 107-115, 1989.
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. arXiv preprint arXiv: 1806.00/87 2018.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier. and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. Proceedings of NAACI-HHLT 2019 Demonstrations, 2019.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Pruceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 2383 2392. Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10. 1/L116663///DD1----1 1264. URL hwtp::/wwwwwwwww....... org/
Stthhology/DD6-1255.
Carl Edward Rasmussen and Zoubin Ghahramani. Occam's razor. Advances in neural information processing systems, pp. 294-300, 2001.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp.
1951.
Victor Sanh, Thomas Wolf, and Alexander M Rush. Movement pruning: Adaptive sparsity by fine-tuning, arXiv preprint arXiv:2005.0768. 2020.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Edinburgh neural machine translation systems for wmt 16. arXiv preprint arXiv: /10.6666662.., 2016.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language processing, pp. 1631-1642, 2013.
Lucas Theis, Iryna Korshunova, Alykhan Tejani, and Ferene Huszár. Faster gaze prediction with dense networks and fisher pruning. arXiv preprint 180/.000000077// 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint
2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint 1rXiv:1905.09418, 2019.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 2018.
Ziheng Wang. Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. arXiv preprint 110.0039. 2019.
Alex Warstadt, Amanpreet Singh, and Samuel Bowman. Neural network acceptability judgments. Transactions of the Association for Computational Linguistics, 7:625-641, 2019.
12

---

Published as a conference paper at ICLR 2022
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technolo- gies, Volume (Lon Papers), pp. 1112-1122, Association for Computational Linguisties, 2018. URL httttttp::/ac: web. org. wwt.boorg/annhhoooor//llll1-111-
Xia Xiao, Zigeng Wang, and Sanguthevar Rajasekaran. Autoprune: Automatic network pruning by regularizing auxiliary parameters. Advances in neural information processing systems, 32, 2019.
Matthew D Zeiler. Adadelta: adaptive learning rate method. arXiv preprint arXiv:1212.5701
2012.
13

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1
NATURAL LANGUAGE UNDERSTANDING
A.1.1 DATA
GLUE is a collection of nine NLU tasks. The benchmark includes question answering Raaiuurknn al. 2016). linguistic acceptability (CoLA, Warstadt et al. 2019). sentiment analysis (SST, Socher
(ee et al. 2013), text similarity (STS-B, Cer et al. 2017), paraphrase detection (MRPC, Dolan & Brockett 2005), ) and natural language inference (RTE Giampiccolo et al. 2007; Bentivogli et al. 2009 Williams et al. 2018) tasks. Details of the GLUE benchmark, including tasks, statistics, and evaluation metrics, are summarized in Table [3
Dolan
MNLI, Dagan al. 2006; Bar-Haim et al. 2006
All the texts were tokenized using wordpieces, and were chopped to spans no longer than 512 tokens.
A.1.2 TRAINING DETAILS
To fine-tune BERT- and RoBERTa-large models on individual tasks, we append a task-specific fully-comnected classification layer them as Devlin
(2018).
Table 7 present the hyper-parameter parameter configurations. We tune this set of hyper- prrameprrameters on a single seed, and report the averaged results obtained with the same configuration over all seeds. For SAGE experiments, We slightly tune Bo within range of 0.1 on different seeds. We apply a linear weight decay rate of 0.01 and gradient norm clipping threshold of for experiments. All experiments are conducted on Nvidia VI00 GPUs.
Hyper-param
Experiment
RTE
MRPC CoLA
SST-2
STS-B
QNLI
QQP
MNLI
BERT BASE, Adam
le-5
le-5
le-5
1c-5
le-5
Ic-5
2c-5
26-5
BERT Adam-SAGE
lc-4
8c-5
8c-5
3c-5
lc-4
8c-5
4c-5
5c-5
BERTRASE, Adamax
lc-4
lc-4
lc-4
5c-5
lc-4
lc-4
lc-4
80-5
Learning Rate
BERTRASE, Adamax-SAGE
3e-4
3e-4
2e-4
2e-4
5e-4
Se-4
3e-4
2e-4
RoBERTaLARGE. Adamax
5e-5
5e-5
3e-5
1e-5
Ie-5
le-4
1e-5
RoBERTaLARGEE Adamax-SAGE
6e-5
2e-4
8e-5
2e-5
8e-5
3e-5
2e-4
8e-5
BERT BASE. Adam-SAGE
0.60
0.80
0.70
0.80
0.60
0.70
0.75
0.70
Po
BERTRASE, Adamax SAGE
Adamax-SAGE
0.65
0,80
0.75
0.70
0.75
0.70
0.75
0.85
RoBERTaLARGE, Adamax-SAGH
0.75
0.65
0.70
0.75
0.80
0.80
0.65
0.60
BERT
16
8
32
32
32
32
32
32
Batch Size
RoBERTa
16
32
32
32
32
32
32
BERT BASE
6
6
6
3
6
3
Epoch
ROBERTALARGE
15
6
10
15
3
BERT
0.1
0.1
0.1
0.1
0.1
0.1
0.0
0.3
Dropout
ROBERTaLARGE
0.1
0.1
0.1
0.1
0.0
0.3
BERT BASE
0.1
0.1
0.1
0.1
0.1
0.1
0.0
0.1
Warmup
RoBERTa
0.1
0.1
0.1
0.1
0.1
0.1
Table 7: Hyper-parameter configurations for GLUE experiments. "Epoch" refers to the total training epochs; we adopt early-stopping strategy practice. "Dropout" refers to classification layer dropout ratio. "Warmup" refers to the ratio of learning rate linear warmup iterations to total training iterations.
A.1.3 EVALUATION RESULTS
Statistics of the dev set results. Table shows the standard deviation of the dev set results.
Average score computation formula. For dev set results, we first obtain a score for each task by averaging the scores of all metrics (e.g., Ace and F1) and test sets (e.g.. MNLI-m and MNLI-mm) within this task, then compute a task-average score. For test set results, we directly averages scores of all reported metrics following Devlin et al. (2018).
14

---

Published as a conference paper at ICLR 2022
Model
Optimizer
RTE
MRPC
CoLA
SST-2
STS-B
QNLI
QQP
MNLI
Adam-SAGE
0.35
0.32
0.85
0.25
0.12
0.06
0.05
0.06
BERT
Adamax-SAGi
0.56
0.69
0.12
0.23
0.03
0.06
0.08
0.10
RoBERTaLARGE
Adamax-SAGE
0.51
0.78
0.50
0.19
0.08
0.00
0.05
0.05
Table 8: Standard deviation of the dev set results.
A.2 NEURAL MACHINE TRANSLATION
A.2.1 DATA
Table shows the number of sentence pairs in each dataset. We use the standard newstest-2013 and newstest-2014 as dev and test set for WMT 16 En-De. We follow Ott et al. (2019) to split the dev/test. sets for IWSLT" 14 De-En.
All datasets are encoded using byte-pair encoding (BPE, (Sennrich et al. (2016)) We preprocess IWSLT' 14 De-En data following fairsed and adopt the preprocessed WMT En-De from Google
Data
Train
Dev
Test
IWSLT' De-En
160K
7283
6750
WM En-De
4.5M
1061 1019
Table 9: The number of parallel sentences in NMT datasets.
A.2.2 TRAINING DETAILS
We adopt the Transformer- base model for both datasets. For IWSLT' 14 De-En, we share the decoder and encoder output embeddings. For WMT' 16 En-De, we share all the embeddings.
Table 10 presents the hyper-parameter configurations for the best models. We apply a linear weight decay rate of 1 10 and label smoothing ratio of 0.1 for all experiments. All experiments are conducted on Nvidia V100 GPUs.
For IWSLT' 14 De-En, we report the BLEU score of the best checkpoint using beam size of 5 and length penalty of 1. For WMT" 16 En- En-De, we report the average of the last 10 checkpoints with beam size of and length penalty of 0.6.
Hyper-param
Experiment
1WSLT'14 De-En
WMT'16 En-De
Adam
Se-4
7e-4
Learning Rate
Adam-SAGE
le-3
2e-3
So
Adam-
0.8
0.4
Batch size
Both
4096
32768
Epoch
Both
60
40
Dropout
Both
Warmup
Both
8000
4000
Table 10: Hyper- parameter configurations for NMT experiments. "Warmup" refers to the learning
rate linear warmup iterations.
15

---

Published as a conference paper at ICLR 2022
A.3
IMAGE CLASSIFICATION
A.3.1
DATA
For CIFAR100, we apply random cropping and random horizontal flipping to the training data.
A.3.2
TRAINING DETAILS
Table 11 present the hyper-parameter configurations for the best models. All experiments
conducted on Nvidia V100 GPUs.
Hyper-param
Experiment
CIFARI00
ImageNet
ViT-B/32, SGD-SAGE
0.02
0.05
Learning Rate
ViT-L/32, SGD-SAGE
0.02
0.08
ViT-B/32, SGD-SAGE
0.95
0.95
VIT-L/32, SGD-SAGE
0.85
0.95
Training Steps
All
10000
20000
Dropout
All
0.0
0.0
Table 11: Hlyee-paarameter configurations for ViT experiments on CIFAR100 and ImageNet.
A.4
SUPPLEMENTS FOR METHOD AND ANALYSIS
A.4.1 ADAM-SAGE ALGORITHM
Algorithm Adam-SAGH (3 denotes Hadamard product and denotes Hadamard division)
Input: Model parameters e E RR;; Data D; Learning rate schedule 77-); Total training iteration T; Moving average coefficient Sg. $2. 1: Initialize 7 (0) (0) = 0 R-
2: for 1, ..... 7 do
3: Sample a minibatch ((() from DD
4: Compute gradient (t) = ( ((((t) (()).
5: Compute sensitivity [(t] = 10(1) @g)))
((t) = (tm (1 - -1))
6:
w( 12r(1-1) + (1 - (0)
7:
7t) = f0 + (1 - (0)
8:
9:
(( = (L)) 101011111 31)
10:
T( /(0/(11 32)
11:
7 710/(1 Bo)
(T()) [7(0) Í(0].
13: Update ((+1) ((1) - ((0)) @ ((0 +c) (Vot) + c)) c g()
14: end for
A.4.2 IMPLEMENTATION DETAILS FOR SECTION 5.1
Figure 2 experiments: Due to the extremely large model size, we only sample 110K parameters per layer (in total 12 x 110K parameters) to calculate the distribution. 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.
Figure 4 experiments: Following previous experiment's practice, we randomly sample 110K parame- ters per layer (in total 12 110K parameters), and for visualization purposes, plot 60 randomly selected iterations. We adopt the learning rate corresponding to the best training performance for both SAGE and the baselines.
16

---

Published as a conference paper at ICLR 2022
A.4.3 IMPLEMENTATION DETAILS FOR SECTION 5.2
Plotting the parameter sensitivity distribution throughout training can be computational expensive. The distribution varies significantly throughout training and often fails to provide a meaningful visualisation. As result, we compute the structured sensitivity score instead of the parameter sensitivity score. Specifically, we compute single sensitivity score for each Transformer weight block O iteration / using the structured counterpart of the parameter sensitivity metric widely adopted in the existing structured pruning literature (Michel al. 2019 Liang et al. (02T). Following common structured pruning practice, we split Transformer models into 12 feed-forward weight modules and 12 multi-head attention weight modules, and plot the average and variance of the sensitivity of these modules' sensitivity scores throughout the training.
We present the results obtained with the hyper-parameter that yield the best generalization perfor- mance on the BERT-base model for both Adamax (Baseline) and Mlamax-SAGE (SAGE).
A.4.4 ABLATION STUDY
To further interpret the role of the parameter sensitivity and the local temporal variation U, conduct an ablation study on these two factors. Specifically, we check five variants of Eq. (4)
Variant 1.
(1) // (T,) + +0) +e)
Variant 2.
(1) (0,0 +)/(U,, +6)
Variant 3. My
(1) ,70(7))) +e)
Variant 4.
7i (1) ,/01/(0) +c)
Variant 5. ni (1)0 +c)
For Variants 1,2 and 3, we aim to check the performance of giving a high/low-sensitive parameter a high/low, instead of low/high learning rate. Specifically, we place (i7 e) in the numerator, that the learning rates increase for the high sensitive parameters and decrease for low sensitive parameters.
For Variants 4 and 5, aim to check the performance of eliminating the influence of one of these factors. Specifically, we fix the local temporal variation term to in Variant and fix the sensitivity term in Variant 5.
A.4.5 HYPER-PARE PARAMETER STUDY
We investigate the influence of hyper- parameters learning rate and 6g on the performance of SAGE (Figure [8. As can be seen, SAGE requires a larger learning rate than the baselines to offset the small scale of the modulation term (the optimal baseline learning rate lies in 5 x 10-5 ~ 1 x 10-4 for MNLI, 10-4 7 x 10-1 for IWSLT 14 De-En and 0.1 ~ 0.2 for CIFARIO). Furthermore, switching to larger learning rate requires lower Bo to maintain the same level of performance.
35.1
92.5
85.2
92.0
85.0
34.9
Ace
91.5
Leurning Rate
34.8
0.3
84.8
Learning Rate
Leaming Rate Se-1
0.4
2e-4
0.5
84.6
0.6
0.8
34.6
0.76
0.78
1o-3
90.0.3
0.4
0.7
.0 MNLI
3h
IWSLT 14 De-Fi
CIFAR 10
Figure Parameter study on learning rate and fg-
All five variants show no clear gain upon the baseline on both RTE and SST-2 datasets after careful dyper-parameted tuning. Specifically, we observe that the Variants and 3 converge very fast the early stage of training, and then quickly start overfitting. In Variants and 4, the training collapses due to gradient explosion vanishing.
17

---

Published as a conference paper at ICLR 2022
Variant Name
Learning Rate Modulating Term
RTE
SST-2
Adam
63.5
92.9
Adam-SAGE
(0 + +c)1(150) +c)
73.3
93.5
Variant 1.
(T. -c)((()) +6)
63.5
91.2
(T()) +c)/(U, +c)
Variant
Unconverged Unconverged
7/0. +e
Variant
63.8
91.1
Variant
1/(T50) +6)
Unconverged
Unconverged
Uf"0+ +e
Variant 5.
63.8
91.1
Table 12: Ablation study on parameter sensitivity and local temporal variations.
Corpus
Task
#Train
#Dev
#Test
#Label
Metries
Single- Sentence Classification (GLUE)
CoLA
Acceptability
8.5k
lk
lk
2
Matthews corr
SST
Sentiment
67k
872
1.8k
Accuracy
Pairwise Text Classification (GLUE)
MNLI
NLI
393k
20k
20k
3
Accuracy
RTE
NLI
2.5k
276
3k
2
Accuracy
QQP
Paraphrase
364k
40k
391k
2
Accuracy/Fi
MRPC
Paraphrase
3.7k
408
1.7k
Accuracy/Fi
QNLI
QA/NLI
108k
5.7k
5.7k
2
Accuracy
Text Similarity (GLUE)
STS-B
Similarity
7k
1.5k
1.4k
Pearson/Spearman corr
Table 13: Summary of the GLUE benchmark.
18


Reviews:
[{"review_id": "wSrrmW5kM6d", "reviewer": "Reviewer_4pzE", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Strengths:\n\n- The idea is novel, straightforward, well motivated, and easy to implement. Computational cost is low.\n- Potential for large impact on model training best practice.\n- The experiments are very thorough, demonstrating gains across different settings, and showing that the technique achieves its goal.\n- The paper is extremely clearly and carefully written.\n\nWeaknesses:\n\n- The scaling formula seems somewhat ad hoc and hard to characterize. In particular, if the sensitivity of some parameter spikes on a given iteration, it will get a larger update than another parameter with the same moving-average sensitivity that has not spiked.\n- There are some hints that the gains might be partly related to regularization (eg, better results on IWSLT than WMT). It would have been good to test the combination with dropout or something similar.\n- The hyper-parameters for SAGE are exhaustively tuned, more so than for the baseline adaptive optimizers, so there’s a potential for bias. Figure 7 counters this possibility, but only by showing heatmaps, so it's hard to gauge the actual numbers.\n\nDetails:\n\n- Table 3 should show results for your implementation of Adamax that’s comparable to Adamax-SAGE, in addition to the Devlin et al numbers.\n- Figure 2: Consider flipping these plots to the standard convention, and stacking them.\n- Figure 4: Why do we care that local temporal variation is lower in SAGE than the baselines?", "summary_of_the_paper": "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.\n", "summary_of_the_review": "A simple adaptive learning-rate formula that seems to work really well, even on top of other optimizers, as demonstrated by very thorough experiments. Downsides are that the formula isn't particularly well justified, and there are a few potential experimental weaknesses.", "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": "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": "wCQ3Ftc1r2W", "reviewer": "Reviewer_sSHP", "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": "Strengths:\n\n- **Writing:** The paper is well written and easy to follow.\n- **Tasks:** The paper evaluates its method with different optimisers and on different tasks in NLP and image classification.\n- **Results:**  The results are quite good; the proposed method surpasses the baseline each time.\n\nWeakness:\n\n- **Architecture:**  The method is evaluated only with transformer architectures. In the context of image classification it would be interesting to evaluate it with CNN and others architecture than ViT.\n\n- **Optimisation:** In image classification, the pre-training procedure used is quite sub-optimal since the paper of Dosovitsky et al. [1] many improvements have been proposed such as the DeiT approach [2]. It would be interesting to see if the proposed method still works when the model is trained with more regularisation and data-augmentation which may also lead to a better use of weights.\n\n- **Image Classification:** Having only results of models pre-trained on ImageNet-21k and fine-tuned with the proposed method on downstream tasks is not very usual in image classification. It would be interesting to have results on ImageNet only where the SAGE method is used during the training.\n\n[1] Dosovitskiy et al, An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, ICLR 2021\n\n[2] Touvron et al, Training data-efficient image transformers & distillation through attention, ICML 2021", "summary_of_the_paper": "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.\n", "summary_of_the_review": "The idea of the paper is interesting and the proposed method seems effective. Nevertheless, more complete experiments in image classification would allow to better evaluate the interest of the proposed method. ", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "qV-wX3wFx1H", "reviewer": "Reviewer_MzBV", "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": "I think that overall the paper is good as it raises and studies an important point: redundancy in parameters is not an axiom we have to accept.\nHowever, after this introduction and motivation, it reads more like a typical optimizer paper introducing a new optimizer based on hand-wavy intuitions. This is unfortunate and not rigorous.\n\nWhat follows are detailed comments:\n1) The use of Taylor approximation is only good close to the operating point for nonlinear functions. However Theta_j,-j may be very large, and thus the approximation may be very bad.\n2) I do not agree that the memory and computational costs of SAGE are \"marginal\". The EMA \"I-hat\" is a full copy of the model since we need one EMA per parameter, especially for large models this is substantial overhead. I would further appreciate timings to verify that indeed the computational overhead is \"marginal\".\n3) It is not clear to me in what scale the quantities U and I-hat are, and hence, the multiplier to the learning-rate. Because of this, it is not clear how one needs to change the original learning-rate \"eta\" when one turns on SAGE.\n4) I am confused by Fig5, it does not look like the two-moon dataset at all to me? Please image-search \"two moon dataset\" and compare.\n\nminor nitpicks:\n5) Fig1: Percent should be [0,100].\n6) I do not agree that one can conclude from the results that SAGE is more effective for small datasets than for large ones, since the \"points\" betwen the datasets/tasks do not live on the same scale.\n7) All experimental evaluation is about fine-tuning of already pre-trained transformer models. This should be more accurately reflected in the title, for example by replacing the word \"training\" by \"fine-tuning\". It would further be interesting to see how well it works when training from scratch; even if it does not work, it is still a valuable fine-tuning method, and stating this may save a lot of people a lot of resources and time.\n8) In Figure 6, since we have per-parameter learning-rates, which one is shown for SAGE?\n9) Typo: \"smoothier\" -> \"smoother\"", "summary_of_the_paper": "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.", "summary_of_the_review": "For an \"optimizer\" paper, it is weak: hand-wavy motivation, no real derivation of update rule or even convergence proofs, and experiments only in one very specific domain: transfer of pre-trained transformer models. I am not convinced that it will be generally useful at all.\n\nHowever, the paper raises a very important point: we should not take redundant parameters as a necessity, and it does propose a method to avoid them. I find this point important enough to still suggest acceptance.", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "--qTFCQKsQr", "reviewer": "Reviewer_NSqH", "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": "Strengths\n1. The paper is well-written and easy to follow. The method section provides helpful intuitions behind the algorithm design.\n2. The idea is interesting since it explores a different direction in dealing with redundant parameters.\n3. Experiments use multiple benchmarks, including both language and vision, and results show noticeable performance improvements.\n4. SAGE is orthogonal to the existing adaptive gradient methods. Jointly using them can bring more gains.\n\nWeaknesses\n1. Table 1 shows that models with different percentages of redundant parameters have similar performance, implying that performance may not be proportional to \"well-trained\" parameters. It means that making more parameters \"well-trained\" does not necessarily improve the performance. This seems not totally in line with the paper's motivation.\n2. Figure 5 is not straightforward to visualize the difference. The curve is drawn only on the right subfigure.\n3. Sufficiently training all parameters seems a double-edged sword. According to Figure 3, models trained with SAGE are susceptible to parameter pruning. In general, we want efficient and compact models in deployment. Will SAGE be harmful to developing efficient deployment models?", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper proposes SAGE, an adaptive learning rate schedule to train redundant parameters more sufficiently to improve model generalization. SAGE, together with existing adaptive optimizers, shows effectiveness in a wide range of downstream tasks. The main concern is whether SAGE has adverse effects in getting efficient deployment models.", "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 paper observes that the number of redundant parameters is a function of the training procedure and proposes a training strategy that encourages all parameters in the model to be trained sufficiently and become useful. The method adaptively adjusts the learning rate for each individual parameter according to its sensitivity (a proxy for the parameter's contribution to the model performance). The approach encourages the use of under fitted parameters while preventing overfitting in the well-fitted ones. Experimental results are presented covering a wide range of tasks and in combination with several optimisers, showing improvements in model generalization.

The paper is very well written and easy to follow (as mentioned by Reviewers NSqH, 4pzE and sSHP). 

The authors provided a strong rebuttal including new experiments, like training using CNN based architectures (as requested by Reviewers sSHP and MzBV). Reviewer sSHP requested these results to be reported with STD, the AC encourages the authors to do so for the camera ready.

Reviewer MzBV points out that the paper could be improved by giving a motivation of the update rule and proving convergence. However, still recommends accepting the paper due to the novelty in the idea of not taking redundant parameters as something inevitable and devising an effective strategy to improve it. This idea was also appreciated by the other reviewers. While the AC agrees that adding these points would improve the work, it takes as valid the point made by the authors. Namely, that the intuition behind the update rule is quite clear, and many other reasonable variants were ablated (in Appendix A.4.4). Furthermore, the empirical evidence shows that the method improves generalization.

Reviewer NSqH points out that while SAGE improves the model’s generalization performance for lightly compressed models, its performance becomes more susceptible to pruning when the model is compressed heavier. While the authors responded with good points, the AC encourages them to follow the reviewer’s advice and incorporate further experiments studying this issue (e.g. other datasets).

In sum, the paper proposes a simple and effective method that is able to improve generalization of large scale models. All four reviewers recommend accepting the paper. The AC agrees and encourages the authors to incorporate the requests mentioned above.
Decision: Accept (Poster)