Original: U NDERSTANDING AND S CHEDULING W EIGHT D ECAY
Revised: U NDERSTANDING AND S CHEDULING W EIGHT D ECAY

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Weight decay is a popular and even necessary regularization technique for training deep neural networks that generalize well.
Revised: Weight decay is a popular and even necessary regularization technique for training deep neural networks that generalize well.

Original: Previous work usually interpreted weight decay as a Gaussian prior from the Bayesian perspective.
Revised: Previous work usually interpreted weight decay as a Gaussian prior from the Bayesian perspective.

Original: However, weight decay sometimes shows mysterious behaviors beyond the conventional understanding.
Revised: However, weight decay sometimes shows mysterious behaviors beyond the conventional understanding.

Original: For example, the optimal weight decay value tends to be zero given long enough training time.
Revised: For example, the optimal weight decay value tends to be zero given long enough training time.

Original: Moreover, existing work typically failed to recognize the importance of scheduling weight decay during training.
Revised: Moreover, existing work typically failed to recognize the importance of scheduling weight decay during training.

Original: Our work aims at theoretically understanding novel behaviors of weight decay and designing schedulers for weight decay in deep learning.
Revised: Our work aims at theoretically understanding novel behaviors of weight decay and designing schedulers for weight decay in deep learning.

Original: This paper mainly has three contributions.
Revised: This paper mainly has three contributions.

Original: First, we propose a novel theoretical interpretation of weight decay from the perspective of learning dynamics.
Revised: First, we propose a novel theoretical interpretation of weight decay from the perspective of learning dynamics.

Original: Second, we propose a novel weight-decay linear scaling rule for large-batch training that proportionally increases weight decay rather than the learning rate as the batch size increases.
Revised: Second, we propose a novel weight-decay linear scaling rule for large-batch training that proportionally increases weight decay rather than the learning rate as the batch size increases.

Original: Third, we provide an effective learningrate-aware scheduler for weight decay, called the Stable Weight Decay (SWD) method, which, to the best of our knowledge, is the ﬁrst practical design for weight decay scheduling.
Revised: Third, we provide an effective learningrate-aware scheduler for weight decay, called the Stable Weight Decay (SWD) method, which, to the best of our knowledge, is the ﬁrst practical design for weight decay scheduling.

Original: In our various experiments, the SWD method often makes improvements over L 2 Regularization and Decoupled Weight Decay.
Revised: In our various experiments, the SWD method often makes improvements over L 2 Regularization and Decoupled Weight Decay.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Weight decay is a popular and even necessary regularization technique for training deep neural networks that generalize well (Krogh and Hertz, 1992).
Revised: Weight decay is a popular and even necessary regularization technique for training deep neural networks that generalize well (Krogh and Hertz, 1992).

Original: People commonly use L 2 regularization as “weight decay” for training of deep neural networks and interpret it as a Gaussian prior over the model weights (David, 1992; Graves, 2011).
Revised: People commonly use L 2 regularization as “weight decay” for training of deep neural networks and interpret it as a Gaussian prior over the model weights (David, 1992; Graves, 2011).

Original: This is true for vanilla Stochastic Gradient Descent (SGD).
Revised: This is true for vanilla Stochastic Gradient Descent (SGD).

Original: However, Loshchilov and Hutter (2018) revealed that, when the learning rate is adaptive, the commonly used L 2 regularization is not identical to the vanilla weight decay proposed by Hanson and Pratt (1989):
Revised: However, Loshchilov and Hutter (2018) revealed that, when the learning rate is adaptive, the commonly used L 2 regularization is not identical to the vanilla weight decay proposed by Hanson and Pratt (1989):

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where λ (cid:48) is the weight decay hyperparameter, η is the learning rate, and L ( θ t − 1 ) is the loss function over one minibatch.
Revised: where λ (cid:48) is the weight decay hyperparameter, η is the learning rate, and L ( θ t − 1 ) is the loss function over one minibatch.

Original: We denote the loss function over the whole training dataset as L ( θ ) .
Revised: We denote the loss function over the whole training dataset as L ( θ ) .

Original: A conventional belief is that, from the Bayesian perspective, weight decay indicates a Gaussian prior over the model parameters, which regularize the complexity of the learned model (Graves, 2011).
Revised: A conventional belief is that, from the Bayesian perspective, weight decay indicates a Gaussian prior over the model parameters, which regularize the complexity of the learned model (Graves, 2011).

Original: However, we argue that this Bayesian interpretation cannot explain the empirical success of weight decay in deep learning.
Revised: However, we argue that this Bayesian interpretation cannot explain the empirical success of weight decay in deep learning.

Original: On one hand, Wenzel et al.
Revised: On one hand, Wenzel et al.

Original: (2020) demonstrated that the conventional
Revised: (2020) demonstrated that the conventional

Original: Bayesian posterior usually does not work well in deep learning.
Revised: Bayesian posterior usually does not work well in deep learning.

Original: This challenged the meaning of weight decay as a proper prior in the Bayesian interpretation.
Revised: This challenged the meaning of weight decay as a proper prior in the Bayesian interpretation.

Original: On the other hand, Lewkowycz and GurAri (2020) and our experimental results (Figure 2) both suggested that weight decay is surprisingly unnecessary for long-time training.
Revised: On the other hand, Lewkowycz and GurAri (2020) and our experimental results (Figure 2) both suggested that weight decay is surprisingly unnecessary for long-time training.

Original: This observation critically challenged the Bayesian interpretation of weight decay in deep learning.
Revised: This observation critically challenged the Bayesian interpretation of weight decay in deep learning.

Original: The Bayesian interpretation believes that the prior knowledge over model parameters truly exist.
Revised: The Bayesian interpretation believes that the prior knowledge over model parameters truly exist.

Original: Thus, the prior knowledge should not depend on the training time of deep neural networks, which contradicts with the experimental results in Figure 2.
Revised: Thus, the prior knowledge should not depend on the training time of deep neural networks, which contradicts with the experimental results in Figure 2.

Original: Recently researchers started to explore other explanations for weight decay.
Revised: Recently researchers started to explore other explanations for weight decay.

Original: Zhang et al.
Revised: Zhang et al.

Original: (2019b) revealed three different roles of weight decay.
Revised: (2019b) revealed three different roles of weight decay.

Original: Some papers (Van Laarhoven, 2017; Zhang et al., 2019b) argued that weigh decay increases the effective learning rate but has no regularization effect
Revised: Some papers (Van Laarhoven, 2017; Zhang et al., 2019b) argued that weigh decay increases the effective learning rate but has no regularization effect

Original: [Table]
Revised: [Table]

Original: when combined with scale-invariant normalization layers
Revised: when combined with scale-invariant normalization layers

Original: [Table]
Revised: [Table]

Original: such as batch normalization (Ioffe and
Revised: such as batch normalization (Ioffe and

Original: Szegedy, 2015)
Revised: Szegedy, 2015)

Original: [Table]
Revised: [Table]

Original: Li et al.
Revised: Li et al.

Original: (2020) deﬁned a new “intrinsic learning rate” parameter that is the product of the vanilla learning rate η and the weight decay factor λ and argued that the convergence time depends on this intrinsic learning rate.
Revised: (2020) deﬁned a new “intrinsic learning rate” parameter that is the product of the vanilla learning rate η and the weight decay factor λ and argued that the convergence time depends on this intrinsic learning rate.

Original: However, people still expect to further explore the theoretical mechanism of weight decay, because weight decay sometimes shows mysterious behaviors beyond the existing theoretical understanding.
Revised: However, people still expect to further explore the theoretical mechanism of weight decay, because weight decay sometimes shows mysterious behaviors beyond the existing theoretical understanding.

Original: As we discussed above, the Bayesian interpretation has been critically challenged by recent empirical ﬁndings (Lewkowycz and Gur-Ari, 2020).
Revised: As we discussed above, the Bayesian interpretation has been critically challenged by recent empirical ﬁndings (Lewkowycz and Gur-Ari, 2020).

Original: The interpretation of increasing the effective learning rate only holds for scale-invariant normalization layers, which is a narrow scope.
Revised: The interpretation of increasing the effective learning rate only holds for scale-invariant normalization layers, which is a narrow scope.

Original: Contributions.
Revised: Contributions.

Original: In this paper, we try to theoretically understand the novel behaviors of weight decay and explore the potential of weight decay.
Revised: In this paper, we try to theoretically understand the novel behaviors of weight decay and explore the potential of weight decay.

Original: Our paper mainly has three contributions.
Revised: Our paper mainly has three contributions.

Original: First, we propose a novel theoretical interpretation of weight decay from the perspective of learning dynamics.
Revised: First, we propose a novel theoretical interpretation of weight decay from the perspective of learning dynamics.

Original: We study how the stability of the stationary points and the convergence time depend on weight decay, which corresponds to two main effects of weight decay.
Revised: We study how the stability of the stationary points and the convergence time depend on weight decay, which corresponds to two main effects of weight decay.

Original: Second, we propose a novel weight-decay linear scaling rule (See Rule 2 in Section 3) for large-batch training.
Revised: Second, we propose a novel weight-decay linear scaling rule (See Rule 2 in Section 3) for large-batch training.

Original: We discover that the existing linear scaling rule (Krizhevsky, 2014; Goyal et al., 2017) for large-batch training that proportionally increases the learning rate may sometimes be harmful to large-batch training due to bad convergence.
Revised: We discover that the existing linear scaling rule (Krizhevsky, 2014; Goyal et al., 2017) for large-batch training that proportionally increases the learning rate may sometimes be harmful to large-batch training due to bad convergence.

Original: Instead, increasing weight decay is a more suitable solution in this case.
Revised: Instead, increasing weight decay is a more suitable solution in this case.

Original: Third, to the best of our knowledge, we are the ﬁrst to design an effective learning-rate-aware scheduler for weight decay which may signiﬁcantly improve test performance.
Revised: Third, to the best of our knowledge, we are the ﬁrst to design an effective learning-rate-aware scheduler for weight decay which may signiﬁcantly improve test performance.

Original: Our work ﬁlls the gap in this direction.
Revised: Our work ﬁlls the gap in this direction.

Original: We note that a trivial weight decay scheduling method recently proposed by
Revised: We note that a trivial weight decay scheduling method recently proposed by

Original: Lewkowycz and Gur-Ari (2020) has limited practical value because it only works well with constant learning rate, as the original paper claimed.
Revised: Lewkowycz and Gur-Ari (2020) has limited practical value because it only works well with constant learning rate, as the original paper claimed.

Original: Inspired by the stability of the stationary points, we show that weight decay should be coupled with the effective learning rate.
Revised: Inspired by the stability of the stationary points, we show that weight decay should be coupled with the effective learning rate.

Original: We call the proposed method
Revised: We call the proposed method

Original: Stable Weight Decay (SWD).
Revised: Stable Weight Decay (SWD).

Original: Adam (Kingma and Ba, 2015) with SWD (AdamS) is displayed in
Revised: Adam (Kingma and Ba, 2015) with SWD (AdamS) is displayed in

Original: Algorithm 2.
Revised: Algorithm 2.

Original: The results in Table 1 fully supports that SWD can signiﬁcantly improve Adam.
Revised: The results in Table 1 fully supports that SWD can signiﬁcantly improve Adam.

Original: Structure of this paper.
Revised: Structure of this paper.

Original: In Section 2, we analyze the dynamics of weight decay.
Revised: In Section 2, we analyze the dynamics of weight decay.

Original: In Section 3, we study how weight decay can improve large-batch training.
Revised: In Section 3, we study how weight decay can improve large-batch training.

Original: In Section 4, we propose SWD.
Revised: In Section 4, we propose SWD.

Original: In Section 5, we empirically study SWD.
Revised: In Section 5, we empirically study SWD.

Original: In Section 6, we conclude our work.
Revised: In Section 6, we conclude our work.

Original:  U NDERSTANDING THE D YNAMICS OF W EIGHT D ECAY
Revised:  U NDERSTANDING THE D YNAMICS OF W EIGHT D ECAY

Original: In this section, we provide a way to understand weight decay from a viewpoint of learning dynamics.
Revised: In this section, we provide a way to understand weight decay from a viewpoint of learning dynamics.

Original: Weight decay should be coupled with the learning rate scheduler.
Revised: Weight decay should be coupled with the learning rate scheduler.

Original: The vanilla weight decay described by Hanson and Pratt (1989) is given by Equation (1).
Revised: The vanilla weight decay described by Hanson and Pratt (1989) is given by Equation (1).

Original: A more popular implementation for vanilla SGD in modern deep learning libraries, such as PyTorch (Paszke et al., 2019) and TensorFlow (Abadi et al., 2016), is given by
Revised: A more popular implementation for vanilla SGD in modern deep learning libraries, such as PyTorch (Paszke et al., 2019) and TensorFlow (Abadi et al., 2016), is given by

Original: [Equation]
Revised: [Equation]

Original: where η t is the learning rate at the t -th step and weight decay is coupled with the learning rate scheduler.
Revised: where η t is the learning rate at the t -th step and weight decay is coupled with the learning rate scheduler.

Original: While previous work did not studied why weight decay should be coupled with the learning rate scheduler, this trick has been adopted by most deep learning libraries.
Revised: While previous work did not studied why weight decay should be coupled with the learning rate scheduler, this trick has been adopted by most deep learning libraries.

Original: We ﬁrst take vanilla SGD as the studied example, where no Momentum is involved.
Revised: We ﬁrst take vanilla SGD as the studied example, where no Momentum is involved.

Original: It is easy to see that vanilla SGD with L 2 regularization λ 2 (cid:107) θ (cid:107) 2 is also given by Equation (2).
Revised: It is easy to see that vanilla SGD with L 2 regularization λ 2 (cid:107) θ (cid:107) 2 is also given by Equation (2).

Original: Suppose the learning rate is ﬁxed in the whole training procedure.
Revised: Suppose the learning rate is ﬁxed in the whole training procedure.

Original: Then Equation (2) will be identical to Equation (1) if we simply choose λ (cid:48)
Revised: Then Equation (2) will be identical to Equation (1) if we simply choose λ (cid:48)

Original: = ηλ .
Revised: = ηλ .

Original: However, learning rate decay is quite important for training of deep neural networks.
Revised: However, learning rate decay is quite important for training of deep neural networks.

Original: So Equation (2) is not identical to Equation (1) in practice.
Revised: So Equation (2) is not identical to Equation (1) in practice.

Original: Which implementation is better?
Revised: Which implementation is better?

Original: In Figure 1, we empirically veriﬁed that the popular Equation-(2)- based weight decay indeed outperforms the vanilla implementation in Equation (1).
Revised: In Figure 1, we empirically veriﬁed that the popular Equation-(2)- based weight decay indeed outperforms the vanilla implementation in Equation (1).

Original: We argue that
Revised: We argue that

Original: Equation-(2)-based weight decay is theoretically better than Equation-(1)-based weight decay in terms of the stability of the stationary points.
Revised: Equation-(2)-based weight decay is theoretically better than Equation-(1)-based weight decay in terms of the stability of the stationary points.

Original: We deﬁne the stability as Deﬁnition 1.
Revised: We deﬁne the stability as Deﬁnition 1.

Original: Deﬁnition 1 (The stability of the stationary point) .
Revised: Deﬁnition 1 (The stability of the stationary point) .

Original: The stationary point is stable if a stationary point is ﬁxed during training.
Revised: The stationary point is stable if a stationary point is ﬁxed during training.

Original: Otherwise, we deﬁne it as an unstable stationary point.
Revised: Otherwise, we deﬁne it as an unstable stationary point.

Original: The training objective at step- t of Equation (1) is L ( θ ) + λ (cid:48)
Revised: The training objective at step- t of Equation (1) is L ( θ ) + λ (cid:48)

Original: η t (cid:107) θ (cid:107) 2 , while that of Equation (2) is
Revised: η t (cid:107) θ (cid:107) 2 , while that of Equation (2) is

Original: L ( θ ) + λ (cid:107) θ (cid:107) 2 .
Revised: L ( θ ) + λ (cid:107) θ (cid:107) 2 .

Original: Thus, the 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.
Revised: Thus, the 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.

Original: Weight decay matters to the stability.
Revised: Weight decay matters to the stability.

Original: As the stability of the stationary points critically depends on the weight decay scheduler, we deﬁne stable weight decay accordingly in Deﬁnition 2.
Revised: As the stability of the stationary points critically depends on the weight decay scheduler, we deﬁne stable weight decay accordingly in Deﬁnition 2.

Original: Deﬁnition 2 (Stable Weight Decay) .
Revised: Deﬁnition 2 (Stable Weight Decay) .

Original: The weight decay is stable if the stationary points are stable during training.
Revised: The weight decay is stable if the stationary points are stable during training.

Original: Otherwise, we deﬁne it as unstable weight decay.
Revised: Otherwise, we deﬁne it as unstable weight decay.

Original: We propose Theorem 1 and prove that even deterministic Gradient Descent (GD) with vanilla weight decay may not converge to any non-zero stationary point due the instability of the stationary points.
Revised: We propose Theorem 1 and prove that even deterministic Gradient Descent (GD) with vanilla weight decay may not converge to any non-zero stationary point due the instability of the stationary points.

Original: Theorem 1 (Non-convergence of GD with unstable weight decay) .
Revised: Theorem 1 (Non-convergence of GD with unstable weight decay) .

Original: Suppose learning dynamics is governed by GD with vanilla weight decay (Equation (1) ) and the learning rate η t ∈ (0 , + ∞ ) holds.
Revised: Suppose learning dynamics is governed by GD with vanilla weight decay (Equation (1) ) and the learning rate η t ∈ (0 , + ∞ ) holds.

Original: If ∃ δ that satisﬁes 0 < δ ≤ | η t − η t +1 | for any t > t 0 , then the learning dynamics cannot converge to any non-zero stationary point satisfying the condition
Revised: If ∃ δ that satisﬁes 0 < δ ≤ | η t − η t +1 | for any t > t 0 , then the learning dynamics cannot converge to any non-zero stationary point satisfying the condition

Original: [Equation]
Revised: [Equation]

Original: Remark.
Revised: Remark.

Original: We leave the proof in Appendix A.1.
Revised: We leave the proof in Appendix A.1.

Original: Even if the gradient is zero at the t -th step, the gradient at the t + 1 -th step will not be zero.
Revised: Even if the gradient is zero at the t -th step, the gradient at the t + 1 -th step will not be zero.

Original: Unstable weight decay causes changing the objective and hence the solution does not converge to any non-zero stationary point.
Revised: Unstable weight decay causes changing the objective and hence the solution does not converge to any non-zero stationary point.

Original: Thus, GD with unstable weight decay has no theoretical convergence guarantee like GD (Wolfe, 1969).
Revised: Thus, GD with unstable weight decay has no theoretical convergence guarantee like GD (Wolfe, 1969).

Original: This explains why weight decay should be coupled with the learning rate to keep the stationary points stable during training.
Revised: This explains why weight decay should be coupled with the learning rate to keep the stationary points stable during training.

Original: We conjecture that the stability of the stationary points should be a desirable property in deep learning.
Revised: We conjecture that the stability of the stationary points should be a desirable property in deep learning.

Original: Dynamics of weight decay.
Revised: Dynamics of weight decay.

Original: Lewkowycz and Gur-Ari (2020) empirically discovered that, if we let the training time be approximately inverse to weight decay, the test error can monotonically decrease as weight decay decreases, which is beyond the existing theoretical understanding.
Revised: Lewkowycz and Gur-Ari (2020) empirically discovered that, if we let the training time be approximately inverse to weight decay, the test error can monotonically decrease as weight decay decreases, which is beyond the existing theoretical understanding.

Original: We will show that it is easy to understand this empirical observation by analyzing the dynamics of weight decay.
Revised: We will show that it is easy to understand this empirical observation by analyzing the dynamics of weight decay.

Original: We ﬁrst state Assumption 1.
Revised: We ﬁrst state Assumption 1.

Original: Assumption 1 which holds near minima are common and useful for analyzing the solution or the convergence behavior in related papers (Mandt et al., 2017; Neyshabur et al., 2017; Xie et al., 2021a;
Revised: Assumption 1 which holds near minima are common and useful for analyzing the solution or the convergence behavior in related papers (Mandt et al., 2017; Neyshabur et al., 2017; Xie et al., 2021a;

Original: Li et al., 2017; Zhang et al., 2019a;
Revised: Li et al., 2017; Zhang et al., 2019a;

Original: Zhou et al., 2020).
Revised: Zhou et al., 2020).

Original: It is known that stochastic optimization can escape multiple bad minima (Xie et al., 2021b; Zhu et al., 2019;
Revised: It is known that stochastic optimization can escape multiple bad minima (Xie et al., 2021b; Zhu et al., 2019;

Original: Kleinberg et al., 2018) and ﬁnally converge to a good minimum during training.
Revised: Kleinberg et al., 2018) and ﬁnally converge to a good minimum during training.

Original: Thus, Assumptioncan be repeatedly applied near each minimum during training.
Revised: Thus, Assumptioncan be repeatedly applied near each minimum during training.

Original: Assumption 1.
Revised: Assumption 1.

Original: The loss function around a minimum θ (cid:63) can be written as
Revised: The loss function around a minimum θ (cid:63) can be written as

Original: [Equation]
Revised: [Equation]

Original: where H is the Hessian at θ (cid:63) .
Revised: where H is the Hessian at θ (cid:63) .

Original: We propose Theorem 2 and demonstrate how the learned model parameters θ may depend on weight decay.
Revised: We propose Theorem 2 and demonstrate how the learned model parameters θ may depend on weight decay.

Original: We leave the proof in Appendix A.2.
Revised: We leave the proof in Appendix A.2.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: With the ﬁxed 200 epochs, the optimal weight decay is about 0 . 0005 .
Revised: With the ﬁxed 200 epochs, the optimal weight decay is about 0 . 0005 .

Original: With 0 .
Revised: With 0 .

Original: 1 λ − 1 epochs, decreasing weight decay monotonically decreases test performance.
Revised: 1 λ − 1 epochs, decreasing weight decay monotonically decreases test performance.

Original: Table 3 further supports that the optimal weight decay is approximately inverse to the number of epochs.
Revised: Table 3 further supports that the optimal weight decay is approximately inverse to the number of epochs.

Original: Theorem 2 (Dynamics of weight decay) .
Revised: Theorem 2 (Dynamics of weight decay) .

Original: Suppose Assumption 1 holds and learning dynamics is governed by SGD with weight decay (Equation (2) ).
Revised: Suppose Assumption 1 holds and learning dynamics is governed by SGD with weight decay (Equation (2) ).

Original: If the learning rate scheduler is constant the expected θ t after t iterations satisﬁes the following:
Revised: If the learning rate scheduler is constant the expected θ t after t iterations satisﬁes the following:

Original: [Equation]
Revised: [Equation]

Original: where I is the identity matrix.
Revised: where I is the identity matrix.

Original: If the learning rate scheduler is given by the sequence { η 1 , η 2 , . . .
Revised: If the learning rate scheduler is given by the sequence { η 1 , η 2 , . . .

Original: , η t } , the expected θ t after t iterations satisﬁes the following:
Revised: , η t } , the expected θ t after t iterations satisﬁes the following:

Original: [Equation]
Revised: [Equation]

Original: Theorem 2 is a useful tool to understand convergence behaviors which may repeatedly happen during searching minima of deep learning.
Revised: Theorem 2 is a useful tool to understand convergence behaviors which may repeatedly happen during searching minima of deep learning.

Original: While Theorem 2 depends on Assumption 1, we will show that it is remarkably consistent with novel empirical behaviors of weight decay and can also reveal novel insights.
Revised: While Theorem 2 depends on Assumption 1, we will show that it is remarkably consistent with novel empirical behaviors of weight decay and can also reveal novel insights.

Original: Theorem 2 indicates that weight decay mainly has two effects: 1) biasing the expected learned solution and 2) accelerating the convergence.
Revised: Theorem 2 indicates that weight decay mainly has two effects: 1) biasing the expected learned solution and 2) accelerating the convergence.

Original: Effect 1.
Revised: Effect 1.

Original: Weight decay biases the stationary points.
Revised: Weight decay biases the stationary points.

Original: Theorem 2 provides us two key insights about weight decay.
Revised: Theorem 2 provides us two key insights about weight decay.

Original: The ﬁrst insight is Corollary 1 derived from Theorem 2 by choosing the long-time limit.
Revised: The ﬁrst insight is Corollary 1 derived from Theorem 2 by choosing the long-time limit.

Original: Corollary 1 indicates that large weight decay may not learn accurate minima well.
Revised: Corollary 1 indicates that large weight decay may not learn accurate minima well.

Original: This insight is consistent with the regularization-based understanding of weight decay (Krogh and Hertz, 1992).
Revised: This insight is consistent with the regularization-based understanding of weight decay (Krogh and Hertz, 1992).

Original: Corollary 1.
Revised: Corollary 1.

Original: Suppose the conditions of Theorem 2 hold and H
Revised: Suppose the conditions of Theorem 2 hold and H

Original: i + λ > 0 , where H i is the i -th eigenvalue of H .
Revised: i + λ > 0 , where H i is the i -th eigenvalue of H .

Original: Then, in the limit of long training time that t → + ∞ , the displacement from the true minimum θ (cid:63) to the expected learned solution along the i -th direction is given by
Revised: Then, in the limit of long training time that t → + ∞ , the displacement from the true minimum θ (cid:63) to the expected learned solution along the i -th direction is given by

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ) where the bias µ = lim t →∞ E [ θ t − θ (cid:63) ] .
Revised: ) where the bias µ = lim t →∞ E [ θ t − θ (cid:63) ] .

Original: Figure 2 suggests that, given long enough training time inverse to weight decay, decreasing the weight decay hyperparameter may consistently improve test performance, which is also veriﬁed by
Revised: Figure 2 suggests that, given long enough training time inverse to weight decay, decreasing the weight decay hyperparameter may consistently improve test performance, which is also veriﬁed by

Original: Lewkowycz and Gur-Ari (2020).
Revised: Lewkowycz and Gur-Ari (2020).

Original: This observation is contradicted with the conventional regularizationbased understanding (Krogh and Hertz, 1992), but fully supports Corollary 1 that weight decay is bad, if we already have good convergence.
Revised: This observation is contradicted with the conventional regularizationbased understanding (Krogh and Hertz, 1992), but fully supports Corollary 1 that weight decay is bad, if we already have good convergence.

Original: Effect 2.
Revised: Effect 2.

Original: Weight decay accelerates convergence.
Revised: Weight decay accelerates convergence.

Original: However, we discover that the long-time posterior cannot explain the convergence behaviors.
Revised: However, we discover that the long-time posterior cannot explain the convergence behaviors.

Original: Figure 2 also suggests that, we need to increase the number of training iterations by 1 λ for maintaining good convergence, which is also discovered by Lewkowycz and Gur-Ari (2020) and Li et al.
Revised: Figure 2 also suggests that, we need to increase the number of training iterations by 1 λ for maintaining good convergence, which is also discovered by Lewkowycz and Gur-Ari (2020) and Li et al.

Original: To explain why the convergence time depends on 1 λ , we need to explore the dynamics of weight decay before approaching the convergence.
Revised: To explain why the convergence time depends on 1 λ , we need to explore the dynamics of weight decay before approaching the convergence.

Original: The second insight is Corollary 2 that estimates the number of iterations for good convergence.
Revised: The second insight is Corollary 2 that estimates the number of iterations for good convergence.

Original: Corollary 2.
Revised: Corollary 2.

Original: Suppose the conditions of Theorem 2 hold and δ i
Revised: Suppose the conditions of Theorem 2 hold and δ i

Original: = H i + λ > 0 , where H i is the i -th eigenvalue of H .
Revised: = H i + λ > 0 , where H i is the i -th eigenvalue of H .

Original: Then, the convergence time in the i -th direction is
Revised: Then, the convergence time in the i -th direction is

Original: [Equation]
Revised: [Equation]

Original: It is well-known that the Hessian in deep learning has a very small number of large eigenvalues and a large number of nearly zero eigenvalues (Sagun et al., 2016; 2017).
Revised: It is well-known that the Hessian in deep learning has a very small number of large eigenvalues and a large number of nearly zero eigenvalues (Sagun et al., 2016; 2017).

Original: Thus, we have H i
Revised: Thus, we have H i

Original: + λ
Revised: + λ

Original: ≈ λ for most directions.
Revised: ≈ λ for most directions.

Original: It means that the convergence time critically depend on the dynamics in the subspace corresponding to those nearly zero eigenvalues:
Revised: It means that the convergence time critically depend on the dynamics in the subspace corresponding to those nearly zero eigenvalues:

Original: [Equation]
Revised: [Equation]

Original: Equation (7) reveals two facts.
Revised: Equation (7) reveals two facts.

Original: First, Equation (7) supports the statement that the initial learning rate and the weight decay hyperparameter λ are linearly coupled for maintaining performance (Loshchilov and Hutter, 2018; Van Laarhoven, 2017; Zhang et al., 2019b; Hoffer et al., 2018).
Revised: First, Equation (7) supports the statement that the initial learning rate and the weight decay hyperparameter λ are linearly coupled for maintaining performance (Loshchilov and Hutter, 2018; Van Laarhoven, 2017; Zhang et al., 2019b; Hoffer et al., 2018).

Original: Second, Equation (7) demonstrates that the number of iterations (not the number of epochs or training time) for good convergence depends on O (cid:0) λ − 1 η − 1 (cid:1) .
Revised: Second, Equation (7) demonstrates that the number of iterations (not the number of epochs or training time) for good convergence depends on O (cid:0) λ − 1 η − 1 (cid:1) .

Original: This not only explains the relation in Figure 2, but also reveals that we should increase the number of iterations instead of the number of epochs by 1 λ .
Revised: This not only explains the relation in Figure 2, but also reveals that we should increase the number of iterations instead of the number of epochs by 1 λ .

Original: We note that the second fact is also reported by Lewkowycz and Gur-Ari (2020), while Lewkowycz and Gur-Ari (2020) did not present formal theoretical analysis like Theorem 2.
Revised: We note that the second fact is also reported by Lewkowycz and Gur-Ari (2020), while Lewkowycz and Gur-Ari (2020) did not present formal theoretical analysis like Theorem 2.

Original: In our theoretical analysis, weight decay mainly has two effects.
Revised: In our theoretical analysis, weight decay mainly has two effects.

Original: The ﬁrst effect is to bias the expected solution in the long-time limit.
Revised: The ﬁrst effect is to bias the expected solution in the long-time limit.

Original: The ﬁrst effect is negative for learning accurate minima.
Revised: The ﬁrst effect is negative for learning accurate minima.

Original: The second effect is to accelerate the convergence in the subspace corresponding to nearly zero eigenvalues of the Hessian.
Revised: The second effect is to accelerate the convergence in the subspace corresponding to nearly zero eigenvalues of the Hessian.

Original: The second effect is positive for accelerating training.
Revised: The second effect is positive for accelerating training.

Original: Tuning weight decay needs to balance the tradeoff between the two effects.
Revised: Tuning weight decay needs to balance the tradeoff between the two effects.

Original:  W EIGHT D ECAY I MPROVES L ARGE -B ATCH T RAINING
Revised:  W EIGHT D ECAY I MPROVES L ARGE -B ATCH T RAINING

Original: In this section, motivated by the theoretical insights in Section 2, we study how weight decay may surprisingly improve large-batch training.
Revised: In this section, motivated by the theoretical insights in Section 2, we study how weight decay may surprisingly improve large-batch training.

Original: Large-batch training can efﬁciently utilize the parallel computation to speedup training of deep networks (Goyal et al., 2017).
Revised: Large-batch training can efﬁciently utilize the parallel computation to speedup training of deep networks (Goyal et al., 2017).

Original: However, large-batch training suffers from the difﬁculty of minimizing the training loss and usually ﬁnd sharp minima that do not generalize well (Hoffer et al., 2017; Keskar et al., 2017).
Revised: However, large-batch training suffers from the difﬁculty of minimizing the training loss and usually ﬁnd sharp minima that do not generalize well (Hoffer et al., 2017; Keskar et al., 2017).

Original: How to perform large-batch training with good generalization is an important problem which has attracted much attention from both academia (Zhang et al., 2019a; Wu et al., 2020; Wen et al., 2020) and industry (You et al., 2017; Goyal et al., 2017).
Revised: How to perform large-batch training with good generalization is an important problem which has attracted much attention from both academia (Zhang et al., 2019a; Wu et al., 2020; Wen et al., 2020) and industry (You et al., 2017; Goyal et al., 2017).

Original: Rule 1 mitigates the small noise problem in large-batch training.
Revised: Rule 1 mitigates the small noise problem in large-batch training.

Original: The famous linear scaling rule, namely Rule 1, for large-batch training is well-known (Krizhevsky, 2014; Goyal et al., 2017).
Revised: The famous linear scaling rule, namely Rule 1, for large-batch training is well-known (Krizhevsky, 2014; Goyal et al., 2017).

Original: Large-batch training can easily get stuck in saddle points and sharp minima due to small gradient noise, because the noise magnitude in SGD dynamics is approximately proportional to the ratio of the learning rate η to batch size B , namely ηB (Mandt et al., 2017; Keskar et al., 2017).
Revised: Large-batch training can easily get stuck in saddle points and sharp minima due to small gradient noise, because the noise magnitude in SGD dynamics is approximately proportional to the ratio of the learning rate η to batch size B , namely ηB (Mandt et al., 2017; Keskar et al., 2017).

Original: Xie et al.
Revised: Xie et al.

Original: (2021b) recently proved that ηB exponentially matters to SGD dynamics.
Revised: (2021b) recently proved that ηB exponentially matters to SGD dynamics.

Original: Thus, the learning-rate linear scaling rule is a natural result of maintaining the similar magnitude of stochastic gradient noise and similar learning dynamics invariant to the batch size.
Revised: Thus, the learning-rate linear scaling rule is a natural result of maintaining the similar magnitude of stochastic gradient noise and similar learning dynamics invariant to the batch size.

Original: When the noise magnitude is the performance bottleneck, Rule 1 is usually an effective solution.
Revised: When the noise magnitude is the performance bottleneck, Rule 1 is usually an effective solution.

Original: Rule 1.
Revised: Rule 1.

Original: When the batch size is multiplied by k , multiply the learning rate by k .
Revised: When the batch size is multiplied by k , multiply the learning rate by k .

Original: For large-batch training, people usually focused on the noise magnitude problem but overlooked the convergence problem (Jain et al., 2018; Ma et al., 2018;
Revised: For large-batch training, people usually focused on the noise magnitude problem but overlooked the convergence problem (Jain et al., 2018; Ma et al., 2018;

Original: Yin et al., 2018).
Revised: Yin et al., 2018).

Original: Sometimes, the noise magnitude is not the performance bottleneck.
Revised: Sometimes, the noise magnitude is not the performance bottleneck.

Original: People usually perform large-batch training with a ﬁxed number of epochs for keeping the total computational costs moderate.
Revised: People usually perform large-batch training with a ﬁxed number of epochs for keeping the total computational costs moderate.

Original: Under the constraint of ﬁxing the number of epochs, obviously, the number of iterations will be multiplied by k − 1 if the batch size is multiplied by k .
Revised: Under the constraint of ﬁxing the number of epochs, obviously, the number of iterations will be multiplied by k − 1 if the batch size is multiplied by k .

Original: Thus, the number of iterations may be too small to achieve good convergence.
Revised: Thus, the number of iterations may be too small to achieve good convergence.

Original: Rule 2 mitigates the bad convergence problem in large-batch training.
Revised: Rule 2 mitigates the bad convergence problem in large-batch training.

Original: If the bad convergence problem is the performance bottleneck, the learning-rate linear scaling rule will even be harmful to large-batch training due to slower convergence.
Revised: If the bad convergence problem is the performance bottleneck, the learning-rate linear scaling rule will even be harmful to large-batch training due to slower convergence.

Original: When Rule 1 causes optimization divergence or bad convergence, we instead propose the weight-decay linear scaling rule as Rule 2 based on Equation
Revised: When Rule 1 causes optimization divergence or bad convergence, we instead propose the weight-decay linear scaling rule as Rule 2 based on Equation

Original: [Figure]
Revised: [Figure]

Original: fully supports the proposed Rule 2.
Revised: fully supports the proposed Rule 2.

Original: Because large weight decay may accelerate the convergence.
Revised: Because large weight decay may accelerate the convergence.

Original: Note that, increasing the learning rate only works in a narrow range, i.e., B ≤ 8192 on ImageNet (Keskar et al., 2017), since too large learning rates may lead to optimization divergence or bad convergence (Keskar et al., 2017; Masters and Luschi, 2018).To the best of our knowledge, we are the ﬁrst to propose the weight-decay-based rule for large-batch training.
Revised: Note that, increasing the learning rate only works in a narrow range, i.e., B ≤ 8192 on ImageNet (Keskar et al., 2017), since too large learning rates may lead to optimization divergence or bad convergence (Keskar et al., 2017; Masters and Luschi, 2018).To the best of our knowledge, we are the ﬁrst to propose the weight-decay-based rule for large-batch training.

Original: Previous papers did not design a rule for the second-type problem.
Revised: Previous papers did not design a rule for the second-type problem.

Original: Rule 2.
Revised: Rule 2.

Original: When the batch size is multiplied by k , multiply the weight decay by k .
Revised: When the batch size is multiplied by k , multiply the weight decay by k .

Original: Empirical Analysis of Large-Batch Training.
Revised: Empirical Analysis of Large-Batch Training.

Original: We conducted experiments to show that, when the learning-rate linear scaling rule is harmful due to the convergence problem, the weight-decay linear scaling rule signiﬁcantly improves large-batch training.
Revised: We conducted experiments to show that, when the learning-rate linear scaling rule is harmful due to the convergence problem, the weight-decay linear scaling rule signiﬁcantly improves large-batch training.

Original: In Figure 3, we trained ResNet18 via SGD on
Revised: In Figure 3, we trained ResNet18 via SGD on

Original: CIFAR-10 with the very large batch size B = 16384 , which is 128 times the common batch size 128 .
Revised: CIFAR-10 with the very large batch size B = 16384 , which is 128 times the common batch size 128 .

Original: We note that, in this common setting on CIFAR-10, even slightly increasing the learning rate is very harmful to large-batch training.
Revised: We note that, in this common setting on CIFAR-10, even slightly increasing the learning rate is very harmful to large-batch training.

Original: Figure 3 demonstrates Rule 1 may seriously cause bad convergence.
Revised: Figure 3 demonstrates Rule 1 may seriously cause bad convergence.

Original: This observation means that, in the setting of training ResNet18 on CIFAR-10, the convergence problem is the main performance bottleneck instead of the noise magnitude.
Revised: This observation means that, in the setting of training ResNet18 on CIFAR-10, the convergence problem is the main performance bottleneck instead of the noise magnitude.

Original: As Rule 2, we observed that multiplying the weight decay hyperparameter by 128 works best.
Revised: As Rule 2, we observed that multiplying the weight decay hyperparameter by 128 works best.

Original: Thus, in this second-type problem of large-batch training, Rule 2 signiﬁcantly outperforms Rule 1.
Revised: Thus, in this second-type problem of large-batch training, Rule 2 signiﬁcantly outperforms Rule 1.

Original:  S TABLE W EIGHT D ECAY
Revised:  S TABLE W EIGHT D ECAY

Original: In this section, we design a weight decay scheduler for adaptive gradient methods, such as Adam.
Revised: In this section, we design a weight decay scheduler for adaptive gradient methods, such as Adam.

Original: Algorithm 1: Adam/AdamW
Revised: Algorithm 1: Adam/AdamW

Original: [Equation]
Revised: [Equation]

Original: Algorithm
Revised: Algorithm

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Decoupled weight decay is unstable weight decay in adaptive gradient methods.
Revised: Decoupled weight decay is unstable weight decay in adaptive gradient methods.

Original: Loshchilov and Hutter (2018)
Revised: Loshchilov and Hutter (2018)

Original: ﬁrst pointed that, when the learning rate is adaptive, the commonly used Lregularization is not identical to weight decay.
Revised: ﬁrst pointed that, when the learning rate is adaptive, the commonly used Lregularization is not identical to weight decay.

Original: In the following analysis, we ignore the effect of Momentum and focus on the effect of Adaptive Learning Rate for simplicity.
Revised: In the following analysis, we ignore the effect of Momentum and focus on the effect of Adaptive Learning Rate for simplicity.

Original: Thus, AdamW (Loshchilov and Hutter, 2018) can be written as
Revised: Thus, AdamW (Loshchilov and Hutter, 2018) can be written as

Original: [Equation]
Revised: [Equation]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Even with similar or higher training losses, AdamS still generalizes better than other Adam variants.
Revised: Even with similar or higher training losses, AdamS still generalizes better than other Adam variants.

Original: We leave the scatter plot on CIFAR-10 in Appendix C.
Revised: We leave the scatter plot on CIFAR-10 in Appendix C.

Original: [Figure]
Revised: [Figure]

Original: by training ResNet34 on CIFAR-
Revised: by training ResNet34 on CIFAR-

Original: [Figure]
Revised: [Figure]

Original: variants.
Revised: variants.

Original: The test performance of other models can be found in Table 1.
Revised: The test performance of other models can be found in Table 1.

Original: [Figure]
Revised: [Figure]

Original: where v t is the exponential moving average of the squared gradients in Algorithm 1 and the power notation of a vector means the element-wise power of the vector.
Revised: where v t is the exponential moving average of the squared gradients in Algorithm 1 and the power notation of a vector means the element-wise power of the vector.

Original: We interpret ηv − 12 t as the effective learning rate for multiplying the gradients.
Revised: We interpret ηv − 12 t as the effective learning rate for multiplying the gradients.

Original: However, we clearly see that decoupled weight decay uses the vanilla learning rate to perform weight decay rather than the effective learning rate.
Revised: However, we clearly see that decoupled weight decay uses the vanilla learning rate to perform weight decay rather than the effective learning rate.

Original: The minimum ˆ θ (cid:63) of the regularized loss function optimized by AdamW at the t -th step is
Revised: The minimum ˆ θ (cid:63) of the regularized loss function optimized by AdamW at the t -th step is

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: There is no guarantee that v t will be constant for t .
Revised: There is no guarantee that v t will be constant for t .

Original: Thus, the regularized loss function optimized by
Revised: Thus, the regularized loss function optimized by

Original: AdamW has unstable minima.
Revised: AdamW has unstable minima.

Original: Scheduling weight decay in adaptive gradient methods.
Revised: Scheduling weight decay in adaptive gradient methods.

Original: An easy solution to ﬁx the unstable weight decay problem is using the following scheduled weight decay:
Revised: An easy solution to ﬁx the unstable weight decay problem is using the following scheduled weight decay:

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: This gives ˆ θ (cid:63) − θ (cid:63) = − λ
Revised: This gives ˆ θ (cid:63) − θ (cid:63) = − λ

Original: ( H + λ )
Revised: ( H + λ )

Original: [Equation]
Revised: [Equation]

Original: 1 θ (cid:63) .
Revised: 1 θ (cid:63) .

Original: Equation (10) can indeed make stationary points of the regularized loss function stable for different t .
Revised: Equation (10) can indeed make stationary points of the regularized loss function stable for different t .

Original: However, if we still use the element-wise scheduler v − 12 t for weight decay, the anisotropic convergence problem will appear in the subspace corresponding to small eigenvalues of the Hessian.
Revised: However, if we still use the element-wise scheduler v − 12 t for weight decay, the anisotropic convergence problem will appear in the subspace corresponding to small eigenvalues of the Hessian.

Original: By Theorem 2, the convergence can be measured by (cid:81) tk =1 ( I − v − 12 k ηλ ) .
Revised: By Theorem 2, the convergence can be measured by (cid:81) tk =1 ( I − v − 12 k ηλ ) .

Original: Due to the element-wise learning rate adaptivity, a common number of iterations may be not enough for achieving the convergence along some dimensions.
Revised: Due to the element-wise learning rate adaptivity, a common number of iterations may be not enough for achieving the convergence along some dimensions.

Original: [Figure]
Revised: [Figure]

Original: Unfortunately, there is no ideal solution to ﬁx the stability problem and the convergence problem at the same time for adaptive gradient methods.
Revised: Unfortunately, there is no ideal solution to ﬁx the stability problem and the convergence problem at the same time for adaptive gradient methods.

Original: We conjecture that this may be an internal fault of all optimizers that use Adaptive Learning Rate.
Revised: We conjecture that this may be an internal fault of all optimizers that use Adaptive Learning Rate.

Original: To balance the tradeoff between the stability and the convergence in presence of Adaptive Learning Rate, we propose a non-element-wise learning-rateaware scheduler for weight decay as
Revised: To balance the tradeoff between the stability and the convergence in presence of Adaptive Learning Rate, we propose a non-element-wise learning-rateaware scheduler for weight decay as

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where ¯ v t is the mean of all elements of the vector v t and p is a hyperparameter to control the adaptivity of the scheduler.
Revised: where ¯ v t is the mean of all elements of the vector v t and p is a hyperparameter to control the adaptivity of the scheduler.

Original: Note that η ¯ v − p t is always isotropic.
Revised: Note that η ¯ v − p t is always isotropic.

Original: We call weight decay with this stability-inspired heuristic scheduler Stable Weight Decay.
Revised: We call weight decay with this stability-inspired heuristic scheduler Stable Weight Decay.

Original: The pseudocode of AdamS is displayed in Algorithm 2.
Revised: The pseudocode of AdamS is displayed in Algorithm 2.

Original: We recommend p = 0 .
Revised: We recommend p = 0 .

Original: 5 as the default value in AdamS unless we specify it, because η ¯ v − p t with p = 0 .
Revised: 5 as the default value in AdamS unless we specify it, because η ¯ v − p t with p = 0 .

Original: 5 may approximate the magnitude of the adaptive learning rate.
Revised: 5 may approximate the magnitude of the adaptive learning rate.

Original: Note that Decoupled Weight
Revised: Note that Decoupled Weight

Original: Decay (Loshchilov and Hutter, 2018; Bjorck et al., 2021) is a special case of Stable Weight Decay with p = 0 .
Revised: Decay (Loshchilov and Hutter, 2018; Bjorck et al., 2021) is a special case of Stable Weight Decay with p = 0 .

Original: 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 (Zhu et al., 2019) and depends on the Hessian (Xie et al., 2021b).
Revised: 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 (Zhu et al., 2019) and depends on the Hessian (Xie et al., 2021b).

Original: In some special cases, such as full-batch training, it is ﬁne to add a small value (i.e., 10 − 8 ) to ¯ v to avoid being zero as a divisor.
Revised: In some special cases, such as full-batch training, it is ﬁne to add a small value (i.e., 10 − 8 ) to ¯ v to avoid being zero as a divisor.

Original:  E MPIRICAL A NALYSIS OF S TABLE W EIGHT D ECAY
Revised:  E MPIRICAL A NALYSIS OF S TABLE W EIGHT D ECAY

Original: In this section, we conducted comprehensive experiments to verify the advantage of the SWD scheduler over Decoupled Weight Decay and L 2 for Adam.
Revised: In this section, we conducted comprehensive experiments to verify the advantage of the SWD scheduler over Decoupled Weight Decay and L 2 for Adam.

Original: Our experiments also include popular
Revised: Our experiments also include popular

Original: Adam variants, including AMSGrad (Reddi et al., 2019), Yogi (Zaheer et al., 2018), AdaBound (Luo et al., 2019), Padam (Chen and Gu, 2018), and RAdam (Liu et al., 2019),.
Revised: Adam variants, including AMSGrad (Reddi et al., 2019), Yogi (Zaheer et al., 2018), AdaBound (Luo et al., 2019), Padam (Chen and Gu, 2018), and RAdam (Liu et al., 2019),.

Original: Models and Datasets.
Revised: Models and Datasets.

Original: We choose Adam as the base optimizer, and train popular deep models, including ResNet18/34/50 (He et al., 2016), VGG16 (Simonyan and Zisserman, 2014),
Revised: We choose Adam as the base optimizer, and train popular deep models, including ResNet18/34/50 (He et al., 2016), VGG16 (Simonyan and Zisserman, 2014),

Original: DenseNet (Huang et al., 2017), GoogLeNet (Szegedy et al., 2015), and Long Short-Term Memory (LSTM) (Hochreiter and Schmidhuber, 1997), on CIFAR-10/CIFAR-100 (Krizhevsky and Hinton, 2009), and
Revised: DenseNet (Huang et al., 2017), GoogLeNet (Szegedy et al., 2015), and Long Short-Term Memory (LSTM) (Hochreiter and Schmidhuber, 1997), on CIFAR-10/CIFAR-100 (Krizhevsky and Hinton, 2009), and

Original: Penn TreeBank (Marcus et al., 1993).
Revised: Penn TreeBank (Marcus et al., 1993).

Original: The implementation details are in Appendix B.
Revised: The implementation details are in Appendix B.

Original: Image Classiﬁcation on CIFAR-10/CIFAR-100.
Revised: Image Classiﬁcation on CIFAR-10/CIFAR-100.

Original: Figure 4 shows the learning curves of AdamS,
Revised: Figure 4 shows the learning curves of AdamS,

Original: AdamW, and Adam on several benchmarks.
Revised: AdamW, and Adam on several benchmarks.

Original: In our experiments, AdamS always leads to lower test errors.
Revised: In our experiments, AdamS always leads to lower test errors.

Original: Figure 5 shows that, even if with similar or higher training losses, AdamS still generalizes signiﬁcantly better than AdamW, Adam, and recent Adam variants.
Revised: Figure 5 shows that, even if with similar or higher training losses, AdamS still generalizes signiﬁcantly better than AdamW, Adam, and recent Adam variants.

Original: Figure 6 displays the learning curves of all adaptive gradient methods.
Revised: Figure 6 displays the learning curves of all adaptive gradient methods.

Original: The test performance of other models can be found in Table 1.
Revised: The test performance of other models can be found in Table 1.

Original: Simply scheduling weight decay for Adam by SWD even outperforms complex Adam variants.
Revised: Simply scheduling weight decay for Adam by SWD even outperforms complex Adam variants.

Original: We also note that most Adam variants surprisingly generalize worse than SGD (See Appendix C).
Revised: We also note that most Adam variants surprisingly generalize worse than SGD (See Appendix C).

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Robustness to the hyperparameters.
Revised: Robustness to the hyperparameters.

Original: Figure 7 further demonstrates that AdamS consistently outperforms Adam and AdamW under various weight decay hyperparameters.
Revised: Figure 7 further demonstrates that AdamS consistently outperforms Adam and AdamW under various weight decay hyperparameters.

Original: According to Figure 7, we also notice that the optimal decoupled weight decay hyperparameter in AdamW can be very different from L 2 regularization and SWD.
Revised: According to Figure 7, we also notice that the optimal decoupled weight decay hyperparameter in AdamW can be very different from L 2 regularization and SWD.

Original: Thus, AdamW requires re-tuning the weight decay hyperparameter in practice, which is time-consuming.
Revised: Thus, AdamW requires re-tuning the weight decay hyperparameter in practice, which is time-consuming.

Original: Figure 8 shows that AdamS is more robust to the learning rate and weight decay than AdamW and Adam.
Revised: Figure 8 shows that AdamS is more robust to the learning rate and weight decay than AdamW and Adam.

Original: AdamS has a much deeper and wider basin than
Revised: AdamS has a much deeper and wider basin than

Original: AdamW and Adam.
Revised: AdamW and Adam.

Original: We leave the experiments with cosine annealing schedulers and warm restarts (Loshchilov and Hutter, 2016) in Appendix E, which also support that AdamS yields superior test performance and converges to lower training losses.
Revised: We leave the experiments with cosine annealing schedulers and warm restarts (Loshchilov and Hutter, 2016) in Appendix E, which also support that AdamS yields superior test performance and converges to lower training losses.

Original: Large-Batch Training.
Revised: Large-Batch Training.

Original: In Figure 9, we again veriﬁed that Rule 2 holds well on all of Adam, AdamW, and AdamS.
Revised: In Figure 9, we again veriﬁed that Rule 2 holds well on all of Adam, AdamW, and AdamS.

Original: We also note that, in the case that applies Rule 2, AdamS generalizes better than AdamW and Adam on large-batch training ( B = 16384 ).
Revised: We also note that, in the case that applies Rule 2, AdamS generalizes better than AdamW and Adam on large-batch training ( B = 16384 ).

Original: Limitations.
Revised: Limitations.

Original: Figure 10 shows that, on ImageNet, the improvement of AdamS over AdamW becomes marginal.
Revised: Figure 10 shows that, on ImageNet, the improvement of AdamS over AdamW becomes marginal.

Original: Figure 11 in Appendix C demonstrates that, for LSTM on Penn TreeBank, L 2 regularization yields better test results than both Decoupled Weight Decay and SWD.
Revised: Figure 11 in Appendix C demonstrates that, for LSTM on Penn TreeBank, L 2 regularization yields better test results than both Decoupled Weight Decay and SWD.

Original:  C ONCLUSION AND D ISCUSSION
Revised:  C ONCLUSION AND D ISCUSSION

Original: Our experiments on large-batch training shows that the learning-rate linear scaling rule sometimes is harmful to large-batch training due to bad convergence.
Revised: Our experiments on large-batch training shows that the learning-rate linear scaling rule sometimes is harmful to large-batch training due to bad convergence.

Original: We propose the weight-decay linear scaling rule for large-batch training as an alternative solution.
Revised: We propose the weight-decay linear scaling rule for large-batch training as an alternative solution.

Original: We note that the results may not be explained by previous papers(Arora et al., 2019; Li et al., 2020), which studied scale-invariant deep loss landscape due to batch normalization(Ioffe and Szegedy, 2015).
Revised: We note that the results may not be explained by previous papers(Arora et al., 2019; Li et al., 2020), which studied scale-invariant deep loss landscape due to batch normalization(Ioffe and Szegedy, 2015).

Original: Instead, the empirical results veriﬁed our theoretical analysis with or without scale-invariant loss functions.
Revised: Instead, the empirical results veriﬁed our theoretical analysis with or without scale-invariant loss functions.

Original: Inspired by the theoretical analysis of weight decay on stability and convergence, we propose a scheduler for weight decay.
Revised: Inspired by the theoretical analysis of weight decay on stability and convergence, we propose a scheduler for weight decay.

Original: Previous work did not studied how to design schedulers for weight decay.
Revised: Previous work did not studied how to design schedulers for weight decay.

Original: Our empirical results demonstrate that SWD often makes improvements over L 2 regularization and
Revised: Our empirical results demonstrate that SWD often makes improvements over L 2 regularization and

Original: Decoupled Weight Decay.
Revised: Decoupled Weight Decay.

Original: The generalization gap between SGD and Adam can be almost closed by a simple weight decay scheduler on CIFAR datasets.
Revised: The generalization gap between SGD and Adam can be almost closed by a simple weight decay scheduler on CIFAR datasets.

Original: Although our analysis mainly focused on Adam,
Revised: Although our analysis mainly focused on Adam,

Original: SWD can be easily combined with other Adam variants, too.
Revised: SWD can be easily combined with other Adam variants, too.

Original: Finally, we discuss the limitations of SWD.
Revised: Finally, we discuss the limitations of SWD.

Original: While SWD makes signiﬁcant improvements in ﬁrst three experiments of Section 5, the experiments on ImageNet and Language Modeling suggest that the improvement of SWD may be marginal for complex loss landscapes.
Revised: While SWD makes signiﬁcant improvements in ﬁrst three experiments of Section 5, the experiments on ImageNet and Language Modeling suggest that the improvement of SWD may be marginal for complex loss landscapes.

Original: SWD as the ﬁrst practical design for weight decay scheduling is not perfect.
Revised: SWD as the ﬁrst practical design for weight decay scheduling is not perfect.

Original: In future, it will be interesting to design weight decay schedulers that can accelerate convergence in the early training phase while learning stable and accurate minima in the late training phase.
Revised: In future, it will be interesting to design weight decay schedulers that can accelerate convergence in the early training phase while learning stable and accurate minima in the late training phase.

Original: A P ROOFS A. P ROOF OF T HEOREM
Revised: A P ROOFS A. P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: We ﬁrst write the regularized loss function corresponding to SGD with vanilla weight decay as
Revised: We ﬁrst write the regularized loss function corresponding to SGD with vanilla weight decay as

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: at t -th step.
Revised: at t -th step.

Original: If the corresponding L 2 regularization λ (cid:48) 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.
Revised: If the corresponding L 2 regularization λ (cid:48) 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.

Original: Suppose we have a non-zero solution θ (cid:63) which is a stationary point of f ( θ, t ) at t -th step and SGD ﬁnds θ t
Revised: Suppose we have a non-zero solution θ (cid:63) which is a stationary point of f ( θ, t ) at t -th step and SGD ﬁnds θ t

Original: = θ (cid:63) at t -th step.
Revised: = θ (cid:63) at t -th step.

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

Original: [Equation]
Revised: [Equation]

Original: It means that (cid:107) g t
Revised: It means that (cid:107) g t

Original: +

Revised: +

Original: (cid:107) 2 =
Revised: (cid:107) 2 =

Original: λ (cid:48) 2 ( η − 1 t − η
Revised: λ (cid:48) 2 ( η − 1 t − η

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: +1 ) 2 (cid:107) θ (cid:63) (cid:107) 2 .
Revised: +1 ) 2 (cid:107) θ (cid:63) (cid:107) 2 .

Original: To achieve convergence, we must have (cid:107) g t +1 (cid:107) 2 = 0 .
Revised: To achieve convergence, we must have (cid:107) g t +1 (cid:107) 2 = 0 .

Original: It requires ( η −
Revised: It requires ( η −

Original: [Equation]
Revised: [Equation]

Original: Theorem 2.2 of Shapiro and Wardi (1996) told us that the learning rate should be small enough for convergence.
Revised: Theorem 2.2 of Shapiro and Wardi (1996) told us that the learning rate should be small enough for convergence.

Original: Obviously, we have η < ∞ in practice.
Revised: Obviously, we have η < ∞ in practice.

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

Original: The proof is now complete.
Revised: The proof is now complete.

Original: A. P ROOF OF T HEOREM
Revised: A. P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: Under Assumption 1, we write the dynamics of SGD with weight decay as
Revised: Under Assumption 1, we write the dynamics of SGD with weight decay as

Original: [Equation]
Revised: [Equation]

Original: Note that the stochastic gradient noise ξ is zero-mean, E [ ξ ] = 0 , where ξ is the difference of the true gradient and the stochastic gradient.
Revised: Note that the stochastic gradient noise ξ is zero-mean, E [ ξ ] = 0 , where ξ is the difference of the true gradient and the stochastic gradient.

Original: The zero-mean property generally holds because SGD randomly selects minibatches from the whole training dataset, which is widely used in related papers (Xie et al., 2021c; Kingma and Ba, 2015;
Revised: The zero-mean property generally holds because SGD randomly selects minibatches from the whole training dataset, which is widely used in related papers (Xie et al., 2021c; Kingma and Ba, 2015;

Original: Reddi et al., 2019).
Revised: Reddi et al., 2019).

Original: The dynamics above is equivalent to minimizing the training loss with L 2 regularization as
Revised: The dynamics above is equivalent to minimizing the training loss with L 2 regularization as

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: Thus, we may rewrite the dynamics of SGD with weight decay as
Revised: Thus, we may rewrite the dynamics of SGD with weight decay as

Original: [Equation]
Revised: [Equation]

Original: where the expectation of the stochastic gradient noise term is zero.
Revised: where the expectation of the stochastic gradient noise term is zero.

Original: As we care about the difference between the expected learned solution and the true minimum θ (cid:63)
Revised: As we care about the difference between the expected learned solution and the true minimum θ (cid:63)

Original: [Equation]
Revised: [Equation]

Original: we may further obtain
Revised: we may further obtain

Original: [Equation]
Revised: [Equation]

Original: In the case of the constant learning rate scheduler
Revised: In the case of the constant learning rate scheduler

Original: [Equation]
Revised: [Equation]

Original: we may simplify the equation above as
Revised: we may simplify the equation above as

Original: [Equation]
Revised: [Equation]

Original: The proof is now complete.
Revised: The proof is now complete.

Original: B E XPERIMENTAL D ETAILS
Revised: B E XPERIMENTAL D ETAILS

Original: Computational environment.
Revised: Computational environment.

Original: The experiments are conducted on a computing cluster with GPUs of
Revised: The experiments are conducted on a computing cluster with GPUs of

Original: NVIDIA R (cid:13) Tesla TM P100 16GB and CPUs of Intel R (cid:13)
Revised: NVIDIA R (cid:13) Tesla TM P100 16GB and CPUs of Intel R (cid:13)

Original: Xeon R (cid:13) CPU E5-2640 v3 @ 2.60GHz.
Revised: Xeon R (cid:13) CPU E5-2640 v3 @ 2.60GHz.

Original: B. I MAGE C LASSIFICATION ON CIFAR-10 AND CIFAR- 
Revised: B. I MAGE C LASSIFICATION ON CIFAR-10 AND CIFAR-

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

Original: 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.
Revised: 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.

Original: We use the default learning rate for adaptive gradient methods in the experiments of Table ??
Revised: We use the default learning rate for adaptive gradient methods in the experiments of Table ??

Original: , while we compared Adam,
Revised: , while we compared Adam,

Original: AdamW, AdanS under various learning rates and batch sizes in the experiments of Figure 7 and
Revised: AdamW, AdanS under various learning rates and batch sizes in the experiments of Figure 7 and

Original: Figure 8.
Revised: Figure 8.

Original: In the experiments on CIFAF-10 and CIFAR-100: η = 0 .
Revised: In the experiments on CIFAF-10 and CIFAR-100: η = 0 .

Original: 1 for SGD and SGDS; η = 0 .for
Revised: 1 for SGD and SGDS; η = 0 .for

Original: Adam, AdamW, AdamW, AMSGrad, AdamW, AdaBound, and RAdam; η = 0 .
Revised: Adam, AdamW, AdamW, AMSGrad, AdamW, AdaBound, and RAdam; η = 0 .

Original: 01 for Padam.
Revised: 01 for Padam.

Original: For the learning rate schedule, the learning rate is divided by 10 at the epoch of { 80 , 160 } for CIFARand { 100 , 150 } for CIFAR-100, respectively.
Revised: For the learning rate schedule, the learning rate is divided by 10 at the epoch of { 80 , 160 } for CIFARand { 100 , 150 } for CIFAR-100, respectively.

Original: The batch size is set to 128 for both CIFAR-10 and
Revised: The batch size is set to 128 for both CIFAR-10 and

Original: CIFAR-100, unless we specify it on large-batch training.
Revised: CIFAR-100, unless we specify it on large-batch training.

Original: The strength of L 2 regularization and SWD is default to 0 .
Revised: The strength of L 2 regularization and SWD is default to 0 .

Original: as the baseline.
Revised: as the baseline.

Original: Considering the linear scaling rule, we choose λ W =
Revised: Considering the linear scaling rule, we choose λ W =

Original: λ
Revised: λ

Original: L 2 η .
Revised: L 2 η .

Original: Thus, the weight decay of AdamW uses λ W = 0 .for
Revised: Thus, the weight decay of AdamW uses λ W = 0 .for

Original: CIFAR-10 and CIFAR-100.
Revised: CIFAR-10 and CIFAR-100.

Original: The basic principle of choosing weight decay strength is to let all optimizers have similar convergence speed.
Revised: The basic principle of choosing weight decay strength is to let all optimizers have similar convergence speed.

Original: We set the momentum hyperparameter β 1 = 0 .
Revised: We set the momentum hyperparameter β 1 = 0 .

Original: 9 for SGD and SGDS.
Revised: 9 for SGD and SGDS.

Original: As for other optimizer hyperparameters, we apply the default hyperparameter settings directly.
Revised: As for other optimizer hyperparameters, we apply the default hyperparameter settings directly.

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

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

Original: B. L ARGE - BATCH TRAINING ON CIFAR- 
Revised: B. L ARGE - BATCH TRAINING ON CIFAR-

Original: Model: we always used the common ResNet18 for large-batch training.
Revised: Model: we always used the common ResNet18 for large-batch training.

Original: Optimizer: we always used Adam and its variant as the optimizer, because adaptive gradient methods often outperform SGD in large-batch training.
Revised: Optimizer: we always used Adam and its variant as the optimizer, because adaptive gradient methods often outperform SGD in large-batch training.

Original: The learning rate scheduler: Suppose the number of epochs is E , the learning rate is divided byat the epoch of { 0 .
Revised: The learning rate scheduler: Suppose the number of epochs is E , the learning rate is divided byat the epoch of { 0 .

Original: 4 E, 0 .
Revised: 4 E, 0 .

Original: 8 E } .
Revised: 8 E } .

Original: The default batch size is 16384 in large-batch training.
Revised: The default batch size is 16384 in large-batch training.

Original: The default weight decay is 0 . 0001 or follows
Revised: The default weight decay is 0 . 0001 or follows

Original: Rule 2.
Revised: Rule 2.

Original: The default learning rate is 0 .
Revised: The default learning rate is 0 .

Original: 001 or follows Rule 1.
Revised: 001 or follows Rule 1.

Original: As for other optimizer hyperparameters, we apply the default hyperparameter setting directly.
Revised: As for other optimizer hyperparameters, we apply the default hyperparameter setting directly.

Original: B. I MAGE CLASSIFICATION ON I MAGE N ET
Revised: B. I MAGE CLASSIFICATION ON I MAGE N ET

Original: Data Preprocessing For ImageNet: For ImageNet, we perform the per-pixel zero-mean unitvariance normalization, horizontal random ﬂip, 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 34 to 43 ) of the original aspect ratio is made.
Revised: Data Preprocessing For ImageNet: For ImageNet, we perform the per-pixel zero-mean unitvariance normalization, horizontal random ﬂip, 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 34 to 43 ) of the original aspect ratio is made.

Original: 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.
Revised: 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.

Original: For the learning rate schedule, the learning rate is divided by 10 at the epoch of { 30 , 60 } .
Revised: For the learning rate schedule, the learning rate is divided by 10 at the epoch of { 30 , 60 } .

Original: We train each model for 90 epochs.
Revised: We train each model for 90 epochs.

Original: The batch size is set to 256 .
Revised: The batch size is set to 256 .

Original: 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 } .
Revised: 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 } .

Original: As for other optimizer hyperparameters, we still apply the default hyperparameter settings directly.
Revised: As for other optimizer hyperparameters, we still apply the default hyperparameter settings directly.

Original: B. L ANGUAGE M ODELING
Revised: B. L ANGUAGE M ODELING

Original: We use a classical language model, Long Short-Term Memory (LSTM) (Hochreiter and Schmidhuber, 1997) 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.
Revised: We use a classical language model, Long Short-Term Memory (LSTM) (Hochreiter and Schmidhuber, 1997) 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.

Original: Note that our baseline performance is better than the reported baseline performance in Zaremba et al.
Revised: Note that our baseline performance is better than the reported baseline performance in Zaremba et al.

Original: The benchmark task is the word-level Penn TreeBank (Marcus et al., 1993).
Revised: The benchmark task is the word-level Penn TreeBank (Marcus et al., 1993).

Original: We empirically compared
Revised: We empirically compared

Original: AdamS
Revised: AdamS

Original: , AdamW, and Adam under the common and same conditions.
Revised: , AdamW, and Adam under the common and same conditions.

Original: Hyperparameter Settings.
Revised: Hyperparameter Settings.

Original: Batch Size: B
Revised: Batch Size: B

Original: [Table]
Revised: [Table]

Original: BPTT Size
Revised: BPTT Size

Original: [Table]
Revised: [Table]

Original: bptt
Revised: bptt

Original: [Table]
Revised: [Table]

Original: Learning Rate
Revised: Learning Rate

Original: [Table]
Revised: [Table]

Original: 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
Revised: 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

Original: dropout probability is set to 0 .
Revised: dropout probability is set to 0 .

Original: We clipped gradient norm to 1 .
Revised: We clipped gradient norm to 1 .

Original: C S UPPLEMENTARY F IGURES AND R ESULTS OF A DAPTIVE G RADIENT M ETHODS
Revised: C S UPPLEMENTARY F IGURES AND R ESULTS OF A DAPTIVE G RADIENT M ETHODS

Original: Popular Adam variants often generalize worse than SGD.
Revised: Popular Adam variants often generalize worse than SGD.

Original: A few Adam variants tried to ﬁx the hidden problems in adaptive gradient methods, including AdamW Loshchilov and Hutter (2018),
Revised: A few Adam variants tried to ﬁx the hidden problems in adaptive gradient methods, including AdamW Loshchilov and Hutter (2018),

Original: AMSGrad (Reddi et al., 2019) and Yogi (Zaheer et al., 2018).
Revised: AMSGrad (Reddi et al., 2019) and Yogi (Zaheer et al., 2018).

Original: A recent line of research, such as
Revised: A recent line of research, such as

Original: AdaBound (Luo et al., 2019), Padam (Chen and Gu, 2018), and RAdam (Liu et al., 2019), believes controlling the adaptivity of learning rates may improve generalization.
Revised: AdaBound (Luo et al., 2019), Padam (Chen and Gu, 2018), and RAdam (Liu et al., 2019), believes controlling the adaptivity of learning rates may improve generalization.

Original: This line of research usually introduces extra hyperparameters to control the adaptivity, which requires more efforts in tuning hyperparameters.
Revised: This line of research usually introduces extra hyperparameters to control the adaptivity, which requires more efforts in tuning hyperparameters.

Original: However, we and Zhang et al.
Revised: However, we and Zhang et al.

Original: (2019c) found that this argument is contradicted with our comparative experimental results (see Table 1).
Revised: (2019c) found that this argument is contradicted with our comparative experimental results (see Table 1).

Original: In our empirical analysis, most advanced
Revised: In our empirical analysis, most advanced

Original: Adam variants may narrow but not completely close the generalization gap between adaptive gradient methods and SGD.
Revised: Adam variants may narrow but not completely close the generalization gap between adaptive gradient methods and SGD.

Original: SGD with a fair weight decay hyperparameter as the baseline performance usually generalizes better than recent adaptive gradient methods.
Revised: SGD with a fair weight decay hyperparameter as the baseline performance usually generalizes better than recent adaptive gradient methods.

Original: The main problem may lie in weight decay.
Revised: The main problem may lie in weight decay.

Original: SGD with weight decay λ = 0 . 0001 , a common setting in related papers, is often not a good baseline, as λ = 0 .
Revised: SGD with weight decay λ = 0 . 0001 , a common setting in related papers, is often not a good baseline, as λ = 0 .

Original: 0005 often shows better generalization on CIFAR-10 and CIFAR-100 (see in
Revised: 0005 often shows better generalization on CIFAR-10 and CIFAR-100 (see in

Original: Figures 19 and 21 of Appendix F).
Revised: Figures 19 and 21 of Appendix F).

Original: We also conduct comparative experiments with λ = 0 .
Revised: We also conduct comparative experiments with λ = 0 .

Original: (see
Revised: (see

Original: Table 2 of Appendix C).
Revised: Table 2 of Appendix C).

Original: 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.
Revised: 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.

Original: For example, for ResNet18 on CIFAR-10, the test error of AdamS is lower than
Revised: For example, for ResNet18 on CIFAR-10, the test error of AdamS is lower than

Original: SGD by nearly one point and no other Adam variant may compare with AdamS.
Revised: SGD by nearly one point and no other Adam variant may compare with AdamS.

Original: Language Modeling.
Revised: Language Modeling.

Original: 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.
Revised: 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.

Original: Figure 11 in Appendix
Revised: Figure 11 in Appendix

Original: C demonstrates that the conventional belief is true that the standard L 2 regularization yields better
Revised: C demonstrates that the conventional belief is true that the standard L 2 regularization yields better

Original: [Figure]
Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Original: [Table] 
Revised: [Table] [Figure] [Table] [Figure]

Revised: [Table]

Original: [Figure] [Table] 
Revised: [Figure] [Table] [Figure]

Original: [Figure] [Table] [Figure] [Table]
Revised: [Table]

Original: [Figure][Table] [Figure] [Table]
Revised: [Figure]

Original: [Figure] [Table] [Figure] [Table]
Revised: [Table]

Original: [Figure][Table] [Figure] [Table]
Revised: [Figure]

Original: [Figure] [Table] [Figure] [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: We report the learning curves of all adaptive gradient methods in Figure 12.
Revised: We report the learning curves of all adaptive gradient methods in Figure 12.

Original: They shows that vanilla
Revised: They shows that vanilla

Original: Adam with SWD can outperform other complex variants of Adam.
Revised: Adam with SWD can outperform other complex variants of Adam.

Original: Figure 13 displays the scatter plot of training losses and test errors during ﬁnal 40 epochs of training
Revised: Figure 13 displays the scatter plot of training losses and test errors during ﬁnal 40 epochs of training

Original: DenseNet121 on CIFAR-100.
Revised: DenseNet121 on CIFAR-100.

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

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

Original: Note that ResNets have scale-invariant loss landscape, while VGGs have no scale-invariant loss landscape.
Revised: Note that ResNets have scale-invariant loss landscape, while VGGs have no scale-invariant loss landscape.

Original: Our theoretical analysis generally holds under Assumption 1, which is independent of the scale-invariant property of the loss landscape.
Revised: Our theoretical analysis generally holds under Assumption 1, which is independent of the scale-invariant property of the loss landscape.

Original: Figure 16 further shows that, with or without normalization layers
Revised: Figure 16 further shows that, with or without normalization layers

Original: [Figure]
Revised: [Figure]

Original: our theoretical result can be empirically observed.
Revised: our theoretical result can be empirically observed.

Original: This is beyond the intrinsic learning rate proposed by Li et al.
Revised: This is beyond the intrinsic learning rate proposed by Li et al.

Original: (2020) which depends on the scale-invariant property of the loss function.
Revised: (2020) which depends on the scale-invariant property of the loss function.

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Table]
Revised: [Table]

Original: D A DDITIONAL A LGORITHMS
Revised: D A DDITIONAL A LGORITHMS

Original: Algorithm 3 is the TensorFlow implementation for SGD.
Revised: Algorithm 3 is the TensorFlow implementation for SGD.

Original: Algorithm 4 is the implementation of Adai, AdaiW, and AdaiS.
Revised: Algorithm 4 is the implementation of Adai, AdaiW, and AdaiS.

Original: As β 3 1 − β 1 is always 1 in Adai, AdaiW is identical to AdaiS.
Revised: As β 3 1 − β 1 is always 1 in Adai, AdaiW is identical to AdaiS.

Original: We note that the implementation of AMSGrad in Algorithm 5 is the popular implementation in
Revised: We note that the implementation of AMSGrad in Algorithm 5 is the popular implementation in

Original: PyTorch.
Revised: PyTorch.

Original: We use the PyTorch implementation in our paper, as it is widely used in practice.
Revised: We use the PyTorch implementation in our paper, as it is widely used in practice.

Revised: Algorithm 3: SGD in TensorFlow

Original: [Equation]
Revised: [Equation]

Original: Algorithm 4: Adai /AdaiS=
Revised: Algorithm 4: Adai /AdaiS=

Original: AdaiW g t = ∇ L ( θ t − 1 ) + λθ t − 1 ; v t = β 2 v t − 1 + (1 − β 2 ) g 2 t ; ˆ v t = v t 1 − β t 2 ; ¯ v t = mean (ˆ v t ); β 1 t
Revised: AdaiW [Equation] ¯ v t = mean (ˆ v t )

Original: [-= (1 − β 
-]
Revised: [-= (1 − β 
-][Equation]

Original: ˆ v t ¯ v t ) .Clip
Revised: v t ¯ v t ) .Clip

Original: (0 , 1 − (cid:15) ) ; m t = β 1 t m t − 1 + (1 − β 1 t ) g t ; ˆ m t = m t 1 − (cid:81) tk =1 β 1 k ; θ t = θ t − 1 − η ˆ m t − ηλθ t − 1 ;
Revised: (0 , 1 − (cid:15) ) [Equation]

Original: Algorithm 5: AMSGrad/AMSGradW g t = ∇ L ( θ t − 1 ) + λθ t − 1 ; m t = β 1 m t − 1 + (1 − β 1 ) g t ; v t = β 2 v t − 1 + (1 − β 2 ) g 2 t ; ˆ m t = m t 1 − β t 1 ; v max = max ( v t , v max ) ; ˆ v t = v max 1 − β t 2 ; θ t = θ t − 1 − η √ ˆ v t + (cid:15) ˆ m t − ηλθ t − 1 ;
Revised: Algorithm 5: AMSGrad/AMSGradW [Equation]

Original: Algorithm 6: AMSGradS g t = ∇ L ( θ t − 1 ) ; m t = β 1 m t − 1 + (1 − β 1 ) g t ; v t = β 2 v t − 1 + (1 − β 2 ) g 2 t ; ˆ m t = m t 1 − β t 1 ; v max = max ( v t , v max ) ; ˆ v t = v max 1 − β t 2 ; ¯ v t = mean (ˆ v t ) ; θ t = θ t − 1 − η √ ˆ v t + (cid:15) ˆ m t − η √ ¯ v t λθ t − 1 ;
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: E S UPPLEMENTARY E XPERIMENTS WITH C OSINE A NNEALING S CHEDULERS AND W ARM R ESTARTS
Revised: E S UPPLEMENTARY E XPERIMENTS WITH C OSINE A NNEALING S CHEDULERS AND W ARM R ESTARTS

Original: 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 (2016).
Revised: 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 (2016).

Original: We set the learning rate scheduler with a recommended setting of Loshchilov and Hutter (2016): T 0 =and T mul = 2 .
Revised: We set the learning rate scheduler with a recommended setting of Loshchilov and Hutter (2016): T 0 =and T mul = 2 .

Original: The number of total epochs is 210 .
Revised: The number of total epochs is 210 .

Original: Thus, we trained each deep network for four runs of warm restarts, where the four runs have 14, 28, 56, and 112 epochs, respectively.
Revised: Thus, we trained each deep network for four runs of warm restarts, where the four runs have 14, 28, 56, and 112 epochs, respectively.

Original: Other hyperparameters and details are displayed in Appendix B.
Revised: Other hyperparameters and details are displayed in Appendix B.

Original: Our experimental results in Figures 17 and 18 suggest that AdamS consistently outperforms AdamW and Adam in the presence of cosine annealing schedulers and warm restarts.
Revised: Our experimental results in Figures 17 and 18 suggest that AdamS consistently outperforms AdamW and Adam in the presence of cosine annealing schedulers and warm restarts.

Original: It demonstrates that, with various learning rate schedulers, the advantage of SWD may generally hold.
Revised: It demonstrates that, with various learning rate schedulers, the advantage of SWD may generally hold.

Original: 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.
Revised: 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.

Original: We noticed that Loshchilov and Hutter (2016) 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.
Revised: We noticed that Loshchilov and Hutter (2016) 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.

Original: There may be two possible causes.
Revised: There may be two possible causes.

Original: First, three-staged piecewise-constant learning rate schedulers, which usually have a longer ﬁrst 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 (2016) are usually not optimal.
Revised: First, three-staged piecewise-constant learning rate schedulers, which usually have a longer ﬁrst 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 (2016) are usually not optimal.

Original: Second, warm restarts may be helpful, while cosine annealing may be not.
Revised: Second, warm restarts may be helpful, while cosine annealing may be not.

Original: The ablation study on piecewise-constant learning rate schedulers with warm restarts is lacked.
Revised: The ablation study on piecewise-constant learning rate schedulers with warm restarts is lacked.

Original: 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.
Revised: 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.

Original: F W EIGHT D ECAY IN M OMENTUM
Revised: F W EIGHT D ECAY IN M OMENTUM

Original: In this section, we rethink weight decay in Momentum.
Revised: In this section, we rethink weight decay in Momentum.

Original: [Figure]
Revised: [Figure]

Original: It is also commonly believed that L 2 regularization is an optimal weight decay implementation when Adaptive Learning Rate is not involved.
Revised: It is also commonly believed that L 2 regularization is an optimal weight decay implementation when Adaptive Learning Rate is not involved.

Original: Almost all deep learning libraries directly use Lregularization as the default weight decay implementation.
Revised: Almost all deep learning libraries directly use Lregularization as the default weight decay implementation.

Original: However, we reveal that L 2 regularization is not identical to decoupled weight decay Loshchilov and Hutter (2018) and sometimes harms performance slightly when Momentum is involved.
Revised: However, we reveal that L 2 regularization is not identical to decoupled weight decay Loshchilov and Hutter (2018) and sometimes harms performance slightly when Momentum is involved.

Original: We take Stochastic Heavy Ball (SHB) (Zavriev and Kostyuk, 1993), which only uses ﬁxed momentum inertia and dampening coefﬁcients, as the studied example in the presence of Momentum, as SHBstyle Momentum methods are widely used by many popular optimizers.
Revised: We take Stochastic Heavy Ball (SHB) (Zavriev and Kostyuk, 1993), which only uses ﬁxed momentum inertia and dampening coefﬁcients, as the studied example in the presence of Momentum, as SHBstyle Momentum methods are widely used by many popular optimizers.

Original: SGD implemented in
Revised: SGD implemented in

Original: PyTorch (Paszke et al., 2019) is actually SHB with default hyperparameters.
Revised: PyTorch (Paszke et al., 2019) is actually SHB with default hyperparameters.

Original: We write SHB with Lregularization and SHB with decoupled weight decay in Algorithm 7.
Revised: We write SHB with Lregularization and SHB with decoupled weight decay in Algorithm 7.

Original: Algorithm 7: SGD/SGDW (SHB/SHBW)
Revised: Algorithm 7: SGD/SGDW (SHB/SHBW)

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Inspired by Theorem 1, we propose SGD with SWD (SGDS) as
Revised: Inspired by Theorem 1, we propose SGD with SWD (SGDS) as

Original: [Equation]
Revised: [Equation]

Original: where β 3 (1 − β t 1 ) 1 − β 1 η is the effective learning rate in SHB, as m t ≈ β 3 (1 − β t 1 ) 1 − β 1 E [ g t ] .
Revised: where β 3 (1 − β t 1 ) 1 − β 1 η is the effective learning rate in SHB, as m t ≈ β 3 (1 − β t 1 ) 1 − β 1 E [ g t ] .

Original: As (1 − β t 1 ) converges to 1 soon, we use the simpliﬁed β 3 1 − β 1 η as the bias correction in practice, where relatively large weight decay in the ﬁrst dozen iterations can work like a model parameter initialization strategy.
Revised: As (1 − β t 1 ) converges to 1 soon, we use the simpliﬁed β 3 1 − β 1 η as the bias correction in practice, where relatively large weight decay in the ﬁrst dozen iterations can work like a model parameter initialization strategy.

Original: It is easy to see that the factor β 3 1 − β 1 for weight decay is exactly the difference between our stable weight decay and decoupled weight decay suggested by Loshchilov and Hutter (2018).
Revised: It is easy to see that the factor β 3 1 − β 1 for weight decay is exactly the difference between our stable weight decay and decoupled weight decay suggested by Loshchilov and Hutter (2018).

Original: The pseudocode of SGDS is displayed in Algorithm 8.
Revised: The pseudocode of SGDS is displayed in Algorithm 8.

Original: It may also be called SHB with SWD (SHBS).
Revised: It may also be called SHB with SWD (SHBS).

Original: Empirical Analysis on SGDS.
Revised: Empirical Analysis on SGDS.

Original: We empirically veriﬁed that the optimal performance of
Revised: We empirically veriﬁed that the optimal performance of

Original: SGDS/SGDW is often slightly better than the widely used vanilla SGD and SGD, as seen in Figures 19 and 21.
Revised: SGDS/SGDW is often slightly better than the widely used vanilla SGD and SGD, as seen in Figures 19 and 21.

Original: We leave more empirical results of SGDS in Appendix G (see Table 4).
Revised: We leave more empirical results of SGDS in Appendix G (see Table 4).

Original: In Figure 19, the optimal performance of SGDS/SGDW is better than the widely used Vanilla SGD and SGD.
Revised: In Figure 19, the optimal performance of SGDS/SGDW is better than the widely used Vanilla SGD and SGD.

Original: For Vanilla SGD, SGD, and SGDS, we may choose λ L 2 = λ S = 0 .
Revised: For Vanilla SGD, SGD, and SGDS, we may choose λ L 2 = λ S = 0 .

Original: to maintain the optimal performance.
Revised: to maintain the optimal performance.

Original: But we have to re-tune
Revised: But we have to re-tune

Original: λ W = 0 .
Revised: λ W = 0 .

Original: 005 for SGDW.
Revised: 005 for SGDW.

Original: Hyperparameter Setting: β 1 = 0 for Vanilla SGD; β 1 = 0 .
Revised: Hyperparameter Setting: β 1 = 0 for Vanilla SGD; β 1 = 0 .

Original: 9 for SGD, SGDW, and SGDS.
Revised: 9 for SGD, SGDW, and SGDS.

Original: We repeat each simulation for three runs.
Revised: We repeat each simulation for three runs.

Original: A similar experimental result for ResNet18 is presented in Appendix G.
Revised: A similar experimental result for ResNet18 is presented in Appendix G.

Original: In Figure 20, we report the test performance of SGDS and SGD for ResNet18 on CIFAR-10 under various learning rates and weight decay.
Revised: In Figure 20, we report the test performance of SGDS and SGD for ResNet18 on CIFAR-10 under various learning rates and weight decay.

Original: SGDS has a deeper blue basin than SGD.
Revised: SGDS has a deeper blue basin than SGD.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: D ATASET
Revised: D ATASET

Original: [Table]
Revised: [Table]

Original: CIFAR-
Revised: CIFAR-

Original: [Table]
Revised: [Table]

Original: CIFARD ENSE N
Revised: CIFARD ENSE N

Original: [Table]
Revised: [Table]

Original: It is not well-known that SGDW often outperforms SGD slightly.
Revised: It is not well-known that SGDW often outperforms SGD slightly.

Original: We believe it is mainly because people rarely know re-tuning λ W based on β 3 1 − β 1 is necessary for maintaining good performance when people switch from SGD to SGDW.
Revised: We believe it is mainly because people rarely know re-tuning λ W based on β 3 1 − β 1 is necessary for maintaining good performance when people switch from SGD to SGDW.

Original: As the effective learning rate of SGD is β 3 1 − β 1 η , the weight decay rate of SGDW is actually R = 1 − β 1 β 3 λ W rather than λ W .
Revised: As the effective learning rate of SGD is β 3 1 − β 1 η , the weight decay rate of SGDW is actually R = 1 − β 1 β 3 λ W rather than λ W .

Original: If we use different settings of βand β 3 in decoupled weight decay, we will undesirably change the weight decay rate R unless we re-tune λ W .
Revised: If we use different settings of βand β 3 in decoupled weight decay, we will undesirably change the weight decay rate R unless we re-tune λ W .

Original: However, people usually directly let λ
Revised: However, people usually directly let λ

Original: W
Revised: W

Original: =
Revised: =

Original: λ
Revised: λ

Original: L 2 for SGDW in practice.
Revised: L 2 for SGDW in practice.

Original: Figureshows that, the optimal λ L 2 and λ S in SGD ( with β 1 = 0 .
Revised: Figureshows that, the optimal λ L 2 and λ S in SGD ( with β 1 = 0 .

Original: 9 and β 3 = 1 ) are both 0 .
Revised: 9 and β 3 = 1 ) are both 0 .

Original: 0005 , while the optimal λ W is 0 .
Revised: 0005 , while the optimal λ W is 0 .

Original: 005 instead.
Revised: 005 instead.

Original: The optimal λ L 2 and λ S are almost same, while the optimal λ W is quite different.
Revised: The optimal λ L 2 and λ S are almost same, while the optimal λ W is quite different.

Original: Thus, the advantage of SGDS over SGDW can save us from re-tuning the weight decay hyperparameter.
Revised: Thus, the advantage of SGDS over SGDW can save us from re-tuning the weight decay hyperparameter.

Original: G S UPPLEMENTARY R ESULTS OF SGD WITH M OMENTUM AND A DAI
Revised: G S UPPLEMENTARY R ESULTS OF SGD WITH M OMENTUM AND A DAI

Original: We compare SGDS, SGDW, vanilla SGD and SGD under various weight decay hyperparameters by training ResNet18 on CIFAR-10.
Revised: We compare SGDS, SGDW, vanilla SGD and SGD under various weight decay hyperparameters by training ResNet18 on CIFAR-10.

Original: We observe that the optimal performance of SGDS/SGDW is better than vanilla SGD and SGD in Figure 21.
Revised: We observe that the optimal performance of SGDS/SGDW is better than vanilla SGD and SGD in Figure 21.

Original: The advantage of SGDS over SGDW is that we do not need to ﬁne-tune the weight decay hyperparameters based on β 3 1 − β 1 .
Revised: The advantage of SGDS over SGDW is that we do not need to ﬁne-tune the weight decay hyperparameters based on β 3 1 − β 1 .

Original: We report the learning curves of SGD and SGDS in Figure 22.
Revised: We report the learning curves of SGD and SGDS in Figure 22.

Original: SGDS compares favorably with SGD.
Revised: SGDS compares favorably with SGD.

Original: Based on a diffusion theoretical framework (Xie et al., 2021b), Xie et al.
Revised: Based on a diffusion theoretical framework (Xie et al., 2021b), Xie et al.

Original: (2020) proposed Adaptive
Revised: (2020) proposed Adaptive

Original: Inertia Estimation (Adai) that uses adaptive momentum inertia instead of Adaptive Learning Rate to help training.
Revised: Inertia Estimation (Adai) that uses adaptive momentum inertia instead of Adaptive Learning Rate to help training.

Original: Adaptive inertia can be regarded as an inertia-adaptive variant of SHB.
Revised: Adaptive inertia can be regarded as an inertia-adaptive variant of SHB.

Original: The previous analysis on SHB can be easily generalized to Adai.
Revised: The previous analysis on SHB can be easily generalized to Adai.

Original: We display Adai with SWD (AdaiS) in Algorithm
Revised: We display Adai with SWD (AdaiS) in Algorithm

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: β 1 =
Revised: β 1 =

Original: [Figure] [Figure]
Revised: [Figure] [Figure]