Why Do We Need Weight Decay in Modern Deep Learning?

The paper studies the user of weight decay (WD) in deep learning. There are three main contributions of the paper:

1. The authors state that WD results in implicit regularization by controlling the learning rate. They provide two conjectures which build on empirical observations they have made.
2. The authors replicate findings from (Hoffman el al. 2022), showing that WD leads to lower final loss in the final stages of training LLMs. The authors argue that WD will modulate the effective LR.
3. The authors observe that for a small GPT style model, WD allows better stability when training in bf16. The authors hypothesize that WD alleviates precision issues for bf16.

This paper studies why we need weight decay in modern neural networks. The paper claims that, different from the regularization role in conventional machine learning, weight decay has novel roles and various roles for over-parameterized neural networks and under-parameterized neural networks. For over-parameterized neural networks (CNNs), weight decay modifies the optimization dynamics and enhancing the ever-present implicit regularization of the SGD noise. For under-parameterized neural networks (LLMs trained with nearly one-pass SGD), weight decay balances the bias-variance tradeoff in stochastic optimization leading to lower training loss. Preventing sudden loss divergences for mixed-precision training of LLMs is also reported for weight decay.

This paper studies the effect of weight decay in modern deep learning tasks, ranging from computer vision to language modeling.

This paper investigates the effect of weight decay in modern deep learning, such as overparameterized models and large language models. The main finding is that, on over-parameterized models, weight decay regularizes the trace of Hessian implicitly, while on LLMs, weight decay mainly benefits the optimization dynamics instead of regularizing the model. Finally, the paper finds the reason for why LLM training has loss spikes from the perspective of the low precision of bfloat16.

The study explores the necessity of weight decay in deep learning, specifically addressing its relevance and effects. Key contributions of the research include:

1. Weight decay is shown to enhance the implicit regularization of SGD noise in overparameterized networks by maintaining the learning trajectory close to that of a process with a regularized Hessian trace.
2. In LLMs trained with nearly one-pass SGD, weight decay does not act as a significant regularizer but adjusts the effective learning rate and improves the bias-variance tradeoff, reducing training loss.
3. Weight decay also serves to prevent abrupt loss divergences during bfloat16 mixed-precision training, critical for training LLMs at scale.

We thank the reviewers for their constructive suggestions and for identifying some inconsistencies in our paper. We have decided to withdraw our submission to further investigate the behavior of the implicit regularization term $\nabla^2 \mathcal{L}(\mathbf w)$ in the conjectures.

Our conjecture is primarily inspired by the existent theoretical works studying implicit regularization of label noise gradient descent (LNGD) on the $*squared*$ loss. Nevertheless, during the rebuttal period, we discovered that the behaviour is different. If our conjecture would hold as formulated, the averaged iterates should have a smaller value of the regularizer for larger values of the learning rate (note that this need not happen for the finetuned iterates). More specifically, although $Tr(\nabla^2)$ is effectively minimized for the finetuned iterates as reported in Figure 3 (c) , $Tr(\nabla^2)$ for the EMA does not necessarily have the correct trend with respect to $\eta$ as posited in our conjecture (see Figure 8 in appendix A.2). Furthermore, the trend seems to depend on the coefficient used in the EMA. We can obtain the expected trend for EMA coefficient $0.95$ but not $0.999$. This made us consider the possibility that $Tr(\nabla^2)$ might not be the correct quantity to explain the improvement in generalization when WD is used in combination with large LRs. This intuition is supported by the observation that $Tr(\nabla^2)$ depends on the value of the training loss, potentially clarifying the apparent disparity between the experiments and the conjecture. Attempting to find a quantity that reveals the implicit regularization effect without being directly affected by the value of the training loss, we conducted some preliminary experiments on the Jacobian norm. This quantity closely relates to $Tr(\nabla^2)$ without carrying the explicit dependence on the value of the loss as we elaborate in Appendix A.2. The preliminary experiments show the correct trend of the EMA for standard ResNet-18 on CIFAR-10. We provide a longer discussion on this in Appendix A.2 of the updated version of the paper, together with more details on why we believe both quantities are at play during the large learning rate phase. However, a thorough investigation of the behaviour of the Jacobian norm is needed (i.e., on multiple architectures and datasets) before updating the conjecture.

We would like to briefly clarify other points about our submission:

• Overparameterization vs. underparameterization. We agree that this is not the best way to present this distinction and we should have called the one-epoch training of LLMs as the “underfitting regime”. What we really meant by (effective) underparameterization is that for one-epoch training, the training and validation losses stay very close to each other. We have updated Figure 4 to more clearly illustrate this for LLMs by plotting the generalization gap directly over training iterations for a GPT-2 small and GPT-2 large models.
• Precise meaning of regularization. By regularization, we meant any technique that restricts the hypothesis class and reduces the generalization gap. As illustrated in Figure 4, weight decay does not noticeably reduce it for LLMs trained for one epoch (unlike for CIFAR-10 trained for multiple epochs). Thus, we conclude that the role of weight decay is different from reducing the generalization gap. We will use more precise wording to explain what we mean.
• Novelty of the bfloat16 observation. We would like to emphasize that bfloat16 has the same dynamic range as float32, so the issue is not in the overflow (which is an issue, for example, for float16 whose maximum possible value is $65\\,519$ and larger values are interpreted as NaN). We will add a more thorough explanation of this to the paper. Moreover, we will investigate the precise mechanism of how weight decay helps to address the precision issue.
• Parameters of the LLM experiments. We used gradient clipping with the default parameter $1.0$ for all experiments so the loss divergence with bfloat16 is not due to the lack of gradient clipping. We used the context length of 256 primarily to speed up the experiments for Section 4.1 but the overall pattern is the same for the context length 1024.

We thank the reviewers again. These comments will help us to formulate our ideas more precisely and collect more comprehensive experimental evidence for the next version of our paper.