Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks

Thank you for the time spent reviewing our work as well as your feedback and suggestions. Below we will try to address the weaknesses you bring up and answer your questions.

Random Walk:

First, we would like to provide further explanation of our random walk setting and why we consider it relevant to real neural networks. We apologize for not making this clearer in the manuscript and will update it to include the following discussion.

Relation between random walk and real training dynamics:

Let's assume we are in the standard empirical risk minimization setting used in deep learning. We have a training dataset $\mathbb{X}$ consisting of points $\\{\mathbf{x}\_{1} \ldots \mathbf{x}\_{N}\\}$, where each point accounts for both an input and a label/target when applicable. Let $\mathbf{\theta}$ be the whole parameter vector of our model, accounting for all layers etc. We can then write our full training loss as: $\mathscr{L}(\mathbb{X}, \mathbf{\theta}) = \frac{1}{N}\sum\_{i=0}^{N} \mathscr{L}(\mathbf{x}\_i, \mathbf{\theta})$ Similarly we can write the loss for a minibatch $\mathbb{B}$ of size $|\mathbb{B}|$ as: $\mathscr{L}(\mathbb{B}, \mathbf{\theta}) = \frac{1}{|\mathbb{B}|}\sum\_{\mathbf{x}\_i \in \mathbb{B}} \mathscr{L}(\mathbf{x}\_i, \mathbf{\theta})$ The full-batch gradient is $\mathbf{g}\_\text{True} = \nabla\_{\mathbf{\theta}}\mathscr{L}(\mathbb{X}, \mathbf{\theta})$ and the mini-batch gradient is $\mathbf{g}\_\mathbb{B}=\nabla\_{\mathbf{\theta}}\mathscr{L}(\mathbb{B}, \mathbf{\theta})$. The noise component of the gradient, $\mathbf{g}\_N$ is then the difference between the two: $\mathbf{g}\_N = \mathbf{g}\_\mathbb{B} - \mathbf{g}\_\text{True}$ Note that this noise component is zero-mean, since the minibatch gradient is an unbiased estimate of the full batch gradient for a randomly sampled $\mathbb{B}$: $\mathbb{E}_{\mathbb{B}}[\mathbf{g}\_\mathbb{B}] = \mathbf{g}\_\text{True}$ Our random walk analysis assumes that the batch gradient is dominated by the noise component, i.e. that: $\mathbf{g}\_\mathbb{B} = \mathbf{g}\_\text{True} + \mathbf{g}\_N \approx \mathbf{g}\_N$ This does of course not hold exactly in practice but may be a reasonable approximation, especially for smaller batch sizes.

We note that the noise component is always zero-mean in expectation which is the main requirement for our random-walk analysis. The presence of normalization also helps ensure that the gradient is orthogonal in expectation even when $\mathbf{g}_\text{True}$ is significant.

Relevance of Random Walk Setting:

We chose this random walk setting for several key reasons:

1. It is a single relatively simple assumption that streamlines the math of the analysis. Without it we would have to make various other approximations or assumptions about the gradients or alternatively introduce unknown quantities describing aspects like the alignment of gradients over time. This would result in a more complicated analysis that is harder to understand as well as equilibrium predictions that may depend on unknown quantities.
2. It allows us to extend the analysis to the unnormalized setting, unlike prior work which has only focused on the fully scale-invariant case. This allows us to provide a broader view of the effects of weight decay in deep learning since most networks are not perfectly scale-invariant in practice.
3. It captures important properties of the gradient unlike simpler models that might for example sample the weight gradient directly from a fixed distribution. Our setup accounts for e.g. the presence of normalization layers so $\mathbf{g}_N$ respects fundamental properties like the orthogonality and the inverse proportionality (Equations 1 and 2 in the paper). For a real optimization problem the distribution of $\mathbf{g}_N$ may change with the progress on $\mathscr{L}(\mathbb{X}, \mathbf{\theta})$, but we believe that if this happens slowly enough compared to the convergence towards equilibrium it won't significantly affect the results (and this would likely be very hard to analyze anyway).

Finally, it is important to highlight that the predictions we derive from analyzing this random walk scenario closely match our measurements of real neural networks in practice. In Appendix G, we explore different scenarios and provide examples illustrating the accuracy of these predictions, showing a close alignment, often within a small percentage margin. We therefore believe our analysis offers valuable insights and predictions applicable to real neural networks, despite being derived for a simplified setting.

We hope this message finds you well. We understand your time is precious and deeply appreciate your efforts in this review process. Since the discussion period is ending soon, we would like to kindly ask for your feedback on our rebuttal. If you feel it addresses at least some of your concerns, we would greatly appreciate it if you updated your score accordingly.

Please let us know if you have any further questions and we will quickly get back to you.

Thank you for taking the time to review our paper and your feedback. We will try to address your concerns and questions below.

On a high level, we believe there may have been some misunderstanding about the main goal of our paper which is understanding the effect of weight decay on the learning dynamics of existing optimizers, not proposing a new SOTA optimizer. Aside from the high level structure of the manuscript, we have tried to further emphasize this in various places e.g. when proposing the RVs "In this section we create an experimental tool to study these phenomena" (referring to the effects of the transient phase and imbalanced rotation) and in the experiment section "Although the RVs have many interesting properties and good performance, we view them primarily as a research tool".

The Advantage of Rotational Optimizers

The main purpose of the RVs in this work is to serve as research tools to further our understanding of weight decay and balanced rotation in deep learning.

We model the RVs to match the equilibrium dynamics of the standard optimizers (with proper normalization) as closely as possible. By design, this leaves no room for changes that could bring significant performance differences except in cases where the original optimizers deviate significantly from the regular "balanced" equilibrium dynamics. When that is the case the rotational variants tend to perform better, see for example the poorly normalized case (imbalanced rotation) and the need for warmup (harmful transient phase) sections of our experiments.

The main conclusions we draw from Table 2 is not that the RVs perform better but rather:

• The transient phase is not needed, the simpler enforced equilibrium dynamics of the RVs suffice for good performance.
• In many cases we can predict what the hyperparameters of the RVs should be (zero-shot transfer). This supports our analytical results since we can predict what the rotational update size should be.

If one wants to use RVs as practical optimizers (not what we propose them for), we believe the main advantage would be simpler and more robust learning dynamics. For example:

• Enforced balanced rotation even in cases with poor / no normalization (more robust to the neural network design).
• Eliminating the transient phase, giving better learning dynamics early in training (likely faster progress, reduced need for warmup).
• Allowing the direct specification of the "effective learning rate schedule" (in terms of the update sizes), otherwise it depends on interactions between the specified schedule with the initialization, weight decay and learning rate used.

On their own these benefits might not be significant enough to warrant switching to the RVs for mainstream, well established / tuned training tasks. Popular architectures are likely to have at least somewhat decent dynamics through "natural selection", even if obtained through trial and error. However, these are desirable properties that may inspire future optimizer design (where there is freedom to deviate arbitrarily from existing optimizers, unlike with the RVs).

Paper Structure

We don't consider this work primarily empirical, all of our experiments aim to validate our analysis or quantify the impact of specific phenomena observed in our analysis such as the transient phase, imbalanced rotation, different roles of learning rate and weight decay etc.

Our analysis reveals a broad range of insights that we empirically explore / validate in our experiments. We believe this structure makes sense for this purpose. If our main goal was to propose a new SOTA optimizer, we agree that this structure would not be optimal and a narrower focus on showing specific advantages would be better. We note that our discussion section further interprets and summarizes our experimental results.

We hope this message finds you well. We understand your time is precious and deeply appreciate your efforts in this review process. Since the discussion period is ending soon, we would like to kindly ask for your feedback on our rebuttal. If you feel it addresses at least some of your concerns, we would greatly appreciate it if you updated your score accordingly.

Please let us know if you have any further questions and we will quickly get back to you.

Thank you for taking the time to review our paper and your feedback. We will try to address your concerns below.

Gap between theory and practice

We believe there may be some misunderstanding here: we find that our theoretical predictions closely match empirical observations and can explain the differences between various optimization algorithms. In the following, we highlight and clarify the results in our paper that support this claim, focusing on Figures 4 and 6.

Figure 4

This figure shows the measured weight norm and angular updates for AdamW and Adam+L2, in a simple system undergoing a random walk with a constant learning rate. We note that these are standard optimizers, not the rotational variants, and will therefore have a transient phase before converging to equilibrium. Our predicted values are derived for the steady state (equilibrium) that occurs after the initial transient phase. The equilibrium values also describe the expected or average behavior, the exact values fluctuate between steps although their expectation (over different random seeds or time) remains constant. We therefore don't expect the depicted predictions to match early in training and the exact instantaneous values in equilibrium can also deviate from them, both are fully consistent with our analysis.

For AdamW all neurons (solid colored lines) and their average (solid black line) converge to the same weight magnitude (left) and angular update (right). In the equilibrium phase, the measurements and our predictions (dashed red line) closely align. In particular, the average (black line) aligns well with our predictions. This is because for AdamW, neurons behave similarly and the black line represents the average expected value.

For Adam+L2 every neuron converges to a different value due to the dependency on the average gradient magnitude. Although not shown in Table 1, Appendix section E thoroughly explains (predicts) this.

We omitted the neuron predictions for Adam + L2 due to their complex dependency on the average gradient norm of each neuron. In a slightly simplified experimental setup, where the gradient distribution is similar across the coordinates (synapses) of each neuron, we can predict the equilibrium values from measured gradient norms.

Figure 6

We make the following two observations.

Left: After extensive tuning, Adam+L2 fails to match AdamW's performance, supporting Loshchilov and Hutter's original findings in their 2019 AdamW paper.

Right: Adam+L2 shows unbalanced average angular updates across layers, unlike AdamW. This aligns with our theory, which predicts such imbalance for Adam+L2 due to varying gradient magnitudes across layers, but not for AdamW. We make this observation for the optimal parameter settings for each optimizer.

General remarks: The angular updates change over time because of the cosine learning rate schedule, the initial transient phase, and towards the end of training, as explained in detail in Section 5 (Scheduling Effects).

We hope this message finds you well. We understand your time is precious and deeply appreciate your efforts in this review process. Since the discussion period is ending soon, we would like to kindly ask for your feedback on our rebuttal. If you feel it addresses at least some of your concerns, we would greatly appreciate it if you updated your score accordingly.

Please let us know if you have any further questions and we will quickly get back to you.

Thank you for taking the time to review our paper and your feedback. We will try to address your concerns below.

AdamW vs Adam+L2

We would like to note that we do analyze the rotational equilibrium of Adam+L2 in Appendix E. Our analysis shows that the rotation depends on the gradient magnitude, resulting in imbalanced rotation when it varies between neurons / layers. This differs from AdamW, where our analysis suggests that the expected average rotation is independent of the gradient norm in equilibrium

Summary of Appendix E: For Adam+L2, the equilibrium weight magnitude depends non-trivially on the average gradient norm. With some simplifying assumptions, we can describe the equilibrium weight magnitude as the roots of a third order polynomial. The closed form solution is quite long and complicated. We were not able to simplify it much in the general case, so we omit it from Table 1. The dependency on the gradient norm does not cancel out when computing the expected angular update in equilibrium. This differs from SGD where the equilibrium weight magnitude also depends on the average gradient norm, but this dependency cancels out when computing the average angular update.

Please let us know if you find this explanation sufficient or if you have any suggestions for how it could be improved. We believe we could also give a more high-level intuitive explanation for the dependency on the gradient magnitude for a specific setting if that would be of interest.

Why does balanced rotation matter?

We note that we do provide some high level intuitions for why balanced updates might be a good heuristic in the intro and why rotation is a natural measurement of an effective update in Section 2.

First we would like to highlight that the standard practice of using a single learning rate across all neurons / layers is a practical but somewhat arbitrary choice. In theory, we could tune a separate learning rate and even a lr schedule for each component (although intractable in practice). We know that a single learning rate is often not optimal, for example AdamW can be seen as SGDM with an adaptive (time varying) learning rate scaling for each coordinate and significantly outperforms SGDM on many problems.

Similarly to the coordinate-wise scaling in AdamW, balanced rotation can be seen as a heuristic that empirically performs well. Despite its performance, we don't believe balanced rotation is likely to be theoretically optimal. Tuning a separate learning rate for each neuron on a given problem seems unlikely to result in perfectly balanced rotation.

We would also like to emphasize that the idea of balanced rotation comes from analyzing the learning dynamics arising from good modern deep learning practices (use of normalization, weight decay etc), it does not come out of nowhere. Our research provides a detailed analysis and empirical evidence to support the effectiveness of balanced rotation. We observe that imbalanced rotation, as seen with methods like Adam+L2 and coarse normalization, often results in poorer performance compared to when rotations are balanced (e.g. through the use of the rotational optimizer variants). We also show that artificially imbalancing the rotation results in worse performance.

We believe that rigorously proving why balanced rotation often outperforms imbalanced rotation is likely to be a challenging theoretical problem. Other forms of coordinate or component wise scaling for example in AdamW remain poorly understood theoretically despite significant effort from the community. Nevertheless, we believe these are valuable observations of practical relevance even if we cannot offer a theoretical explanation for why (perfectly) balanced rotation performs well. Our observations may inspire further theoretical research to explore why and when balanced rotation helps.

We hope this message finds you well. We understand your time is precious and deeply appreciate your efforts in this review process. Since the discussion period is ending soon, we would like to kindly ask for your feedback on our rebuttal. If you feel it addresses at least some of your concerns, we would greatly appreciate it if you updated your score accordingly.

Please let us know if you have any further questions and we will quickly get back to you.

Thank you, we find your work very interesting as well and will make sure to discuss/cite it in a future version of our manuscript.

Please note that we have uploaded a new, revised version of our manuscript. In this version, we have highlighted the changes from the last update in light red for easier identification.