$\boldsymbol{\mu}\mathbf{P^2}$: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling

Thank you for carefully reading our paper, providing detailed feedback and your help in improving the presentation of our results. We take your concerns about clarity and presentation seriously as we would like that a large audience is able to appreciate our results. If we are able to address some of your concerns and manage to improve clarity, we kindly ask you to consider updating your score as you judged the content overall positively as ‘original’ and ‘significant’.

Lack of figures. We will include a figure that visualizes the phases of all bcd-parameterizations. After fixing the initialization and learning rate scalings, there is indeed a simple way to visualize all unstable, effective SGD, perturbation non-trivial and effective perturbation regimes by drawing a quadrant in a 2D plane (see the pdf attached to the global response).

Familiarity with Tensor Programs. As our theory relies on Tensor Program (TP) theory, including it is hard to avoid. However we agree that readers unfamiliar with it should also be able to intuitively understand the results. We will try to reduce the focus on TPs in the main body of the updated manuscript, and highlight spectral and intuitive conditions more (see also the last paragraph in this response, the answer to reviewer egBf and our global response).

Distinction between non-vanishing and effective perturbations. For the reader’s convenience, we will also include a spectral condition on the weights $||\varepsilon^l ||\_\ast/||W^l ||_\ast =\Theta(1)$ for all l, that essentially states that the effect of the weight perturbation on a layer’s output should be of the same scale as the original output of that layer. We will highlight that non-vanishing perturbations are achieved if and only if at least one layer is effectively perturbed. However we want all layers to be effectively perturbed, otherwise the layer’s perturbation could be set to 0, which would save computation. Hence non-vanishing perturbations is not the correct notion to measure but effective perturbations are, if we want to achieve the best layerwise perturbation scaling for SAM. Effective perturbations really measure an individual layer’s perturbation effect on the output. We hope this resolves the confusion.

Lacking definitions and notation table. We agree that we should provide a complete list of definitions used in the paper, and will do so in the updated version of the manuscript. We also love the idea and will include a table that summarizes the notation. Just like the proofs have to be moved to the appendix, the full formal definitions are quite spacious and can therefore not be included in the main paper, whereas we try to give a shorter and more intuitive presentation of the relevant terms in the main text.

How to use $\mu P^2$ in practice. We hope the following changes facilitate the understanding and adoption of $\mu P^2$ in practice. We will:

(1) explain how to set layerwise learning rates and perturbation radii to achieve $\mu P^2$ in the main text,

(2) make Pytorch code to reproduce all of our experiments publicly available upon acceptance,

(3) refer to the pseudocode in Appendix E.8 and rewrite that appendix for improved clarity to provide alternative implementations and perspectives, using the mup-package and the spectral perspective.

Concretely, concerning (1), each layer either behaves input-like, hidden-like or output-like. At width $n$, $\mu P$ is implemented by scaling the learning rate in each layer to $\eta\cdot n^{-c_l}$, where $n^{-c_l}$ can be read off from the table below. In addition, our results show that for SAM in $\mu P^2$, the global perturbation radius should be scaled as $\rho\cdot n^{-½}$, and in SAM’s weight perturbation step, the layerwise gradients should be multiplied by the scalar $n^{-d_l}$, where $n^{-d_l}$ can be read off from the table below (as was provided in Table 1 (right) for several variants of SAM).

Now to apply $\mu P^2$, we will explain the following steps in the main text:

1. Parameterize the network and the SAM update rule according to $\mu P^2$ as explained above.
2. Tune the learning rate $\eta$ and perturbation radius $\rho$ on a small model.
3. Train the large model only once using the optimal learning rate and perturbation radius from the small model.

More intuitive way to present $\mu P^2$. We find ourselves unable to replace Tensor Programs in our analysis, but the spectral perspective as alluded to by Reviewer egBf provides a complementary perspective to understand our results. However, the analysis of SAM is necessarily more complicated than that of SGD or ADAM due to layer coupling through the joint gradient normalization in SAM’s denominator. See our response to Reviewer egBf for more details. Related literature has shown that a careful signal propagation analysis both forward and backward is necessary to not get stuck to an analysis at or close to initialization like the NTK (Jacot et al., 2018), because of scaling mismatches. While writing out the computations for SAM learning using TP rules is quite technical, a TP analysis is conceptually a simple and reliable way to find the correct scalings: The TP framework intuitively just states that a matrix vector product either introduces a factor $n^{½}$ when matrix and vector are sufficiently independent, or a factor $n$ when matrix and vector are sufficiently correlated. Wrong update or perturbation scalings can then be corrected with layerwise scalings of the learning rate and perturbation radius.

References:

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks, NeurIPS 2018.

Thank you for carefully reading our paper and providing detailed feedback. We are delighted about your overall positive evaluation of our work. If we are able to address some of your concerns, we kindly ask you to consider updating your score as you positively alluded to both our rigorous theoretical analysis that extends the community’s understanding of SAM and improves over standard SAM as well as the extensive experiments we conducted.

Discussion of abc-parameterizations. The footnote on page 5 discussed $abc$-parameterizations and referred to Appendix E.7, where we provide a more detailed comparison. $abc$-parameterizations only differ from $bcd$-parameterizations in that they do not consider perturbation scalings $d$, but introduce additional layerwise weight multipliers $a$ in the architecture. These weight multipliers result in equivalence classes of abc-parameterizations, out of which we pick the representative with $a=0$ in all layers. In this way, $bc$-parameterizations without the perturbations effectively recover all $abc$-parameterizations, reducing them to their essence: Each $abc$-parameterization is effectively just a layerwise initialization and learning rate scaling. In the updated version of the manuscript, we will explain this in the main text and refer to Appendix E.7 for more details. We believe that omitting weight multipliers $a$ in the definition of $bcd$-parameterizations improves clarity, as we already have to introduce an unavoidable complication of introducing perturbation scalings $d$.

Role of batch size. As for SGD, fixed batch size is covered by our theory. Since small batches are particularly useful for SAM, we do not see additional value in studying the limit $m\to\infty$. In the updated version of the manuscript, we will include a comment that fixed batch size is covered by our theory. In our experiments we make sure to always use batch size 64 on CIFAR10. As our focus is width-scaling, changing the batch size would introduce confounding effects. By achieving width-independent SAM dynamics, we expect low batch size to also be beneficial for generalization in $\mu P^2$ at large width whenever it is at small width, but a systematic analysis and understanding constitutes an interesting avenue for future work.

Layerwise perturbation scaling is not novel, but its rigorous understanding is. We do not claim to be the first to propose layerwise perturbation scaling. Instead we aim to provide the first rigorous infinite-width theory that informs practice how layerwise perturbations should be scaled without having to tune all layers individually, which would be much more costly. This paper rigorously resolves the question how exactly layerwise perturbations should be scaled as we scale up model size. We are not aware of other work that makes meaningful progress in this regard, but would be very interested in further related work.

In bcd-parameterizations, what distinguishes SAM from SGD? In terms of understanding SAM, our scaling analysis has shown that standard scaling (1) becomes unstable if $\rho$ is held fixed if we scale up the width, (2) can at most perturb the last layer in wide neural networks and hence could instead be replaced by SGD (=set perturbations to 0) in all previous layers, and (3) can also recover width-independent perturbation dynamics (=effective perturbations) with the correct layerwise adjustment.

Generalization. We consciously analyze the SAM update rule without alluding too much to sharpness, as contributing to the discussion about the connection between sharpness and generalization is not our goal. Indeed we do not make claims about generalization just like other Tensor Program theory. Yang and Hu (2021) show that $\mu P$ is necessary to achieve maximal stable updates in all layers in the infinite-width limit. If feature learning is necessary to achieve optimal generalization, then $\mu P$ will outperform other parameterizations at large width. Similarly, our goal are width-independent perturbation scalings as a necessary requirement for effective SAM dynamics in the infinite-width limit. If SAM enhances generalization over its base optimizer, then $\mu P^2$ can be expected to outperform other parameterizations at large width. However, as we mentioned in the ‘Future work’ section, we agree that further insights into generalization are very relevant and interesting.

Depth transfer. While for SGD and Adam, depth transfer has been achieved in several papers with a simple $1/\sqrt{`depth`}$-scaling of the residual connections and adapted learning rate scaling (see the extended related work appendix A), depth transfer for SAM remains open and is a question that we are currently working on, that lies beyond the scope of this paper. We will mention this question in the ‘Future work’ section.

References:

Greg Yang and Edward Hu. Feature learning in infinite width neural networks, ICML 2021.

Thank you for your detailed response. While my original assessment and score remain unchanged, I encourage the authors to incorporate this discussion into the revision.

We sincerely appreciate your thoughtful review of our paper. We are delighted about your positive evaluation of our work and are grateful for your insights which have been invaluable in enhancing the accessibility and clarity of our work.

Improving the presentation. In the main paper, after line 271, we provide an intuitive condition that perturbations should scale like updates under $\mu P$. From this condition, it follows that a layer is effectively perturbed iff the spectral condition holds for perturbations (in TP architectures). We agree with you that emphasizing the spectral perspective will improve the accessibility of the paper and we will highlight it in the updated paper (see global response).

However, unlike for SGD or Adam, deriving correct layerwise perturbation scalings for SAM from the spectral condition is complex. This complexity arises from the gradient normalization (in Frobenius norm) in the perturbation's denominator, which couples all layer scalings and has been shown to be practically relevant [1]. To simplify the analysis, we require the perturbation's denominator to be $\Theta(1)$ in the definition of bcd-parameterizations (bcd-params). This allows the numerator to be scaled just like updates, under layer coupling constraints. Further updates to enhance clarity in terms of practical use can be found in our response to Reviewer dLWn.

Relevance of TP theory. While we agree that the spectral condition provides a useful and accessible perspective, note that a rigorous justification of the spectral condition [3] crucially relies on NE⊗OR⊤ program (TP) theory. We would also like to point out that a recent ICML 2024 workshop spotlight paper “On Feature Learning in Structured SSMs” shows that the spectral scaling condition fails to achieve feature learning in Mamba layers, which cannot be represented as TPs. This highlights that the spectral scaling condition is not universal and requires further validation of its underlying assumptions.

The significance of our TP theory extends beyond $\mu P^2$. It enables a comprehensive characterization of SAM's scaling behavior and allows for the analytical derivation of infinite-width limits for all bcd-params, including standard SAM. Furthermore, any scaling analysis for SAM introduces additional complexities over SGD/Adam, not only because of the layer coupling due to SAM's denominator. A priori, it is unclear how perturbations and updates interact even to provide conditions for stable learning. Interestingly, our findings reveal that the conditions on perturbation scalings are largely independent of initialization and learning rate scalings, a result that only becomes apparent in retrospect.

Given these challenges, our TP analysis is conceptually simple: we write SAM learning (two forward and backward passes for each update) as a TP to understand how evaluating the gradient on gradient-perturbed weights affects the activations and output function and rigorously track all update and perturbation scalings. This allows us to derive rigorous statements under weak and simple assumptions.

Uniqueness claim. We study infinite-width limits of bcd-params under fixed depth and training time. In this setting, Theorem 11 indeed shows that $\mu P^2$ is the unique stable bcd-param (up to smaller last layer initialization) that achieves both maximal stable updates and effective perturbations in all layers in the infinite-width limit. If the concern is about non-uniqueness related to equivalences (e.g., using weight multipliers), we addressed this in the paper: weight multipliers are covered in the footnote on page 5 and explained in detail in Appendix E.7. We will be sure to further emphasize this in the main text. If the reviewer has knowledge of other parameterizations to achieve effective perturbations, we are very interested in learning about them.

Extensions to other layer types. We will omit the word “straightforward” and just refer to the respective appendix, mentioning that many common layers either behave like input, hidden or readout. In the respective appendix, we will also explain in more detail why our theory and derived scalings extend to common layer types even under SAM's layer coupling.

Missing related work. Thank you for the missing reference. We will discuss it in the updated manuscript. The ideas of compositionality, perturbations and automatic update scaling are intriguing and related. While the analytical strategies might become helpful in developing simpler analyses in the future, we do not see an immediate path to derive SAM width-scaling analysis using these ideas. In particular, the assumption that perturbations are full rank is violated for gradient-based perturbations on small mini-batches (as used for SAM), and the condition number of random matrices explodes with increasing width already at initialization [2], which may render the upper bound in Theorem 1 quite loose at large width. In contrast, the assumptions required for our analysis are easy to understand.

Minor comments on Figures. We will correct the unclear legend for $\mu P$-global in Figure 1. Concerning Figure 2, we thought that the Frobenius norm tending to 0 is an even stronger statement: The effect of the perturbations on the activations vanishes even if we accumulate the perturbations over all directions. Arguably, this is unsurprising, as the two norms are equivalent because the gradient on a mini-batch of size 64 is always low-rank. We also found it striking that there is a width-independent limit for last-layer perturbations, even in Frobenius norm (explained in Appendix G.1).

References:

[1] Monzio Compagnoni et al. An SDE for modeling SAM: Theory and insights, 2023.

[2] Alan Edelman. Eigenvalues and condition numbers of random matrices, 1988.

[3] Yang et al. A spectral condition for feature learning, 2023.

Thank you for carefully reading our paper. We are delighted about your overall positive evaluation of our work. As other reviewers have pointed out, this paper is already quite dense and contains extensive experiments. Hence we would like to defer experiments on further SAM variants to future work.

We thank the reviewer for the engaged discussion and appreciate their acknowledgement of our efforts to improve the accessibility of our paper. If it was not clear in our original response, we want to emphasize that we like the spectral perspective, which is why we will further highlight it in the revision. We believe that it has a potential for simplification and generalization.

Improving the exposition. In addition to the changes promised in our original global response, the following changes aim to further improve the clarity of presentation:

• In Section 4.2, we will clearly discuss the spectral perspective and how to achieve effective perturbations with spectral arguments in more detail.
• We will emphasize in the setup that TP theory relies on the assumption that width dominates training time and batch size.
• When discussing the simpler-to-analyse version of SAM without layer coupling (the gradient normalization applied layerwise instead of jointly), we will add ablations to study its $\rho$-transfer and generalization properties for preliminary insights whether this version has practical potential.

Before we respond to the remaining comments, we aim to first clarify the primary objective of our paper.

Theoretical understanding vs algorithm design. Our main goal in this work is to analyze the infinite-width behavior of existing SAM update rules as they are implemented in practice. Specifically, we aimed to analytically derive infinite-width limits and identify desirable/undesirable behaviors of these rules. For example, we use TP theory to prove statements such as "muP with global perturbation scaling yields vanishing perturbations in all hidden layers in the infinite-width limit (Proposition 2)"

While the development of $\mu P^2$ is an interesting and extremely valuable outcome of this theoretical investigation, it is not our sole or primary objective. From the reviewer's comments, we infer that they may have assumed that our primary aim was to derive $\mu P^2$ or a scaling rule that allows width-transferable perturbation radius.

In line with our previous response, we still maintain that deriving $\mu P^2$/the correct width-dependent perturbation scaling is possible via spectral scaling conditions potentially supported by simpler analyses, that first scaling SAM's denominator to $\Theta(1)$ simplifies the analysis, and that in a version of SAM without layer coupling each layer can be individually scaled like updates (there would just be layerwise $d_l$ outside of the normalization and no need for additional global perturbation scaling $d$). We will make sure that we further clarify these points in the revision.

However, we respectfully emphasize that our focus is primarily on advancing theoretical understanding which we strongly believe to be inherently valuable. For example, our analysis can be used as a starting point to understand generalization properties of infinitely wide linear neural networks trained under SAM. Practical improvements like $\mu P^2$ emerge as secondary, albeit welcome, consequences of our analysis.

We hope this clarification helps better contextualize our work and the remainder of our response.

Statement about rigor. We are confused by the reviewer's comments on our statement about rigor. First, we want to emphasize that we have never stated or implied anywhere in our response or in any other form that either spectral normalization or algorithm design based on majorization constitutes a non-rigorous method. Second, could the reviewer kindly clarify which analysis with upper bounds (without matching lower bounds) are being referred to in this context? The spectral scaling paper does not use any upper bounds without matching lower bounds. It first uses some elementary arguments to intuitively motivate the claim that spectral condition implies feature learning in the infinite-width limit and then uses TP theory to formally prove that muP is the unique abc-parameterization that satisfies the spectral condition, thereby formally connecting it to the infinite-width theory established in the TP framework. Therefore, we are very confused about which analysis the reviewer is referring to here.

To further clarify the statement in our original response, in the context of our paper, where the goal is to understand infinite-width behavior, we are interested in proving the following: “$\mu P^2$/spectral scaling implies effective perturbations in the infinite-width limit.” Our response aimed to state that the spectral condition paper [3, Appendix B] utilizes TP theory to make a similar statement about feature learning in the infinite-width limit and therefore is rigorously justified (in the aforementioned context). We want to emphasize that the essence of our response was to highlight the utility of TP theory. We did not claim that it is superior to some another analysis that motivates the spectral condition.

dear egBf,

The dispute here is definitely a scientific one, no worries.

1/ Regarding the weakness you are heavily relying on now, would it not possible that a paper written on optimization for ML using variational inequalities may assume optimization knowledge on variational inequalities?

On this vein, we already provided concrete ways of improving accessibility along spectral perspective, even promising to define bigO and Omega notations. Is this not sufficient?

2/ Can you please respond to our comment if you are still not sure on what concrete basis we believe the spectral perspective is non-rigorous in absence of TP or infinite-width theory?

"the reviewer's statement that equations (4) and (5) from the spectral paper have no matching lower bounds appears incorrect. The paper makes simplifying assumptions that the terms in the expansion do not cancel each other and under this assumption provide matching lower bounds in eq. (8) for eq. (4), and, citing from the same paper: ‘Combining Equation (10) with our matching upper bound from Equation (5), we conclude that $||\Delta h_l(x)||_2=\Theta(\sqrt{n_l})$ as desired.’ Also in the appendix, via TP theory, these lower-bounds are further supported."

In particular, please explain how we can support the assumption "the terms in the expansion do not cancel each other" of the spectral perspective in our minimax/SAM context without the TP theory support?

3/ Regarding our uniqueness claim, we can qualify it under restricted stability constraints, which address your concern.

Dear authors,

First of all, thank you for the feedback. I take your concerns seriously and will separately contact the AC for feedback on my conduct. I think that since the reviews are all recorded, I would expect that if there has been a conduct violation it will be caught by the AC and dealt with.

I think that what might help is for me to explain why I believe the question of the perceived rigor of TP arguments versus the spectral perspective is central to my evaluation of your paper. First of all, both these topics are discussed in your paper. But second, both myself and reviewer dLWn found the reliance of your paper on Tensor Programs arguments to be a serious weakness. Reviewer dLWn wrote as a weakness that "The paper is written in a way which assumes that the readers are very familiar with Tensor program" and cited presentation issues as their main reason for a negative review. Therefore, if there is a simpler way to derive your results, it is of critical importance whether or not that approach is rigorous in the absence of Tensor Programs. After all this discussion, I am still not sure on what concrete basis you believe the spectral perspective is non-rigorous in absence of TP or infinite-width theory.

I also want to note that I raised other issues with your technical results, for instance the uniqueness claim. When pushed by the authors, I provided a counterexample to the claim. Upon providing the counterexample, I then felt that the authors dismissed my concern rather than taking it seriously.

Although I praised the authors' work for some positive aspects, I made it clear in my review and repeated engagement with the authors that I would be prepared to either raise or lower my score depending on whether the authors could address my concerns. Due to the message that the authors were still "confused" and "baffled" (after I thought I had made my final assessment) I realized that the discussion period was not a success in terms of bringing consensus between this reviewer and the authors, I lost trust that the authors could make the necessary changes, and therefore I felt it right to lower my score.

I do feel that the dispute here is a scientific one and not a personal one. I sincerely wish the authors all the best, and good luck

Dear Reviewer egBf,

I would like to flag what the authors are perceiving from these discussions:

1. Power play (you agree with me or else) on a related paper, which we must evaluate in its spirit and not what it actually contains. This kind of inflexible behavior is terrifying to new authors.

2. Clear personal conflict conflated into our work/your review.

Dear Reviewer egBf,

we sincerely regret this discussion to have become so heated as both sides seem to be very passionate about finding and understanding the correct way to scale neural network optimization. Misunderstandings seem to be mostly caused by differing attitudes to achieve this goal. Our goal was to analyze past and current SAM practices, while the reviewer has more visionary perspectives of designing the optimization and scaling procedures of the future.

For our analysis, we have decided to use TP theory, which comes with its limitations which we will point out very clearly, but still seems to us like a legitimate approach to take, as our results demonstrate. We certainly agree that there is a lot of follow up work to be done; there is a lot to do in terms of understanding long training dynamics.

In our own interest, we will try to clarify the used TP framework as much as possible, as well as highlighting the practical and more intuitive ways to arrive at the corrected scaling, but we tried to argue that some of the difficulties in our analysis stem from the layer coupling in the original SAM algorithm, which can only be resolved by changing the optimization algorithm, which again points to our differences in algorithm design versus analysis.

We may also have differing views along the theory versus practice dimension. Our statement on rigor was only related to the infinite-width theory, we are doing in this paper. We embrace a multitude of perspectives and agree that the approach of making assumptions that are empirically verifiable followed by logical arguments may yield further insights into long training dynamics than TP theory can, as previously stated with a potential for simplification and generalization. We just disagree that TP theory is not a valid approach for our analysis.

To engage with your comment about the spectral paper, we agree that all assumptions should be made crystal clear. The spectral paper does this and the original TP papers also state the underlying assumptions but are perhaps hard to parse. To us, it appears that the spectral paper makes strictly stronger assumptions than the TP papers, including width dominates training time, batch size, and depth. The spectral paper does not assume that the activation function is smooth (required by TP) but instead assumes that co-ordinate wise non-linearities do not affect the scaling of vectors (Assumption 2).

Regarding whether the act of making assumptions itself makes things non-rigorous: we believe the answer to this question to be quite nuanced. In the extreme case, one can simply assume the result and then prove that the result holds. To our understanding of the word, this does not make it rigorous. Our statement “a rigorous justification of the spectral condition [3] crucially relies on NE⊗OR⊤ program (TP) theory” meant that the simplifying assumptions like Assumption 1 and line above eq. (10) of the spectral paper (used for providing the lower bounding arguments) can be proven using TP theory. This does not mean that these assumptions are not useful to understand what is essentially happening. On the contrary, we strongly believe that the argumentation in the spectral paper to arrive at the correct scalings from the assumptions made is indeed logical and useful.

Overall, we hope that in a year both parties will be able to look back at this discussion with empathy for each other, as we follow the same goal and seem to appreciate each other’s efforts. Please feel free to score the paper as you genuinely feel the paper and promised changes deserve.