SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zeke Xie $^{1,2}$ , Zhiqiang Xu $^{3}$ , Jingzhao Zhang $^{4}$ , Issei Sato $^{1}$ , and Masashi Sugiyama $^{2,1}$
<sup>1</sup>The University of Tokyo
2RIKEN Center for AIP
3MBZUAI
$^{4}$ Tsinghua University
Correspondence: zekexie16@gmail.com, Zhiqiang.Xu@mbzuai.ac.ae
Weight decay is a simple yet powerful regularization technique that has been very widely used in training of deep neural networks (DNNs). While weight decay has attracted much attention, previous studies fail to discover some overlooked pitfalls on large gradient norms resulted by weight decay. In this paper, we discover that, weight decay can unfortunately lead to large gradient norms at the final phase (or the terminated solution) of training, which often indicates bad convergence and poor generalization. To mitigate the gradient-norm-centered pitfalls, we present the first practical scheduler for weight decay, called the Scheduled Weight Decay (SWD) method that can dynamically adjust the weight decay strength according to the gradient norm and significantly penalize large gradient norms during training. Our experiments also support that SWD indeed mitigates large gradient norms and often significantly outperforms the conventional constant weight decay strategy for Adaptive Moment Estimation (Adam).
Deep learning [25] has achieved great success in numerous fields over the last decade. Weight decay is the most popular regularization technique for training deep neural networks that generalize well [24]. In deep learning, there exist two types of "weight decay":  $L_{2}$  regularization and weight decay.
People commonly use the first type of  $L_{2}$  regularization for training of deep neural networks as  $\theta_{t} = \theta_{t - 1} - \eta_{t}\nabla (L(\theta_{t - 1}) + \lambda_{L_{2}}\| \theta_{t - 1}\|^{2})$ , where  $\nabla$  is the gradient operator,  $\| \cdot \|$  is the  $L_{2}$ -norm,  $t$  is the index of iterations,  $\theta$  denotes the model weights,  $\eta_{t}$  is the learning rate,  $L(\theta)$  is the training loss, and  $\lambda_{L_2}$  is the  $L_{2}$  regularization hyperparameter. In principle, the  $L_{2}$  regularization strength is chosen to be constant during training of DNNs and people do not need to schedule  $L_{2}$  regularization like learning rate decay. The second type of weight decay was first proposed by Hanson and Pratt [13] to directly regularize the weight norm:
$$
\theta_ {t} = \left(1 - \lambda^ {\prime}\right) \theta_ {t - 1} - \eta_ {t} \frac {\partial L \left(\theta_ {t - 1}\right)}{\partial \theta}, \tag {1}
$$
where  $\lambda^{\prime}$  is the (vanilla) weight decay hyperparameter.
While the two types of "weight decay" seem similar to Stochastic Gradient Descent (SGD), Loshchilov and Hutter [32] reported that they can make huge differences in adaptive gradient methods, such as Adam [22]. Adam often generalizes worse and finds sharper minima than SGD [40, 56, 46] for popular convolutional neural networks (CNNs), where the flat minima have been argued to be closely related with good generalization [16, 17, 14, 51, 21]. Understanding and bridging the generalization
Table 1: Test performance comparison of  $L_{2}$  regularization, Decoupled Weight Decay, and SWD for Adam, which corresponds to Adam, AdamW, and the proposed AdamS, respectively. We report the mean and the standard deviations (as the subscripts) of the optimal test errors. We make the lowest two test errors bold for each model. The SWD method even enables Adam to generalize as well as SGD and even outperform complex Adam variants for the five popular models.
DATASET	MODEL	ADAMS	ADAM	ADAMW	SGD	AMSGRAD	ADABOUND	PADAM	YOGI	RADAM
CIFAR-10	RESNet18	4.910.04	6.960.02	5.080.07	5.010.03	6.160.18	5.650.08	5.120.04	5.870.12	6.010.10
VGG16	6.090.11	7.310.25	6.590.13	6.420.02	7.140.14	6.760.12	6.150.06	6.900.22	6.560.04
CIFAR-100	RESNet34	21.760.42	27.160.55	22.990.40	21.520.37	25.530.19	22.870.13	22.720.10	23.570.12	24.410.40
DENSENET121	20.520.26	25.110.15	21.550.14	19.810.33	24.430.09	22.690.15	21.100.23	22.150.36	22.270.22
GOOGLENET	21.050.18	26.120.33	21.290.17	21.210.29	25.530.17	23.180.31	21.820.17	24.240.16	22.230.15
gap between SGD and Adam have been a hot topic recently. Loshchilov and Hutter [32] argued that directly regularizing the weight norm is more helpful for boosting Adam than  $L_{2}$  regularization, and proposed Decoupled Weight Decay as
$$
\theta_ {t} = \left(1 - \eta_ {t} \lambda_ {W}\right) \theta_ {t - 1} - \eta_ {t} \frac {\partial L \left(\theta_ {t - 1}\right)}{\partial \theta}, \tag {2}
$$
where  $\lambda_W$  is the decoupled weight decay hyperparameter. We display Adam with  $L_{2}$  regularization and Adam with Decoupled Weight Decay (AdamW) in Algorithm 1. To avoid abusing notations, we strictly distinguish  $L_{2}$  regularization and weight decay in the following discussion.
While fine analysis of weight decay has attracted much attention recently [52, 39, 4], some pitfalls of weight decay have still been largely overlooked. This paper analyzes the currently overlooked pitfalls of (the second-type) weight decay and studies how to mitigate them by a gradient-norm-aware scheduler. The main contributions can be summarized as follows.
First, supported by theoretical and empirical evidence, we discover that, weight decay can lead to large gradient norms especially at the final phase of training. The large gradient norm problem is significant in the presence of scheduled or adaptive learning rates, such as Adam. From the perspective of optimization, a large gradient norm at the final phase of training often leads to poor convergence. From the perspective of generalization, penalizing the gradient norm is regarded as beneficial regularization and may improve generalization [27, 3, 2, 55]. However, previous studies fail to discover that the pitfalls of weight decay may sometimes significantly hurt convergence and generalization.
Second, to the best of our knowledge, we are the first to demonstrate the effectiveness of scheduled weight decay for mitigating large gradient norms and boosting performance. Inspired by the gradient-norm generalization measure and the stability analysis of stationary points, we show that weight decay should be scheduled with the gradient norm for adaptive gradient methods. We design a gradient-norm-aware scheduler for weight decay, called Scheduled Weight Decay (SWD), which significantly boosts the performance of weight decay and bridges the generalization gap between Adam and SGD. Adam with SWD (AdamS) is displayed in Algorithm 2. The empirical results in Table 1 using various DNNs support that SWD often significantly surpasses unscheduled weight decay for Adam.
Algorithm 1: Adam/AdamW
$$
\begin{array}{l} g _ {t} = \nabla L \left(\theta_ {t - 1}\right) + \lambda \theta_ {t - 1}; \\ m _ {t} = \beta_ {1} m _ {t - 1} + (1 - \beta_ {1}) g _ {t}; \\ v _ {t} = \beta_ {2} v _ {t - 1} + (1 - \beta_ {2}) g _ {t} ^ {2}; \\ \hat {m} _ {t} = \frac {m _ {t}}{1 - \beta_ {t} ^ {2}}; \\ \hat {v} _ {t} = \frac {v _ {t}}{1 - \beta_ {2} ^ {t}}; \\ \theta_ {t} = \theta_ {t - 1} - \frac {\eta}{\sqrt {b _ {t} + \epsilon}} \hat {m} _ {t} - \eta \lambda \theta_ {t - 1}; \\ \end{array}
$$
Algorithm 2: AdamS
$$
\begin{array}{l} g _ {t} = \nabla L \left(\theta_ {t - 1}\right); \\ m _ {t} = \beta_ {1} m _ {t - 1} + (1 - \beta_ {1}) g _ {t}; \\ v _ {t} = \beta_ {2} v _ {t - 1} + (1 - \beta_ {2}) g _ {t} ^ {2}; \\ \hat {m} _ {t} = \frac {m _ {t}}{1 - \beta_ {t} ^ {2}}; \\ \hat {v} _ {t} = \frac {v _ {t}}{1 - \beta_ {2} ^ {\prime}}; \\ \bar {v} _ {t} = \text {m e a n} (\hat {v} _ {t}) (\text {g r a d i e n t - n o r m - a w a r e}); \\ \theta_ {t} = \theta_ {t - 1} - \frac {\eta}{\sqrt {\hat {v} _ {t} + \epsilon}} \hat {m} _ {t} - \frac {\eta}{\sqrt {\hat {v} _ {t} + \epsilon}} \lambda \theta_ {t - 1}; \\ \end{array}
$$
Structure of this paper. In Section 2, we analyze the overlooked pitfalls of weight decay, namely large gradient norms. In Section 3, we discuss why large gradient norms are undesired. In Section 4, we propose SWD to mitigate the large gradient norm and boosting performance of Adam. In Section 5, we empirically demonstrate the effectiveness of SWD. In Section 6, we conclude our work.
2 Overlooked Pitfalls of Weight Decay
In this section, we report the overlooked pitfalls of weight decay, on convergence, stability, and gradient norms.
Weight decay should be coupled with the learning rate scheduler. The vanilla weight decay described by Hanson and Pratt [13] is given by Equation (1). A more popular implementation for vanilla SGD in modern deep learning libraries, such as PyTorch and TensorFlow, is given by Equation (2), where weight decay is coupled with the learning rate scheduler. While existing studies did not formally study why weight decay should be coupled with the learning rate scheduler, this trick has been adopted by most deep learning libraries.
We first take vanilla SGD as a studied example, where no momentum is involved. It is easy to see that vanilla SGD with  $L_{2}$  regularization  $\frac{\lambda}{2}\| \theta \| ^2$  is also given by Equation (2). Suppose the learning rate is fixed in the whole training procedure. Then Equation (2) will be identical to Equation (1) if we simply choose  $\lambda^{\prime} = \eta \lambda$ . However, learning rate decay is quite important for training of deep neural networks. So Equation (2) is not identical to Equation (1) in practice.
Which implementation is better? In Figure 1, we empirically verified that the popular Equation-(2)-based weight decay indeed outperforms the vanilla implementation in Equation (1). We argue that Equation-(2)-based weight decay is theoretically better than Equation-(1)-based weight decay in terms of the stability of the stationary points. While we use the notations of "stable"/"unstable" stationary points for simplicity, we formally term "stable"/"unstable" stationary points as strongly/weakly stationary points in Definition 1. This clarifies that there can exist a gap between true updating and pure gradient-based updating for stationary points.
Definition 1 (Strongly/Weakly Stationary Points). Suppose  $\mathcal{A}:\Theta \to \Theta$  is a function that updates model parameters during training where  $\Theta$  is the parameter space. We define  $\theta^{\star}$  a strongly stationary point for  $\mathcal{A}$  if  $\| \theta^{\star} - \mathcal{A}(\theta^{\star})\| = 0$  holds during the whole training process. We define  $\theta^{\star}$  a weakly stationary point for  $\mathcal{A}$  if  $\| \theta^{\star} - \mathcal{A}(\theta^{\star})\| = 0$  holds only for some iterations.
The training objective at step-  $t$  of Equation (1) is  $L(\theta) + \frac{\lambda'}{\eta_t} \| \theta \|^2$ , while that of Equation (2) is  $L(\theta) + \lambda \| \theta \|^2$ . Thus, stationary points of the regularized loss given by Equation (2) are stable during training, while that given by Equation (1) is unstable due to the dynamic training objective.
Figure 1: We compared Equation(1)-based weight decay and Equation(2)-based weight decay by training ResNet18 on CIFAR-10 via vanilla SGD. In the presence of a popular learning rate scheduler, Equation-(2)-based weight decay shows better test performance. It supports that the form  $-\eta_t\lambda \theta$  is a better weight decay implementation than  $-\lambda' \theta$ .
Weight decay increases gradient norms at the final
phase of deterministic optimization and stochastic optimization. Following the standard convergence analysis of Gradient Descent (GD) [6] with the extra modification of weight decay, we prove Theorem 1 and show that even GD with vanilla weight decay may not converge to any non-zero stationary point due to missing stability of the stationary points.
Theorem 1 (Non-convergence due to vanilla weight decay and scheduled learning rates). Suppose learning dynamics is governed by  $GD$  with vanilla weight decay (Equation (1)) and the learning rate  $\eta_t \in (0, +\infty)$  holds. If  $\exists \delta$  such that satisfies  $0 < \delta \leq |\eta_t - \eta_{t+1}|$  for any  $t > 0$ , then the learning dynamics cannot converge to any non-zero stationary point satisfying the condition
$$
\max  (\| \nabla f _ {t} (\theta_ {t}) \| ^ {2}, \| \nabla f _ {t} (\theta_ {t + 1}) \| ^ {2}) \geq \frac {\lambda^ {\prime 2} \sigma^ {2} \| \theta^ {\star} \| ^ {2}}{\eta_ {t} \eta_ {t + 1}} > 0,
$$
where  $f_{t}(\theta) = L(\theta) + \frac{\lambda^{\prime}}{\eta_{t}}\| \theta \|^{2}$  is the regularized loss.
We leave the proof in Appendix A.1. Theorem 1 means that even if the gradient norm is zero at the  $t$ -th step, the gradient norm at the  $t + 1$ -th step will not be zero. Thus, the solution does not
(a) AdamW
(c) SGD
(d) AdamW
(e) Adam
(f) SGD
Figure 2: Weight decay increases the gradient norms and the squared gradient norms ResNet18 for various optimizers. Top Row: the gradient norms. Bottom Row: the squared gradient norms.
converge to any non-zero stationary point. GD with unstable minima has no theoretical convergence guarantee like GD [41]. More importantly, due to unstable stationary points, the gradient norms (of the regularized/original loss) naturally also do not converge to zero. This may explain why weight decay should be coupled with the learning rate to improve stability of stationary points during training. This non-convergence problem is especially significant when  $\delta$  is relatively large in the presence of scheduled or adaptive learning rates, such as Adam.
Existing studies [48, 12] proved that the gradient norm bound of SGD in convergence analysis depends on the properties of the training objective function and gradient noise. We further prove Theorem 2 and reveal how the gradient norm bound in convergence analysis depends on the weight decay strength  $\lambda$  when weight decay is involved. We leave the proof in Appendix A.2.
Theorem 2 (Convergence analysis on weight decay). Assume that  $L(\theta)$  is an  $\mathcal{L}$ -smooth function $^1$ ,  $L$  is lower-bounded as  $L(\theta) \geq L^{\star}$ ,  $L(\theta, X)$  is the loss over one minibatch  $X$ ,  $\mathbb{E}[\nabla L(\theta, X) - \nabla L(\theta)] = 0$ ,  $\mathbb{E}[\| \nabla L(\theta, X) - \nabla L(\theta) \|^2] \leq \delta^2$ , and  $\| \nabla L(\theta) \| \leq G$  for any  $\theta$ . Let SGD optimize  $f(\theta) = L(\theta) + \frac{\lambda}{2} \| \theta \|^2$  for  $t + 1$  iterations. If  $\eta \leq \frac{C}{\sqrt{t + 1}}$ , we have
$$
\min  _ {k = 0, \dots , t} \mathbb {E} [ \| \nabla f (\theta_ {k}) \| ^ {2} ] \leq \frac {1}{\sqrt {t + 1}} \left[ C _ {1} + C _ {2} \right], \tag {3}
$$
$$
C _ {1} = \frac {L \left(\theta_ {0}\right) + \frac {\lambda}{2} \left\| \theta_ {0} \right\| ^ {2} - L ^ {\star}}{C}, \tag {4}
$$
$$
C _ {2} = C (\mathcal {L} + \lambda) ((G + \sup  (\lambda \| \theta \|)) ^ {2} + \sigma^ {2}), \tag {5}
$$
and  $\sup (\| \theta \|) = \sup_{k = 0,\dots ,t}\| \theta_k\|$  is the maximum  $L_{2}$  norm of  $\theta$  over the iterations.
Theorem 1 shows the non-convergence problem with improper weight decay. Theorem 2 further shows that, even with proper convergence guarantees and stable minima, the gradient norm upper bound at convergence still monotonically increases with the weight decay strength  $\lambda$ . Obviously, Theorem 2 theoretically supports that weight decay can result in large gradient norms. The phenomenon can also be widely observed in the following empirically analysis.
It means that  $\| \nabla L(\theta_a) - \nabla L(\theta_b)\| \leq \mathcal{L}\| \theta_a - \theta_b\|$  holds for any  $\theta_{a}$  and  $\theta_{b}$
(a) VGG16 with BatchNorm
(b) VGG16 without BatchNorm
(c) VGG16 with BatchNorm
(d) VGG16 without BatchNorm
Figure 3: Weight decay increases the gradient norms and the squared gradient norms of VGG16 with BatchNorm (scale-invariant) and VGG16 without BatchNorm (scale-variant) on CIFAR-10. Top Row: the gradient norms. Bottom Row: the squared gradient norms.
Empirical evidence of large gradient norms resulted by weight decay. Theorem 1 suggests that both unstability and strength of weight decay may increase the gradient norm even in full-batch training. While weight decay is helpful for regularizing the model weights' norm, we clearly observed in Figure 2 that weight decay significantly increases the gradient norms in training of DNNs as we expect. Note that the displayed weight decay of AdamW is rescaled by the factor 0.001. The observation generally holds for various optimizers, such as AdamW, Adam and SGD. Figure 3 further shows that the observation holds for both scale-invariant models (VGG16 with BatchNorm) and scale-variant models (VGG16 without BatchNorm). Moreover, the observation still holds for gradient norms with or without regularization included, because the numerical difference due to regularization is smaller than their gradient norms by multiple orders of magnitude in the experiments. Zhang et al. [52] argued that weight decay increases the "effective learning rate" due to smaller weight norms resulted by weight decay. However, the "effective learning rate" interpretation may not explain large gradient norms, because a small weight norm does not necessarily indicate a large gradient norm. Interestingly, Figure 3 shows that the gradient curves of VGG16 without BatchNorm and weight decay sometimes exhibit strange spikes, which can be wiped out by weight decay.
3 Large Gradient Norms Are Undesired
In this section, we demonstrate that the large gradient norm is a highly undesirable property in deep learning because large gradient norms correspond to multiple pitfalls.
First, small gradient norms are required for good convergence, which is desired in optimization theory [36]. Popular optimizers also expect clear convergence guarantees [8]. The optimizers with no convergence guarantees may cause undesirable training behaviors in practice.
Second, regularizing the gradient norm has been recognized as an important beneficial regularization effect in recent papers [28, 3, 11, 55, 54]. Li et al. [28] and Geiping et al. [11] expressed implicit regularization of stochastic gradients as
$$
R _ {\mathrm {s g}} = \frac {\eta}{4 | \mathcal {B} |} \sum_ {X \in \mathcal {B}} \| \nabla L (\theta , X) \| ^ {2} = \frac {\eta}{4} \mathcal {G} (\theta), \tag {6}
$$
where  $L(\theta, X)$  indicate the loss over one data minibatch  $X \in \mathcal{B}$ ,  $|\mathcal{B}|$  is the number of minibatches per epoch, and  $\mathcal{G}$  is the expected squared gradient norms over each minibatch. Recent studies, Zhao et al. [55] and Zhang et al. [54], also successfully improved generalization by explicitly adding the gradient-norm penalty into stochastic optimization. Obviously, penalizing the gradient norm is considered as an important and beneficial regularizer in deep learning.
Third, large gradient norms may theoretically indicate poor generalization according to the gradient-norm generalization measures [27, 2]. For example, Theorem 11 of Li et al. [27] derived a gradient-norm-based generalization error upper bound that uses the squared gradient norms, written as
$$
\operatorname {S u p} \operatorname {G e n} = \mathcal {O} \left(\frac {L ^ {\star}}{n} \sqrt {\mathbb {E} \left[ \frac {\eta_ {t} ^ {2}}{\sigma_ {t} ^ {2}} \mathcal {G} \left(\theta_ {t}\right) \right]}\right), \tag {7}
$$
where  $n$  is the training data size,  $\sigma_t^2$  is the gradient noise variance at the  $t$ -th iteration, and  $\mathcal{G}(\theta_t) = \frac{1}{|\mathcal{B}|}\sum_{X\in \mathcal{B}}\| \nabla L(\theta ,X)\| ^2$  is the squared gradient norm with the batch size as 1. An et al. [2] presented more empirical evidences for supporting the gradient-norm-based generalization bound (7). Moreover, large gradient norms are also believed to be closely related to minima sharpness [20, 57, 43, 47, 42, 46, 7], while sharp minima often indicate poor generalization [16, 17, 14, 51, 21]. As our analysis focuses on the large gradient norm at the final phase (or the terminated solution) of training, our conclusion is actually not contradicted to the viewpoint [1, 29, 10] that the gradient-noise regularization effect of large initial learning rates at the early phase of training is beneficial.
We also report that the gradient norm and the squared gradient norm show very similar tendencies in our experiments. For simplicity, we mainly use  $\mathcal{G}$  as the default gradient norm measure in this paper, unless we specify it otherwise. As we discussed above, the squared gradient norm  $\mathcal{G}$  is more closely related to implicit gradient regularization, generalization bounds, and SWD than other gradient norms. Overall, the theoretical results above supported that the small gradient norm is a desirable property in deep learning. However, our theoretical analysis suggests that weight decay can lead to large gradient norms and hence hurt generalization. We suggest that gradient-norm-centered poor convergence and generalization can lead to the serious but overlooked performance degradation of weight decay.
4 Gradient-Norm-Aware Scheduled Weight Decay
In this section, we design the first practical scheduler to boost weight decay by mitigating large gradient norm.
We focus on scheduled weight decay for Adam in this paper for three reasons. First, Adam is the most popular optimization for accelerating training of DNNs. Second, as the adaptivity of Adam updates the learning rate every iteration, Adam suffers more than SGD from the pitfalls of weight decay due to adaptive learning rates. Third, SGD usually employs  $L_{2}$  regularization rather than weight decay.
Adam with decoupled weight decay is often called AdamW [32], which can be written as
$$
\theta_ {t} = (1 - \eta \lambda) \theta_ {t - 1} - \eta v _ {t} ^ {- \frac {1}{2}} m _ {t}, \tag {8}
$$
where  $v_{t}$  and  $m_{t}$  are the exponential moving average of the squared gradients and gradients in Algorithm 1, respectively, and the power notation of a vector means the element-wise power of the vector. We interpret  $\eta v_{t}^{-\frac{1}{2}}$  as the effective learning rate for multiplying the gradients. However, we clearly see that decoupled weight decay couples weight decay with the vanilla learning rate rather than the effective learning rate. Due to the differences of weight decay and  $L_{2}$  regularization, AdamW and Adam optimize different training objectives. The minimum  $\hat{\theta}^{\star}$  of the regularized loss function optimized by AdamW at the  $t$ -th step is not constant during training. Thus, the regularized loss function optimized by AdamW has unstable minima, which may result in large gradient norms
(a) Norm  $\mathcal{G}$
(b) Scheduler
(c) Training Loss
(d) Norm  $\mathcal{G}$
(e) Scheduler
(f) Training Loss
Figure 4: The curves of the weight decay scheduler,  $\mathcal{G}$ , and training losses of SWD and constant weight decay. Top Row: ResNet18 on CIFAR-10. Bottom Row: ResNet34 on CIFAR-100. SWD mitigates large gradient norms.
suggested by Theorem 1. This may be an internal fault of all optimizers that use adaptive learning rates and weight decay at the same time.
To mitigate the pitfall of large gradient norms, we propose a gradient-norm-aware scheduler for weight decay as
$$
\theta_ {t} = \left(\mathbf {1} - \eta \bar {v} _ {t} ^ {- \frac {1}{2}} \lambda\right) \theta_ {t - 1} - \eta v _ {t} ^ {- \frac {1}{2}} m _ {t}, \tag {9}
$$
where  $\bar{v}_t$  is the mean of all elements of the vector  $v_t$ . Moreover, as  $v_t$  is the exponential moving average of squared gradient norms, the factor  $\bar{v}_t$  is expected to be equivalent to the squared gradient norms divided by model dimensionality. While the proposed scheduler cannot give perfectly stable minima for Adam, it can penalize large gradient norms by dynamically adjusting weight decay. We call weight decay with this gradient-norm-aware scheduler Scheduled Weight Decay (SWD). The pseudocode is displayed in Algorithm 2. The code is publicly available at GitHub. In Section 5, we empirically demonstrate the advantage of the gradient-norm-aware scheduler. Note that AdamS is slightly slower than AdamW but slightly faster than Adam, while the cost difference is nearly ignorable (usually less than  $5\%$ ).