Title: On the Overlooked Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm Perspective
Venue: NeurIPS 2023
Accepted: yes
Keywords: Weight Decay

Regularization

Optimization

Deep Learning
Abstract: Weight decay is a simple yet powerful regularization technique that has been very widely used in training of deep neural networks (DNNs). While weight decay has attracted much attention, previous studies fail to discover some overlooked pitfalls on large gradient norms resulted by weight decay. In this paper, we discover that, weight decay can unfortunately lead to large gradient norms at the final phase (or the terminated solution) of training, which often indicates bad convergence and poor generalization. To mitigate the gradient-norm-centered pitfalls, we present the first practical scheduler for weight decay, called the Scheduled Weight Decay (SWD) method that can dynamically adjust the weight decay strength according to the gradient norm and significantly penalize large gradient norms during training. Our experiments also support that SWD indeed mitigates large gradient norms and often significantly outperforms the conventional constant weight decay strategy for Adaptive Moment Estimation (Adam).

On the Overlooked Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm Perspective

Abstract

Weight decay is a simple yet powerful regularization technique that has been very widely used in training of deep neural networks (DNNs). While weight decay has attracted much attention, previous studies fail to discover some overlooked pitfalls on large gradient norms resulted by weight decay. In this paper, we discover that, weight decay can unfortunately lead to large gradient norms at the final phase (or the terminated solution) of training, which often indicates bad convergence and poor generalization. To mitigate the gradient-norm-centered pitfalls, we present the first practical scheduler for weight decay, called the Scheduled Weight Decay (SWD) method that can dynamically adjust the weight decay strength according to the gradient norm and significantly penalize large gradient norms during training. Our experiments also support that SWD indeed mitigates large gradient norms and often significantly outperforms the conventional constant weight decay strategy for Adaptive Moment Estimation (Adam).

1 Introduction

Deep learning [25] has achieved great success in numerous fields over the last decade. Weight decay is the most popular regularization technique for training deep neural networks that generalize well [24]. In deep learning, there exist two types of 'weight decay': L 2 regularization and weight decay.

People commonly use the first type of L 2 regularization for training of deep neural networks as θ t = θ t -1 -η t ∇ ( L ( θ t -1 ) + λ L 2 ∥ θ t -1 ∥ 2 ) , where ∇ is the gradient operator, ∥ · ∥ is the L 2 -norm, t is the index of iterations, θ denotes the model weights, η t is the learning rate, L ( θ ) is the training loss, and λ L 2 is the L 2 regularization hyperparameter. In principle, the L 2 regularization strength is chosen to be constant during training of DNNs and people do not need to schedule L 2 regularization like learning rate decay. The second type of weight decay was first proposed by Hanson and Pratt [13] to directly regularize the weight norm:

where λ ′ is the (vanilla) weight decay hyperparameter.

While the two types of 'weight decay' seem similar to Stochastic Gradient Descent (SGD), Loshchilov and Hutter [32] reported that they can make huge differences in adaptive gradient methods, such as Adam [22]. Adam often generalizes worse and finds sharper minima than SGD [40, 56, 46] for popular convolutional neural networks (CNNs), where the flat minima have been argued to be closely related with good generalization [16, 17, 14, 51, 21]. Understanding and bridging the generalization

gap between SGD and Adam have been a hot topic recently. Loshchilov and Hutter [32] argued that directly regularizing the weight norm is more helpful for boosting Adam than L 2 regularization, and proposed Decoupled Weight Decay as

where λ W is the decoupled weight decay hyperparameter. We display Adam with L 2 regularization and Adam with Decoupled Weight Decay (AdamW) in Algorithm 1. To avoid abusing notations, we strictly distinguish L 2 regularization and weight decay in the following discussion.

While fine analysis of weight decay has attracted much attention recently [52, 39, 4], some pitfalls of weight decay have still been largely overlooked. This paper analyzes the currently overlooked pitfalls of (the second-type) weight decay and studies how to mitigate them by a gradient-norm-aware scheduler. The main contributions can be summarized as follows.

First, supported by theoretical and empirical evidence, we discover that, weight decay can lead to large gradient norms especially at the final phase of training. The large gradient norm problem is significant in the presence of scheduled or adaptive learning rates, such as Adam. From the perspective of optimization, a large gradient norm at the final phase of training often leads to poor convergence. From the perspective of generalization, penalizing the gradient norm is regarded as beneficial regularization and may improve generalization [27, 3, 2, 55]. However, previous studies fail to discover that the pitfalls of weight decay may sometimes significantly hurt convergence and generalization.

Second, to the best of our knowledge, we are the first to demonstrate the effectiveness of scheduled weight decay for mitigating large gradient norms and boosting performance. Inspired by the gradientnorm generalization measure and the stability analysis of stationary points, we show that weight decay should be scheduled with the gradient norm for adaptive gradient methods. We design a gradient-norm-aware scheduler for weight decay, called Scheduled Weight Decay (SWD), which significantly boosts the performance of weight decay and bridges the generalization gap between Adam and SGD. Adam with SWD (AdamS) is displayed in Algorithm 2. The empirical results in Table 1 using various DNNs support that SWD often significantly surpasses unscheduled weight decay for Adam.

Structure of this paper. In Section 2, we analyze the overlooked pitfalls of weight decay, namely large gradient norms. In Section 3, we discuss why large gradient norms are undesired. In Section 4, we propose SWD to mitigate the large gradient norm and boosting performance of Adam. In Section 5, we empirically demonstrate the effectiveness of SWD. In Section 6, we conclude our work.

2 Overlooked Pitfalls of Weight Decay

In this section, we report the overlooked pitfalls of weight decay, on convergence, stability, and gradient norms.

Weight decay should be coupled with the learning rate scheduler. The vanilla weight decay described by Hanson and Pratt [13] is given by Equation (1). A more popular implementation for vanilla SGD in modern deep learning libraries, such as PyTorch and TensorFlow, is given by Equation (2), where weight decay is coupled with the learning rate scheduler. While existing studies did not formally study why weight decay should be coupled with the learning rate scheduler, this trick has been adopted by most deep learning libraries.

We first take vanilla SGD as a studied example, where no momentum is involved. It is easy to see that vanilla SGD with L 2 regularization λ 2 ∥ θ ∥ 2 is also given by Equation (2). Suppose the learning rate is fixed in the whole training procedure. Then Equation (2) will be identical to Equation (1) if we simply choose λ ′ = ηλ . However, learning rate decay is quite important for training of deep neural networks. So Equation (2) is not identical to Equation (1) in practice.

Which implementation is better? In Figure 1, we empirically verified that the popular Equation-(2)-based weight decay indeed outperforms the vanilla implementation in Equation (1). We argue that Equation-(2)-based weight decay is theoretically better than Equation-(1)-based weight decay in terms of the stability of the stationary points. While we use the notations of 'stable'/'unstable' stationary points for simplicity, we formally term 'stable'/'unstable' stationary points as strongly/weakly stationary points in Definition 1. This clarifies that there can exist a gap between true updating and pure gradient-based updating for stationary points.

Definition 1 (Strongly/Weakly Stationary Points) . Suppose A : Θ → Θ is a function that updates model parameters during training where Θ is the parameter space. We define θ ⋆ a strongly stationary point for A if ∥ θ ⋆ - A ( θ ⋆ ) ∥ = 0 holds during the whole training process. We define θ ⋆ a weakly stationary point for A if ∥ θ ⋆ -A ( θ ⋆ ) ∥ = 0 holds only for some iterations.

The training objective at stept of Equation (1) is L ( θ ) + λ ′ η t ∥ θ ∥ 2 , while that of Equation (2) is L ( θ ) + λ ∥ θ ∥ 2 . Thus, stationary points of the regularized loss given by Equation (2) are stable during training, while that given by Equation (1) is unstable due to the dynamic training objective.

0

25

50

75

100

125

150

175

200

Epochs

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Test Error

Eq1

=5×10 6

Eq1

=5×10 5

Eq1

=5×10 4

Eq2

=5×10 4

Weight decay increases gradient norms at the final phase of deterministic optimization and stochastic optimization. Following the standard convergence analysis of Gradient Descent (GD) [6] with the extra modification of weight decay, we prove Theorem 1 and show that even GD with vanilla weight decay may not converge to any non-zero stationary point due to missing stability of the stationary points.

Theorem 1 (Non-convergence due to vanilla weight decay and scheduled learning rates) . Suppose learning dynamics is governed by GD with vanilla weight decay (Equation (1) ) and the learning rate η t ∈ (0 , + ∞ ) holds. If ∃ δ such that satisfies 0 < δ ≤ | η t -η t +1 | for any t > 0 , then the learning dynamics cannot converge to any non-zero stationary point satisfying the condition

where f t ( θ ) = L ( θ ) + λ ′ η t ∥ θ ∥ 2 is the regularized loss.

We leave the proof in Appendix A.1. Theorem 1 means that even if the gradient norm is zero at the t -th step, the gradient norm at the t + 1 -th step will not be zero. Thus, the solution does not

0

25

50

75

100

125

150

175

200

Epochs

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

Gradient Norm

(a) AdamW

0

25

50

75

100

125

150

175

200

Epochs

0.000

0.005

0.010

0.015

0.020

0.025

0.030

0.035

0.040

Gradient Norm

Adam WD=1e-3

(b) Adam

0

25

50

75

100

125

150

175

200

Epochs

0.000

0.005

0.010

0.015

0.020

0.025

0.030

0.035

0.040

Gradient Norm

SGD WD=1e-3

SGD WD=5e-4

SGD WD=1e-4

SGD WD=0

(c) SGD

0

25

50

75

100

125

150

175

200

Epochs

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Squared Gradient Norm

(d) AdamW

0

25

50

75

100

125

150

175

200

Epochs

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Squared Gradient Norm

0

25

50

75

100

125

150

175

200

Epochs

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Squared Gradient Norm

(f) SGD

converge to any non-zero stationary point. GD with unstable minima has no theoretical convergence guarantee like GD [41]. More importantly, due to unstable stationary points, the gradient norms (of the regularized/original loss) naturally also do not converge to zero. This may explain why weight decay should be coupled with the learning rate to improve stability of stationary points during training. This non-convergence problem is especially significant when δ is relatively large in the presence of scheduled or adaptive learning rates, such as Adam.

Existing studies [48, 12] proved that the gradient norm bound of SGD in convergence analysis depends on the properties of the training objective function and gradient noise. We further prove Theorem 2 and reveal how the gradient norm bound in convergence analysis depends on the weight decay strength λ when weight decay is involved. We leave the proof in Appendix A.2.

Theorem 2 (Convergence analysis on weight decay) . Assume that L ( θ ) is an L -smooth function 1 , L is lower-bounded as L ( θ ) ≥ L ⋆ , L ( θ, X ) is the loss over one minibatch X , E [ ∇ L ( θ, X ) -∇ L ( θ )] = 0 , E [ ∥∇ L ( θ, X ) - ∇ L ( θ ) ∥ 2 ] ≤ δ 2 , and ∥∇ L ( θ ) ∥ ≤ G for any θ . Let SGD optimize f ( θ ) = L ( θ ) + λ 2 ∥ θ ∥ 2 for t +1 iterations. If η ≤ C √ t +1 , we have

and sup( ∥ θ ∥ ) = sup k =0 ,...,t ∥ θ k ∥ is the maximum L 2 norm of θ over the iterations.

Theorem 1 shows the non-convergence problem with improper weight decay. Theorem 2 further shows that, even with proper convergence guarantees and stable minima, the gradient norm upper bound at convergence still monotonically increases with the weight decay strength λ . Obviously, Theorem 2 theoretically supports that weight decay can result in large gradient norms. The phenomenon can also be widely observed in the following empirically analysis.

where

0

25

50

75

100

125

150

175

200

Epochs

0.000

0.025

0.050

0.075

0.100

0.125

0.150

0.175

0.200

Gradient Norm

(a) VGG16 with BatchNorm

0

25

50

75

100

125

150

175

200

Epochs

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Gradient Norm

(b) VGG16 without BatchNorm

0

25

50

75

100

125

150

175

200

Epochs

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Squared Gradient Norm

(c) VGG16 with BatchNorm

0

25

50

75

100

125

150

175

200

Epochs

0

1

2

3

4

5

6

7

8

Squared Gradient Norm

(d) VGG16 without BatchNorm

Empirical evidence of large gradient norms resulted by weight decay. Theorem 1 suggests that both unstability and strength of weight decay may increase the gradient norm even in full-batch training. While weight decay is helpful for regularizing the model weights' norm, we clearly observed in Figure 2 that weight decay significantly increases the gradient norms in training of DNNs as we expect. Note that the displayed weight decay of AdamW is rescaled by the factor 0 . 001 . The observation generally holds for various optimizers, such as AdamW, Adam and SGD. Figure 3 further shows that the observation holds for both scale-invariant models (VGG16 with BatchNorm) and scale-variant models (VGG16 without BatchNorm). Moreover, the observation still holds for gradient norms with or without regularization included, because the numerical difference due to regularization is smaller than their gradient norms by multiple orders of magnitude in the experiments. Zhang et al. [52] argued that weight decay increases the 'effective learning rate' due to smaller weight norms resulted by weight decay. However, the 'effective learning rate' interpretation may not explain large gradient norms, because a small weight norm does not necessarily indicate a large gradient norm. Interestingly, Figure 3 shows that the gradient curves of VGG16 without BatchNorm and weight decay sometimes exhibit strange spikes, which can be wiped out by weight decay.

3 Large Gradient Norms Are Undesired

In this section, we demonstrate that the large gradient norm is a highly undesirable property in deep learning because large gradient norms correspond to multiple pitfalls.

First, small gradient norms are required for good convergence, which is desired in optimization theory [36]. Popular optimizers also expect clear convergence guarantees [8]. The optimizers with no convergence guarantees may cause undesirable training behaviors in practice.

Second, regularizing the gradient norm has been recognized as an important beneficial regularization effect in recent papers [28, 3, 11, 55, 54]. Li et al. [28] and Geiping et al. [11] expressed implicit regularization of stochastic gradients as

where L ( θ, X ) indicate the loss over one data minibatch X ∈ B , |B| is the number of minibatches per epoch, and G is the expected squared gradient norms over each minibatch. Recent studies, Zhao et al. [55] and Zhang et al. [54], also successfully improved generalization by explicitly adding the gradient-norm penalty into stochastic optimization. Obviously, penalizing the gradient norm is considered as an important and beneficial regularizer in deep learning.

Third, large gradient norms may theoretically indicate poor generalization according to the gradientnorm generalization measures [27, 2]. For example, Theorem 11 of Li et al. [27] derived a gradientnorm-based generalization error upper bound that uses the squared gradient norms, written as

where n is the training data size, σ 2 t is the gradient noise variance at the t -th iteration, and G ( θ t ) = 1 |B| ∑ X ∈B ∥∇ L ( θ, X ) ∥ 2 is the squared gradient norm with the batch size as 1. An et al. [2] presented more empirical evidences for supporting the gradient-norm-based generalization bound (7). Moreover, large gradient norms are also believed to be closely related to minima sharpness [20, 57, 43, 47, 42, 46, 7], while sharp minima often indicate poor generalization [16, 17, 14, 51, 21]. As our analysis focuses on the large gradient norm at the final phase (or the terminated solution) of training, our conclusion is actually not contradicted to the viewpoint [1, 29, 10] that the gradient-noise regularization effect of large initial learning rates at the early phase of training is beneficial.

We also report that the gradient norm and the squared gradient norm show very similar tendencies in our experiments. For simplicity, we mainly use G as the default gradient norm measure in this paper, unless we specify it otherwise. As we discussed above, the squared gradient norm G is more closely related to implicit gradient regularization, generalization bounds, and SWD than other gradient norms. Overall, the theoretical results above supported that the small gradient norm is a desirable property in deep learning. However, our theoretical analysis suggests that weight decay can lead to large gradient norms and hence hurt generalization. We suggest that gradient-norm-centered poor convergence and generalization can lead to the serious but overlooked performance degradation of weight decay.

4 Gradient-Norm-Aware Scheduled Weight Decay

In this section, we design the first practical scheduler to boost weight decay by mitigating large gradient norm.

We focus on scheduled weight decay for Adam in this paper for three reasons. First, Adam is the most popular optimization for accelerating training of DNNs. Second, as the adaptivity of Adam updates the learning rate every iteration, Adam suffers more than SGD from the pitfalls of weight decay due to adaptive learning rates. Third, SGD usually employs L 2 regularization rather than weight decay.

Adam with decoupled weight decay is often called AdamW [32], which can be written as

where v t and m t are the exponential moving average of the squared gradients and gradients in Algorithm 1, respectively, and the power notation of a vector means the element-wise power of the vector. We interpret ηv -1 2 t as the effective learning rate for multiplying the gradients. However, we clearly see that decoupled weight decay couples weight decay with the vanilla learning rate rather than the effective learning rate. Due to the differences of weight decay and L 2 regularization, AdamW and Adam optimize different training objectives. The minimum ˆ θ ⋆ of the regularized loss function optimized by AdamW at the t -th step is not constant during training. Thus, the regularized loss function optimized by AdamW has unstable minima, which may result in large gradient norms

1.6

0

25

50

75

100

125

150

175

200

Epochs

0.2

0.4

0.6

0.8

1.0

1.2

1.4

Squared Gradient Norm

AdamS

AdamW

(a) Norm

G

0

25

50

75

100

125

150

175

200

Epochs

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

AdamS

AdamW

0

25

50

75

100

125

150

175

200

Epochs

10

4

10

3

10

2

10

1

10 0

Training Loss

AdamS

AdamW

(c) Training Loss

0

25

50

75

100

125

150

175

200

Epochs

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Squared Gradient Norm

AdamS

AdamW

(d) Norm

G

0

25

50

75

100

125

150

175

200

Epochs

0.5

1.0

1.5

2.0

2.5

AdamS

AdamW

0

25

50

75

100

125

150

175

200

Epochs

10

3

10

2

10

1

10 0

Training Loss

AdamS

AdamW

(f) Training Loss

suggested by Theorem 1. This may be an internal fault of all optimizers that use adaptive learning rates and weight decay at the same time.

To mitigate the pitfall of large gradient norms, we propose a gradient-norm-aware scheduler for weight decay as

where ¯ v t is the mean of all elements of the vector v t . Moreover, as v t is the exponential moving average of squared gradient norms, the factor ¯ v t is expected to be equivalent to the squared gradient norms divided by model dimensionality. While the proposed scheduler cannot give perfectly stable minima for Adam, it can penalize large gradient norms by dynamically adjusting weight decay. We call weight decay with this gradient-norm-aware scheduler Scheduled Weight Decay (SWD). The pseudocode is displayed in Algorithm 2. The code is publicly available at GitHub. In Section 5, we empirically demonstrate the advantage of the gradient-norm-aware scheduler. Note that AdamS is slightly slower than AdamW but slightly faster than Adam, while the cost difference is nearly ignorable (usually less than 5% ).

We note that ¯ v is not expected to be zero at minima due to stochastic gradient noise, because the variance of the stochastic gradient is directly observed to be much larger than the expectation of the stochastic gradient at/near minima and depends on the Hessian [43]. In some cases, such as full-batch training, it is fine to add a small value (e.g., 10 -8 ) to ¯ v to avoid being zero as a divisor.

Lewkowycz and Gur-Ari [26] recently proposed a simple weight decay scheduling method which has a limited practical value because it only works well with a constant learning rate, as the original paper claimed. Bjorck et al. [4] empirically studied early effects of weight decay but did not touch mitigating the pitfalls of weight decay. In contrast, our experiments suggest that SWD indeed effectively penalizes the large gradient norm and hence improve convergence and generalization.

Figure 4 verifies multiple advantages of SWD. First, SWD dynamically adjusts the weight decay strength inverse to G as we expect. Second, obviously, SWD can effectively penalize large gradient norms compared with constant weight decay during not only the final training phase but also the nearly whole training procedure. Third, SWD can often lead to lower training losses than constant weight decay. Figure 5 clearly reveals that, to choose proper weight decay strength, SWD dynamically

0

25

50

75

100

125

150

175

200

Epochs

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

AdamS WD=1e-3

AdamS WD=5e-4

AdamS WD=1e-4

0

25

50

75

100

125

150

175

200

Epochs

0.20

0.22

0.24

0.26

0.28

0.30

0.32

0.34

Test Error

AdamS

AdamW

Adam

AdaBound

Padam

RAdam

0

25

50

75

100

125

150

175

200

Epochs

0.06

0.07

0.08

0.09

0.10

0.11

0.12

Test Error

AdamS

AdamW

Adam

(a) VGG16 on CIFAR-10

0

25

50

75

100

125

150

175

200

Epochs

0.20

0.22

0.24

0.26

0.28

0.30

0.32

0.34

Test Error

AdamS

AdamW

Adam

(b) ResNet34 on CIFAR-100

0

25

50

75

100

125

150

175

200

Epochs

0.20

0.22

0.24

0.26

0.28

0.30

0.32

0.34

Test Error

AdamS

AdamW

Adam

(c) DenseNet121 on CIFAR-100

0

25

50

75

100

125

150

175

200

Epochs

10

4

10

3

10

2

10

1

10 0

Training Loss

AdamS

AdamW

Adam

0

25

50

75

100

125

150

175

200

Epochs

10

3

10

2

10

1

10 0

Training Loss

AdamS

AdamW

Adam

0

25

50

75

100

125

150

175

200

Epochs

10

3

10

2

10

1

10 0

Training Loss

AdamS

AdamW

Adam

employs a stronger scheduler when the initially selected weight decay hyperparameter λ is smaller and a weaker scheduler when λ is larger compared to the optimal weight decay choice, respectively.

5 Empirical Analysis

In this section, we first empirically verify that weight decay indeed affects gradient norms and further conducted comprehensive experiments to demonstrate the effectiveness of the SWD scheduler in mitigating large gradient norms and boosting performance.

Models and Datasets. We train various neural networks, including ResNet18/34/50 [15], VGG16 [37], DenseNet121 [19], GoogLeNet [38], and Long Short-Term Memory (LSTM) [18], on CIFAR10/CIFAR-100 [23], ImageNet[9], and Penn TreeBank [34]. The optimizers include popular Adam variants, including AMSGrad [35], Yogi [49], AdaBound [33], Padam [5], and RAdam [30]. See Appendix B and C for more details and results.

0

20

40

Eigenvalue Rank

10

2

10

1

10 0

10 1

10 2

Eigenvalue

0

20

40

Eigenvalue Rank

10 1

10 2

Eigenvalue

Adam

AdamW

AdamS

Generalization. Figure 7 shows the learning curves of AdamS, AdamW, and Adam on several benchmarks. In our experiments, AdamS always leads to lower test errors. Figure 19 in Appendix shows that, even if with similar or higher training losses, AdamS still generalizes significantly better than AdamW, Adam, and recent Adam variants. Figure 6 displays the learning curves of all adaptive gradient methods. The test performance of other models can be found in Table 1. The results on ImageNet (See Figure 18 in Appendix) shows that, AdamS can also make improvements over AdamW. Simply scheduling weight decay for Adam by SWD even outperforms complex Adam variants. We report that most Adam variants surprisingly generalize worse than SGD (See Appendix C).

Minima sharpness. Previous studies [20, 57, 43, 45, 44, 7] reported that the minima sharpness measured by the Hessian positively correlate to the gradient norm. Note that a number of popular minima sharpness depends on the Hessian at minima [21]. In Figures 8 and 9, we visualize the top Hessian eigenvalues to measure the minima sharpness. Figure 8 shows that stronger weight decay leads to sharper minima. Figure 9 suggests that AdamS learns significantly flatter minima than AdamW. Both large gradient norms and sharp minima correspond to high-complexity solutions and often hurts generalization [16]. This is a surprising finding for weight decay given the conventional belief that weight decay as a regularization technique should encourage low-complexity solutions.

The Robustness of SWD. Figure 11 shows that AdamS is more robust to the learning rate and weight decay than AdamW and Adam. AdamS has a much deeper and wider basin than AdamW and Adam. Figure 10 further demonstrates that AdamS consistently outperforms Adam and AdamW under various weight decay hyperparameters. According to Figure 10, we also notice that the optimal decoupled weight decay hyperparameter in AdamW can be very different from L 2 regularization and SWD. Thus, AdamW requires re-tuning the weight decay hyperparameter in practice, which is timeconsuming. The robustness of selecting hyperparameters can be supported by Figure 5, where the proposed scheduler also depends on the initial weight decay strength. It means that the proposed gradient-norm-aware scheduler exhibit the behavior of negative feedback to stabilize the strength of weight decay during training.

10

7

10

6

10

5

10

4

10

3

10

2

Weight Decay

0.05

0.06

0.07

0.08

0.09

0.10

0.11

0.12

Test Error

AdamS

AdamW

Adam

Cosine learning rate schedulers and warm restarts. We conducted comparative experiments on AdamS, AdamW, and Adam in the presence of cosine annealing schedulers and warm restarts proposed by Loshchilov and Hutter [31]. We set the learning rate scheduler with a recommended

10

6

10

5

10

4

10

3

Weight Decay

10

5

10

4

10

3

10

2

Learning Rate

AdamS

4.50

4.72

4.94

5.16

5.38

5.60

5.82

6.04

6.26

6.48

10

6

10

5

10

4

10

3

Weight Decay

10

5

10

4

10

3

10

2

Learning Rate

AdamW

4.50

4.72

4.94

5.16

5.38

5.60

5.82

6.04

6.26

6.48

10

6

10

5

10

4

10

3

Weight Decay

10

5

10

4

10

3

10

2

Learning Rate

Adam

4.50

4.72

4.94

5.16

5.38

5.60

5.82

6.04

6.26

6.48

setting of Loshchilov and Hutter [31]. Our experimental results in Figure 12 suggest that AdamS consistently outperforms AdamW and Adam in terms of both training losses and test errors. It demonstrates that, with complex learning rate schedulers, the advantage of SWD still holds.

0

25

50

75

100

125

150

175

200

Epochs

0.05

0.10

0.15

0.20

0.25

0.30

Test Error

AdamS

AdamW

Adam

(a) ResNet18

0

25

50

75

100

125

150

175

200

Epochs

0.05

0.10

0.15

0.20

0.25

0.30

Test Error

AdamS

AdamW

Adam

(b) VGG16

0

25

50

75

100

125

150

175

200

Epochs

10

4

10

3

10

2

10

1

10 0

Training Loss

AdamS

AdamW

Adam

(c) ResNet18

0

25

50

75

100

125

150

175

200

Epochs

10

4

10

3

10

2

10

1

10 0

Training Loss

AdamS

AdamW

Adam

(d) VGG16

Limitations. While the proposed scheduler boosts weight decay in most experiments, the improvement of SWD may be limited sometimes. SWD does not work well particularly when L 2 regularization even yields better test results than weight decay (See the experiments on Language Modeling Figure 13 in Appendix C). Our paper does not touch the pitfall of L 2 regularization and how to schedule its strength. We leave scheduling L 2 regularization as future work.

6 Conclusion

While weight decay has attracted much attention, previous studies failed to recognize the problems of unstable stationary points and large gradient norm caused by weight decay. In this paper, we theoretically and empirically studied the overlooked pitfalls of weight decay from the gradient-norm perspective. To mitigate the pitfalls, we propose a simple but effective gradient-norm-aware scheduler for weight decay. Our empirical results demonstrate that SWD can effectively penalize large gradient norms and improve convergence. Moreover, SWD often significantly improve generalization over constant schedulers. The generalization gap between SGD and Adam can be almost closed by such a simple weight decay scheduler for CNNs.

We emphasize that the proposed method is the first scheduled weight decay method rather than a novel optimizer. Although our analysis mainly focused on Adam, SWD can be easily combined with other optimizers, including SGD and Adam variants. To the best of our knowledge, this is the first formal touch on the art of scheduling weight decay to boost weight decay. We believe that scheduled weight decay as a new line of research may inspire more theories and algorithms that will help further understanding and boosting weight decay in future. Our work just made the first step along this line.

References

Achille, A., Rovere, M., and Soatto, S. (2017). Critical learning periods in deep networks. In International Conference on Learning Representations .

An, H., Xiong, H., Li, X., Li, X., Dou, D., and Zhu, Z. (2020). Can we use gradient norm as a measure of generalization error for model selection in practice?

Barrett, D. and Dherin, B. (2020). Implicit gradient regularization. In International Conference on Learning Representations .

Bjorck, J., Weinberger, K. Q., and Gomes, C. (2021). Understanding decoupled and early weight decay. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 6777-6785.

Chen, J. and Gu, Q. (2018). Closing the generalization gap of adaptive gradient methods in training deep neural networks. arXiv preprint arXiv:1806.06763 .

Chong, E. K., Lu, W.-S., and ˙ Zak, S. H. (2023). An introduction to optimization . John Wiley & Sons.

Daneshmand, H., Kohler, J., Lucchi, A., and Hofmann, T. (2018). Escaping saddles with stochastic gradients. In International Conference on Machine Learning , pages 1155-1164.

Défossez, A., Bottou, L., Bach, F., and Usunier, N. (2020). On the convergence of adam and adagrad. arXiv preprint arXiv:2003.02395 .

Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. (2009). Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee.

Frankle, J., Schwab, D. J., and Morcos, A. S. (2020). The early phase of neural network training. In International Conference on Learning Representations .

Geiping, J., Goldblum, M., Pope, P., Moeller, M., and Goldstein, T. (2021). Stochastic training is not necessary for generalization. In International Conference on Learning Representations .

Ghadimi, S. and Lan, G. (2013). Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization , 23(4):2341-2368.

Hanson, S. J. and Pratt, L. Y . (1989). Comparing biases for minimal network construction with back-propagation. In Advances in neural information processing systems , pages 177-185.

Hardt, M., Recht, B., and Singer, Y. (2016). Train faster, generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning , pages 1225-1234.

He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778.

Hochreiter, S. and Schmidhuber, J. (1995). Simplifying neural nets by discovering flat minima. In Advances in neural information processing systems , pages 529-536.

Hochreiter, S. and Schmidhuber, J. (1997a). Flat minima. Neural Computation , 9(1):1-42.

Hochreiter, S. and Schmidhuber, J. (1997b). Long short-term memory. Neural computation , 9(8):1735-1780.

Huang, G., Liu, Z., Van Der Maaten, L., and Weinberger, K. Q. (2017). Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 4700-4708.

Jastrzkebski, S., Kenton, Z., Arpit, D., Ballas, N., Fischer, A., Bengio, Y., and Storkey, A. (2017). Three factors influencing minima in sgd. arXiv preprint arXiv:1711.04623 .

Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. (2019). Fantastic generalization measures and where to find them. In International Conference on Learning Representations .

Kingma, D. P. and Ba, J. (2015). Adam: A method for stochastic optimization. 3rd International Conference on Learning Representations, ICLR 2015 .

Krizhevsky, A. and Hinton, G. (2009). Learning multiple layers of features from tiny images.

Krogh, A. and Hertz, J. A. (1992). A simple weight decay can improve generalization. In Advances in neural information processing systems , pages 950-957.

LeCun, Y., Bengio, Y., and Hinton, G. (2015). Deep learning. nature , 521(7553):436.

Lewkowycz, A. and Gur-Ari, G. (2020). On the training dynamics of deep networks with l _ 2 regularization. Conference on Neural Information Processing Systems .

Li, J., Luo, X., and Qiao, M. (2019a). On generalization error bounds of noisy gradient methods for non-convex learning. In International Conference on Learning Representations .

Li, Q., Tai, C., et al. (2017). Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference on Machine Learning-Volume 70 , pages 2101-2110. JMLR. org.

Li, Y., Wei, C., and Ma, T. (2019b). Towards explaining the regularization effect of initial large learning rate in training neural networks. In Advances in Neural Information Processing Systems , pages 11669-11680.

Liu, L., Jiang, H., He, P., Chen, W., Liu, X., Gao, J., and Han, J. (2019). On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations .

Loshchilov, I. and Hutter, F. (2016). Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983 .

Loshchilov, I. and Hutter, F. (2018). Decoupled weight decay regularization. In International Conference on Learning Representations .

Luo, L., Xiong, Y., Liu, Y., and Sun, X. (2019). Adaptive gradient methods with dynamic bound of learning rate. 7th International Conference on Learning Representations, ICLR 2019 .

Marcus, M., Santorini, B., and Marcinkiewicz, M. A. (1993). Building a large annotated corpus of english: The penn treebank.

Reddi, S. J., Kale, S., and Kumar, S. (2019). On the convergence of adam and beyond. 6th International Conference on Learning Representations, ICLR 2018 .

Shapiro, A. and Wardi, Y. (1996). Convergence analysis of gradient descent stochastic algorithms. Journal of optimization theory and applications , 91(2):439-454.

Simonyan, K. and Zisserman, A. (2014). Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556 .

Szegedy, C., Liu, W., Jia, Y., Sermanet, P., Reed, S., Anguelov, D., Erhan, D., Vanhoucke, V., and Rabinovich, A. (2015). Going deeper with convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 1-9.

Wan, R., Zhu, Z., Zhang, X., and Sun, J. (2021). Spherical motion dynamics: Learning dynamics of normalized neural network using sgd and weight decay. Advances in Neural Information Processing Systems , 34.

Wilson, A. C., Roelofs, R., Stern, M., Srebro, N., and Recht, B. (2017). The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems , pages 4151-4161.

Wolfe, P. (1969). Convergence conditions for ascent methods. SIAM review , 11(2):226-235.

Xie, Z., He, F., Fu, S., Sato, I., Tao, D., and Sugiyama, M. (2021a). Artificial neural variability for deep learning: On overfitting, noise memorization, and catastrophic forgetting. Neural Computation , 33(8).

Xie, Z., Sato, I., and Sugiyama, M. (2021b). A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations .

Xie, Z., Tang, Q., Sun, M., and Li, P. (2023). On the overlooked structure of stochastic gradients. In Thirty-seventh Conference on Neural Information Processing Systems .

Xie, Z., Tang, Q.-Y., Cai, Y., Sun, M., and Li, P. (2022a). On the power-law spectrum in deep learning: A bridge to protein science. arXiv preprint arXiv:2201.13011 .

Xie, Z., Wang, X., Zhang, H., Sato, I., and Sugiyama, M. (2022b). Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 24430-24459.

Xie, Z., Yuan, L., Zhu, Z., and Sugiyama, M. (2021c). Positive-negative momentum: Manipulating stochastic gradient noise to improve generalization. In International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 11448-11458. PMLR.

Yan, Y., Yang, T., Li, Z., Lin, Q., and Yang, Y. (2018). A unified analysis of stochastic momentum methods for deep learning. In IJCAI International Joint Conference on Artificial Intelligence .

Zaheer, M., Reddi, S., Sachan, D., Kale, S., and Kumar, S. (2018). Adaptive methods for nonconvex optimization. In Advances in neural information processing systems , pages 9793-9803.

Zaremba, W., Sutskever, I., and Vinyals, O. (2014). Recurrent neural network regularization. arXiv preprint arXiv:1409.2329 .

Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. (2017). Understanding deep learning requires rethinking generalization. In International Conference on Machine Learning .

Zhang, G., Wang, C., Xu, B., and Grosse, R. (2019a). Three mechanisms of weight decay regularization. International Conference on Learning Representations .

Zhang, J., He, T., Sra, S., and Jadbabaie, A. (2019b). Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Representations .

Zhang, X., Xu, R., Yu, H., Zou, H., and Cui, P. (2022). Gradient norm regularizer seeks flat minima and improves generalization.

Zhao, Y., Zhang, H., and Hu, X. (2022). Penalizing gradient norm for efficiently improving generalization in deep learning. In International Conference on Machine Learning .

Zhou, P., Feng, J., Ma, C., Xiong, C., HOI, S., et al. (2020). Towards theoretically understanding why sgd generalizes better than adam in deep learning. arXiv preprint arXiv:2010.05627 .

Zhu, Z., Wu, J., Yu, B., Wu, L., and Ma, J. (2019). The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In ICML , pages 7654-7663.

A Proofs

A.1 Proof of Theorem 1

Proof. We first write the regularized loss function corresponding to SGD with vanilla weight decay as

at t -th step.

If the corresponding L 2 regularization λ ′ 2 η t is unstable during training, the regularized loss function f t ( θ ) will also be a time-dependent function and has no non-zero stable stationary points.

Suppose we have a non-zero solution θ ⋆ which is a stationary point of f ( θ, t ) at t -th step and SGD finds θ t = θ ⋆ at t -th step.

Even if the gradient of f t ( θ ) at t -step is zero, we have the gradient at ( t +1) -th step as

To achieve convergence, we must have ∥ g t +1 ∥ 2 = 0 .

Theorem 2.2 of Shapiro and Wardi [36] told us that the learning rate should be small enough for convergence. Obviously, we have η < ∞ in practice.

As η t = η t +1 does not hold, SGD cannot converging to any non-zero stationary point.

The proof is now complete.

A.2 Proof of Theorem 2

Before formally completing the proof, we first introduce a useful Lemma 1, which is a specialized case of Theorem 1 in Yan et al. [48] with β = s = 0 .

Lemma 1 (Convergence of SGD) . Assume that L ( θ ) is an L -smooth function 2 , L is lower bounded as L ( θ ) ≥ L ⋆ , E [ ∇ L ( θ, X ) -∇ L ( θ )] = 0 , E [ ∥∇ L ( θ, X ) -∇ L ( θ ) ∥ 2 ] ≤ δ 2 , ∥∇ L ( θ ) ∥ ≤ G for any θ . Let SGD optimize L for t +1 iterations. If η ≤ C √ t +1 , we have

Proof. Given the conditions of L ( θ ) in Lemma 1, we may obtain the resulted conditions of f ( θ ) = L ( θ ) + λ 2 ∥ θ ∥ 2 .

As L ( θ ) is an L -smooth function, we have

holds for any θ a and θ b . It shows that f ( θ ) is an ( L + λ ) -smooth function.

As L is lower bounded as L ( θ ) ≥ L ⋆ , we have

It means that

As E [ ∇ L ( θ, X ) -∇ L ( θ )] = 0

As E [ ∥∇ L ( θ, X ) -∇ L ( θ ) ∥ 2 ] ≤ δ 2 , we have

As ∥∇ L ( θ ) ∥ ≤ G , we have

where ∥ θ ∥ max is the maximum L 2 norm of any θ .

Introducing the derived conditions Eq. (12) - (16) for f into Lemma 1, we may treat f as the objective optimized by SGD. Then we have

Obviously, the gradient norm upper bound in convergence analysis monotonically increases as the weight decay strength λ .

The proof is complete.

B Experimental Details

Computational environment. The experiments are conducted on a computing cluster with GPUs of NVIDIA ® Tesla ™ P100 16GB and CPUs of Intel ® Xeon ® CPU E5-2640 v3 @ 2.60GHz.

B.1 Image Classification on CIFAR-10 and CIFAR-100

Data Preprocessing For CIFAR-10 and CIFAR-100: We perform the common per-pixel zero-mean unit-variance normalization, horizontal random flip, and 32 × 32 random crops after padding with 4 pixels on each side.

Hyperparameter Settings: We select the optimal learning rate for each experiment from { 0 . 0001 , 0 . 001 , 0 . 01 , 0 . 1 , 1 , 10 } for non-adaptive gradient methods. We use the default learning rate for adaptive gradient methods in the experiments of Table 1, while we also compared Adam, AdamW, AdamS under various learning rates and batch sizes in other experiments. In the experiments on CIFAF-10 and CIFAR-100: η = 0 . 1 for SGD and SGDS; η = 0 . 001 for Adam, AdamW, AdamS, AMSGrad, Yogi, AdaBound, and RAdam; η = 0 . 01 for Padam. For the learning rate schedule, the learning rate is divided by 10 at the epoch of { 80 , 160 } for CIFAR-10 and { 100 , 150 } for CIFAR-100, respectively. The batch size is set to 128 for both CIFAR-10 and CIFAR-100.

The strength of L 2 regularization and SWD is default to 0 . 0005 as the baseline. Considering the linear scaling rule, we choose λ W = λ L 2 η . Thus, the weight decay of AdamW uses λ W = 0 . 5 for CIFAR-10 and CIFAR-100. The basic principle of choosing weight decay strength is to let all optimizers have similar convergence speed.

We set the momentum hyperparameter β 1 = 0 . 9 for SGD and SGDS. As for other optimizer hyperparameters, we apply the default hyperparameter settings directly.

We repeated each experiment for three times in the presence of the error bars.

We leave the empirical results with the weight decay setting λ = 0 . 0001 in Appendix C.

B.2 Image classification on ImageNet

Data Preprocessing For ImageNet: For ImageNet, we perform the per-pixel zero-mean unitvariance normalization, horizontal random flip, and the resized random crops where the random size (of 0.08 to 1.0) of the original size and a random aspect ratio (of 3 4 to 4 3 ) of the original aspect ratio is made.

Hyperparameter Settings for ImageNet: We select the optimal learning rate for each experiment from { 0 . 0001 , 0 . 001 , 0 . 01 , 0 . 1 , 1 , 10 } for all tested optimizers. For the learning rate schedule, the learning rate is divided by 10 at the epoch of { 30 , 60 } . We train each model for 90 epochs. The batch size is set to 256 . The weight decay hyperparameter of AdamS, AdamW, Adam are chosen from { 5 × 10 -6 , 5 × 10 -5 , 5 × 10 -4 , 5 × 10 -3 , 5 × 10 -2 } . As for other optimizer hyperparameters, we still apply the default hyperparameter settings directly.

B.3 Language Modeling

We use a classical language model, Long Short-Term Memory (LSTM) [18] with 2 layers, 512 embedding dimensions, and 512 hidden dimensions, which has 14 million model parameters and is similar to the 'medium LSTM' in Zaremba et al. [50]. Note that our baseline performance is better than the reported baseline performance in Zaremba et al. [50]. The benchmark task is the word-level Penn TreeBank [34]. We empirically compared AdamS, AdamW, and Adam under the common and same conditions.

Hyperparameter Settings. Batch Size: B = 20 . BPTT Size: bptt = 35 . Learning Rate: η = 0 . 001 . We run the experiments under various weight decay selected from { 10 -4 , 5 × 10 -5 , 10 -5 , 5 × 10 -6 , 10 -6 , 5 × 10 -7 , 10 -7 } .The dropout probability is set to 0 . 5 . We clipped gradient norm to 1 .

C Supplementary Figures and Results of Adaptive Gradient Methods

Popular Adam variants often generalize worse than SGD. A few Adam variants tried to fix the hidden problems in adaptive gradient methods, including AdamW Loshchilov and Hutter [32], AMSGrad [35] and Yogi [49]. A recent line of research, such as AdaBound [33], Padam [5], and RAdam [30], believes controlling the adaptivity of learning rates may improve generalization. This line of research usually introduces extra hyperparameters to control the adaptivity, which requires more efforts in tuning hyperparameters. However, we and Zhang et al. [53] found that this argument is contradicted with our comparative experimental results (see Table 1). In our empirical analysis, most advanced Adam variants may narrow but not completely close the generalization gap between adaptive gradient methods and SGD. SGD with a fair weight decay hyperparameter as the baseline performance usually generalizes better than recent adaptive gradient methods. The main problem may lie in weight decay. SGD with weight decay λ = 0 . 0001 , a common setting in related papers, is often not a good baseline, as λ = 0 . 0005 often shows better generalization on CIFAR-10 and CIFAR-100. We also conduct comparative experiments with λ = 0 . 0001 . Under the setting λ = 0 . 0001 , while some existing Adam variants may outperform SGD sometimes due to the lower baseline performance of SGD, AdamS shows superior test performance. For example, for ResNet18 on CIFAR-10, the test error of AdamS is lower than SGD by nearly one point and no other Adam variant may compare with AdamS.

10

7

10

6

10

5

10

4

Weight Decay

65

70

75

80

85

90

95

100

105

Test Perplexity

AdamS

AdamW

Adam

0

25

50

75

100

125

150

175

200

Epochs

0.04

0.05

0.06

0.07

0.08

0.09

0.10

0.11

0.12

Test Error

AdamS

AdamW

Adam

AdaBound

Padam

RAdam

(a) ResNet18

0

25

50

75

100

125

150

175

200

Epochs

0.06

0.07

0.08

0.09

0.10

0.11

0.12

Test Error

AdamS

AdamW

Adam

AdaBound

Padam

RAdam

(b) VGG16

0

25

50

75

100

125

150

175

200

Epochs

0.20

0.22

0.24

0.26

0.28

0.30

0.32

0.34

Test Error

AdamS

AdamW

Adam

AdaBound

Padam

RAdam

(c) DenseNet121

0

25

50

75

100

125

150

175

200

Epochs

0.20

0.22

0.24

0.26

0.28

0.30

0.32

0.34

Test Error

AdamS

AdamW

Adam

AdaBound

Padam

RAdam

(d) GoogLeNet

Language Modeling. It is well-known that, different from computer vision tasks, the standard Adam (with L 2 regularization) is the most popular optimizer for language models. Figure 13 in Appendix C demonstrates that the conventional belief is true that the standard L 2 regularization yields better test results than both Decoupled Weight Decay and SWD. The weight decay scheduler suitable for language models is an open problem.

We report the learning curves of all adaptive gradient methods in Figure 14. They shows that vanilla Adam with SWD can outperform other complex variants of Adam.

Figure 15 displays the scatter plot of training losses and test errors during final 40 epochs of training DenseNet121 on CIFAR-100.

Figure 16 displays the test performance of AdamS, AdamW, and Adam under various weight decay hyperparameters of ResNet34 on CIFAR-100.

0.04

0.05

0.06

0.07

0.08

0.09

0.10

Test Errors

10

3

10

2

10

1

Training Loss

AdamS

AdamW

Adam

0.20

0.22

0.24

0.26

0.28

0.30

Test Errors

10

3

10

2

10

1

Training Loss

AdamS

AdamW

Adam

10

7

10

6

10

5

10

4

10

3

10

2

Weight Decay

0.175

0.200

0.225

0.250

0.275

0.300

0.325

0.350

0.375

Test Error

AdamS

AdamW

Adam

0

100

200

300

400

500

600

700

800

900

Epochs

0.04

0.06

0.08

0.10

0.12

0.14

Test Error

AdamS

AdamW

Adam

0

10

20

30

40

50

60

70

80

90

Epochs

0.24

0.26

0.28

0.30

0.32

0.34

0.36

0.38

0.40

Test Error

AdamS

AdamW

Adam

We train ResNet18 on CIFAR-10 for 900 epochs to explore the performance limit of AdamS, AdamW, and Adam in Figure 17.

Figure 18 in Appendix shows that, AdamS can make marginal improvements over AdamW.

0.20

0.22

0.24

0.26

0.28

0.30

0.32

0.34

Test Errors

10

3

10

2

10

1

10 0

Training Loss

AdamS

AdamW

Adam

AdaBound

Padam

RAdam

D Additional Algorithms

We note that the implementation of AMSGrad in Algorithm 3 is the popular implementation in PyTorch. We use the PyTorch implementation in our paper, as it is widely used in practice.

Algorithm 3: AMSGrad/AMSGradW

Algorithm 4: AMSGradS

E Supplementary Experiments with Cosine Annealing Schedulers and Warm Restarts

In this section, we conducted comparative experiments on AdamS, AdamW, and Adam in the presence of cosine annealing schedulers and warm restarts proposed by Loshchilov and Hutter [31]. We set the learning rate scheduler with a recommended setting of Loshchilov and Hutter [31]: T 0 = 14 and T mul = 2 . The number of total epochs is 210 . Thus, we trained each deep network for four runs of warm restarts, where the four runs have 14, 28, 56, and 112 epochs, respectively. Other hyperparameters and details are displayed in Appendix B.

We conducted comparative experiments on AdamS, AdamW, and Adam in the presence of cosine annealing schedulers and warm restarts proposed by Loshchilov and Hutter [31]. We set the learning rate scheduler with a recommended setting of Loshchilov and Hutter [31] Our experimental results in Figures 12 and 20 suggest that AdamS consistently outperforms AdamW and Adam in the presence of cosine annealing schedulers and warm restarts. It demonstrates that, with various learning rate schedulers, the advantage of SWD may generally hold.

Moreover, we did not empirically observe that cosine annealing schedulers with warm restarts may consistently outperform the common piecewise-constant learning rate schedulers for adaptive gradient methods. We noticed that Loshchilov and Hutter [31] empirically compared four-staged piecewiseconstant learning rate schedulers with cosine annealing schedulers with warm restarts, and argue that cosine annealing schedulers with warm restarts are better. There may be two possible causes. First, three-staged piecewise-constant learning rate schedulers, which usually have a longer first stage and decay learning rates by multiplying 0 . 1 , are the recommended settings, while the four-staged piecewise-constant learning rate schedulers in Loshchilov and Hutter [31] are usually not optimal. Second, warm restarts may be helpful, while cosine annealing may be not. The ablation study on piecewise-constant learning rate schedulers with warm restarts is lacked. We argued that how to choose learning rate schedulers may still be an open question, considering the complex choices of schedulers and the complex loss landscapes.

10

6

10

5

10

4

10

3

10

2

Weight Decay

0.04

0.05

0.06

0.07

0.08

0.09

0.10

0.11

0.12

Test Error

AdamS

AdamW

Adam

(a) ResNet18

10

6

10

5

10

4

10

3

10

2

Weight Decay

0.04

0.05

0.06

0.07

0.08

0.09

0.10

0.11

0.12

Test Error

AdamS

AdamW

Adam

(b) VGG16

Peer reviews:
comments: SUMMARY:
The authors propose Scheduled Weight Decay (SWD), a method that mitigates the large gradient norm issue caused by constant weight decay factors. The authors demonstrate that SWD can improve the generalization performance of Adam and outperform other adaptive optimizers on CIFAR-10/100 datasets.

STRENGTHS:
1. The proposed method is simple and easy to implement. Moreover, it addresses the long-standing generalization gap between adaptive optimizers and SGD on certain tasks.
2. The authors provide a theoretical analysis of the problem of unstable stationary points.
3. The authors conduct extensive experiments to validate their claims and compare their method with other optimizers.

WEAKNESSES:
1. The gradient analyzed in Theorem 2 seems to include the gradient of L2 regularization, which is not included in the empirical analysis (e.g. Fig. 2).
2. Table 1 in appendix and table 1 in the main text show very different results, especially for SGD. The reason for this difference should be made more clear.

QUESTIONS:
See above.

LIMITATIONS:
The authors have discussed the limitations of the proposed method on datasets other than CIFAR-10/100.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: SUMMARY:
This paper studies the overlooked pitfalls of weight decay, a regularization technique used in deep neural networks (DNNs). The authors discovered that weight decay can lead to large gradient norms, particularly at the final phase of training, often indicating poor convergence and generalization. To address this issue, the authors propose a novel method called Scheduled Weight Decay (SWD), which dynamically adjusts weight decay strength according to the gradient norm and penalizes large gradient norms during training. The paper concludes that the SWD approach outperforms the conventional constant weight decay strategy, especially for the Adaptive Moment Estimation (Adam) optimization algorithm.

STRENGTHS:
1. SWD dynamically adjusts weight decay, which can help to avoid large gradient norms that can lead to poor convergence and generalization. The paper demonstrates that SWD outperforms the traditional constant weight decay, especially in Adam optimization, potentially leading to better model performance.
2. The concept is simple to understand and can be easily implemented in most deep learning frameworks.

WEAKNESSES:
1. The proposal is specific to weight decay regularization and might not generalize to other regularization techniques.
2. SWD's performance might depend heavily on how well the scheduling function is chosen, which can be a complex task.
3. The paper doesn't discuss how SWD might affect the training time, and adjusting weight decay dynamically could potentially increase computational complexity.

QUESTIONS:
How does SWD affect the training time, given the additional complexity introduced by dynamically adjusting the weight decay?

LIMITATIONS:
Not applicable.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the role of weight decay and its connection with large gradient norms in deep learning settings. In particular, the paper highlights differences in variants of weight decay and also the effect of weight decay on gradient norm in the final phase. Based on the observation that weight decay yields large gradient norms, the authors propose a scheduler for weight decay, called the Scheduled Weight Decay, which dynamically adjusts the weight decay strength.\n\nSTRENGTHS:\nThe paper is reasonably well-written and the motivation for studying the effect of weight decay is both important and clear. The approach is easy to implement and the results on some smaller scale benchmarks seem encouraging.\n\nWEAKNESSES:\n(1) The theoretical results are weak and not very interesting. Theorem 2 follows from standard convergence rates for SGD. I could not find any novelty in this result. The authors should highlight the novelty of the result and proof technique in the main paper.\n\n(2) While Theorem 1 seems interesting, it seems somewhat weak. The lower bound on the difference in learning rate is somewhat artificial and unreasonable. For instance, consider decreasing learning rates of SGD typically used in machine learning settings i.e., eta_t >= eta_{t+1}. Then Theorem 1 assumes delta <= eta_t - eta_{t+1}. This implies eta_0 >= delta . t for any t > 0. This makes sense if delta is going to 0 as t -> infinity. Maybe I am missing something but would really appreciate it if the authors can comment about this. \n\n(3) Theorem 1 & 2 seem to analyze gradients of two different functions, which doesn't seem like a proper comparison. More remarks regarding this will be valuable. Furthermore, I would like to get clarification if the y-axis in plots of Figure 2 are gradient norms with regularization included.\n\n(4) Definition 1 is very unclear. The authors should provide more remarks around definition 1 to help readers understand the definition better.\n\n(5) The dependence of C_2 on sup norm of theta is somewhat weird in Theorem 2. Are you assuming ||theta|| is bounded? (otherwise rewrite the theorem statement without this dependence) .\n\n(6) I think the theoretical basis for Gradient-Norm-Aware Scheduled Weight Decay is weak as presented in the paper. While I understand the intuition, it is not clear why this particular proposal is meaningful. Since the theoretical basis for this is lacking, I would expect much more comprehensive experiments to support the method. Unfortunately, the empirical analysis in the paper is somewhat limited.\n\nLine 20: eta -> eta_t is the learning rate & Line 125: E[||∇L(θ, X) − ∇L(θ)||^2] ≤ \\sigma^2?\n\n\nPost-rebuttal\n==========\nI think the theoretical analysis is still not convincing but given that these observations are interesting, I am slightly increasing the score.\n\nQUESTIONS:\nPlease refer to the concerns raised in the weakness section.\n\nLIMITATIONS:\nNot applicable.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-31"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nI would divide this paper's contributions into two parts.  The first part is an algorithm (a variant of Adam) which the authors argue generalizes better than Adam/AdamW and is easier to tune.  The second part is the justification for the effectiveness of that algorithm.\n\n**The algorithm itself** The proposed algorithm, AdamS, is similar to AdamW, except that the weight decay strength is divided by the mean of the current squared gradient EMA.  This is effectively a schedule for weight decay.  It has the effect of penalizing the overall weights less strongly during phases of training when the overall gradient is large, and penalizing the overall weights more strongly during phases of training when the overall gradient is small.   For example, Figure 3 shows gradient norm and weight decay strength during the training of ResNet-34 on CIFAR-10; after the first learning rate drop, squared gradient norm grows, which causes the weight decay strength to shrink; after the second learning rate drop, squared learning rate plummets, which causes weight decay to rise.\n\n**Experimental evaluation** Table 1 and section 5 experimentally evaluate AdamS and argue that the algorithm is better generalizing / easier to tune than AdamW.\n\n**Justification**: The authors justify their algorithm along the following lines (which I disagree with; see more below):\n\n1. When the learning rate is potentially changing, gradient descent with weight decay has trouble converging to stationary points - in particular, the authors argue (Theorem 1) that in this setting, the gradient norm will be lower bounded near stationary points.  Therefore, the authors argue that weight decay causes the gradient to be large, especially at the end of training.\n2.  The authors argue that large gradients are bad for generalization.\n3. As a consequence of #1 and #2, AdamW-style weight decay is bad for generalization (since it causes large gradients near stationary points, and large gradients are bad for generalization)\n\nThe authors acknowledge that weight decay on networks with normalization layers has been shown in the literature to control the effective learning rate, but insist that their paper's mechanism is unrelated to this, and propose an entirely different mechanism.\n\n---\n\n**post-rebuttal update**\n\nAfter discussion with the authors, I'm raising my score from 3 (\"reject\") to 6 (\"weak accept\").    Firstly, it's clear that my original review was wrong; I had argued that the proposed algorithm's benefits are probably due to the interaction between weight decay and normalization layers, but the authors have correctly pointed out that the same phenomena occur on networks without normalization layers, and I have verified this myself.   The reason for 6 as opposed to 7 is that I believe that the proposed theoretical explanation for why weight decay would cause large gradients is not the real explanation (reviewer xE1B also didn't like this theoretical portion of the paper).  That said, it's entirely plausible that weight decay causing large gradients is one of the *many* aspects of deep learning that is just theoretically unexplainable at the present time.  So, my issue with the submission is less that the paper fails to include a correct theoretical justification, and more that it includes a faulty theoretical justification; I think it is better to have no justification at all than to have a faulty justification.  \n\nThat being said, on the positive side, it seems there is a good chance that the proposed algorithm is a worthwhile addition to the deep learning toolbox (though I am not an expert at evaluating this type of contribution).\n\nSTRENGTHS:\nI'm not an expert in experimentally evaluating the performance of neural network training algorithms, but from the experiments in section 5 it does seem plausible to me that AdamS is indeed better / easier to tune than AdamW.\n\nWEAKNESSES:\nWhile I can accept that the algorithm may be a good idea, I strongly disagree with the proposed justification given in the paper.  I believe the effectiveness of the algorithm is probably unrelated to the justification given in the paper, and is instead closely linked to the implicit effects of weight decay on effective learning rate when the network has normalization layers.\n\nFirst, I would note that all of the experiments in the paper are on networks with normalization layers.  The authors argue in Figure 2 that VGG-16 is \"not scale invariant\" but the VGG-16 architecture does have many BatchNorm layers, even though it is not fully scale-invariant, so I believe that the literature on scale invariance is still quite related to what is going on with the VGG-16.\n\nFor networks with normalization layers, there is a clear mechanism by which weight decay causes large gradients.  If $L$ is the loss for a scale-invariant network, then $\\nabla L(c \\theta) = \\frac{1}{c} \\nabla L(\\theta)$, i.e. scaling down the weights $(c < 1)$ will automatically scale up the gradients (see Lemma 1 here: http://www.offconvex.org/2020/04/24/ExpLR1/).  The large gradients have been reported to cause the 'effective learning rate' of training to be higher, which is thought to be good for generalization.  My intuition for the mechanism in this paper is: if the gradient norm is too large, our effective learning rate will be too large for the algorithm to converge, so we need to decrease weight decay so the gradient norms come back down; if the gradient norm is too small, our effective learning rate will be too small (which makes convergence slow and is bad for generalization) so we need to increase weight decay so that the gradient norm moves back up.\n\nBy contrast, the proposed justification doesn't really make sense to me.  Figure 2 shows that higher weight decay causes growth in the gradient norm _in the middle of training_, not at convergence.  This is much more easily explained by scale-invariance (for the VGG-16, the scale invariance of certain layers) than by Theorems 1 and 2.\n\nFigure 5 shows that weight decay increases the eigenvalue spectrum at convergence.  I believe this is related to cross-entropy loss, and wouldn't happen with e.g. squared loss.  For cross-entropy loss, the Hessian is small when the margins are large.  Adding weight decay prevents the margins from becoming large and therefore the Hessian from shrinking.\n\nQUESTIONS:\nHow do the authors respond to my critique above?\n\nLIMITATIONS:\ndiscussed above", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-08-01"}