We thank the reviewer for the detailed comments.

1. 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.

2. See points (1,2,3) in the main response.

3. See (4) in the main response.

4. The conjecture has been formulated for clarity of exposition and summarize the main message of this part of the paper.

5. 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.

6. 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.

7. 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).

8. 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.

Questions:

1. 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.

2. 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.

3. 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.

4. Yes, Fig. 4 also reports EMA results. Notice the axis labels ("Test Error EMA"), we will clarify in the caption.

5. 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

We thank the reviewer for their helpful comments. Below, we address the concerns raised by the reviewer in detail

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".

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.

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.

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).

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.

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.

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?

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.

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.

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.

• 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.
• 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.

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.

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.

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. For 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)$

Similarly 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.

We 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.

Questions:

1. 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.

2. 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.

3. 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.

We thank the reviewer for their helpful comments, which will allow us to improve the manuscript.

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.

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)$

Similarly 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.

Questions:

1. 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.

2. 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.

3. 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.

We appreciate the reviewer's suggestion to enhance the theoretical explanations in our paper. We understand the importance of providing thorough theoretical insights.

• 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.

• 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.

• 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.

To further validate our conjecture, we performed the experiment suggested by reviewer tBnp. Specifically, within the same ResNet18 on CIFAR-10 setting as in our main experiments, we created snapshot ensembles 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 (TVD) between the softmax outputs of the ensemble and the EMA on the Test set $TVD = \frac{1}{2N} \sum_{i=1}^{N} \sum_{j=1}^{C} \left| p^{(i)}_{\text{ensemble}, j} - p^{(i)}\_{EMA,j} \right|$ . The results in Table 1 show a strong alignment in test accuracies, while those in Table 2 indicate a low Total Variation across all combinations. Together, these findings offer further validation for our conjecture.

We believe these additional experiments contribute to strengthening our contributions. While time constraints limited us to running these experiments on CIFAR-10, we plan to extend them to TinyImageNet and include results from both datasets in the revised version of the manuscript.

Table 1: Test Error for Snapshot ensemble and EMA for different values of LR and WD

Table2: Total variation Distance between softmax output of Ensemble and EMA.

[17] Huang, Gao, et al. "Snapshot ensembles: Train 1, get m for free."