Why Do We Need Weight Decay in Modern Deep Learning?
Authors: Francesco D'Angelo
Maksym Andriushchenko
Aditya Varre
Nicolas Flammarion
Venue: NeurIPS 2024 poster (2024)
Keywords: Weight decay
overparameterization
implicit regularization
large language models
optimization dynamics
Primary area: optimization_for_deep_networks

Abstract: Weight decay is a broadly used technique for training state-of-the-art deep networks from image classification to large language models. Despite its widespread usage and being extensively studied in the classical literature, its role remains poorly understood for deep learning. In this work, we highlight that the role of weight decay in modern deep learning is different from its regularization effect studied in classical learning theory. For deep networks on vision tasks trained with multipass SGD, we show how weight decay modifies the optimization dynamics enhancing the ever-present implicit regularization of SGD via the *loss stabilization mechanism*. In contrast, for large language models trained with nearly one-epoch training, we describe how weight decay balances the *bias-variance tradeoff* in stochastic optimization leading to lower training loss and improved training stability. 
Overall, we present a unifying perspective from ResNets on vision tasks to LLMs: weight decay is never useful as an explicit regularizer but instead changes the training dynamics in a desirable way.

Why Do We Need Weight Decay in Modern Deep Learning?
Francesco D' Angelo, Maksym Andriushchenko, Aditya Varre, Nicolas Flammarion
Theory of Machine Learning Lab EPFL, Lausanne, Switzerland
(francesco. dangelo, makayn anwrishhheenoo, aditya. varre www.ea nicolas. (ommmarronnS@ppfff
Abstract
Weight decay is a broadly used technique for training state-of-the-art deep networks from image classification to large language models. Despite its widespread usage and being extensively studied in the classical literature, its role remains poorly understood for deep learning. this work, we highlight that the role of weight decay in modern deep learning is different from its regularization effect studied in classical learning theory. For deep networks on vision tasks trained with multipass SGD, we show how weight decay modifies the optimization dynamics enhancing the ever-present implicit regularization of SGD via the loss stabilization mechanism. In contrast, for large language models trained with nearly one-epoch training, we describe how weight decay balances the bias-variance tradeoff' in stochastic optimization leading to lower training loss and improved training stability. Overall, we present a unifying perspective from ResNets on vision tasks to LLMs: weight decay is never useful an explicit regularizer but instead changes the training dynamics in a desirable way. The code is available at https: //gi thub. com/ tml-epf 1/why-/wwwwa/
werchh-wweieee wew
Introduction
The training of modern neural networks broadly falls into two regimes: over-training. which involves multiple passes through dataset and necessitates effective regularization strategies to avoid overfitting; and umder-training, characterized by fewer passes due large amounts of training data and computational constraints (Hoff- mann al., 2022). Modern deep learning unequivocally embodies both training regimes: ResNet architectures on computer vision tasks (He et al., 2016) serve as quintessential examples of the wver-training regime, while the training of large language models (Brown et al., 2020) stands hallmark of the under-training regime. Despite their differences, both regimes extensively adopt weight decay as a regularization technique, though its effectiveness and role remain subjects of ongoing debate. For the first regime, Zhang et al. (2016) 6) showed that even when using weight decay, neural networks can still fully memorize the data, thus questioning its regularization prop- erties. For the second, regularization is inherently unnecessary as the limited number of passes already prevents overfitting. These considerations raise important questions about the necessity and purpose of weight decay, introducing uncertainty about its widespread usage. To illustrate the effect of weight decay in the two regimes,
0.35
Tan WD
0.30
Torn
AN M117
0.25
0.20
0.15
0.30
0.05
0.00
50k
250k
5m
Dataset
Figure 1: Test error dataset size CIFAR-IO-5m for fixed number of training iteration. Weight decay is helpful in both: the over-training and the under- training, one-pass regime.
"quual contribution
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

we conduct a simple experiment. We train a 8 on subsets of the CIFAR-5m dataset (Nakkiran et al., 2020) with sizes from 10 000 to 5 mln. The computational budget of each training session is fixed to 5 mln iterations, which amounts to a range of passes between 500 and one. In the over-training regime (left in Fig. 1), weight decay does not prevent the models from achieving zero training error, but its presence still improves the test error. Attempting to explain this generalization benefit, recent works (Li & Arora, 2019; Li et al., 2020) bring forth the hypothesis that it is inadequate to think about weight decay as a capacity constraint since it still bears an effect on the training of scale-invariant models. As result, understanding the effect of weight decay on the optimization dynamics becomes effective
crucial to understanding generalization. Nevertheless, this line of work heavily relies learning rate (ELR) which only emerges as a consequence of inveriiiveeeanee and therefore does not apply general architectures. In the under-training regime (right in Fig. 1), where the generalization gap vanishes, weight decay seem to facilitate faster training for slightly better accuracy. However, a characterization of the mechanisms through which weight decay impacts the training speed in this regime remains underexplored.
Our work delves into the mechanisms underlying the benefits of weight decay by training established machine learning models in both regimes: ResNet on popular vision tasks (over-training) and Transformer on text data (under-training). Towards this goal, we make the following contributions:
In the over-training regime, we unveil the mechanism by which weight decay effectively reduces the generalization gap. We demonstrate that combining weight decay with large learning rates enables vanissiin SGD noise, which through its implicit regularization controls the norm of the Jacobian leading to improved performance. Moreover, our investigation offers a thorough explanation for the effectiveness of employing exponential moving average and learning rate decay in combination with weight decay.
In the under-training regime, particularly for LLMs trained with one-pass Adam, we confirm experimentally that weight decay does not bring any regularization effect and is simply equivalent a modified ELR. We explain the training curves commonly observed with weight decay: through this ELR, weight decay better modulates the bias-variance trade-off, resulting in lower loss. Additionally, we show that weight decay has another important practical benefit: enabling stable training with the bfloat16 precision.
1.1 Related work
The concept of employing le weight penalty traces back to studies on the stability of solutions for ill-posed problems (Tikhonov, 1943). It has since been extensively explored in statistics (Foster, 1961; Hoerl, 1962; Hoerl & Kennard, 1970). Krogh & Hertz (1991) present one of the earliest systematic studies on weight decay tailored for neural networks. Generalization bounds, such as those by Shalev-Shwartz & Ben-David (2014), suggest that weight decay can be sufficient for generalization, although not necessary, e.g., due to the implicit regularization of gradient methods (Soudry et al., 2018). Zhang et al. (2016) argue that while weight decay improves test accuracy, the improvement is not substantial (== 1- 1-2% on ImageNet), indicating the key role of implicit regularization. Loshchilov & Hutter (2019) highlight the distinct effects of weight decay and l2 regularization, particularly for Adam, suggesting that Adam combined with weight decay (AdamW) leads to superior regularization and simpler hyperparameter tuning. For GPT-3 training, Brown et al. (2020) suggest that they include weight decay to provide a small amount of regularization, although we believe it is not the primary reason as we discuss in Sec.
Multiple works have focused on weight decay as a tool influencing optimization dynamics. Van Laarhoven (2017) emphasizes that weight decay's impact on scale-invariant networks is primarily seen in terms of an effective learning rate. Zhang et al. (2018) propose three mechanisms of weight decay regularization: (1) increasing the effective learning rate for scale-invarianl networks, although as we discuss, the same holds for networks beyond scale invariance (2) approximating the regulariza- tion of the input Jacobian for an optimizer inspired by second-order methods, (3) inducing a specific dampening effect in this optimizer. Li & Arora (2019); Li et al. (2020) explore the optimization properties of scale-invariant deep networks for which the effective learning rate can be derived. Lewkowycz & Gur-Ari (2020) suggest that the best generalization is achieved with the smallest X although it necessitates longer training. Additionally, Lewkowyez (2021) propose a criterion for detecting when to decay the learning rate based on the evolution of the weight norm. Bjorck et al. (2021) explore the effect of decoupling weight decay, especially in the early stage of training. Li et al. (2022a) make BERT architecture inle-invarriatt to enhance training stability and make it more

---

Resnet18 CIFARIO
Resnet18 CIFARIO
Resnet18 CIFARIO
Resnet18 on Tiny ImageNet
100%
1eat S3re
6010
5004
024
Epochs
Epochs
Epochs
Epochs
(a)
(d)
Figure 2: Training with and w/o weight decay. We report the test error for Resnet18 8 on CIFAR-10 10 (2a) and
Imageee (2d) trained with and without weight decay and with small and large learning rates, We also include the correspondent EMA, represented by dashed lines. After the first 250 epochs the learning rate decayed to 1 = 10 for all the curves. We report also the L2 norm of the parameters (2c) and Train CE (2b) which after the decay converges to the same value for all the runs with the same X.
compatible with standard SGD. Recently, Kosson et al. (2023) show a mechanism through which weight decay balances rotational updates across different layers that motivates
optimizer.
The seminal paper of Krizhevsky et al. (2012) that introduced AlexNet suggest that weight decay serves not only as regularizer but also reduces the model's training error, functioning as an optimization tool. In recent work, Hoffmann et al. (2022) briefly observe that weight decay enhances the training performance of Adam for training LLMs, but only after x 80% of the total iterations. However, they do not provide an explanation for this behavior, a point we delve into in Sec. 3.
2 Weight decay in the over-training regime
In this section, we delve into the influence of weight decay in the over- over-training regime, with specific focus on image classification tasks. We focus on the training of ResNet models (He et al., 2016) using SGD on Tiny-ImageNet (Wu et al., 2017) and report additional experiments in appendix C for VGG, Resnet32 and scale-invariant Resnet architectures on CIFAR10 and CIFARI00 (Krizhevsky & Hinton, 2009).
Notations and setup. Let (xr,, "= be the training inputs and labels where v, E Rd,, yi E IR. and is number classes. Let h IP x IR + IR be the hypothesis class of neural network and for any parameter R2 where the function hw (.) .Rd > R represents the network predictions. The training loss L and the fz-regularized loss Lx, for À 2 O, are given by:
A 2
((w)= N Σ e (ve, (x:))
(x(w)= = L(w)
w/
where (((-)): R R + R denotes the cross-entropy (CE) loss function. With ~ U([N]), the SGD algorithm on Lx (here with batch size and with replacement) with a learning rate (LR) 1) is
W++1 = W - V/ w (yi, h(w,, xz,)) 7//w.
(I)
Along the training we track three different iterates: (1) large-L. denoted by w, which use a large constant LR to exploit the SGD noise, (2) fine-tuning w, which starting from w, use a small LR and (3) the exponential moving average (EMA) w, along the large-LR iterates.
2.1 Loss stabilization and weight decay
To understand whether minimizing the regularized objective LA alone ensures optimal generalization, we compare test errors in Fig. 2a across various settings. Although both large and small LRs minimize the regularized objective, the evidence that optimal performance is achieved exclusively with large I.Rs indicates that the objective alone is insufficient to explain the benefits of WD or ensure generalization. This experiment reaffirms the widely acknowledged consensus that implicit regularization induced by the LR is crucial (Keskar et al., 2016; Li et al., 2019; Andriushchenko et al.,
Thh red, blue and green curves have the same CE 2b and norm 2c, hence LA but different test error 2a.
3

---

2023). Despite revealing an interplay between weight decay and large initial LR, the understanding of the corresponding dynamics remains limited. In this section, our goal is to comprehensively understand these dynamics, particularly to elucidate the difference in generalization between training with and without weight decay and using different learning rates, as observed in Fig. 2a. Given the regularization of the l2 norm of parameters, it is natural to wonder whether weight decay's improvement primarily stems from its ability to control the norm of the trained model. The experiment in Fig. 2c clearly illustrates that distinct training trajectories, while resulting in the same final Ez norm for parameters, can yield different levels of generalization stating that the fg-norm of the learned model's parameters is inconsequential. This observation suggests that once the norm is constrained by weight decay, the critical factor influencing the model's generalization is the subsequent choice of LR. We should note that neural networks can be explicitly made layers and small architectural changes. L. & Arora (2019) have used this setting to reveal that the training dynamics has effective learning rate which, depending reduces the effect of WD to merely a scheduler for the learning rate. Our analysis presents a broader perspective that does not depend on scale invariance. At the core of our examination are the unique properties of exponentially tailed loss functions, such as CE: when the data is separable and WD not applied, the infimum of the loss norm (Ji & Telgarsky, 2019; Soudry et al., 2018). The application of WD, by inhibiting this growth, prevents the decrease of CE loss, which in turn, significantly alters the optimization dynamics.
by means of normalization
the 1.2-norm of the parameters,
infinity, leading to the unbounded growth of the weight
Indeed, examining the parameter norm evolution in Fig. 2c, we notice how it rapidly decreases to stabilize within small, approximately constant interval. Similarly, the Train CE in Fig. 2b displays a stabilization effect beyond which it cannot decrease without a reduction in the learning rate. We hypothesize that WD enables an optimization dynamic akin to that on the surface of sphere of certain radius thus triggering the following essential mechanism:
Constraining the parameter norm hinders the decrease of the CE, thereby enabling non-vanishing noise in SGD. This allows SGD implicit regularization to unfold and steer the optimization trajectory.
Next, we empirically characterize this implicit regularization mechanism.
2.2 The noise driven process
The long-held belief that the implicit regularization property of SGD is pivotal to the generalization capabilities of Neural Networks has been a cornerstone in the field of deep learning (Keskar et al., 2016). Many theoretical studies (Blanc et al., 2020; Li et al., 2021b; Damian et al., 2021), attempting to understand this phenomenon, draw upon the essential finding that, in the case of regression and when Gaussian noise is added to the labels, the shape of the covariance of the stochastic gradients matches the shape of the Hessian. This allows Damian et al. (2021) and Pillaud- Vivien et al. (2022) to show that the trajectory of the SGD iterates closely tracks the solution of a regularized problem. In our analysis, we conjecture a similar result; the dynamics of SGD with CE, closely track a regularized process. The important difference in our statement is that unlike Blanc et al. (2020) and Damian al. (2021), we do not need add noise to the labels at each iteration. Instead, weight decay, in combination with large-LR induces a label noise-like behavior via loss stabilization (Andriushchenko al., 2023). To better understand the interplay between weight decay, loss stabilization and the noise of SGD, it is convenient to consider the binary classification case where y/s E (0, 13. We also define the Jacobian of the network as J JZz, w) Vh (a:)) E I2" and its norm averaged across the dataset: [[[ N EN Tr hh (zx) hhw (x;) Denoting the noise of the gradient by 9: = [(w))) - Vw (yi. h(W,, Iz. )), the SGD update in equation becomes:
W++1 = (1 - n)) ((W)) + 199t
Furthermore, we consider Gaussian approximation of the SGD noise, matching the first and second
moment of gt- A substantial body of research has built upon this approximation and verified its
validity (Li et al., 2020, 2021a; Smith et al., 2020; Xie et al., 2020; Li et al., 2021b). In particular Li
et al. (2021a) demonstrated how modelling the SGD noise by a Gaussian sufficient to understand
its generalization. The mean is zero due to the unbiased mini-batch gradients: g(w.) == E[gd] = 0

---

Test Error EMA
[[
0.05
0.15
6,025
00225
0.
8
0,03
6.005
0025
0:80
0.001
00-
0002
6003
000
001
008
63
015
c3-
025
0.003
000
007
099
0.1
62-
0.35
Figure 3: Resnet18 Imay-Im Heatmap of the error and Jacobian for the EMA for all the different combinations 7) and X.
whereas for the second moment:
N
Eww
Vli(w,)Ve (w:) - VL((,)
2
Ve, w: Vl, (wa)
N
N
N Σ ([((w))) (1x) Vlo (I2)
~
1 N n.
V Vlw, (2:a Vhw, (x:)
(3)
In equation 3, consider VL(w,) (wr) negligible compared to (w, (wr) the gradient noise variance dominates the square of the gradient noise mean. This fact has been used previous works (Mori et al., 2022: Zhu al., 2018: Jastrzebski al., 2017) and in particular Jastrzebski et al. (2017) and Saxe et al. (2019) empirically verify it. Finally, we assume that the first derivative is approximately constant across all datapoints: ((ww) Rs (((w)) Vi.j and denote this common quantity as x (w.). This last approximation is referred to as "decoupling approximation" and has been empirically verified for classification (Mori et al., 2022). Furthermore, in App. C.6 we performed additional experiments to verify the decoupling approximation by comparing the spectrum of the SGD covariance with and without this approximation during the large LR phase.
Altogether, these considerations lead to the following formulation of the SGD update:
Wx+ 2 (1 ndww, --V--v.W))))) N A
(w:)) Σ Vhw, (x:) E where E ~ N(0 (4)
1=1
This series of approximations, allows us to define the quantity Jg (We). which has fundamental role in the characterization of the training dynamics. We refer as the scale of the noise because regulates its intensity. Although 1) and À influence the noise scale indirectly through the trajectory ofw,, we explicitly highlight this dependence to emphasize our objective: characterize the influence of WD and LR on the stochastic dynamics of SGD. We develop this characterization building upon the connection between the Jacobian of the network and the covariance of the SGD noise, which motivates the introduction of the following implicit regularization mechanism:
Conjecture 1. Consider the algorithm in Eq. with wo initialized from distribution μr (R(()). For any input x, let w., h(w.. be the random variables that denote the iterate at time and its functional value. The stochastic process (h (h(we,)) 2)EN converges to the stationary distribution
pe with mean (n,)(x) h ( for which wi, is a first-order stationary point of the following regularized loss:
[x(w)= = [x(w) no? ]
(5)
The conjecture illustrates how varying noise levels correspond distinct processes, wherein the mean of each solves a unique regularized problem. Moreover, it describes how each noise level determines the strength of the regularization. Using the mean formulate the conjecture is natural choice: even at stationarity," the values of the loss ((w,) and the regularizer [[w)]] would be dominated
Assuming the existence of stationary distribution, the iterates we eventually realizations from it.

---

14
al Sol oduunge
go
1Vtellle
Train
02
1A
(a)
(b)
(c)
Figure 4: Resnet18 on Tiny-ImageNet. Training for 200 epochs with different 7 and X; the scale of the noise monotonically increases with the train loss and x À Fig. 4a, 4b. The test error instead, presents optimal value of 7 x Fig. 4e while the Jacobian norm decreases monotonically Fig. 4d.
by the noise. To unveil the existence of an implicit regularization and to analyze the evolution of the distribution, we need look at its summary statistics, in this case. the mean. While insights from Langevin dynamics suggest employing learning rate annealing to converge towards the mean of the stationary distribution, this approach introduces additional complexities which we discuss in Section 2.3. We instead consider an exponential moving average (EMA) (We)e20 of the SGD iterates 0.999.
with parameter
The most important implication of the conjecture is that the strength of the regularization Gx,A depends on both the LR and the WD parameter À. Our experiments in Fig. 3, provide empirical validation for this conjecture. When trained with different combinations of and X, the EMA converges to models with different test performances. When fixing there exists an optimal value of learning rate which gives the best test performance while the Jacobian norm monotonically increases. A similar picture can be drawn when fixing the learning rate.
Therefore, given two solutions A and A for which Test (mror(mm.)))) < Test arror(pn...o the difference in their performances can be explained with the difference in their regularization strengths T.A- The solution Ar benefits from better regularization and therefore endows better generalization properties. Furthermore, the heatmap for test error in Fig. 3 indicates that the minimum error is not achieved by a single combination of and X, but rather along a contour where their product 17 x À appears to be constant. This observation suggests an optimal trade-off between the learning rate and weight decay parameter, characterized by a curve in the parameter space where their product remains constant. Characterizing this relationship might reveal useful tool which practitioners can adopt to optimally select values of weight decay and learning rate. Fig. 4c, 4d confirm that the product ε/λ is the quantity controlling the regularization; combinations of v) and A with the same product show similar test performances and Jacobian norm. For iny-Immgenet, the test error exhibits an optimal value for yd ~ 0.005 where increases beyond this point lead to over-regularizztion and decreases result in reder-regularization. Simultaneously, the Jacobian norm |||| consistently exhibits a monotonically decreasing trend.
To better understand the properties of the noise scale during training, we can observe in Fig. 4a, 4b how higher values of the training loss given by larger 17λ correspond to higher values of the noise scale .. The latter is measured by computing the Frobenius norm of the first derivative of the loss with respect the predictions* averaged over all training datapoints scale of the noise and therefore the strength of the regularization vanish when L(w) = which happens for small values of njλ. Therefore, WD combined with a large LR stabilizes the CE to a larger value, preventing the noise from vanishing and regulating the implicit regularization. This fact further confirms Conjecture 1, demonstrating that with smaller noise scales, the mean (EMA) tends towards point where the Jacobian's norm is higher, compared to trajectories with larger noise scales. To further validate our conjecture, we created snapshot ensembles in the ResNet18 on CIFAR-10 setting, by averaging in function space along the SGD trajectory every 10 epochs for the combinations of learning rate (LR) and weight decay (WD) considered. To assess whether the mean of the stationary distribution in function space aligns closely with the EMA, where the Jacobian norm is regularized, we compared the performance of snapshot ensembles with that of the EMA. Additionally, we computed the Total Variation Distance between the softmax outputs of the ensemble and the EMA. The results are reported in App. C.7 and show strong alignment.
Ci=/lee( [| The
In the culti-class setting (w:)) E RR; to quantify the scale, compute its Frobenius norm.

---

The role of the dynamics of the norm. As discussed at the end of previous sub-section, after rapid initial decrease of norm, the optimization resembles the dynamics of SGD projected onto a sphere. We stress that this is the crucial phase in training and the implicit regularization induced by SGD during this spherical optimization leads to better generalization, To validate this observation and isolate it from the initial norm drop, we investigate the behavior of SGD on a sphere with scale-invariant networks (Li & Arora, 2019). Scale invariance is chosen for its ease of LR tuning and for comparison with existing works (Kodryan et al., 2022; Li et al., 2020). Fig. 15 depicts a similar phenomenon as Fig. 3, where the test error vs. LR exhibits a U-shaped curve. While Kodryan et al. (2022) makes a similar observation, they do not provide explanation. We demonstrate that the implicit regularization of the Jacobian norm is the key factor, elucidating its dependence on LR.
Effective learning rate vs. high training loss. Zhang et al. (2018); Van Laarhoven (2017) explored
the relationship between LR and WD, introducing the concept of effective LR 1/2 - -n//|| w I|2 These studies highlight that WD, preventing unbounded growth of the norm, enables the training process to evolve with a higher effective LR. This hypothesis is justified only with scale-invariance (which does not hold for general architecture). Furthermore, the underlying mechanism by which higher LR enhances generalization understood only limited settings (Li et al., 2019). We propose that high LR, combined with WD, leads to an increase in yT This hypothesis allows us to fully characterize and understand the mechanism through which generalization enhanced.
Mixing in the function space. A simpler conjecture could have been stated in terms of the mixing of the iterates (we) towards a solution of the regularized objective w However, Li et al. (2020) shows empirical evidence against mixing in the parameter space, emphasizing the necessity of considering the function space. Hence, our conjecture is formulated to capture stationarity in function space.
On the benefit of normalization. Our conjecture characterizes the mixing distribution but does not delve into the speed of the mixing process. In our experiments, we observe that normalization layers enable faster mixing. Li et al. (2020) observes similar phenomenon in the case of scale-invariant networks, specifically the fast equilibrium conjecture, which is addressed by Li et al. (2022b). We note that this phenomenon persists even when the models are not exactly scale-invariant.
2.3 EMA and Fine-tuning
The large-LR phase sets the stage for SGD's inherent biases to emerge but to actually exploit such bias, reducing the stochastic noise is necessary. This can be attained in two different ways: averaging (EMA) or decaying the learning rate (fine-tuning), both strategies are widely adopted in practice. This section illustrates the relation between the two and highlights the benefits of using one over the other and the implications for our analysis. From a practical standpoint, implementing EMA is more efficient than LR-decay because it does not require additional gradient iterations or hyperparameter tuning. While both methods enhance performance, their effectiveness is contingent on being combined with loss stabilization, supporting the hypothesis that the noisy dynamics is the underlying factor for their success. Although EMA shows only a slight advantage, our experiments in Fig. 5b, 2a, 2d demonstrate that it consistently outperforms learning rate decay in various settings.
When empirically validating Conjecture the EMA is useful to characterize the limit point (i.c.. t ++ oo) but cannot adequately capture the dynamics throughout the entire trajectory. This limitation arises because different points along the trajectory are at different loss values, making the comparison of any relevant regularized quantities inconsistent. An approach to overcome this inconsistency is to project the iterate w, onto a manifold of constant loss. This can be achieved via early-stopped gradient flow (Li et al., 2021b) (SGD with small LR) on the CE loss with À 0 where w, is projected to nearby point w. such that L(W,) ~ const., Vt. In practice, this corresponds to fine-tuning via LR-decay. Since after fine-tuning L(w)) L(w,), Vt, t' see Fig. 5a, we can compare [|J(W)) and ||J( and understand its evolution. In the experiments detailed in Fig. 5c, we report ||]] along the fine-tuned iterates W, and observe decreasing trend, i.e., the sequence [1))())) ) is
decreasing. This fact empirically validates that the entire trajectory of the iterates (w:) closely following the trajectory of the fine-tuned iterates (W,)r20. bias the model towards a regularized solution that might enhances generalization. This also explains why learning rate schedules, such as step-decay, which starts with a large value and then decrease it, can enhance generalization.
7

---

Itais Cross Entropy
M-will
Epocha
Epccw
Epocks
Spccw
Figure 5: EMA vs Fine-tuning, Training of standard Resnet18 on CIFAR-10 10 for 100 epochs fixing À = 0.0125 and varying the learning rate. In Fig. 5a we report different levels of loss stabilization, in Fig. 5b report the test errors and Fig. 5c and Fig. 5d the norm of the Jacobian and the weights respectively. The quantities are measured for the SGD iterates, the EMA and the fine-tuning. The latter is performed for 100 epochs every with r = 10
Despite providing a straightforward methodology to analyze the trajectory, LR-decay introduces additional complexities that cause deviations from the conjecture. Indeed, in Fig. 5c we observe that the final points of the nnnet iterates report the opposite trend compared to the EMA (smaller 1) lead to larger || |||||> p)) This discrepancy is potentially due to the stte-dependent nature of the SGD noise covariance in equation 3; decreasing the LR and removing WD can alter the stationary distribution and the regularized objective, leading to a different solution than anticipated by Conjecture 1.
3 Weight decay in the under-training regime
In this section, we investigate how WD enhances optimization in the under-training regime. Although the phenomenon is more general, we focus on LLMs for which one-epoch training is typically used.
Two key effects of weight decay in the under-training regime. WD is widely used in training state-of-the-art LLMs like GPT-3, Chinchilla, and Llama (Brown et al., 2020; Hoff- mann et al.. 2022; Touvron et al., 2023). While Brown et al. (2020) suggest that WD offers "a small amount of regulariza- tion." its necessity remains unclear in the context of one-pass training where the population loss is directly minimized. As sanity check, in Fig. 19 in Appendix, we verify that the generalization gap is close to zero even for models trained without WD. Instead of the regularization effect, we suggest that the two most crucial effects of WD in the under-training regime are (1) better optimization of the training loss as briefly observed by Hoffmann et al. (2022), (2) prevention of loss divergences under the bfloat16 weight precision. We repro- duce this phenomenon at a smaller scale with 124M parameters in Fig. 6: the final training loss is smaller for X equal to 0.1 and 0.3 compared to O. We study both mechanisms which stand in contrast to the over-training regime of Sec. 2, where the primary concerns are not optimization and stability, but rather generalization.
Awe in O.O, My
Avo t O.1,
Ave 0.3,
Training lass.
Avo
3.5
Ave
20000
30000
40202
50000
50300
lteration
Figure 6: GPT-2-124M on OpenWeb- Text. We reproduce the improvement from WD as in Hoffmann et al. (2022) but at much smaller scale. Performing fine-tuning with a tiny LR reveals that a higher starting training loss can still be a better point in terms of the final loss after fine-tuning.
Experimental setup. We use the NanoGPT repository (Karpathy, 2023) for training GPT-2 models (Radford et al., 2019) on Webxee., We train a 124M parameter model (known as GPT-2-Small) for 50 000 iterations with a batch size of 256. For most experiments, we reduce the default context length from 1024 to 256 to ensure practicality within an academic budget. Alternatively, we could have reduced the number of training iterations or batch size, but this would lead to insufficiently trained models. Unless mentioned otherwise, we train with AdamW using the default LR 0.0006, short 400-iteration LR warmup, gradient clipping with the thrthshhold 1.0, and 10x cosine LR decay. We keep all other hyperparameters at their default values (see App. B).
Better optimization with WD is reproducible at smaller scale. The findings from Hoffmann et al. (2022) (Fig. A7 therein) indicate that WD in AdamW leads to lower training loss ( 0.02 lower), primarily towards the end of training. The reduction of training loss directly translates to better downstream performance and makes this observation practically relevant. Additionally, performing fine-tuning with tiny LR reveals that a higher starting training loss can still be a better starting point
8

---

0.04
Effective learning rate
B000 App
Training loss
0.000
0,000
0.000
Iteration
Steration
Reration
Figure 7: GPT-2-124M on Open WebTex. Left: The effective LR np. /I wallz for the models reported in Fig. 6. Middle: The LR schedule that matches the effective LR //el l| the runs with weight decay 0.1 and 0.3. Right: Matching the effective LR is sufficient to match the whole training dynamics of the loss if we avoid the loss spikes by using full precision (float32 instead of bfloat16).
in terms of the final loss after fine-tuning. Moreover, in Fig. 21, we show that decoupling WD, as advocated by Loshchilov & Hutter (2019), is not necessary to achieve this effect: a simple & penalty added to the loss suffices. Lastly, in Fig. 22, we show that similar improvement in training loss is observed for SGD with momentum suggesting that adaptive LRs are not key for this phenomenon.
Effective LR induced by weight decay in Adam We hypothesize that the use of WD for LLM training results in an increased effective LR, even in the absence of scale invariance of the training loss for modern transformer architectures. Here we show that WD in combination with sign SGD utilized as surrogate for Adam (Balles & Hennig, 2018) is equivalent to projected SGD on the sphere, with an effective LR 7Jeff 2//lw:lla. similarly to Van Laarhoven (2017). Consider the update rule of sign SGD on loss with WD:
Wx+1 = (1 - ():))) W. - sign (V((w:)) (1 :1dde)) ||well2
W
Ne sign (Ve(We))
[well
(1 De) W/
Considering the evolution of the direction w =- =://|wll,
7/ sign( (gem((e((.)))))
W X
W. Ao) well2
When sign( (Ve(W))) is determined solely by the direction W., the evolution of the direction of weights becomes the primary matter. This scenario occurs when the function é is scale-invariant or homogeneous. Observing the trend of the gradient norm and the parameter norm |wall from Fig. 26, we note an inverse relationship, i.e., the gradient norm is higher when the parameter norm is lower. This behavior is reminiscent of
networks for which V((aw) = 1 V((ow), for any f Thus, controlling parameter norms with WD allows implicit changes to the LR schedule which we verify experimentally next.
Matching the effective LR without weight decay. To verify the key role of the ELR nr/lww, la. we train a model without WD but with ELR corresponding to models trained with À E (0.1, and report results in Fig. We observe that the whole training dynamics is matched which confirms our hypothesis. This fully explains the observation from Hoffmann et al. (2022) about the advantage of AdamW W over Adam: there exists an LR schedule (albeit non-standard one) shown in Fig. (middle) that leads to the same loss profile as the original A run. However, we note that this holds only for full float32 precision, and models trained with bfloat16 precision diverge in the middle of training. This experience suggests that WD is still necessary in practice to prevent loss divergence. We also note that matching the ELR "t/lww.le derived above for sign SGD instead of n/1 w for plain SGD (Zhang et al., 2018; Hoffer et al., 2018) is critical for AdamW. Otherwise, the runs diverge very early in training, even with float32 parameter precision.
Explaining the training dynamics of A.amW. Classical optimization theory suggests that con- vergence of SGD-based algorithms primarily depends on two factors: the bias term that influences the rate at which initial conditions are forgotten and the variance term that results from noise in the gradient estimates (Moulines & Bach, 201 1). We argue that these two factors, together with the observation about higher ELR for WD, can explain the loss profiles from Fig. 6. If we consider the simple case of SGD with a constant LR r) applied to linear least-squares problem, the expected excess risk after iterations can be bounded as a sum of a bias and variance terms:
Excess Risk $ (1 nu)" Wo w. + n
9

---

where a is uniform bound on the variance of the noise of gradient estimates, Je lower bound on the objective function's 's Hessian, Wo the initial point and w, the optimum. For linear models, is well-established that a larger L.R accelerates the contraction of the bias term but has a detrimental impact on the variance term, ultimately leading the variance term to dominate. Coming back to the dynamics in Fig. 6, with a large ELR at the start, the convergence becomes primarily bottlenecked by the high variance term proprtional to the learning rate, leading to higher loss values in the presence of WD. Conversely, towards the end of training, when ELR and the variance term are reduced, we see that WD catches up and performs better at the end, thanks to its relatively higher ELR throughout the training and thus better bias contraction. This perspective sheds light on the observation that EMA for LLMs is most advantageous when employed with large LRs (Sanyal et al., 2023) as we also illustrate in Fig. 23. As the variance dominates in this case, variance reduction of the averaging helps.
Experiments with bfloat16. Training in reduced precision essential for speeding up training and reducing GPU memory requirements (Kalamkar et al., 2019). We further elaborate on the fact that WD is not fully equivalent to higher ELR and remains necessary for stable bfloat16 training. While Scao et al. (2022) observe that usage of float16 can cause loss divergences, bfloat16 is considered much more stable and is de-facto standard in LLM training. Although bfloat16 shares the same floating-point exponent size as float32 (thus, the range of possi- ble values is the same), it offers lower precision, with only 7 bits for the fraction instead of 23. We observe that even more stable bfloat16 data type can still exhibit late-training spikes that irreparably harm model per- formance in standard practical settings, such as with larger context length (e.g.. 1024 instead of 256 as in the previous experiments). Therefore, experiments shown in Fig. 8. Runs with a moderate LR 0.001 (the default LR of Adam) without WD exhibit late-training divergence for all random seeds when using 16. By comparison, training with float32 remains entirely stable. Importantly, we observe that the model does not recover after the loss spikes which contrasts with the loss spikes described in the Edge of Stability phenomenon (Cohen et al., 2021, 2022). We emphasize that all these runs use gradient clipping with the standard fz-threshold. Finally, we observe that divergences can be prevented by reducing the LR, e.g., from 0.001 to 0. 0.0006. However, this adjustment leads to slower training, as illustrated in Fig. 24 in the Appendix. Instead, the most effective approach is to use higher LR of 0.001 with WD, which enables stable bfloat16 training and yields better final training loss.
staata LtwO 001
3:35
stsatle 13 0000 Ano 0.2 stzitoo Lt- DOL, Aap Seatzz Jtw A.0
2:33
Training loss
3.25
focus on this configuration for the
3.00
39383
2030
sreen
ceres
50002
Iteration
Figure 8: GPT-2-124M on OpenWeb- Text with context length 1024. Weight decay prevents divergence for LR 0.001 and enables stable bfloatt training. The three random seeds are denoted with lines.
Conclusions
In this paper, we demonstrate how weight decay, single hyperparameter, can manifest three distinct effects across different training regimes: it offers regularization when combined with stochastic noise, improves optimization of the training loss, and guarantees stability in low-precision training environments. In the over-training regime, the scale of the noise a. is the fundamental quantity governing the implicit regularization strength of SGD. Weight decay combined with large LR enables the noisy dynamics to evolve by maintaining the scale at a constant level. Techniques such as EMA or fine-tuning work by reducing noise, thereby allowing for the effective exploitation of the accumulated hidden regularization. Coming to the under-training regime, Adam W (Loshchilov & Hutter, 2019) was introduced as a regularization method. Instead, we argue that it is effective as it modulates the ELR to attune the bias ttas-variance tradeoff. In addition, it also improves the stability of training. In summary, weight decay is seldom valuable as an explicit regularizer: instead, its widespread adoption can be attributed to its ability to induce desirable changes in optimization dynamics. We also acknowledge limitations of our work: given our limited computational resources, we do not conduct truly large-scale experiments. Moreover, we do not prove new theoretical results. Instead, we strive to provide a clear experimental picture and formulate general explanations for the effectiveness of weight decay in different training regimes
10

---

Acknowledgements
We thank Atli Kosson and Alex Damian for fruitful discussions and suggestions. M.A. was supported by the Google Fellowship and Open Phil AI Fellowship. A. V. was supported by the Swiss Data Science Center Fellowship. F.D. was supported by the Swiss National Science Foundation (grant number 212111)
References
Andriushchenko, M., Varre, A., Pillaud- Vivien, L.,, and Flammarion, N. SGD with large step sizes learns sparse features. In International Conference Machine Learning, 2023.
Balles, I.. and Hennig, P. Dissecting adam: The sign, magnitude and variance of stochastic gradients.
In Dy, J. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 404-413 PMLR, 10-15 Jul 2018.
Bjorck, J., Weinberger, K. Q., and Gomes, C. Understanding decoupled and early weight decay. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 6777-6785, 2021,
Blanc, G., Gupta, N., Valiant, G., and Valiant, P. Implicit regularization for deep neural networks driven by an OrnsteinUUllenbeck like process. In Conference on Learning Theory, 2020.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P. Sastry, G., Askell, A., al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
Cohen, J. M., Kaur, S., Li, Y., Kolter, J. Z., and Talwalkar, A. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2021.
Cohen, J. M.. Ghorbani, B., Krishnan, S., Agarwal. N., Medapati. S., Badura, M., Suo, D., Cardoze, D.. Nado, Z., Dahl, E., al. Adaptive gradient methods at the edge of stability, arXiv preprint 2PXi::2207.14484 2022.
Damian, A., Ma. T., and Lee, J. D. Label noise sgd provably prefers flat global minimizers. In Advances in Neural Information Processing Systems, volume 34, pp. 27449-2744 2021.
Foster, M. An application of the wiener- koomoo-oo smoothing theory matrix inversion. Journal of the Society for Industrial and Applied Mathematics, 93):3877-32,, 1961.
Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. Openwebtext corpus, 2019. http: //Skylion007 github..hh..... 1o/OpenweeTextcorpss..
He, K., Zhang, X., Ren, S., and Sun, Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. ICCV, 2015.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In CVPR, 2016.
Hoerl, A. F. and Kennard, R. W. Ridge regression: Biased estimation for nonorthogonal problems. Technometrics, 12(1):55-67, 1970.
Hoerl, A. R. Application of ridge analysis to regression problems. Chemical Engineering Progress, 58:54-59, 1962.
Hoffer, E., Banner, R., Golan, I., and Soudry, D. Norm matters: efficient and accurate normalization schemes deep networks. Advances Neural Information Processing Systems, 31, 2018.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models, In Advances in Neural Information Processing Systems, 2022.
Huang, G., Li, Y., Pleiss, G., Liu, Z., Hoperoft, J. E., and Weinberger, Q. Snapshot ensembles: Train get for free. arXiv preprint arXiv: 2017.

---

Jastrzebski, S., Kenton, Z., Arpit, D., Ballas, N., Fischer, A., Bengio. Y., and Storkey, A. Three factors influencing minima sgd. arXiv preprint arXiv:111.004233, 2017.
Ji, Z. and Telgarsky, M. The implicit bias of gradient descent on nonseparable data. In Conference on Learning Theory, pp. 1772-1798 PMLR, 2019.
Kalamkar, D., Mudigere, D., Mellempudi, N., Das, D., Banerjee, K., Avancha, S., Vooturi, D. T., Jammalamadaka, N., Huang, J. Yuen, H., et al. A study of 6 for deep learning training. arXiv preprint arXiv: /2322, 2019.
Karpathy, A. Taposit 2023. httpsttss///ittie htwpaa/......aatttaaaaanaooooo
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. P. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2016.
Kodryan, M., Lobacheva, E., Nakhodnov. M., and Vetrov, D. P. Training scale-invaricn neural networks on the sphere can happen in three regimes. Advances Neural Information Processing Systems, 35:14058 35:14058-14070, 2022.
Kosson, A., Messmer, B., and Jaggi, M. Rotational equilibrium: How weight decay balances learning across neural networks. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning. 2023.
Krizhevsky, A. and Hinton, G. Learning multiple layers of features from tiny images. Technical Report, 2009.
Krizhevsky. A., Sutskever, I., and Hinton, G. Imagenet classification with deep convolutional neural networks. In Advances in Neural Information Processing Systems, 2012.
Krogh, A. and Hertz, J. A simple weight decay can improve generalization. Advances in Neural Information Processing Systems, 1991.
Lewkowycz, A. How to decay your learning rate. arXiv preprint arXiv:2103.12683 2021.
Lewkowycz, A. and Gur-Ari, G. On the training dynamics of deep networks with 1_2 regularization. In Advances in Neural Information Processing Systems, volume 33, pp. 4790-4799, 2020.
Li, Y., Wei, C., and Ma, T. Towards explaining the regularization effect of initial large learning rate in training neural networks. In Advances in Neural Information Processing Systems, 2019.
Li, Z. and Arora, An exponential learning rate schedule for deep learning. arXiv preprint arXiv:1910.074 2019.
Li, Z., Lyu, K., and Arora, S. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate. In Advances in Neural Information Processing Systems, volume 33, pp. 14544-14555, 2020.
Li, Z., Malladi, S., and Arora, S. On the validity modeling sgd with stochastic differential equations (sdes). In Advances in Neural Information Processing Systems, 2021a.
Li, Z., Wang, T., and Arora, S. What happens after sgd reaches zero loss?-a a mathematical framework. arXiv preprint arXiv:2110.0690 2021
Li, 7., Bhojanapalli, S., Zaheer, M., Reddi, S., and Kumar, S. Robust training of neural networks using scale invariant architectures. In Proceedings of 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 12656-12684, PMLR. 17-23 Jul 2022a.
Li, 7., Wang, T., and Yu, D. Fast mixing of stochastic gradient descent with normalization and weight decay. In Advances in Neural Information Processing Systems, 2022b.
Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
12

---

Mori, T., Liu, Z., Liu, K., and Ueda, M. Power-law escape rate of SGD. In International Conference on Machine Learning. PMLR, 2022.
Moulines, E. and Bach, F. Non-asymptotic analysis of stochastic approximation algorithms for machine learning. In Advances in Neural Information Processing Systems, volume 24, 2011.
Nakkiran, P., Kaplun, G., Bansal, Y., Yang, T., Barak, B., and Sutskever, I. Deep double descent: Where bigger models and more data hurt. In International Conference on Learning Representations, 2020.
Papyan, V. The full spectrum of deepnet hessians scale: Dynamics with sgd training and sample size. arXiv preprini arXiv: /81 I. 1.07062, 2018.
Pillaud- Pillaud-Vivien, L., Reygner, J., and Flammarion, N. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Conference on Learning Theory, 2022.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Sagun, L., Evci, U., Guney, V. U., Dauphin, Y., and Bottou, L. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint arXiv:1706.04454 2017.
Sanyal, S., Kaddour, J., Kumar, A., and Sanghavi, S. Understanding the effectiveness of early weight averaging for training large language models. arXiv preprint arXiv:2306.0322 03241, 20233.
Saxe, A. M., Bansal, Y., Dapello, J., Advani, M., Kolchinsky, A., Tracey, B. D., and Cox, D. D. On the information bottleneck theory of deep learning. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):124000, 2019.
Scao, T. L., Fan, A., Akiki, C., Pavlick, E., Ilié, S., Hesslow, D., Castagné, R., Luccioni, S., Yvon, F., Gallé, M., et al. Bloom: A robb-parameter open-access multilingual language model. arXiv preprint arXiv:2211.05100 2022.
Shalev-Shwarze S. and Ben-David, S. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
Smith, S., Elsen, E., and De, S. On the generalization benefit of noise in stochastic gradient descent. In International Conference on Machine Learning, pp. 9058-9067 PMLR, 2020.
Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., and Srebro, N. The implicit bias gradient descent separable data. The Journal of Machine Learning Research, 19(1):2222-2878 2018.
Tikhonov, A. N. On the stability of inverse problems. In Dokl. Akad. Nauk SSSR, volume 39, pp- 195-198, 1943.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint ar 2xx.:23322.19977
2023.
Van Laarhoven, T. L2 regularization versus batch and weight normalization. arXiv preprint 1rXiv:1706.05350 2017.
Wojtowytsch, S. Stochastic gradient descent with noise machine learning type. Part I: Discrete time analysis. arXiv preprint arXiv:2105.01660 2021.
Wu, J., Zhang, Q., and Xu, G. Tiny imagenet challenge. Technical report, 2017.
Xie, Z.. Sato, I., and Sugiyama, M. A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima. arXiv preprint arXiv:2002.08 004.00....344,.... 2020.
Yao, Z., Gholami, A., Keutzer, K., and Mahoney, M. W. Pyhessian: Neural networks through the lens of the hessian. 2020 IEEE international conference on big data (Big data), pp. 581-590. IEEE, 2020.
13

---

Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2016.
Zhang, G., Wang, C., Xu, B., and Grosse, R. Three mechanisms of weight decay regularization. In International Conference Learning Representations, 2018.
Zhu, Z., Wu, J., Yu, B., Wu, L.,, and Ma, J. The anisotropic noise in stochastic gradient descent: Its be- havior of escaping from sharp minima and regularization effects. arXiv preprint arXiv:/803.00/99 2018.
14

---

A An additional comparison with related works
Our focus in Section 2 is on an empirical illustration of the implicit regularization phenomenon, hence we refrain from attempting to prove this general conjecture, which we believe a challenging task. The existing theoretical works Blanc et al. (2020); Li et al. (2021b); Damian et al. (2021) present two major weaknesses; they are essentially limiting analysis and as such fail at capturing the entire optimization trajectory and they primarily target regression tasks. The powerful mathematical framework for scale-invariant networks developed by Li & Arora (2019); Li et al. (2020) allows them to study in detail the benefits of normalization and its interplay with weight decay. By means of this framework, they state a fast equilibrium conjecture, which gives qualitative guarantees for the speed of convergence of the stochastic process to the stationary distribution in function space. They disentangle the evolution of the norm and the direction of the parameters and show how the evolution of the direction only depends on the intrinsic LR X; n1. However, a qualitative description of the stationary distribution, its dependence on this intrinsic LR and the relationship with generalization is missing Li et al. (2020). We attempt to fill this gap by providing qualitative depiction of the stationary distribution and its dependence on the intrinsic LR shading some light towards understanding the relationship with generalization. The work of Kodryan et al. (2022) reports a similar observation, where the best test loss is achieved at LR where the loss neither converges nor diverges but does not provide any explanation.
Table I: Comparison of our work with elosely related works regression and implicit regularization phc- nomenon induced by noise the algorithm. Paper
Loss function
Algorithm
Implicit regularization
Damian et al. (2021) &
Squared loss &
Label noise GD
Trace Hessian
Li et al. (2021b)
CE label smoothing
Blanc et al. (2020)
Squared loss
Label noise GD
Jacobian norm
Li et al. (2020)
svale-ivvrriant loss
SGD
Andriushchenko et al. (2023)
Squared loss
SGD with large LR SGD with large LR
Jacobian norm
Our work
Regularized CE
Jacobian norm
B Training details
Full experimental details are available in our public repository https://gitiee wttpsollleomm//e/eeeeeeee epfl/ why-weight www.httteeeeeeeeeeeeee decay but we also list the main training details here. All the experiments are conducted for 3 different random seeds, the error-bars report one standard deviation. CIFAR-10/100 experi- ments. We train a VGG network without BatchNorm and preactivation ResNet-18 on CIFAR-10 10 and ResNet-34 on CIFAR- 100 without data augmentations. We use standard SGD without momentum for all experiments. We note that &g regularization and weight decay are exactly the same in this case. We use the standard He initialization (He al., 2015) for all parameters. To make ResNets scale-invariant, we follow the approach of Li et al. (2020) consisting of fixing the last layer, removing the learnable parameters of the normalization layers and adding a normalization layer in the skip connection. For the experiments in Fig. 11, VGG is trained with LR = 0.1 and LR = 0.01 and weight 0.008. The ResNet-18 trained with LR- = 0.08 and LR = 0.001 and x= = 0.0 or X= = 0.0125. The ResNet-34 is trained with LR 0.15 and LR- 0.001 and weight decay parameter À= = 0.0 or X = 0.01. The total number of epochs 1000 in all experiments in Fig. 11 and all the I.R are decayed at epoch 500 to 0.0001. For the experiments in Fig. 15 we use scale-invariam R and project the SGD iterates on the unitary sphere. We test the following LRs in the large-LR phase (0.0001, 0.0005.00 0.00075, 0.001, 0.002, 0.003, 0.004, 0.005) to show different generalization performance. After 100 epochs all the learning rates are decayed to the same value 0.0001. In Fig. 15 we fine-tune every 2 epochs for 100 additional epochs with LR=0.0001. To measure the Norm of the Jacobian or the Trace of the Hessian we use a subset of 5000 training datapoints. Each run requires approximately 2 GPU hours
decay parameter is fixed to be either X = 0.0 or
Nvidia A100 GPU.
Tiny- ImageNet experiments. We train Resnet-18 without data augmentation. We use stan- dard SGD without momentum in all our experiments. We use the following learning rates (0.0005, 0.00TO, 0.0050, 0.0100, 0.0500, 1000, O. 15000 0.2000, 0. 0.2500) and weight decay param- eter (0.0200, 0.0150, 0.0125, 0.0100, 0.0075, 0.0050, 0.0025, 0.0010, 0.0005, 0.0000). To measure
15

---

the norm of the Jacobian we use subset of the training data of 2500 examples. Each run requires approximately 5GPU hour on A100.
LLM experiments. We use the NanoCPT repository (Karpathy, 2023) for training GPT-2 models (Radford al., 2019) Open WeeTet (Gokaslan al., 2019). All training documents con- catenated in single stream from which a new batch is sampled with replacement on every iteration of training. We train 124M parameter model known GPT-2-small for 50 000 iterations instead of the default 600 000 to make grid searches over the learning rate and weight decay parameters more accessible within an academic budget. We mostly use the context length of 256 for faster experiments except for Fig. 8 where we use the context length of 1024 since we observed that larger context length is crucial to observe loss divergences with moderate learning rates (such as 0.001 for Adam). We train with Adam (Loshchilov & Hutter, 2019) using batch size 256, default LR 0.0006 (unless mentioned otherwise). 81 0.9, 82 = 0.95, short 400-iteration LR warmup, and 10x cosine LR decay. For the runs with SGD with momentum, we use the learning rate 0.3 and momentum parameter 0.9 using the same LR schedule for Adam W. We initialize all parameters with the standard deviation equal to 0.02. We keep all other hyperparameters at their default values as in the NanoGPT repository. We perform all experiments on A100 Nvidia GPUs that support fast bf 1fl training. Each training run of GPT-2-small for 50 000 iterations takes around 12 hours on a single GPU.
C Weight decay for overparametrized deep networks: additional experiments and details
C.1 graphical illustration of the fine-tuning phase
Here, plot an illustrative graphic in Fig. 9 to give an idea of what happens during the fine-tuning phase.
Figure 9: A graphical illustration of the fine-tuning phase.
C.2 Supporting derivations
Here prove that the scale of noise is well approximated by training loss in the of binary classification instead of classification in the case of multiple classes. The proof follows the lines of Wojtowytsch (2021).
Proposition 2. Assume |w||
[a, b), for any E D, (w,,)) E [m, M] holds. For n sufficiently large, there exists constants C1, 02 such that
[|](w)]]] s egC(w)
Proof. The noise the case when the gradient is computed at (I: y:)
g(w) (v,, [h(w
ve' Vh(w
16

---

Taking the expectation over uniform sampling over i, we have,
E||a|| = Σ (Cw..hww,a (17/(w,,))]1
Σ V (Wt, ((,, m;)l|² (6)
n
Upper bound:
Using the self-bounding property of the binary CE, i.e., (ea) S I and
VV (w,,), < M
E||o|| < M
Σ T2
Cly, h(w,))( = M²C(w).
Comment on the Lower bound: Since the iterates are bound, we can assume there exists constant such that (E²²) 2 cl. as the second term in 6 is decreasing with O(n ?), we can assume that the first term is dominating and relevant and can lower bound the first term as,
E||g|| z cm
lly, (yr, =mm ((w).
C.3 Additional figures for the over-training regime
In this section, report additional experimental results related to Section 2 in the main text. In Fig. 10 we report analogous results for the jacobian norm and test error of the EMA for Resnet on CIFARIO. In Fig. and 2 we report the train CE for VGG and ResNet18 on CIFAR-I0 and ResNet34 trained on CIFAR-100. We can observe how when weight decay is used in combination with large LR, the train CE stabilizes at some approximately constant level.
Test Error EMA
|
303
Erv EMA
(a)
(b)
epochs with different v Fig. 10a reports a heatmap of the test crror and Jacobian norm for the EMA for all the different combinations of parameters. The test error presents an optimal value of 7 when is fixed and, consistently with conjecture 1, the Jacobian norm decreases monotonically. More over, Fig. 10b shows how the optimality might depend only the product yu for which the test error has U-shape and the Jacobian
Figure 10: Resnet18 on CIFARIO train Resnet18 on CIFARIO
decreases monotonically.
CIFAR10
Resnet CIFAR10
on CIFAR100
49%
+23
Erer
39%
49%
33%
30%
Epochs
Epochs
Epochs
(b)
Figure 11: Training with and w/o weight decay. We report the test error for VGG (1la) and ResNet (1lb, 11c) trained CIFAR-1O// with and without weight decay and with small and large learning rates. After the first 500 epachs the learning rate decayed to 10- for all the curves.
17

---

Resnet CIFARID
Resnet CIFARI00
VGG on CIFARIO
201
10-11
1E-4
Cress
10-4
1200
Epochs
Epochss
Epochs
(b)
(c)
Figure 12: Training with and w/o weight decay. We report the train CE for VGG (12a) and ResNet (12b, 12c) trained CIFAR- 10/-00001000000 with and without weight decay and with small and large learning rates. After the first 500 epochs the learning rate is decayed to 10- for all the curves.
Connection between SGD covariance and Hessian. Much of the literature related to implicit bias relies on the assumption that the covariance of the noise of SGD is strictly related the hessian the loss function as discussed in Sec 2. Denoting the Hessian H(w) == V L( w) we can write it as the so-called Gauss-Newton decomposition (Sagun al., 2017: Papyan, 2018) H(w) G(w) + E(w).
To measure the cosine similarity (CS) between H(w) and the covariance E, compute
CS = EE ccos (H( ww, E,v)
where is sampled from the Gaussian distribution in R. and cos(u, v) (u.=)/||=] |l|- The results are reported in Fig. 13.
Cosine similarity
0.86
0.021
0.0244
Cosine similarity between hessian and Noise covariance: we compute the cosine similarity between the hessian and the covariance of the SGD noise for a scale-invariant ResNet after one epoch with large Ir 1 0.005. The results show how the two matrices are correlated and in particular how the SGD noise covariance is highly correlated with G(w).
Figure 13:
18

---

C.4 Trace of Hessian and Jacobian Norm
As reported in Table previous works related with label noise, theoretically derived the connection between the trajectory of SGD and trajectory which regularizes either the trace of the Hessian or the Jacobian norm. The Hessian of a loss function L can be decomposed as:
N
V2W(W)
Vh(zn w) (Mll
+
Vnl(h(Za Vllth(::
G.(w)
F0g (w)
Many works demonstrated empirically that the G2 is the dominant part of the Hessian decomposition and ((W( ~ E, G2 (Sagun et al., 2017). The Jacobian (J) norm instead is defined as:
[[(W)]
N Σ Tr (Vh
in the case of square loss, VRl = I where is the identity matrix. Hence, Tr (V2L(W)) ~ 1 The similarity is an exact equality at an interpolating solution since Val(h(a;;; w)) 0 and therefore not much ambiguity is left regarding which quantity to study. However, in the case of classification, this fact does not hold. In partieular, the trace of V I can significantly deviate from the identity matrix and varies depending on the value of the training loss. Consequently, although the two quantities seem closely related even when using the CE, we opt for analyzing ||||| Fr This choice is motivated by its lack of explicit dependence on the training loss, enabling straightforwan comparisons between different solutions. the following, we report an empirical comparison of the two quantities along the training trajectory. In Fig. (14) we report a comparison between the trace of the Hessian and the Jacobian norm for both the fine-tuned iterates and the EMA. We can observe that for the fine- tuned iterates, both quantities display decreasing trend; nevertheless, the ranking appears to be inverted to what expected from previous theoretical works, i.e. larger LRs should lead to a stronger regularization effect. If we compare the EMA instead, we can see that both quantities are still decreasing along the iterations but for the Hessian we don't observe any meaningful ranking for the final solutions whereas the Jacobian norm is lower for higher LRs.
X107
20.09
Neat 5314
36.6
1M
Epocha
Epochs
Foocks
(b)
(c)
(d)
Figure 14: Trace of Hessian and Jacobian Norm We train standard Resnet18 on CIFAR-10 for 100 epachs fixing x = 0.0125 25 and varying the learning rate. Wc report the EMA and we finetune for 100 epochs every 3 with y 10
C.5 Experiments with scale-invariant Resnet the sphere
In order to isolate the implicit regularization mechanism from the large initial drop and even small fluctuations in the dynamics of the &g norm, we consider a simplified setting. We train scale-invariant networks (Li & Arora, 2019; Li et al., 2020) with projected SGD on the unitary sphere S(p-1), We project the SGD iterates on the unitary sphere g(p-1) within the context of scale-invariant ResNet architectures (Li & Arora, 2019; Li et al., 2020). This setup is helpful for two reasons: (a) it simplifies the selection the LR and hence tremendously reduces the experimental overhead (b) scale-invariant networks have been extensively studied in previous works Li & Arora (2019); Li et al. (2020); Kodryan et al. (2022). Moreover, in the context of scale invariance, the optimization on the sphere is the natural object to study the evolution of the direction is the only quantity that matters. The
19

---

projected SGD update writes as
W++1 Nggo 1) (W: - VV (Yir
(8)
where
ligre
w + w/ w//lw//
(9)
The training framework still consists of two phases separated by LR decay. The primary insight from our experiments on the sphere is depicted in Fig. 15: the test performance achieved in the fine-tuning phase depends on the LR used in the large-LR phase and. moreover, there is an optimal value. The work of Kodryan et al. (2022) reports similar observation, where the best test loss achieved at a LR where the loss neither converges too fast nor diverges but doesn't 't provide any explanation. Once again, our investigation reveals that the key to understand this behaviour and the dependence on the LR lies in the noisy dynamics in the large LR phase which closely tracks a regularized process. To summarize this idea postulate conjecture similar to the reported in Section 2.2.
Conjecture 3. Consider the algorithm Eq. with wo initialized from a distribution ps (S(P-1)) For any input x, let Wt, hlw,, x) be the random variables that denote the iterate at time t and its functional value. The stochastic process (h( Me (x) with mean Ba (x) for which w is a first-order stationary point of the following regularized loss:
will converge to a stationary distribution
((w) L(w) + ]())l
(10)
20.0%
10.0%
18.0%
Testt Error fine turiin
Error tuning
15.0%
40%
14.0%
24.0%
13.0%
20%
12.0%
11.0%
10.0%
40
Large-LR phase
Epochs
Epochs
(b)
(c)
x10?
X30
Train Cross Entropy
2.5
20
2.0
[|/[W||F
0.5
Epochs
Epachs
Figure 15: Training scale-invariant ResNets on the sphere. We train on CIFAR-10 with different large LR for the first 100 epochs and decay it to 10-4 afterwards. Fig. (15a) reports the test error with respect to different LRs in the first phase showing the existence of an optimal value. Fig. (15b) reports the test SGD iterates (-) and the EMA (--). Figure (15c) reports the decreasing trend the test error after fine-tunin fine-tuning for 100 epochs with = 10 4 every 2 epochs. Finally, Fig. (15e) reports the norm of the Jacobian for the EMA and the fine-tuned iterates and Figure (15f) reports comparison with the trace of the Hessian for the fine tuning iterates.
for
20

---

C.6 Empirical validation of Covariance approximation
We empirically verify the validity of the "decoupling approximation" introduced in Sec. 2.2 we use Stochastic Lanczos Quadrature (Yao et al., 2020) to estimate the empirical spectral density of the SGD covariance with and without the decoupling approximation, during the large-LR phase. The experiments are performed for ResNet- 88 trained on the cifar 10 dataset for different combinations of learning rate and weight decay which are used in the manuscript. The results Figures 16, 17, and 18 illustrate a substantial overlap in the two spectra which serves as a validation of the reliability of our approximation.
23n0
-0.75
-0.0
0.7
1.00
-02
92
(a) 1r=0000 wd=0.0005
(b) 1r=00001 wd=0.001
(c) 1r=00001 wd=0.0025
wde6.005
001
10.00
055
32
00
04
06
0.
10
24
03
07
PA
(d) 1r=0.001 wd=0.005
(e) 1r=0,001 wd=0.0075
(1) 1r=00001 wd=0.01
Spectrum of S&D couadinne r-0.035
Szertmun SCD rovariance 5.55055 wd-33.03
Steemim 5G0 cuadance
2047
-28
13
25
-352 -22>
110
321
3.7
1.00
-2.33
-0.3
5.0
025
5.0
0.0
140
(g) |r=0.005 wd=0.0005
(h) 1==0005 wd=0.001
1r=00005 wd=0.0025
Sacutrum of 5GD conatianus 1e0.002 0000
Soctirum of SGD conoliance awd nowd 2015
Soctirum of 5GD constlance 1w2 ndwo
23
12
30
-5.50
-925
633
8.2
3.50
0.7
1.00
-0.5
-259
-2.25
023
3.5
8.50
3.5
2.50
-575
-058
-2.25
5.6
0.50
653
035
1.0
(i) 1r=0.005 wd=0.005
(k) 1r=00005 wd=0.0075
(1) 1==00005 wd=0.01
Figure 16: We use Stochastic Lanczos Quadrature to estimate the empirical spectral density of the SGD covariance, with and without the decoupling approximation, during the large-LR phase. Experiments with ResNet-18 on the CIFAR-10 10 dataset, varying learning rate and weight decay, show substantial overlap in the two spectra, validating our approximation.
21

---

Spectrum cocar anca 1--0.22 300000
Spectrum im-
covatfrrce m-0
540
562
30
2/s
Egerisia
(a) 1r=0.01 wd=0.0005
(b) Ir=0.01 wd=0.001
(c) 1r=0.01 wd=0.0025
Spectium o1502 covorterce Iwo
Scectrum SGD covarincee 1-2.02 wowo 0075
Scectrum SGD ance 2.02 wowo
520
1344
Speriebs
Openmin
(d) lr=0.01 wd=0.005
(c) l==000 wd=0.0075
(1) 1r=0.01 wd=0.01
Spuctum conatiance 0000
Scecbrum 3GD k-2.025 No
Seccirum conalanee Kdwo
10mm
0.5
(g) 1r=0.025 wd=0.0005
(h) 1r=0.025 wd=0.001
(i) 1r=0.025 wd=0.0025
31-
car corappiic
33
10th
12
11
0.35
456
2.25
335
6.55
066
0.5
(j) 1r=0.025 wd=0.005
(k) 1r=0.025 wd=0.0075
(1) 1r=0.025 wd=0.01
We Stochastic Lanczos Quadrature to estimate the empirical spectral density of the SGD covariance, with and without the decoupling approximation, during the large-LR phase. Experiments with ResNet-18 on the CIFAR-10 dataset, varying learning rate and weight decay, show a substantial overlap in two spectra, validating our approximation.
Figure 17:
22

---

Spectrum cocar anza 1--0.25 wdr2 3.000
Spectrum
wdrd.001
covatfrrce
540
62
Sperisia
(a) Ir=0.05 wd=0.0005
(b) Ir=0.05 wd=0.001
(c) Ir=0.05 wd=0.0025
Spectium o15 covur
Scectrum 2GD covarincee 1w2.02 wowd 0075
Scectrum SGD ance 1-2002 wowd
281
11-4
11
10
Speriebs
Openmin
(d) 1r=0.05 wd=0.005
(c) 1r=0.05 wd=0.0075
(1) 1r=0.05 wd=0.01
Specirum 5G2 Lovur ence 0003
Spodirum covar unoo 1-o knocl-w0
Specirum nOwO
10-1
apanc
72-
|r=0.1 wd=0.0005
(h) 1r=0.1 wd=0.001
Ir=0.1 wd=0.0025
33
375 332 4.25
6.35
4350
125
(j) 1r=0.1 wd=0.005
(k) 1r=0.1 wd=0.0075
(I) 1r=0.1 wd=0.01
We Stochastic Lanczos Quadrature to estimate the empirical spectral density of the SGD covariance, with and without the decoupling approximation, during the large-LR phase. Experiments with ResNet-18 on the CIFAR-10 dataset, varying learning rate and weight decay, show a substantial overlap in two spectra, validating our approximation.
Figure 18:
23

---

C.7 Empirical verification of the conjecture through snapshot ensembles
To further validate our conjecture, we performed additional experiments. Specifically, within the same ResNet18 on CIFAR-10 10 setting in our main experiments, we created snapshot ensembles (Huang et al., 2017) by averaging in function space along the SGD trajectory every 10 epochs for the combinations of learning rate (LR) and weight decay (WD) considered in the paper.
To assess whether the mean of the stationary distribution in function space aligns closely with the EMA, where the Jacobian norm is regularized, we compared the performance of snapshot ensembles with that of the EMA. Additionally, we computed the Total Variation Distance Drv between the softmax outputs of the ensemble and the EMA on the Test set
C
DTV
ΣΣ 2N
1
(1) Pensemble.
(R) PEMAJ
The results in Table 2 show strong alignment in test accuracies, while those in Table 3 indicate low Total Variation across all combinations. Together, these findings offer further validation for our conjecture.
Table Test Error Snapshot Ensemble and EMA different values learning rate (LR) weight decay (WD).
WD
LR=0.001
LR=0.005
LR=0.01
LR=0.025
LR=0.05
LR=0.1
LR=0.15
EMA
ENS
EMA
ENS
EMA
ENS
EMA
ENS
EMA
EMA
ENS
EMA
0.0000
0.32
0.33
0.27
0.26
0.24
0.17
0.17
0.17
0.17
0.13
0.13
0.13
0.13
0.0005
0.32
0.32
0.29
0.29
0.24
0.24
0.18
0.17
0.13
0.16
0.13
0.13
0.11
0.13
0.0010
0.32
0.33
0.25
0.27
0.21
0.21
0.13
0.19
0.10
0.13
0.10
0.11
0.11
0.11
0.0015
0.32
0.34
0.23
0.22
0.22
0.25
0.11
0.14
0.10
0.10
0.10
0.09
0.09
0.0025
0.30
0.30
0.22
0.22
0.19
0.20
0.09
0.10
0.10
0.11
0.10
0.10
0.10
0.11
0.0075
0.35
0.37
0.15
0.16
0.10
0.11
0.11
0.11
0.09
0.08
0.10
0.09
0.12
0.10
0.0100
0.34
0.13
0.15
0.11
0.11
0.10
0.10
0.111
0.10
0.11
0.09
0.13
0.13
Table 3: Total Variation Distance between softmax output of Ensemble and EMA.
WD
LR=0.001
LR=0.005
LR=0.01
LR=0.025
LR=0.05
LR=0.1
LR=0.15
0,0000
0.03
0,02
0.01
0.01
0.01
0.01
0.01
0,0005
0.04
0.02
0.02
0.03
0.09
0.08
0.07
0.0010
0.04
0.05
0.04
0.10
0.09
0.07
0.07
0.0015
0.04
0.07
0.07
0.08
0.08
0.07
0.07
0.0025
0.04
0.11
0.10
0.06
0.08
0.08
0.09
0.0050
0.06
0.15
0.11
0.09
0.08
0.10
0.11
0.0075
0.12
0.10
0.10
0.10
0.11
0.12
0,0100
0.10
0.09
0.10
0.10
0.11
0.12
0.13
24

---

D Weight decay for large language models: additional figures and details
We present the following additional figures related to the LLM experiments. We show that the validation loss of a GPT-2-124M model is determined by the training loss and not influenced by in Fig. 19. We also show that the generalization gap stays close zero throughout training for different À for both 124M and 774M parameter models. We show the results for models trained weight decay on LayerNorm weights in Fig. 20. We see that penalizing all parameters in weight decay (i.e., including the LayerNorm parameters) leads to the same effect for smaller À (like 0.1) but underperforms on larger (like 0.3). Note that when WD is applied on all weights, this changes the optimal value of the objective. In Fig. 21, we train models with &g regularization instead of decoupled weight decay as in AdamW (Loshchilov & Hutter, 2019). We observe that & regularization instead of weight decay leads to the same effect as decoupled weight decay (Loshchilov & Hutter, 2019). We train models using SGD with momentum and show the results in Fig. 22. We see that weight decay leads to similar improvement in training loss for SGD with momentum well. We show multiple metrics in Fig. 26 for the models shown in Fig. 6: gradient variance, gradient norm, and weight norm plots that complement Fig. 7 in the main part. In Fig. 23, we show results of weight averaging that suggests the suboptimality gap between runs with different À is much smaller than what the loss at wy suggests. However, weight averaging is still less effective than fine-tuning with a tiny LR in Fig. 6. Finally, in Fig. 25, we show results of models trained context length 1024. We see that the training loss over iterations for models trained with range of LR and WD (all are bfloat16). All runs with LR smaller than 0.001 successfully converge but the final training loss is higher than for LR 0.001. In addition, we observe that lower learning rates prevent the weights from growing too much.
Awe
param, Avo
Ase
774M param, Aap 000
3.7-
Awp 0.3
Generalization gap
1 24M param, Ava 0.0
0.06
14M 0.3
Validation loss
124M param, Aso
0.04
3.5
0.00
3.3
-0.02
3.5
3.7
10000
2,0000
30000
40000
50000
Training loss
Iteration
Figure 19: Left: The validation loss of a GPT-2- model is determined by the training loss and not influenced by A. Right: The generalization gap stays close to throughout training for different X for both 124M and 774M parameter models.
3.7
Auo 0.0
Avo 0.03
App
Avo 0,15
Training loss
Avo 0.3
3.5
3.3
15000 20000 25000 30000 35000 40000 45000 50000
Iteration
Figure 20: GPT-2-124M on Text with weight decay on LayerNorm weights. Penalizing all
parameters in weight decay (i.c., including the LayerNorm parameters) leads to the same effect for smaller (like 0.1) but underperforms on larger X (like 0.3). Note that when WD is applied on all weights, this changes the optimal value of the objective.
25

---

3.7
Avo 0.0
Awe 11-10
3-10
Training loss
3.4
10000 15000 20000
25000 30000 35000 40000 45000 50000
Iteration
Figure 21: GPT-2-124M WpenWe with & regularization. We observe that &2 regularization instead of weight decay leads the effect decoupled weight decay (Loshchilov & Hutter, 2019).
Ano 0.0
11.10
Awo 3-10-
Training loss
3.7
3.6
3.5
15000
35000 45000 50000
lteration
Figure 22: GPT-2-124M Open Web trained with SGD with momentum. Weight decay leads to similar improvement training loss SGD with momentum as (all other experiments are done with AdamW).
AdamW, decay
decay
dol -il
Jva 2.0
=1
Training loss,
E.D. adr
12"
Training
3.a
Training
1-10- wn
Auy 011
100K 1530 2000
2000 2005 2583 1000 1200 5000
20033
1300 2508 2000 2536 2008 2500 4389 4080 5000 Iteration
Iteration
Itaration
Figure 23: Weight averaging for GPT-2-124N TetnWeb Weight averaging (" shows that the suboptimality gap between with different is much smaller than what the loss We suggests. However, weight averaging is still less effective than fine-tuning with a tiny as Fig.
26

---

340
bleadte, 0.0003, Ase
bloatte.
Asc
Training loss
3.5
325
320
335
332
13000
20002
30000
50000
Iteration
Figure 24: GPT-2-124M on OpenWebTee with context length 1024. The model trained with moderate
LR 0.001 diverges for bfloat16 but not for float32; lowering the LR prevents the divergence but leads worse loss.
Ut 0.000000...4 Aaz 0.0
UK 0-00003..Au Aaz
0R-000003AAu
Are
Training loss
Weight norm
100
LR 0.000066 Axo -00
LR 00.00006, Aue m.
LR 0.0000 Aag
LR 0.3
UR 0.001, Awo 0.0
UIl 0.001,
0.1
60
-800.003. Anc
3.35
3.10
3.05
3x
3.00
10000
20000
30000
40000 50000
20000
200000
30000
40000
50000
Iteration
Iteration
Figure 25: GPT-2-124M on OpenWebTeew with context length 1024. (Left) The training loss over iterations for models trained with range of LR and WD (all are bfloat16). All runs with LR smaller 0.001 successfully converge but the final training loss is higher than for LR 0.001. (Right) Weight norms for LR in 0.0003, 0.0006, 0.001 for À = 0.1 which does not diverge. Lower learning rates prevent the weights from growing much.
101
ÁwD 0.0
Awd 0.0
Gradient norm
Awp = 0.1
Awo 0.1
Awp 0.3
Weight nom
Awo 0.3
10a
4x
3x
10000
20000
30000
40000
S0000
Iteration
Iteration
Figure 26: GPT-2-124M on OpenWeb The gradient norm and weight norm plots for the models reported
Fig. 7.
27

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims in the abstract and introduction accurately reflect all our contributions presented in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not perceived well by the reviewers.
The claims made should match thcoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We mention the limitations of our work at the end of Section 4.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g., independence assumptions, noiseless settings, model sell-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
28

---

Justification: We do not have new standalone theoretical results like theorems or lemma, but we provide short derivations presented in Section 2.2. Section and Appendix C.2 that support our experimental results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material. Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide detailed settings of our experiments in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided not.
If the contribution a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c)
If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, may be that access to the model is limited some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
29

---

The authors should provide scripts to reproduce all experimental results for the new reproducible, they Answer: [Yes]
Justification:
The code is uploaded in supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/ for more details.
criides/CodeSubmisstt
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips.. wwpplic/guudes/////epppoppaaeaep //ode/Sumissiioool///
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
proposed method and baselines. If only subset of experiments
should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized ver- sions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We describe them in the corresponding sections (Section 2 and Section 3) and provide further details in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide error bars in figures wherever it feasible (i.e., most experiments in Section 2), but we omit them in Section 3 since LLM experiments are much more computationally demanding.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, statistical significance tests, least for the experiments that support the main
claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar the standard deviation or the standard error of the mean.
30

---

It is OK to report 1- I-sigma enror bars, but one should state it. The authors should preferably report 2-sigma enror bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars reported in tables or plots. The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide these details in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http///eerrpppeeeeee cc/ppblicc/hiccetuueeeeeeeee wtw.hhiiciieeeeeeeeeee.
Answer: [Yes]
Justification:
We confirm that our submissions complies with the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics. If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consideration due laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: our opinion, this work does not directly lead to socictal impacts since it focuses general improvement of our understanding of deep learning optimization methods.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, legitimate to point out that an improvement in the quality of generative models could be used to generate
31

---

deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mecha- nisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release data models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions access the model implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the code, data, and models we use our work, most of them directly when we discuss the corresponding experimental results (Section 2 and Section 3) and some of them in the appendix (Appendix B).
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g., website), the copyright and terms of service
of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi th choode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are pa-ckagedd, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
32

---

Answer:
[NA]
Justification:
We do not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their sub-
missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose assel is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We do not conduct studies with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then much detail possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our work does not require IRB approval.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "cNiujUN6KI", "reviewer": "Reviewer_YyLF", "summary": "The paper explores the role of weight decay (WD) in training modern deep neural networks, focusing on its impact on optimization dynamics in two distinct training regimes: over-training (multiple passes through the dataset) and under-training (limited passes due to large dataset sizes and computational constraints). \n\nIn the over-training regime, the authors demonstrate that WD does not act as a capacity constraint but influences optimization dynamics by enhancing implicit regularization of SGD, as well as combines with large learning rates to maintain non-vanishing SGD noise, improving generalization by controlling the norm of the Jacobian. \n\nFor large language models (LLMs) trained with limited data passes, the authors discuss the relationship between WD and the effective learning rate, and how WD can prevent divergence in training with reduced precision, such as bfloat16.", "questions": "Please refer to weaknesses.", "limitations": "N/A", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This paper studies the role of weight decay in deep learning, moving beyond traditional views of regularization to consider its impact on optimization dynamics. By introducing the new perspectives of weight decay, the paper aids practitioners in better configuring their models(like the combination of learning rate and weight decay rate), potentially improving both training stability and model performance.\n\n2. Extensive experiments on ResNet architectures and large language models on both multi-pass training and one-pass training validate the theoretical claims, demonstrating WD's practical benefits in both over-training and under-training contexts. The experimental results well validate the hypotheses introduced by the authors.", "weaknesses": "1. While the practical implications are well-explored, the theoretical explanations could be helpful. Some claims about implicit regularization and optimization dynamics are not fully explained.\n\n2. The study focuses on specific architectures like ResNet-18. It would be beneficial to see similar analyses on a broader range of models to strengthen the generality of the conclusions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's suggestion to enhance the theoretical explanations in our paper. We understand the importance of providing thorough theoretical insights.\n\n- **Conjecture:** We acknowledge that our results regarding implicit regularization are presented as conjectures. This decision was made consciously, as we believe it is challenging to rigorously prove these results within the scope of our current work. However, we included this conjecture to summarize our intuition and facilitate the understanding of our main message providing a conceptual framework that could guide future research.\n\n\n\n- **Optimization dynamics:** We have largely built upon the existing literature on label noise, as detailed in Section 2.2 of our paper. Previous studies have observed implicit regularization of the Jacobian in the presence of label noise. Building on this, we hypothesize that weight decay with large learning rates induces a similar implicit regularization effect, without any additional label noise but through loss stabilization mechanisms. We believe this perspective offers a novel viewpoint on how weight decay influences the optimization dynamics and enhances SGD implicit regularization.\n\n\n\n- **Additional architectures:** Regarding a broader range of models, in the appendix, we present similar experiments showing loss stabilization and better generalization due to weight decay for VGG and ResNet-34 architectures, along with fully scale-invariant ResNet architectures. For language models, the transformer architecture completely dominates the research field, therefore we considered it to be sufficient.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HLs8QfRq0c", "reviewer": "Reviewer_6pLz", "summary": "This paper studies the effect of weight decay (WD) in the over-training and under-training regimes. It hypothesizes and empirically verifies different roles that WD plays in these scenarios. Specifically, in the over-training regime, WD primarily induces the implicit regularization of SGD via controlling the noise scale $\\sigma_{\\eta,\\gamma}$. In the under-training regime, WD serves to stabilize and improve the optimization of the training loss as well as enabling better stability in low-precision training settings.", "questions": "1. The statement that EMA or fine-tuning “allows for the effective exploitation of the accumulated hidden regularization” is not entirely clear to me. Does generalization require a smaller scale of the stochastic noise and why? I feel this point is not adequately addressed or relevant prior work is not cited. In addition, if this is true, the current message seems to suggest that a large LR, combined with WD, is required to induce the stochastic noise but its scale has to be tuned down for better generalization. Is it not possible to induce the noise of appropriate scale at the very beginning with a careful selection of LR and WD? \n2. Related to the previous question, some works actually use a warm-up schedule for the WD, together with LR decay(e.g. DINO: https://arxiv.org/abs/2104.14294). This would lead to a constant noise scale $\\sigma_{\\eta,\\gamma}$. This seems to contradict with the previous statement about the necessity of a smaller noise scale. I am interested in learning about the authors’ thoughts about this. \n3. The current work studies the training of networks from scratch. Would these conclusions/observations carry over to the fine-tuning of networks (e.g. LoRA)?", "limitations": "The authors note the limitations of the current work as the experiments reported are of smaller scale due to computational constraints.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper studies WD, an important hyperparameter in the modern machine learning practice. Although widely used, understanding of its effectiveness is limited. This paper provides a rather comprehensive study on this problem. \n\n2. By treating the over-training and under-training regimes separately, the paper uncovers the different roles of WD in model training. The results are insightful and could lead to future work. \n\n3. The presentation is structured and mostly clear.", "weaknesses": "1. Some figures could be improved and(or) better explained. For example, Figure 2c presents results on the evolution of norm of weights and the correlation between the norm and the test error. It is quite difficult to cross reference which lines correspond to which test errors. It would be better if the authors could pinpoint specific lines in the figure or in the caption which lines one should look at. This would improve readability of the experimental results. Similarly for figures such as Figure 4.\n2. EMA and finetuning would introduce additional hyperparameters (i.e. the momentum parameter and the smaller lr respectively). In the paper and appendix, I did not find their specific settings. The effect of these hyperparameters on the optimization is not, but should be, examined in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their helpful comments, which will allow us to improve the manuscript. \n\n**Improved figures:** We thank the reviewer for suggesting improvements to our figures. We will make sure to reference specific lines in the text to improve the readability of the results. \n\n**Additional hyperparameters:** Some of the secondary hyperparameters are fixed because they do not alter the conclusions of our experiments. We fixed the cofficient $\\beta$ of the exponential moving average $EMA_t = \\beta \\cdot EMA_{t-1} + (1-\\beta) \\cdot w_t$ to be $\\beta = 0.999$. For $\\beta$ closer to 1, the effect is to make the EMA change slowly and smoother. This means that the EMA will incorporate a longer history of the parameter values and will be less influenced by short-term fluctuations therefore better approximating the mean of the stationary distribution. In our experiments we do not see any significant difference for $\\beta \\in [0.9,1)$ \n\nSimilarly for the fine-tuning, in our experiments, we observed that as long as the learning rate used in the fine-tuning is smaller than $10^{-3}$, we converge to solutions which have the same test error but for smaller learning rates longer time is needed. We will better discuss these observations and provide further details regarding the choice of these hyperparameters in the manuscript. \n\n\n\n**Questions**: \n\n1. The meaning behind the statement \"EMA or fine-tuning allows for the effective exploitation of the accumulated hidden regularization\" is discussed in more detail in Section 2.3. Specifically, we do not imply that generalization requires a smaller scale of noise throughout the entire training process. As shown in Figure 4, there is an optimal value for the scale of the noise. Our point is that to fully leverage the generalization benefits accrued from the noisy process, it is necessary to subsequently reduce the variance. To further clarify this point we can consider the red and green lines in Figure 2a. Without reducing the noise after a certain number of iterations, the test error remains high, around 35%. However, when using EMA or fine-tuning, we can observe the benefits of having a large noise level during the initial phase. This is because the implicit regularization of the Jacobian regulated by the noise scale affects the mean of the distribution and not the current interates. To realize these benefits, the noise level needs to be reduced to converge to the mean after this initial transient phase. Therefore, it is not possible to have an appropriate noise scale from the beginning that remains constant. The generalization benefits arise from starting with a large noise scale, which then needs to be appropriately reduced to converge to the mean.\n\n\n1. After reviewing the implementation details of the referenced work, it appears that they indeed combine LR warm-up and decay with a cosine schedule but also apply decay to weight decay using a cosine schedule. Given that both learning rate (LR) and weight decay are decreasing along the training trajectory, this implies, according to our findings, that the scale of the noise should also decrease over time. Moreover, this aligns with our intuition that after the initial phase with a large noise scale, one should apply some variance reduction to fully exploit the generalization benefits. \n\n\n1. Fine-tuning typically uses only a few epochs and LoRA restricts the model's capacity, so we believe its training regime is closer to undertraining. Thus, we believe that our main findings should transfer to the fine-tuning regime. However, we are not aware of any work that would notice a crucial effect of weight decay specifically for fine-tuning (assuming that the learning rate is carefully tuned), which is why we have not covered this topic.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Vn6N70kjSs", "reviewer": "Reviewer_S7kS", "summary": "In the paper the authors investigate WD in current deep learning practices. The authors argue that for deep networks used in vision tasks, WD enhances implicit regularisation by modifying the optimisation through loss stabilisation. They also argue that for LLMs, WD helps balance the bias-variance trade-off, leading to improved training stability. The study offers a unified perspective on the role of WD across different training regimes, suggesting that WD’s benefits are more related to changes in training dynamics than to explicit regularization effects.", "questions": "- There is an assumption that the first derivative is approximately constant across all datapoints (eq 3), what is this assumption based on? Will this hold across all types of data and data complexity, including multimodal data?\n- Related to the previous question, does the data have to be homogenous? \n- In the sections discussing the benefits of weight decay with bfloat16 precision, could the authors elaborate on the mechanisms through which weight decay interacts with lower precision calculations to enhance stability?", "limitations": "The discussion on limitations is severely lacking.\n\nThere is one mention of a limitation of computational resources which hinder large scale experiments but I cannot find any other mention of limitations. This is not sufficient. The authors make a number of assumptions in their experimental setups and in the interpretation of the findings (I mention only a few in my weakness list). The authors must provide a more thorough discussion of the limitations of their work.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper has a number of strengths:\n- It presents new insights into how it impacts optimisation dynamics rather than merely acting as a regulariser\n- It offers practical recommendations for tuning WD and learning rates, providing actionable insights for practitioners\n- The authors set out a theoretical basis for the empirical observations, linking them to the underlying optimization principles", "weaknesses": "The weaknesses are as follows:\n- The study primarily focuses on ResNets and GPT-2. It would have been interesting to see how the results scale to larger models, such as the open source Llama-3. This fact makes it harder to generalise the results.\n- No comparison with other regularisation methods, like batch norm or dropout. I'm not saying this is necessary for the paper to be accepted but it would be a considerable added strength to have at least some of the experiments on other regularising techniques for comparison\n- The study assumes fixed hyperparameters in many experiments, which may not reflect the dynamic tuning often required in real-world training\n- Lack of justification and explanation around assumptions (see questions)\n\nMinor comments:\n- Font in figures is far too small and thin, making it very hard to read\n- The use of \"weight decay\" vs \"WD\" in the main text is mixed and inconsistent", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments. We acknowledge the lack of a detailed discussion of the limitations of our work. We did mention the main limitations in the conclusions (no truly large-scale experiments, no proofs of new theoretical results), we will improve upon this by adding a paragraph discussing the limitations due to our approximations. We also elaborate on them below.\n\n**Assumptions and approximations in SGD:** We would like to clarify that our series of approximations does not need to hold along the entire trajectory of SGD but only in the stabilization phase. This is an important detail that we will clarify in the manuscript. Moreover, we acknowledge that we do not correctly cite the related works to justify our approximations, we will add such justifications in the manuscript.\n\n - **Gaussian approximation:** A substantial body of research has built upon this approximation and verified its validity ([1], [5], [10], [11], [12]). In particular [5] demonstrated how modelling the SGD noise by a Gaussian *is sufficient* to understand its generalization properties which is the main interest of our work. \n - **First derivative constant:** This approximation allows us to disentangle the first derivative of the loss and the product of the gradients isolating the scale of the noise, this corresponds to the decoupling approximation introduced and empirically verified in [2] (Figure 1, page 10). To further verify this assumption in our setting, we perform additional experiments comparing the spectrum of the SGD covariance with our approximations at stabilization. The results illustrate a substantial overlap in the spectrum (see attached pdf and point (3) in the main response) confirming the validity of our approximation.\n\n**Training Llama-3:** We agree that our study primarily focuses on small models like ResNets and GPT-2. Nevertheless, the primary aim of our work is to investigate the impact of weight decay when training models *from scratch*, for this reason, and given our computational resources, it would be infeasible to extend our experiments to significantly larger models such as Llama-3.\n\n**Other regularization methods:** Despite comparing with other regularization methods could give better guidelines for practitioners, the focus of this manuscript is not on understanding which regularization method works better but rather on providing a specific understanding of the mechanisms through which weight decay improves the performance of models trained in two main settings. Therefore, we did not compare with other regularization methods such as dropout because it would not provide further insights regarding the functioning of weight decay. Regarding batch normalization, we would like to point out that most of the experiments are conducted with architectures that have a normalization layer (BatchNorm for ResNets and LayerNorm for GPT2). Nevertheless, removing the normalization layers would be impractical, leading to exponential growth of the output variance (see [14]), therefore requiring additional modifications which would alter the consistency of the comparison. Moreover, in Appendix C.4, we report some experiments with *smaller* VGG networks without batch normalization and observe a similar dynamics when using weight decay. \n\n**Fixed hyperparameters:** For all main experiments, we have performed a thorough grid search over the learning rate and weight decay parameters (which are the main hyperparameters tuned in practice). E.g., see Figure 3 for the over-training regime and Figure 22 for the under-training regime. However, some other secondary hyperparameters are fixed because they do not alter the conclusions.\nFor example, in the over-training setting, we fixed the coefficient $\\beta$ of the EMA $EMA_t = \\beta \\cdot EMA_{t-1} + (1-\\beta) \\cdot w_t$ to be $\\beta = 0.999$. For $\\beta$ closer to 1, the effect is to make the EMA change slowly and smoother. This means that the EMA will incorporate a longer history of the parameter values and will be less influenced by short-term fluctuations therefore better approximating the mean of the stationary distribution. In our experiments we didn't observe any significant difference for $\\beta \\in [0.9,1)$ \n\n\nSimilarly for the fine-tuning, we observed that as long as the learning rate used in this phase is smaller than $10^{-3}$, we converge to solutions which have the same test error but for smaller learning rates longer time is needed. We are happy to provide further details about the effect of other hyperparameters. \n\nWe thank the reviewer for pointing out the font size issue in the figures and the inconsistency in using \"weight decay\" versus \"WD\" in the main text. We will address these issues and make the necessary corrections in the manuscript.\n\n# Questions: \n1. We answered the first part of the question above in the paragraph \"**First derivative constant**\". Given that we focus on standard classification tasks, could the reviewer elaborate on the meaning of multimodal data in this context? Generally, the i.i.d. assumption for the dataset is necessary for our approximations to hold. \n\n2. In our work, we are considering multi-class classification over image datasets, such as CIFAR-10 and Tiny ImageNet. These datasets are inherently heterogeneous, comprising multiple classes. We hope this addresses your concern otherwise we would be happy to provide further details.\n\n3. We suspect that the numerical instability comes from `bfloat16`'s limited precision, which can occur when different components in the network with varying scales are added together (e.g., values in the residual stream). Since the `bfloat16` data type has the same exponent size as `float32`, high-weight norm alone should not pose a problem. However, the precision of `bfloat16` is very limited, e.g., `bfloat16(256.0) + bfloat16(1.0)` is equal to `256.0` instead of `257.0`. We suspect this precision limitation is the primary problem in `bfloat16` runs without weight decay.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iFIqhJFoQk", "reviewer": "Reviewer_g6ok", "summary": "This paper investigates the role of weight decay (WD) in modern deep learning, differentiating its effects in over-training and under-training regimes. Through experiments with ResNets on vision tasks (over-training) and Transformers on text data (under-training), the authors demonstrate that WD’s primary benefit is not explicit regularization, but rather its influence on optimization dynamics.  In the over-training regime, WD with large learning rates leads to non-vanishing SGD noise, controlling the Jacobian norm and improving generalization. In the under-training regime, WD functions as a modified effective learning rate, improving training loss and stability, particularly with bfloat16 precision.", "questions": "From Figures 3 and 10, but also my own experience: A meaningful improvement in test-performance often requires carefully tuned WD, within at most a factor of 2. Would you agree with that observation? Many of the experiments show ablations with significantly larger differences in WD strength. Is that an issue for any of the conclusions described in the paper?", "limitations": "See above.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The paper provides a valuable review of prior work investigating the role of weight decay in modern neural networks, highlighting the need for a deeper understanding of its impact on optimization dynamics.\n* A key strength is the clear distinction made between over-training and under-training regimes,  demonstrating that WD's function can vary drastically depending on the training context.\n* The paper effectively alternates between theoretical analysis (despite its highly approximate nature) and empirical experiments. Notably, the presented experiments are well-designed to confirm predictions derived from the paper's conjectures and hypotheses, demonstrating a strong consistency between predictions and observations.", "weaknesses": "* Although the paper provides a compelling narrative connecting WD's influence to optimization dynamics, it doesn't offer concrete, actionable guidelines for practitioners on how to tune WD for different model architectures or datasets. Insights partially boil down to \"instead of weight decay, other regularizer that avoid loss-collapse would have a similar effect\".  \n* The theoretical analysis, while providing some interesting insights, rely heavily on approximations and lacks formal mathematical proofs for the proposed conjectures. This leaves room for inaccuracies not covered by the simplified theoretical framework.\n* Especially the language experiments in the under-training regime use tiny models by today's standard. One can indeed hope though, that the results generalize to bigger architectures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their helpful comments. Below, we address the concerns raised by the reviewer in detail\n\n>Although the paper provides a compelling narrative connecting WD's influence to optimization dynamics, it doesn't offer concrete, actionable guidelines for practitioners on how to tune WD for different model architectures or datasets. Insights partially boil down to \"instead of weight decay, another regularizer that avoid loss-collapse would have a similar effect\".\n\n**Practical takeaways:** While our paper primarily focuses on providing a conceptual understanding of how weight decay influences optimization dynamics and improves generalization, we believe our findings also offer practical takeaways. In practice, both weight decay and learning rate are often tuned through extensive grid searches considering combinations of both parameters. However, our results in Figures 3 and 4 demonstrate that the key parameter is actually the product of the two, $\\eta \\times \\lambda$. This insight suggests that a 2D grid search is not necessary. The existence of a connected region of good performance where the product is constant implies that practitioners can simplify their tuning process. I.e., it is not needed to tune both at the same time but it is sufficient to fix $\\eta$ and find the best $\\lambda$ or vice-versa. \n\n\n>The theoretical analysis, while providing some interesting insights, rely heavily on approximations and lacks formal mathematical proofs for the proposed conjectures. This leaves room for inaccuracies not covered by the simplified theoretical framework.\n\n**Conjecture and approximations:** We acknowledge that our results regarding implicit regularization are presented as conjectures. This decision was made consciously, as we believe it is challenging to rigorously prove these results. However, we included these conjectures to summarize our intuition and facilitate the understanding of our main message, providing a conceptual framework that could guide future research. Regarding the approximations in our analysis, we would like to clarify a few points. First, our series of approximations do not need to hold throughout the entire trajectory of SGD, but only during the stabilization phase. This is an important detail that we will clarify in the manuscript. Additionally, we will better justify our approximations in the revised manuscript in the context of the related literature. In particular, the work of [1] demonstrated that modelling the SGD noise by a Gaussian is sufficient to understand its generalization properties. Moreover, assuming the first derivative constant across datapoints was introduced and empirically verified in [2] (Figure 1, page 10). To further verify the latter assumption in our setting, we performed additional experiments comparing the spectrum of the SGD covariance with that of our approximation at stabilization. The results illustrate a substantial overlap in the spectrum, as detailed in the supplementary material (see the attached PDF in the main response).\n\n> Especially the language experiments in the under-training regime use tiny models by today's standard. One can indeed hope though, that the results generalize to bigger architectures.\n\n**Scaling challenges:** We agree that GPT-2-scale models are considered tiny now, and we acknowledge this limitation of our work. It is obviously challenging for a typical research lab to train significantly larger language models, especially when also having to perform grid searches over multiple hyperparameters.\n\n> From Figures 3 and 10, but also my own experience: A meaningful improvement in test performance often requires carefully tuned WD, within at most a factor of 2. Would you agree with that observation? Many of the experiments show ablations with significantly larger differences in WD strength. Is that an issue for any of the conclusions described in the paper?\n\n**Question regarding tuning weight decay:** In our preliminary experiments, we explored an even finer grid for weight decay values than those reported in Figure 3. However, we did not observe any substantial differences in test performance for these finer values, at least in the overtraining regime. The only difference would have been a higher resolution for the heatmap that would lead to the same conclusions. Furthermore, as shown in Figure 3, despite the larger differences in $\\lambda$ as highlighted by the reviewer, we found that similar test performances can be achieved with proper tuning of the LR. This aligns with our main practical takeaway: one can fix one of the parameters (either weight decay or LR) and tune only the other to achieve optimal performances.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3gXfwTKquI", "reviewer": "Reviewer_tBnp", "summary": "In this work, the role of weight decay (WD) in contemporary deep learning is investigated. The authors note that interpreting weight decay as a classical $L_2$ regularizer seems less relevant for modern deep learning. Instead, it is suggested to separately regard WD from positions of over- and under-training. In the over-training regime, typical for common image classification with ResNets on CIFAR-like datasets, weight decay paired with learning rate acts as an implicit regularizer on the Jacobian of the model, controlling the scale of the noise defined by the derivative of the loss. Optimal combination of weight decay and learning rate thus yields the best generalization due to a balanced regularization. In the under-training regime, attributed to a single-pass LLMs training, WD takes a different role controlling the effective learning rate of the training process and saving it from numerical instabilities.", "questions": "1. Figure 2: How come the huge spikes of the dark blue line in the training loss and test error are absolutely not reflected in the weight norm?\n2. Equation 4: How exactly does this formula approximate the SGD update in eq. (2)? Where is the WD factor $(1 - \\eta \\lambda)$ before $w_t$? Why is a single $g_t$ term replaced with a sum of normally distributed values? \n3. According to the formulation of Conjecture 1, it seems more natural to perform MCMC ensembling of the models in the training trajectory to approximate the mean value $\\bar{\\mu}_{\\eta, \\lambda} (x)$ rather than doing EMA or fine-tuning (the authors discuss limitations of these methods in Sec. 2.3). I would suggest to conduct such experiment to more accurately validate the conjecture.  \n4. Does Fig. 4 also report the EMA results? \n5. How does the authors derivation of the effective learning rate for Adam compare with a similar prior result obtained for Adam in the scale-invariant case [1]? I acknowledge that the authors result could be more general, albeit it still requires independence of the gradient sign from the parameters scale akin to the scale-invariant case. \n\n[1] Roburin, Simon, et al. \"Spherical perspective on learning with normalization layers.\" Neurocomputing 487 (2022): 66-74.", "limitations": "The authors partly discuss the limitations in the conclusion of the work.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "Despite a variety of literature examining the widely used weight decay technique, there still lacks precise understanding of its mechanisms in the context of deep neural networks training. The idea to consider its role from the dual perspective of over- and under-training regimes is interesting and indeed relevant for modern deep learning. The paper is well written and structured accordingly, its claims are stated clearly and justified by multiple experiments. The mutual role of weight decay and learning rate acting as an implicit regularizer is intriguing while intuitive, especially in light of prior work studying the effective learning rate in scale-invariant architectures and promoting the significance of large step sizes for better generalization. It is remarkable that in the context of LLM training, WD provides other benefits: it implicitly controls the pace of optimization, as shown both theoretically by examining a simplified Adam update and experimentally, and empirically prevents training instabilities.", "weaknesses": "**Over-training regime**\n\nMy personal expertise is the over-training regime, hence most questions and weaknesses are attributed to the first part of the work.\n1. The mechanism, by which WD helps in the over-training regime via constraining the optimization on a sphere, as described in lines 137-138, is unconvincing. Would similar reasoning be applicable to other loss functions not requiring infinite increase of parameters norms, such as MSE, or fully scale-invariant models, which also benefit from the use of WD during training but do not depend on the parameters norm by definition?\n2. Equation 3 involves a series of unrealistic approximations. First of all, where did the total loss gradient go after the first approximation? According to the previous assumption, the loss is separated from its minimum value by WD, thus its gradient must be substantially non-zero. Second, the assumption that the first derivative is approximately constant across all data points is too stringent as the loss on individual objects has a heavy-tailed distribution in practice, where the mean and median values may be very distinct.\n3. The Gaussian approximation of the SGD noise is questionable [1-3] and must be justified experimentally or at least by citing related work.\n4. Conjecture 1 is formulated in a concrete form, however only proxy values are measured to justify it. (see Question 3 below for further detail). \n5. It is unclear how regularization of model Jacobian and generalization are related in the first place, as trends plotted in Fig. 3 and 4 could be mere correlations caused by other intrinsic factors [4]. In particular, in Figure 5, the authors note that the Jacobian norm trends for EMA and fine-tuning are inverse, making it hard to conclude whether the Jacobian norm is regularized or actually fostered during the large-LR stage of training.\n6. Dependence of the training dynamics on the product of learning rate and weight decay is not novel and was particularly studied, for example, in [5] where this value is called intrinsic learning rate.\n\n**Under-training regime**\n\nOverall, I like this part of the work more, especially the derivation of the effective learning rate for Adam and its empirical examination. However, I still have a few concerns.\n\n7. The informal explanation of the AdamW training dynamics using bias-variance reasoning is vary vague and not supported by any experiments. I suppose that similar linear least-squares arguments could be similarly applied to the over-training regime, where WD is expected to play a different role according to the claims of the paper.\n8. Although empirically validated, no explanation is offered as to how the use of WD allows avoiding numerical instabilities in LLM training.\n\n[1] Simsekli, Umut, Levent Sagun, and Mert Gurbuzbalaban. \"A tail-index analysis of stochastic gradient noise in deep neural networks.\" International Conference on Machine Learning. PMLR, 2019.\n\n[2] Panigrahi, Abhishek, et al. \"Non-Gaussianity of stochastic gradient noise.\" arXiv preprint arXiv:1910.09626 (2019).\n\n[3] Zhou, Pan, et al. \"Towards theoretically understanding why sgd generalizes better than adam in deep learning.\" Advances in Neural Information Processing Systems 33 (2020): 21285-21296.\n\n[4] Andriushchenko, Maksym, et al. \"A Modern Look at the Relationship between Sharpness and Generalization.\" International Conference on Machine Learning. PMLR, 2023.\n\n[5] Li, Zhiyuan, Kaifeng Lyu, and Sanjeev Arora. \"Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate.\" Advances in Neural Information Processing Systems 33 (2020): 14544-14555.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed comments.\n\n1\\. The main focus of our work is classification because, at the core of our examination, are the unique properties of exponentially tailed loss functions, such as cross-entropy (CE), as explained in Section 2.1. Therefore, our reasoning does not extend to MSE. Regarding scale-invariant models, [1] shows that training with SGD and WD also leads to the norm stabilizing. Additionally, in Appendix C.6, we investigate training scale-invariant models on the sphere as a proxy for WD. Our empirical results show analogous phenomena in this setting: stabilization of the loss function, a decrease in the Jacobian norm with a correspondent improvement in test error.  \n\n2\\. See points (1,2,3) in the main response.\n\n3\\. See (4) in the main response.\n\n4\\. The conjecture has been formulated for clarity of exposition and summarize the main message of this part of the paper.  \n\n5\\. Previous works ([6], [7], [8]) already show that explicitly regularizing the Hessian or Jacobian of the network is beneficial for generalization. Regarding the discrepancy between fine-tuning and EMA, we would like to point out that in both cases, the Jacobian norm decreases during the large LR phase and is never fostered. The difference arises with different LRs, where the final Jacobian norm is expected to be smaller for larger LR, as seen with EMA. As explained in Section 2.3, fine-tuning is not ideal for studying dynamics with large LR because it may deviate from the mean process, potentially optimizing a different objective. Therefore, we believe that the results with fine-tuning do not compromise the validity of our conclusions regarding the mean process and EMA.\n\n6\\. We agree that the role of the product of LR and WD is not novel, nevertheless in [1] this product is not related to generalization but only to the effective speed of learning and to the equilibrium distribution. Therefore, the mechanism through which different intrinsic LRs affect the generalization of the correspondent equilibrium distribution remains an open question which we instead try to address. Moreover, the work heavily relies on having fully scale-invariant models, which is not needed in our framework therefore making it more general. \n\n7\\. Our bias-variance decomposition applies only to single-pass SGD, as multi-pass SGD lacks sample independence, invalidating the decomposition in the overtraining regime. We verify the bias-variance tradeoff through fine-tuning (FT) at various iterations (Fig. 6), identifying phases where bias or variance dominates. Early FT, where variance is high, reduces the variance revealing a sufficient contraction of the bias, aiding later phases. In contrast, late FT has minimal impact as the variance is already low. This tradeoff also explains the peculiar loss profiles with high ELR, showing higher initial loss and lower final loss (lines 341-346).\n\n8\\. We suspect that the numerical instability comes from `bfloat16`'s limited precision, which can occur when different components in the network with varying scales are added together (e.g., values in the residual stream). Since the `bfloat16` data type has the same exponent size as `float32`, high-weight norm alone should not pose a problem. However, the precision of `bfloat16` is very limited, e.g., `bfloat16(256.0) + bfloat16(1.0)` is equal to `256.0` instead of `257.0`. We suspect this precision limitation is the primary problem in `bfloat16` runs without WD. \n\n## Questions: \n1. The spikes in the training loss do not necessarily need to be reflected in the L2 norm of the weights. These spikes can occur due to the trajectory passing through sharp regions of the landscape, where the loss value changes significantly even with small changes in the weights. Consequently, the L2 norm may remain relatively stable while the loss exhibits large fluctuations.\n\n2. We thank the reviewer for pointing out the missing factor $(1-\\eta \\lambda)$. The fact that $g_t$ can be written as a sum of random vectors stems from our Gaussian approximation and the structure of SGD's covariance, where its noise is spanned by the gradients [9]. We can consider the random vector $z_t := \\frac{1}{\\sqrt{N}} \\sigma_{\\eta,\\lambda}(w_t) \\sum_{i = 1}^{N} \\nabla h_{w_t}(x_i) \\, \\xi_i^t$  with  $\\xi_i^t \\sim \\mathcal{N}(0,1)$ and verify that $z \\sim \\mathcal{N}(0,\\Sigma_{w_t})$. First of all it is clear that $\\mathbb{E}[z_t] = 0$, furthermore we can compute the covariance by  $\\mathbb{E}[z_t z_t^\\top] = \\frac{1}{N} \\sigma^2_{\\eta,\\lambda}(w_t) \\sum_{i = 1}^{N} \\sum_{j = 1}^{N} \\nabla h_{w_t}(x_i) \\,  \\nabla h_{w_t}(x_j)^\\top \\mathbb{E}[  \\xi_i^t \\xi_j^t ]$   the latter double summation is non zero only when $i=j$ therefore  we recover the covariance of SGD. \n\n3. Sampling from the stationary distribution of SGD would be interesting, but we lack the analytical form of this distribution, making MCMC sampling inapplicable. We argue that the SGD trajectory at stationarity corresponds to samples from a stationary distribution, whose mean is a stationary point of the regularized loss in Equation 5. Although we do not provide an analytical expression for it, one approach could be MCMC sampling from the Gibbs measure using the Jacobian-regularized loss as an energy function. However, there is no guarantee that this would be the stationary distribution of SGD, and this method is computationally demanding due to the prohibitive cost of calculating the Jacobian for the entire dataset.\n\n4. Yes, Fig. 4 also reports EMA results. Notice the axis labels (\"Test Error EMA\"), we will clarify in the caption.\n\n5. We derive the effective LR of Adam by leveraging its (asymptotic) similarity with SignGD, unlike [16] which directly derived it for Adam. This leads to a *simpler* estimate of the ELR, which shows excellent empirical agreement (Fig. 7). Our estimate still requires independence from the scale, which applies with scale-invariance or when the norm remains approximately constant", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their suggestions and feedback on our manuscript, which will help improve the quality and clarity of our work. We will incorporate your comments into the revised version of the manuscript.

We acknowledge that we did not adequately justify our approximations by citing the related work that introduces these approximations. Specifically, we approximate the noise of SGD with a Gaussian and make certain approximations on the covariance. 
These approximations are widely recognized and employed in a substantial body of research. We would like to clarify the following points and provide additional experiments, which can be found in the attached pdf, to support our approach.

1. **Assumptions:** our series of approximations does not need to hold along the entire trajectory of SGD but only in the stabilization phase. This detail is important and we will clarify it in the manuscript. 

2. **Negligible total gradient:** The term $L(\theta) \nabla L^\top(\theta)$ being negligible in equation 3 comes from the observation that the gradient noise variance dominates the square of the gradient noise mean. This fact has been used in previous works [2], [3], [4] and in particular [3] and [13] empirically verify it. 

3. **Constant First Derivative:** The assumption of the first derivative being approximately constant across all data points, sometimes referred to as "decoupling approximation" [2], has been empirically verified for classification in [2]. Furthermore, we performed additional experiments to verify the decoupling approximation in our setting by comparing the spectrum of the SGD covariance with our approximations. In particular, we use Stochastic Lanczos Quadrature [15] to estimate the empirical spectral density of the SGD covariance with and without the decoupling approximation, during the large-LR phase. The experiments are performed for ResNet-18 trained on the cifar10 dataset for different combinations of learning rate and weight decay which are used in the manuscript. The results (reported in the attached pdf) illustrate a substantial overlap in the two spectra which serves as a validation of the reliability of our approximation. 


4. **Gaussian approximation:** We agree that the validity of the Gaussian approximation for the noise of SGD has been long questioned. Nevertheless, a substantial body of research has built upon this approximation and verified its validity ([1], [5], [10], [11], [12]). In particular [5] demonstrated how modelling the SGD noise by a Gaussian *is sufficient* to understand its generalization properties which is the main interest of our work. 

We will ensure these justifications and the additional experiments are added to the manuscript, with particular attention to the following references, which we report here to use in the responses below.

[1] Li, Z., Lyu, K., and Arora, S. "Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate."

[2] Mori, T., et al. "Power-law escape rate of SGD."

[3] Zhu, Z., et al. "The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects."

[4] Jastrzębski, S., et al. "Three factors influencing minima in SGD."

[5] Li, Z., Malladi, S., and Arora, S. "On the validity of modeling SGD with stochastic differential equations (SDEs)."

[6] Liu, Y., Yu, S., and Lin, T. "Regularizing deep neural networks with stochastic estimators of Hessian trace."

[7] Hoffman, J., Roberts, D.A., and Yaida, S. "Robust learning with Jacobian regularization."

[8] Sokolić, J., et al. "Robust large margin deep neural networks."

[9] Andriushchenko, M., et al. "SGD with large step sizes learns sparse features."

[10] Smith, S., Elsen, E., and De, S. "On the generalization benefit of noise in stochastic gradient descent."

[11] Xie, Z., Sato, I., and Sugiyama, M. "A diffusion theory for deep learning dynamics: Stochastic gradient descent exponentially favors flat minima."

[12] Li, Z., Wang, T., and Arora, S. "What happens after SGD reaches zero loss?--A mathematical framework."

[13] Saxe, A.M., et al. "On the information bottleneck theory of deep learning."

[14] Zhang, H., Dauphin, Y.N., and Ma, T. "Fixup initialization: Residual learning without normalization."

[15] Yao, Zhewei, et al. "Pyhessian: Neural networks through the lens of the hessian." 

[16] Roburin, Simon, et al. "Spherical perspective on learning with normalization layers."

Meta-review: This paper studies the role of weight decay in the modern learning regimes of over-training computer vision models and under-training language models. The authors show that the weight decay does not act as a simple explicit regularizer, but as a component that affects the training dynamics in different ways for over-training and under-training. 

This paper provides new insights that contribute to the ongoing efforts to explain the deep learning practice. The authors base their claims by extensive experiments and their explanations use nonrigorous mathematical formulations of the SGD dynamics. All of these, sufficiently support the authors’ claims. 


The reviews reflect that the contributions of this paper are of sufficient interest and significance. The provided reviews are thorough, and the authors succeeded to improve their paper during the rebuttal and discussion with the reviewers. Accordingly, Reviewers S7kS and tBnp increased their review scores. It is worth to elaborate here on the following points that appeared in the Authors-Reviewers discussion and later in the AC-Reviewers discussion: 
1. *Actionable insights and guidelines:*  This paper adds to the foundational understanding of the deep learning practice. Therefore, the understanding gained from this paper can help experts in machine learning and optimization to further improve the practical guidelines related to weight decay. In addition, the authors' results add to the existing literature by further emphasizing the role of the product of the weight decay and learning rate hyperparameters, which is relevant for tuning these hyperparameters in practice. Following my discussion with the reviewers, I believe that these contributions are of sufficient interest for publication. 
2. *Assumptions and approximations:*  While the reviews considered some of the mathematical approximations as potential weaknesses of this work, the authors emphasize in their rebuttal and discussion that the approximations need to hold only in the stabilization phase and not throughout the entire SGD optimization trajectory. This seems adequate, considering that this paper does not intend to provide a rigorous theory but to empirically explain the practice.


Based on reading the paper, the reviews, the authors’ rebuttal and their discussion with the reviewers, as well as my discussion with the reviewers, I find this paper of sufficient quality and significance for publication. My recommendation is therefore to accept this paper.
Decision: Accept (poster)