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

Summary

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

Strengths

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

Weaknesses

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

Questions

See above.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors have discussed the limitations of the proposed method on datasets other than CIFAR-10/100.

Summary

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

Strengths

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

Weaknesses

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

Questions

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

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Not applicable.

Summary

This paper studies the role of weight decay and its connection with large gradient norms in deep learning settings. In particular, the paper highlights differences in variants of weight decay and also the effect of weight decay on gradient norm in the final phase. Based on the observation that weight decay yields large gradient norms, the authors propose a scheduler for weight decay, called the Scheduled Weight Decay, which dynamically adjusts the weight decay strength.

Strengths

The paper is reasonably well-written and the motivation for studying the effect of weight decay is both important and clear. The approach is easy to implement and the results on some smaller scale benchmarks seem encouraging.

Weaknesses

(1) The theoretical results are weak and not very interesting. Theorem 2 follows from standard convergence rates for SGD. I could not find any novelty in this result. The authors should highlight the novelty of the result and proof technique in the main paper. (2) While Theorem 1 seems interesting, it seems somewhat weak. The lower bound on the difference in learning rate is somewhat artificial and unreasonable. For instance, consider decreasing learning rates of SGD typically used in machine learning settings i.e., eta_t >= eta_{t+1}. Then Theorem 1 assumes delta <= eta_t - eta_{t+1}. This implies eta_0 >= delta . t for any t > 0. This makes sense if delta is going to 0 as t -> infinity. Maybe I am missing something but would really appreciate it if the authors can comment about this.  (3) Theorem 1 & 2 seem to analyze gradients of two different functions, which doesn't seem like a proper comparison. More remarks regarding this will be valuable. Furthermore, I would like to get clarification if the y-axis in plots of Figure 2 are gradient norms with regularization included. (4) Definition 1 is very unclear. The authors should provide more remarks around definition 1 to help readers understand the definition better. (5) The dependence of C_2 on sup norm of theta is somewhat weird in Theorem 2. Are you assuming ||theta|| is bounded? (otherwise rewrite the theorem statement without this dependence) . (6) I think the theoretical basis for Gradient-Norm-Aware Scheduled Weight Decay is weak as presented in the paper. While I understand the intuition, it is not clear why this particular proposal is meaningful. Since the theoretical basis for this is lacking, I would expect much more comprehensive experiments to support the method. Unfortunately, the empirical analysis in the paper is somewhat limited. Line 20: eta -> eta_t is the learning rate & Line 125: E[||∇L(θ, X) − ∇L(θ)||^2] ≤ \sigma^2? Post-rebuttal ========== I think the theoretical analysis is still not convincing but given that these observations are interesting, I am slightly increasing the score.

Questions

Please refer to the concerns raised in the weakness section.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

1 poor

Limitations

Not applicable.

Summary

I would divide this paper's contributions into two parts. The first part is an algorithm (a variant of Adam) which the authors argue generalizes better than Adam/AdamW and is easier to tune. The second part is the justification for the effectiveness of that algorithm. **The algorithm itself** The proposed algorithm, AdamS, is similar to AdamW, except that the weight decay strength is divided by the mean of the current squared gradient EMA. This is effectively a schedule for weight decay. It has the effect of penalizing the overall weights less strongly during phases of training when the overall gradient is large, and penalizing the overall weights more strongly during phases of training when the overall gradient is small. For example, Figure 3 shows gradient norm and weight decay strength during the training of ResNet-34 on CIFAR-10; after the first learning rate drop, squared gradient norm grows, which causes the weight decay strength to shrink; after the second learning rate drop, squared learning rate plummets, which causes weight decay to rise. **Experimental evaluation** Table 1 and section 5 experimentally evaluate AdamS and argue that the algorithm is better generalizing / easier to tune than AdamW. **Justification**: The authors justify their algorithm along the following lines (which I disagree with; see more below): 1. When the learning rate is potentially changing, gradient descent with weight decay has trouble converging to stationary points - in particular, the authors argue (Theorem 1) that in this setting, the gradient norm will be lower bounded near stationary points. Therefore, the authors argue that weight decay causes the gradient to be large, especially at the end of training. 2. The authors argue that large gradients are bad for generalization. 3. As a consequence of #1 and #2, AdamW-style weight decay is bad for generalization (since it causes large gradients near stationary points, and large gradients are bad for generalization) The authors acknowledge that weight decay on networks with normalization layers has been shown in the literature to control the effective learning rate, but insist that their paper's mechanism is unrelated to this, and propose an entirely different mechanism. --- **post-rebuttal update** After discussion with the authors, I'm raising my score from 3 ("reject") to 6 ("weak accept"). Firstly, it's clear that my original review was wrong; I had argued that the proposed algorithm's benefits are probably due to the interaction between weight decay and normalization layers, but the authors have correctly pointed out that the same phenomena occur on networks without normalization layers, and I have verified this myself. The reason for 6 as opposed to 7 is that I believe that the proposed theoretical explanation for why weight decay would cause large gradients is not the real explanation (reviewer xE1B also didn't like this theoretical portion of the paper). That said, it's entirely plausible that weight decay causing large gradients is one of the *many* aspects of deep learning that is just theoretically unexplainable at the present time. So, my issue with the submission is less that the paper fails to include a correct theoretical justification, and more that it includes a faulty theoretical justification; I think it is better to have no justification at all than to have a faulty justification. That being said, on the positive side, it seems there is a good chance that the proposed algorithm is a worthwhile addition to the deep learning toolbox (though I am not an expert at evaluating this type of contribution).

Strengths

I'm not an expert in experimentally evaluating the performance of neural network training algorithms, but from the experiments in section 5 it does seem plausible to me that AdamS is indeed better / easier to tune than AdamW.

Weaknesses

While I can accept that the algorithm may be a good idea, I strongly disagree with the proposed justification given in the paper. I believe the effectiveness of the algorithm is probably unrelated to the justification given in the paper, and is instead closely linked to the implicit effects of weight decay on effective learning rate when the network has normalization layers. First, I would note that all of the experiments in the paper are on networks with normalization layers. The authors argue in Figure 2 that VGG-16 is "not scale invariant" but the VGG-16 architecture does have many BatchNorm layers, even though it is not fully scale-invariant, so I believe that the literature on scale invariance is still quite related to what is going on with the VGG-16. For networks with normalization layers, there is a clear mechanism by which weight decay causes large gradients. If $L$ is the loss for a scale-invariant network, then $\nabla L(c \theta) = \frac{1}{c} \nabla L(\theta)$, i.e. scaling down the weights $(c < 1)$ will automatically scale up the gradients (see Lemma 1 here: http://www.offconvex.org/2020/04/24/ExpLR1/). The large gradients have been reported to cause the 'effective learning rate' of training to be higher, which is thought to be good for generalization. My intuition for the mechanism in this paper is: if the gradient norm is too large, our effective learning rate will be too large for the algorithm to converge, so we need to decrease weight decay so the gradient norms come back down; if the gradient norm is too small, our effective learning rate will be too small (which makes convergence slow and is bad for generalization) so we need to increase weight decay so that the gradient norm moves back up. By contrast, the proposed justification doesn't really make sense to me. Figure 2 shows that higher weight decay causes growth in the gradient norm _in the middle of training_, not at convergence. This is much more easily explained by scale-invariance (for the VGG-16, the scale invariance of certain layers) than by Theorems 1 and 2. Figure 5 shows that weight decay increases the eigenvalue spectrum at convergence. I believe this is related to cross-entropy loss, and wouldn't happen with e.g. squared loss. For cross-entropy loss, the Hessian is small when the margins are large. Adding weight decay prevents the margins from becoming large and therefore the Hessian from shrinking.

Questions

How do the authors respond to my critique above?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

discussed above

Grateful thanks and Response to Additional Questions

We gratefully thank the reviewer for the prompt reply and carefully re-evaluate our work. Your responsibility will be highly appreciated and is exactly what the whole community eagerly expects. We respond to your additional questions as follows. Q6: on networks with no batch normalization, does SGD with weight decay also cause high gradient norm mid-training? (That is, does mid-training gradient norm get higher as weight decay strength gets higher?) Or is this just something that Adam does? A6: Yes, SGD with weight decay also significantly cause high gradient norm during training. The observation is quite general. It does not only happen to Adam. Q7: Additionally, could you point me towards code for a VGG-16 sized for CIFAR-10 with no batch norm? A7: Yes, of course. The following github repo (https://github.com/kuangliu/pytorch-cifar/blob/master/models/vgg.py) also contains a very population implementation of VGG for CIFAR-10/100. Moreover, the torchvision.models.vgg16 is also a standard VGG-16 without BatchNorm (usually for ImageNet). The VGG-16 with BatchNorm refers to as torchvision.models.vgg16_bn.

More Discussion?

We hope our responses could address your concerns. If there are any further questions, we are very glad to continue the discussion!